Add examples for MoE models - Mixtral in TE

[faradawn]

Summary

This PR adds a complete tutorial for integrating HuggingFace Mixtral (MoE) with Transformer Engine, addressing the gap identified in #2573.

What's included

• te_mixtral.py — Drop-in TEMixtralSparseMoeBlock that replaces HF's loop-over-experts with TE's GroupedLinear (batched GEMM) + moe_permute/moe_unpermute. Includes replace_moe_block context manager, TEMixtralForCausalLM with HF weight loading, and replace_params for expert weight packing.
• utils.py — Data loading, BF16/FP8 model init, Accelerate wrapping, fine-tuning loop — mirrors te_llama/utils.py style.
• requirements.txt — Pinned dependencies matching the Llama/Gemma tutorials.
• Tutorial notebook — Full tutorial matching the quality bar of te_llama and te_gemma, covering:
  1. Architecture overview: Transformer → Mixtral MoE, HF bottleneck, TE approach
  2. Unit-test cell verifying output shape/dtype against the HF block
  3. [Baseline] HF Mixtral in BF16
  4. [Improvement 1] TE GroupedLinear MoE in BF16
  5. [Improvement 2] TE GroupedLinear MoE in FP8
  6. Expert routing considerations with mixed precision (m_splits, per-expert FP8 scaling, aux loss passthrough)
  7. Generalisation guide for other MoE architectures (DeepSeek, Grok-1, etc.)

Bug fix

Corrected the m_splits calculation flagged by the automated review:

# Before (wrong): double-counts tokens by reducing with .any() then multiplying by top_k
expert_mask = (selected_experts == expert_idx).any(dim=-1)
m_splits.append(expert_mask.sum().item() * self.top_k)

# After (correct): count the actual number of (token, top_k_slot) pairs per expert
m_splits = [(selected_experts == i).sum().item() for i in range(self.num_experts)]

Scope

Covers all topics requested in #2573:

• How to wrap MoE layers with TE modules ✓
• FP8 training configuration for MoE ✓
• Expert routing considerations with mixed precision ✓
• Generalisation to arbitrary MoE architectures ✓

[greptile-apps]

Greptile Summary

This PR adds a complete Transformer Engine tutorial for HuggingFace Mixtral (MoE), covering BF16 GroupedLinear, FP8, and MXFP8 expert-parallel fine-tuning, and addresses many previously flagged issues (wrong map_type, m_splits double-count, missing flash-attn/nvidia-cudnn-frontend deps, hard-coded warmup steps, wandb dep, sorted() key for numeric expert-weight keys, and DP gradient sync).

• te_mixtral.py + te_moe_dispatch.py: Two separate AllToAllTokenDispatcher classes implement BF16 and MXFP8 dispatch/combine paths; the BF16 class uses moe_permute(map_type=\"index\") + moe_unpermute, while the MXFP8 class uses moe_permute_and_pad_with_probs aligned to the 256-token CuTeDSL kernel requirement.
• utils.py: Handles EP-aware DistributedSampler, manual DP gradient all-reduce, Accelerate wrapping, and the monkey-patch shim that enables the fused MXFP8 grouped-MLP kernel without recompiling TE.

Confidence Score: 4/5

Close to mergeable — the major routing, weight-loading, and distributed-training correctness issues from earlier rounds are all addressed.

The critical dispatch bugs (map_type, m_splits, deprecated probs kwarg, wrong num_out_tokens) are fixed, the DP sampler and gradient sync look correct, and dependency gaps are closed. Remaining findings are quality and performance suggestions that do not break correctness.

utils.py (dead total_loss accumulator, foreach=False optimizer regression for EP>1); te_mixtral.py (duplicated sort-index helper vs te_moe_dispatch.py)

Important Files Changed

Filename	Overview
docs/examples/te_mixtral/te_mixtral.py	Core BF16 TE Mixtral model; AllToAllTokenDispatcher correctly uses moe_permute(map_type=index) and moe_unpermute with matching map_type. Previously flagged m_splits, probs kwarg, and DTensor issues are addressed.
docs/examples/te_mixtral/te_mixtral_mxfp8.py	MXFP8 fused-MLP model; decode guard now keyed solely on seq-len==1, OrderedDict import fixed, variables initialised before conditional pack block.
docs/examples/te_mixtral/te_moe_dispatch.py	MXFP8 dispatcher using moe_permute_and_pad_with_probs; correctly pads to 256 for CuTeDSL kernel, _build_expert_sort_indices is torch-compiled and fullgraph-safe.
docs/examples/te_mixtral/utils.py	Training utilities; DP gradient sync, EP-aware sampler, and warmup-step accounting now correct. foreach=False in build_adamw for EP>1 is overly conservative, and total_loss is accumulated but never reported.
docs/examples/te_mixtral/hf_to_te_weights.py	Weight-mapping from HF Mixtral to BF16 and MXFP8 TE state dicts; numeric sort key for weight{i} keys fixes the previous lexicographic-ordering bug.
docs/examples/te_mixtral/collator.py	THD sequence-packing collator now ships with the PR, resolving the prior ImportError from bionemo_mixtral.
docs/examples/te_mixtral/requirements.txt	All previously missing deps (flash-attn, nvidia-cudnn-frontend>=1.23.0) are now present; wandb removed.
docs/examples/te_mixtral/test_accuracy.py	Accuracy parity test between HF and TE (BF16 + MXFP8); uses strict=False with explicit unexpected/missing-key checks.

Sequence Diagram

sequenceDiagram
    participant HF as HF Mixtral weights
    participant INIT as init_te_mixtral_model
    participant MODEL as TEMixtralForCausalLM
    participant DISP as AllToAllTokenDispatcher
    participant TE as TE GroupedLinear

    INIT->>HF: "from_pretrained device=cpu"
    INIT->>MODEL: ForCausalLM(te_config).to(cuda)
    INIT->>MODEL: replace_params(hf_sd, te_sd)
    INIT->>MODEL: "load_state_dict strict=False"
    INIT->>MODEL: set_ep_groups(ep_group)
    MODEL->>DISP: set_ep_group(ep_group)
    Note over MODEL,DISP: Forward pass training
    MODEL->>DISP: dispatch(hidden, selected_experts, weights)
    DISP->>DISP: "moe_permute map_type=index"
    DISP-->>DISP: "all_to_all_single EP>1"
    DISP->>DISP: _build_expert_sort_indices
    DISP-->>MODEL: DispatchOutput
    MODEL->>TE: _experts_ffn_op(expert_input, split_sizes)
    TE-->>MODEL: expert_output
    MODEL->>DISP: combine(expert_output, handle)
    DISP-->>DISP: "reverse all_to_all EP>1"
    DISP->>DISP: "moe_unpermute map_type=index"
    DISP-->>MODEL: combined output N H

Loading

_{Reviews (52): Last reviewed commit: "fix none mask" | Re-trigger Greptile}

[greptile-apps]

_{1 file reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[sudhakarsingh27]

Thanks, @faradawn! Also adding @sbhavani to the discussion. Compared to other llama/gemma tutorials, this one seems a quite barebones and looks more like a code example than a tutorial. @sbhavani do you think in its current form, it covers the scope as you requested in #2573?

[faradawn]

Hi @sudhakarsingh27 can you check if this addresses your comments? Tested in 2x H100.

[sbhavani]

Thanks, @faradawn! Also adding @sbhavani to the discussion. Compared to other llama/gemma tutorials, this one seems a quite barebones and looks more like a code example than a tutorial. @sbhavani do you think in its current form, it covers the scope as you requested in #2573?

Agreed! I think any example should show some perf gain and include the whole weight mapping so the user can run the example.

[pggPL]

Documentation build is not working, if you fix it please ping me and I'll review.

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.

[faradawn]

Hi @pggPL I've fixed the doc build and refined the content - can you help check?

[pggPL]

I was not reviewing it in detail, I would like to discuss on high-level flow of the tutorial first.
I see some issues:

1. We use GroupedLinear with EP=8 and 8 experts, so 1 GroupedLinear is essentially having 1 expert, so it does not differ from using standard linear. Can we do EP=2 with 4 experts per each GPU or something like that.
2. We use different types of parallelism in different tiers and it was not clear to me why. Maybe we can also use EP in HF if it is supported.
3. The gains from Tier 1 -> Tier 2 are very huge, but Tier 2 -> Tier 3 -> Tier 4 are very small. Do you know why?

Apart of that it would be nice to prepare documentation (not tutorial) explaining why grouped linear is needed and which permute kernels we provide. I will work on that in different PR and it will land in Features section in our docs.

[greptile-apps]

attention_mask is None crash inside InferenceParams branch

attention_mask is an optional parameter with no None-guard before line 852. In the decode step of model.generate() the mask is frequently omitted, and attention_mask.shape[:2] will raise AttributeError: 'NoneType' object has no attribute 'shape', crashing every auto-regressive generation call.

if isinstance(past_key_values, InferenceParams):
    _ref = input_ids if input_ids is not None else inputs_embeds
    lengths = (
        attention_mask.sum(dim=1).tolist()
        if attention_mask is not None and attention_mask.shape[:2] == _ref.shape[:2]
        else [1] * _ref.shape[0]
    )
    past_key_values.pre_step(OrderedDict(zip(list(range(len(lengths))), lengths)))

[jberchtold-nvidia]

Thanks! Rebase is now fixed

It is my own review and requested changes were resolved. However, the PR itself contains PyTorch usage so I should not give full approval on the PR

[sudhakarsingh27]

Lgtm. Thanks!