[mxfp8 moe training] wrap 3d quantize tensor in custom ops and integrate it

[danielvegamyhre]

Stacked PRs:

• [mxfp8 moe training] update 3d quant colwise scaling kernel to use single input/output TMA descriptors #3034
• [mxfp8 moe training] remove mxfp8_gemms.py #3033
• ->[mxfp8 moe training] wrap 3d quantize tensor in custom ops and integrate it #3004
• [mxfp8 moe training] add CUDA kernel to quantize 3d tensor colwise #3002

---

[mxfp8 moe training] wrap 3d quantize tensor in custom ops and integrate it

Torchtitan Llama4 e2e training benchmarks

Llama4 debug model

• dim=5120 (standard)
• num_layers=2, num_experts=2 (to allow for higher seq len and avoid OOM)
• FSDP=2
• compile=True

Note there are typically 1-8 experts per device depending on EP degree (for llama4 and DSV3, so these tests simulate 2 and 8 experts per device. We can do real tests using EP once pytorch/torchtitan#1651 is resolved).

seq_len=8192, experts per device = 2

tl;dr:

• mxfp8 dense only: 1.05x over bf16
• mxfp8 moe + dense: 1.27x speedup over bf16

Config:

    "debugmodel": TransformerModelArgs(
        dim=5120,
        n_layers=2,
        n_heads=40,
        n_kv_heads=8,
        ffn_dim_multiplier=1.2,
        multiple_of=2048,
        rope_theta=500000,
        max_seq_len=10485760,
        moe_args=MoEArgs(num_experts=2),
        interleave_moe_layer_step=1,
    ),

BF16

rm -rf /tmp/torchinductor_danvm; TORCHTITAN_ROOT=/home/danvm/torchtitan CUDA_VISIBLE_DEVICES="2,3,4,5" NGPU=2 EXTRA_ARGS="--parallelism.data_parallel_shard_degree=2 --parallelism.tensor_parallel_degree=1 --model.print-after-conversion --metrics.log_freq=10 --training.steps=100 --compile.enable --training.seq_len=8192" ./llama4.sh
Median Tokens/Second (excluding step 1): 66828.5
Max Memory Usage: 113.41 GiB

=========================================================================

MXFP8 DENSE ONLY

rm -rf /tmp/torchinductor_danvm; TORCHTITAN_ROOT=/home/danvm/torchtitan CUDA_VISIBLE_DEVICES="2,3,4,5" NGPU=2 EXTRA_ARGS="--parallelism.data_parallel_shard_degree=2 --parallelism.tensor_parallel_degree=1 --model.converters="mx" --mx.recipe_name="mxfp8_cublas" --mx.filter_fqns="output,router.gate,wk,wv" --model.print-after-conversion --metrics.log_freq=10 --training.steps=100 --compile.enable --training.seq_len=8192" ./llama4.sh 
Median Tokens/Second (excluding step 1): 70277.0
Max Memory Usage: 113.35 GiB

=========================================================================

MXFP8 MOE + DENSE

seq_len=8192 -> total_M=65600

(torch) [[email protected] ~/ao/benchmarks/float8/training (mx-moe-compile)]$ rm -rf /tmp/torchinductor_danvm; TORCHTITAN_ROOT=/home/danvm/torchtitan CUDA_VISIBLE_DEVICES="2,3,4,5" NGPU=2 EXTRA_ARGS="--parallelism.data_parallel_shard_degree=2 --parallelism.tensor_parallel_degree=1 --model.converters="mx" --mx.recipe_name="mxfp8_cublas" --mx.filter_fqns="output,router.gate,wk,wv" --mx.moe_fqns_prototype="experts" --model.print-after-conversion --metrics.log_freq=10 --training.steps=100 --compile.enable --training.seq_len=8192" ./llama4.sh 
Median Tokens/Second (excluding step 1): 85169.0
Max Memory Usage: 112.35 GiB

seq_len=2048, experts per device = 8

tl;dr:

• mxfp8 dense only: 1.08x over bf16
• mxfp8 moe + dense: 1.19x speedup over bf16

Llama4 debug model

• dim=5120 (standard)
• num_layers=8, num_experts=8 (to avoid OOM)
• FSDP=4
• compile=True

Config:

    "debugmodel": TransformerModelArgs(
        dim=5120,
        n_layers=8,
        n_heads=40,
        n_kv_heads=8,
        ffn_dim_multiplier=1.2,
        multiple_of=2048,
        rope_theta=500000,
        max_seq_len=10485760,
        moe_args=MoEArgs(num_experts=8),
        interleave_moe_layer_step=1,
    ),

BF16

rm -rf /tmp/torchinductor_danvm; TORCHTITAN_ROOT=/home/danvm/torchtitan CUDA_VISIBLE_DEVICES="2,3,4,5" NGPU=4 EXTRA_ARGS="--parallelism.data_parallel_shard_degree=4 --parallelism.tensor_parallel_degree=1 --model.print-after-conversion --metrics.log_freq=10 --training.steps=100 --compile.enable" ./llama4.sh 

Median Tokens/Second (excluding step 1): 24753.0
Max Memory Usage: 89.77 GiB

======================================================================

MX dense only

rm -rf /tmp/torchinductor_danvm; TORCHTITAN_ROOT=/home/danvm/torchtitan CUDA_VISIBLE_DEVICES="2,3,4,5" NGPU=4 EXTRA_ARGS="--parallelism.data_parallel_shard_degree=4 --parallelism.tensor_parallel_degree=1 --model.converters="mx" --mx.recipe_name="mxfp8_cublas" --mx.filter_fqns="output,router.gate,wk,wv" --model.print-after-conversion --metrics.log_freq=10 --training.steps=100 --compile.enable" ./llama4.sh 

Median Tokens/Second (excluding step 1): 26778.5
Max Memory Usage: 89.47 GiB

Speedup:  ~1.08x speedup over bf16

======================================================================

MX MoE + dense

rm -rf /tmp/torchinductor_danvm; TORCHTITAN_ROOT=/home/danvm/torchtitan CUDA_VISIBLE_DEVICES="2,3,4,5" NGPU=4 EXTRA_ARGS="--parallelism.data_parallel_shard_degree=4 --parallelism.tensor_parallel_degree=1 --model.converters="mx" --mx.recipe_name="mxfp8_cublas" --mx.filter_fqns="output,router.gate,wk,wv" --mx.moe_fqns_prototype="experts" --model.print-after-conversion --metrics.log_freq=10 --training.steps=100 --compile.enable" ./llama4.sh 

Median Tokens/Second (excluding step 1): 29373.5
Max Memory Usage: 82.39 GiB

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 a68b868 with merge base f75b251 ():
💚 Looks good so far! There are no failures yet. 💚

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