fix the outdated end2end training examples of moe+torchtitan

[rakkit]

as titled. the example here uses some outdated TorchTitan APIs. This PR fixes them and aligns the example with other torchtitan end-to-end usages.

# usage for test fp8_rowwise (default use fp8_rowwise)
python ./torchao/prototype/moe_training/examples/simple_moe_layer.py 
# usage for test fp8_rowwise
python ./torchao/prototype/moe_training/examples/simple_moe_layer.py  --scaling_type fp8_rowwise

# usage for test mxfp8
python ./torchao/prototype/moe_training/examples/simple_moe_layer.py  --scaling_type mxfp8

results:

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3242

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit bf18374 with merge base 03c2d28 ():

NEW FAILURE - The following job has failed:

• PR Label Check / Check PR Labels (gh)
  Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[danielvegamyhre]

Thanks for updating this! Did you confirm it is successfully runnable via copy/paste?

[rakkit]

I made a few changes to set the random seed and allow it to pass scaling_type from CLI.

[rakkit]

@danielvegamyhre yes, the outpout should be

> python ./torchao/prototype/moe_training/examples/simple_moe_layer.py 
step 0 loss: 2656.0
step 1 loss: 2624.0
step 2 loss: 2592.0
step 3 loss: 2560.0
step 4 loss: 2528.0
step 5 loss: 2512.0
step 6 loss: 2480.0
step 7 loss: 2448.0
step 8 loss: 2432.0
step 9 loss: 2416.0

i have already revised the PR message and put the screenshoot there