make mxfp8 dim1 cast kernel configurable

[danielvegamyhre]

Stacked PRs:

• ->make mxfp8 dim1 cast kernel configurable #1427

---

make mxfp8 dim1 cast kernel configurable

Summary

• We recently added a new CUDA kernel for the mxfp8 dim1 cast which is ~1.4x faster than the existing Triton kernel or torch.compile, and using it results in an e2e training speedup of +1.5-2.5% TPS with Llama3 8b using FSDP=4/8 (Add CUDA kernel for MXFP8 dim1 casting ao#2513). The integration work for composability with torch.compile + FSDP is complete as well: integration of new mxfp8 casting cuda kernel ao#2564
• This PR updates the mxfp8 user facing API to replace the boolean flag "--mx.use_triton_for_dim1_cast=[true|false] to mxfp8_dim1_cast_kernel_choice=[triton|cuda|torch]

Test plan

• Triton: NGPU=8 CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh --training.steps=100 --model.converters="mx" --mx.recipe_name="mxfp8" --training.compile --mx.mxfp8_dim1_cast_kernel_choice="triton"
• Cuda: NGPU=8 CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh --training.steps=100 --model.converters="mx" --mx.recipe_name="mxfp8" --training.compile --mx.mxfp8_dim1_cast_kernel_choice="cuda"
• Torch: NGPU=8 CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh --training.steps=100 --model.converters="mx" --mx.recipe_name="mxfp8" --training.compile --mx.mxfp8_dim1_cast_kernel_choice="torch"

Limitations

• TP is currently not supported yet, as both the Triton kernel and CUDA kernel are affected by an issue: RuntimeError: Attempting to use FunctionalTensor on its own. Instead, please use it with a corresponding FunctionalTensorMode(). This is a known issue we were talking to Brian about, will continue following up on it.

[tianyu-l]

what's the benefit of letting user choose?

[danielvegamyhre]

1. It makes it easy for torchao developers to benchmark torch.compile vs triton vs cuda implementations as we work on perf improvements, especially on improving torch.compile performance for casts.
2. Users may be using a python only torchao installation that doesn't include the CUDA kernel. This is probably not common but still worth considering.

[danielvegamyhre]

@vkuzo @tianyu-l this is ready for review

CI error is unrelated to this change:

Exception: Integration test failed, flavor : 2D eager, command : TORCH_TRACE="artifacts-to-be-uploaded/2d_eager/compile_trace" CONFIG_FILE=./torchtitan/models/llama3/train_configs/debug_model.toml NGPU=4 LOG_RANK=0,1,2,3 ./run_train.sh --job.dump_folder artifacts-to-be-uploaded/2d_eager --parallelism.tensor_parallel_degree 2

[vkuzo]

lgtm, I'll let someone from titan team accept

[tianyu-l]

TP is currently not supported yet, as both the Triton kernel and CUDA kernel are affected by an issue: RuntimeError: Attempting to use FunctionalTensor on its own. Instead, please use it with a corresponding FunctionalTensorMode(). This is a known issue we were talking to Brian about, will continue following up on it.

Do you think we can error out in mx.py, since you do have JobConfig on tp info?

Please rebase before merge.

[danielvegamyhre]

confirmed error is CUDA driver error in async TP, which is not related to this change. as an aside, it looks like an error at the C++ symm mem level, or perhaps the containerized env running the test had dependencies updated?

[danielvegamyhre]

Do you think we can error out in mx.py, since you do have JobConfig on tp info?

Sure, done