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 TE Mixtral tutorial with NVMixtralSparseMoeBlock (TE GroupedLinear + moe_permute/moe_unpermute), an AllToAllTokenDispatcher for expert parallelism, optional DeepEP fused dispatch, FP8 support, and a fine-tuning harness. Several issues flagged in prior rounds have been fixed (class renames, wandb dep, collator shipping), but EXPERIMENT_LOG.md — a personal benchmark log containing NVIDIA-internal cluster paths and a developer username — has been added to the public repo and should be removed or sanitised before merge.

Confidence Score: 3/5

Not safe to merge: EXPERIMENT_LOG.md exposes internal NVIDIA cluster details in a public repo, and several critical runtime bugs in te_mixtral.py flagged in earlier review rounds remain unaddressed.

The P1 finding (EXPERIMENT_LOG.md with internal paths) pulls the ceiling to 4, but the overall state of te_mixtral.py — which still carries unresolved P1 bugs from prior rounds (wrong DTensor isinstance check in _sync_expert_views, null-pointer crashes in the InferenceParams branch, incorrect padding slot assignment) — warrants a lower score.

docs/examples/te_mixtral/EXPERIMENT_LOG.md (internal cluster details must be removed), docs/examples/te_mixtral/te_mixtral.py (multiple unresolved runtime bugs from prior review rounds).

Important Files Changed

Filename	Overview
docs/examples/te_mixtral/te_mixtral.py	Core TE Mixtral implementation: NVMixtralSparseMoeBlock, AllToAllTokenDispatcher, weight mapping. Multiple previously-flagged P1 bugs remain unresolved (DTensor check in _sync_expert_views, null-pointer crashes in InferenceParams branch, incorrect token padding slot assignment).
docs/examples/te_mixtral/EXPERIMENT_LOG.md	Personal benchmark log containing NVIDIA-internal cluster paths (SLURM account, developer username, lustre paths). Should not be committed to a public repository.
docs/examples/te_mixtral/utils.py	Training utilities: model init, Accelerate wrapping, fine-tuning loop. Previously-flagged issues (hardcoded warmup steps, missing wandb dep, broken class name import) appear to be addressed in current version.
docs/examples/te_mixtral/collator.py	THD sequence-packing collator (previously shipped as bionemo_mixtral, now local). Logic is correct; replaces the previously-broken external dependency.
docs/examples/te_mixtral/fused_token_router.py	DeepEP-backed TokenDispatcher with correct map_type="mask" for the multihot path and merging_probs kwarg.
docs/examples/te_mixtral/requirements.txt	Pinned HF stack dependencies; wandb (previously missing) no longer needed since log_with was removed from Accelerator call.
docs/examples/te_mixtral/install_hybridep.sh	DeepEP build script with GPU arch detection and NVSHMEM-disable patch. Contains an internal NVIDIA bug tracker URL (nvbugspro.nvidia.com) in a public-facing comment.
docs/examples/te_mixtral/test_accuracy.py	Weight-mapping parity test between HF and TE models. Clean and minimal; correctly filters _extra_state keys from missing list.

Sequence Diagram

sequenceDiagram
    participant Input as Input [N, H]
    participant Gate as Router (gate)
    participant Dispatch as TokenDispatcher<br/>(AllToAll / FusedDeepEP)
    participant GU as GroupedLinear<br/>gate_up [E, 2*F, H]
    participant Act as SiLU x up
    participant Down as GroupedLinear<br/>down [E, H, F]
    participant Combine as Combine<br/>(moe_unpermute + all_to_all)
    participant Output as Output [N, H]

    Input->>Gate: router_logits [N, num_experts]
    Gate->>Gate: softmax -> topk -> normalize
    Gate->>Dispatch: hidden_states, selected_experts [N, top_k], routing_weights
    Dispatch->>Dispatch: moe_permute(index) + all_to_all
    Dispatch->>GU: expert_input [T, H], m_splits
    GU->>Act: gate_up_output [T, 2F]
    Act->>Down: intermediate [T, F]
    Down->>Combine: expert_output [T, H]
    Combine->>Combine: moe_unpermute(index) + reverse all_to_all
    Combine->>Output: weighted sum [N, H]

Loading

_{Reviews (20): Last reviewed commit: "docs(te_mixtral): add architecture diagr..." | 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)))

[greptile-apps]

Wrong DTensor check — GroupedLinear receives DTensor views under EP

self.experts_gate_up_weight is an nn.Parameter. After set_ep_group() wraps it with nn.Parameter(DTensor.from_local(...)), the parameter's type is still nn.Parameter (not DTensor), so isinstance(gate_up_w, DTensor) is always False. The guard never fires, to_local() is never called, and every subsequent gate_up_w[i] slice handed to object.__setattr__ is a DTensor. TE's GroupedLinear CUDA kernels cannot accept DTensor views and will either raise a type error or silently use incorrect data.

The _expert_ffn loop path already applies the correct pattern with .data:

if isinstance(gate_up_w.data, DTensor):
    gate_up_w = gate_up_w.data.to_local()

Suggested change

	"""
	gate_up_w = self.experts_gate_up_weight
	if isinstance(gate_up_w, DTensor):
	gate_up_w = gate_up_w.to_local()
	for i in range(self.num_local_experts):
	object.__setattr__(self.experts_gate_up, f"weight{i}", gate_up_w[i])
	down_w = self.experts_down_weight
	if isinstance(down_w, DTensor):
	down_w = down_w.to_local()
	gate_up_w = self.experts_gate_up_weight
	if isinstance(gate_up_w.data, DTensor):
	gate_up_w = gate_up_w.data.to_local()
	for i in range(self.num_local_experts):
	object.__setattr__(self.experts_gate_up, f"weight{i}", gate_up_w[i])
	down_w = self.experts_down_weight
	if isinstance(down_w.data, DTensor):
	down_w = down_w.data.to_local()
	for i in range(self.num_local_experts):
	object.__setattr__(self.experts_down, f"weight{i}", down_w[i])

[greptile-apps]

Padding appended to the last expert's slot regardless of which expert has remaining tokens

When orig_num_tokens % 32 != 0, pad_count tokens are appended to expert_input and tokens_per_expert[-1] is incremented by pad_count. This is only correct if the padded tokens happen to logically belong to the last expert. If the last expert already has 0 tokens on this rank (a routine EP scenario), GroupedLinear will schedule a GEMM for pad_count zero-padded tokens against that expert's weights, inflating compute and potentially triggering a misaligned-size error rather than a clean no-op. A safer approach is to distribute the padding tokens as a trailing dummy entry (tokens_per_expert.append(pad_count)) or to pad at batch-preparation time before dispatch so padding is always aligned with a real expert slot.

[faradawn]

Will wait on this PR: #2923

[greptile-apps]

Internal cluster details in a public repository

EXPERIMENT_LOG.md is a personal working-notes file that contains NVIDIA-internal infrastructure details that should not be committed to a public GitHub repo:

• A developer username embedded in paths (faradawny): /lustre/fsw/coreai_prod_infbench/faradawny/docker/pytorch-25.12-py3.sqsh
• An internal SLURM account: coreai_prod_infbench
• Internal shared storage paths: /lustre/share/coreai_prod_infbench/exp3_nvtx_per_phase/

This file looks like a benchmark log intended for internal coordination, not documentation for external tutorial users. Consider removing it entirely or replacing it with a trimmed BENCHMARKS.md that shows only public results (median step times, speedup numbers) without cluster-specific details.