mxfp8 training: add TP sharding strategy for dim1 kernel

[vkuzo]

Summary:

Enables mxfp8 training with the dim1 triton kernel and TP, in eager mode. In detail:

1. write a sharding strategy for the triton kernel custom op
2. properly wrap the outputs of the custom op in DTensor
3. add a test which enables the triton custom op, and modify the test model sizes since this kernel needs block_size 32

Note that compile does not work yet, seeing https://www.internalfb.com/phabricator/paste/view/P1851219639

Test Plan:

// TP test passes (with eager mode)
./test/prototype/mx_formats/test_mx_dtensor.sh

// unit tests still pass
pytest test/prototype/mx_formats

// torchtitan debug model loss goes down (in eager)
with-proxy CONFIG_FILE="torchtitan/models/llama3/train_configs/debug_model.toml " ./run_train.sh --model.print_after_conversion --training.steps 50 --model.converters mx --mx.recipe_name "mxfp8" --parallelism.tensor_parallel_degree=2 --mx.use_fp8_dim1_cast_triton_kernel

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• -> mxfp8 training: add TP sharding strategy for dim1 kernel #2436

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit dc0803c with merge base 32599be ():
💚 Looks good so far! There are no failures yet. 💚

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

[vkuzo]

need to uncomment this and run ./test/prototype/mx_formats/test_mx_dtensor.sh to reproduce