TE EP MXFP8 fails with last_dim % MXFP8_BLOCK_SIZE == 0

[faradawn]

Summary

transformer_engine.pytorch.ops.GroupedLinear fails with MXFP8BlockScaling in a Mixtral MoE expert-parallel run when per-expert token splits are not divisible by 32.

Repro context

• Model: Mixtral-8x7B
• GPUs: 8x B300
• EP size: 2
• Batch size: 8
• Sequence length: 8192
• Precision recipe: MXFP8BlockScaling
• Expert FFN path: transformer_engine.pytorch.ops.GroupedLinear
• Dispatcher: NCCL all-to-all
• Tutorial/example PR: Add examples for MoE models - Mixtral in TE #2642

Repro command:

cd /lustre/fsw/coreai_prod_infbench/faradawny/TransformerEngine/docs/examples/te_mixtral

torchrun --standalone --nproc_per_node=8 run_finetune_ep.py \
  --improvement 8 \
  --ep-size 2 \
  --batch-size 8 \
  --max-seq-length 8192 \
  --warmup-steps 5 \
  --train-steps 10 \
  2>&1 | tee logs/sweep_seq8k_ep2_8gpus_sequential_ops/seq8k_batch8_ep2_tier8_sequential_ops_mxfp8.log

Error

[rank0]:   File "/lustre/fsw/coreai_prod_infbench/faradawny/TransformerEngine/docs/examples/te_mixtral/te_mixtral.py", line 687, in _expert_ffn
[rank0]:     gate_up_output = self.experts_gate_up(tokens, split_sizes)
[rank0]:                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/ops/op.py", line 522, in forward
[rank0]:     return OperationFuser([self])(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/ops/basic/grouped_linear.py", line 739, in fuser_forward
[rank0]:     xs = tex.split_quantize(x, split_sizes_int, input_quantizers)
[rank0]:          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: RuntimeError: /workspace/TransformerEngine/transformer_engine/pytorch/csrc/quantizer.cpp:1668 in function get_scale_shape: Assertion failed: last_dim % MXFP8_BLOCK_SIZE == 0 && (numel / last_dim) % MXFP8_BLOCK_SIZE == 0. MXFP8 requires tensor dims that are divisible by 32 (got shape=(2283,4096))

Other ranks fail similarly with shapes like (1441,4096), (1178,4096), and (1225,4096).

Expected

The Sequential Ops grouped path should either handle MXFP8 padding internally per split, or provide a clear documented requirement/workaround for MoE token splits whose per-expert token counts are not multiples of 32.

[faradawn]

This is fixed by using. permute with padding, padding to 32. Yet, a new error occurred for division of 128.

[rank0]: Traceback (most recent call last):
[rank0]:   File "/lustre/fsw/coreai_prod_infbench/faradawny/TransformerEngine/docs/examples/te_mixtral/run_finetune_ep.py", line 159, in <module>
[rank0]:     main()
[rank0]:   File "/lustre/fsw/coreai_prod_infbench/faradawny/TransformerEngine/docs/examples/te_mixtral/run_finetune_ep.py", line 155, in main
[rank0]:     run_te_mixtral_finetune(hp)
[rank0]:   File "/lustre/fsw/coreai_prod_infbench/faradawny/TransformerEngine/docs/examples/te_mixtral/utils.py", line 432, in run_te_mixtral_finetune
[rank0]:     finetune_model(model, hyperparams, accelerator, train_dataloader, optimizer, lr_scheduler)
[rank0]:   File "/lustre/fsw/coreai_prod_infbench/faradawny/TransformerEngine/docs/examples/te_mixtral/utils.py", line 355, in finetune_model
[rank0]:     outputs = model(**batch)
[rank0]:               ^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1789, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/lustre/fsw/coreai_prod_infbench/faradawny/TransformerEngine/docs/examples/te_mixtral/te_mixtral.py", line 1170, in forward
[rank0]:     outputs: BaseModelOutputWithPast = self.model(
[rank0]:                                        ^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1789, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/lustre/fsw/coreai_prod_infbench/faradawny/TransformerEngine/docs/examples/te_mixtral/te_mixtral.py", line 1040, in forward
[rank0]:     hidden_states = decoder_layer(
[rank0]:                     ^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1789, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/lustre/fsw/coreai_prod_infbench/faradawny/TransformerEngine/docs/examples/te_mixtral/te_mixtral.py", line 876, in forward
[rank0]:     hidden_states = self.mlp(hidden_states)
[rank0]:                     ^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1789, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/lustre/fsw/coreai_prod_infbench/faradawny/TransformerEngine/docs/examples/te_mixtral/te_mixtral.py", line 793, in forward
[rank0]:     expert_output = self._expert_ffn(expert_input, tokens_per_expert)
[rank0]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/lustre/fsw/coreai_prod_infbench/faradawny/TransformerEngine/docs/examples/te_mixtral/te_mixtral.py", line 686, in _expert_ffn
[rank0]:     gate_up_output = self.experts_gate_up(tokens, split_sizes)
[rank0]:                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1789, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/ops/op.py", line 523, in forward
[rank0]:     return OperationFuser([self])(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/ops/fuser.py", line 517, in __call__
[rank0]:     return forward_func(*args)
[rank0]:            ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/autograd/function.py", line 596, in apply
[rank0]:     return super().apply(*args, **kwargs)  # type: ignore[misc]
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/ops/fuser.py", line 131, in forward
[rank0]:     x, fused_op_extra_outputs = op.fuser_forward(
[rank0]:                                 ^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/ops/basic/grouped_linear.py", line 946, in fuser_forward
[rank0]:     out, tensors_to_save = self._fuser_forward_grouped_tensor(
[rank0]:                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/ops/basic/grouped_linear.py", line 1185, in _fuser_forward_grouped_tensor
[rank0]:     grouped_x = tex.group_quantize(x, input_quantizer, num_groups, split_sizes)
[rank0]:                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: RuntimeError: /lustre/fsw/coreai_prod_infbench/faradawny/TransformerEngine/transformer_engine/common/cast/dispatch/../mxfp8/group_quantize_mxfp8.cuh:842 in function group_quantize: Assertion failed: first_logical_dim % 128 == 0. First logical dimension of a grouped tensor must be divisible by 128.

[faradawn]

Changing change moe_permute_and_pad_with_probs 's padding from 32 to 128 solved the issue!

[sbhavani]

Closing as resolved. The Mixtral EP MXFP8 path now pads per-expert splits and reporter confirmed padding to 128 resolves the repro.