diff --git a/docs/examples/te_mixtral/collator.py b/docs/examples/te_mixtral/collator.py
new file mode 100644
index 0000000000..b9a53cf542
--- /dev/null
+++ b/docs/examples/te_mixtral/collator.py
@@ -0,0 +1,148 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Data collator for THD sequence packing (variable-length flash attention)."""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from transformers import DataCollatorForLanguageModeling
+
+
+logger = logging.getLogger(__name__)
+
+
+def _pt_flatten_collate(features: list[dict[str, list[int]]], return_position_ids: bool = False):
+ """Flatten a list of tokenized samples into a single packed batch with cumulative sequence lengths."""
+ is_labels_provided = "labels" in features[0]
+ sample_lengths = [len(sample["input_ids"]) for sample in features]
+
+ batch = {}
+ batch["max_length_q"] = batch["max_length_k"] = max(sample_lengths)
+ batch["input_ids"] = torch.tensor(
+ [[token for sample in features for token in sample["input_ids"]]], dtype=torch.int64
+ )
+ if is_labels_provided:
+ batch["labels"] = torch.tensor(
+ [[label for sample in features for label in sample["labels"]]], dtype=torch.int64
+ )
+ cu_seq_lens = torch.zeros(len(features) + 1, dtype=torch.int32)
+ cu_seq_lens[1:] = torch.cumsum(torch.tensor(sample_lengths), dim=0, dtype=torch.int32)
+ batch["cu_seq_lens_q"] = batch["cu_seq_lens_k"] = cu_seq_lens
+ if "attention_mask" in features[0]:
+ batch["attention_mask"] = torch.tensor(
+ [[v for sample in features for v in sample["attention_mask"]]], dtype=torch.int64
+ )
+ if return_position_ids:
+ batch["position_ids"] = torch.hstack(
+ [torch.arange(sample_len, dtype=torch.int64) for sample_len in sample_lengths]
+ ).unsqueeze(0)
+
+ return batch
+
+
+def _pt_pad_to_multiple_of(
+ batch: dict[str, Any], pad_to_multiple_of: int, token_pad: int, label_pad: int
+):
+ """Pad a batch to a multiple of ``pad_to_multiple_of`` by appending a mock sequence."""
+ remainder = -batch["input_ids"].numel() % pad_to_multiple_of
+ if remainder == 0:
+ return batch
+
+ batch["input_ids"] = torch.cat(
+ [batch["input_ids"], torch.full((1, remainder), token_pad, dtype=batch["input_ids"].dtype)],
+ dim=1,
+ )
+ if "labels" in batch:
+ batch["labels"] = torch.cat(
+ [batch["labels"], torch.full((1, remainder), label_pad, dtype=batch["labels"].dtype)],
+ dim=1,
+ )
+ if "cu_seq_lens_q" in batch:
+ batch["cu_seq_lens_q"] = torch.cat(
+ [
+ batch["cu_seq_lens_q"],
+ torch.tensor(
+ [batch["cu_seq_lens_q"][-1] + remainder], dtype=batch["cu_seq_lens_q"].dtype
+ ),
+ ],
+ dim=0,
+ )
+ batch["cu_seq_lens_k"] = batch["cu_seq_lens_q"]
+ if "max_length_q" in batch:
+ batch["max_length_q"] = max(batch["max_length_q"], remainder)
+ batch["max_length_k"] = batch["max_length_q"]
+ if "attention_mask" in batch:
+ batch["attention_mask"] = torch.cat(
+ [
+ batch["attention_mask"],
+ torch.zeros((1, remainder), dtype=batch["attention_mask"].dtype),
+ ],
+ dim=1,
+ )
+ if "position_ids" in batch:
+ batch["position_ids"] = torch.cat(
+ [
+ batch["position_ids"],
+ torch.arange(remainder, dtype=batch["position_ids"].dtype).unsqueeze(0),
+ ],
+ dim=1,
+ )
+
+ return batch
+
+
+@dataclass
+class DataCollatorWithFlattening:
+ """Data collator that flattens variable-length sequences into a single packed tensor for flash attention.
+
+ Wraps a ``DataCollatorForLanguageModeling`` and produces THD-format batches with
+ ``cu_seq_lens_q`` / ``cu_seq_lens_k`` metadata for TE's fused attention kernels.
+
+ Args:
+ collator: The base collator for MLM/CLM masking.
+ pad_to_multiple_of: If set, pads the total token count to be divisible by this number.
+ separator_id: Label value inserted at sequence boundaries (typically -100 for causal LM).
+ """
+
+ collator: DataCollatorForLanguageModeling
+ pad_to_multiple_of: int | None = None
+ separator_id: int | None = None
+
+ def __call__(self, features, return_tensors=None):
+ """Pack features into a single THD batch with flash-attention metadata."""
+ if return_tensors is not None and return_tensors != "pt":
+ raise NotImplementedError(
+ f"Only return_tensors='pt' is supported, got '{return_tensors}'"
+ )
+
+ bshd_batch = self.collator(features, return_tensors=return_tensors)
+ packed_batch = _pt_flatten_collate(features)
+
+ masked_input_ids = bshd_batch["input_ids"][bshd_batch["attention_mask"].bool()].unsqueeze(0)
+ masked_labels = bshd_batch["labels"][bshd_batch["attention_mask"].bool()].unsqueeze(0)
+
+ if self.separator_id is not None:
+ masked_labels[:, packed_batch["cu_seq_lens_q"][1:-1]] = self.separator_id
+
+ packed_batch["input_ids"] = masked_input_ids
+ packed_batch["labels"] = masked_labels
+
+ if self.pad_to_multiple_of is not None:
+ pad_token_id = self.collator.tokenizer.pad_token_id
+ if not isinstance(pad_token_id, int):
+ logger.warning(
+ f"tokenizer.pad_token_id is not an integer, using 1 instead: {pad_token_id}"
+ )
+ pad_token_id = 1
+ packed_batch = _pt_pad_to_multiple_of(
+ packed_batch,
+ self.pad_to_multiple_of,
+ token_pad=pad_token_id,
+ label_pad=-100,
+ )
+
+ return packed_batch
diff --git a/docs/examples/te_mixtral/hf_to_te_weights.py b/docs/examples/te_mixtral/hf_to_te_weights.py
new file mode 100644
index 0000000000..86f89d7abc
--- /dev/null
+++ b/docs/examples/te_mixtral/hf_to_te_weights.py
@@ -0,0 +1,426 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""HuggingFace Mixtral -> Transformer Engine state-dict mapping.
+
+Two top-level entry points share the same top-level / attention / layernorm
+plumbing and only differ in how they place the expert MoE weights:
+
+ * :func:`replace_params_bf16` — used by ``te_mixtral.py`` (Improvements 1/2,
+ BF16). Expert weights land in stacked ``mlp.experts_{gate_up,down}_weight``
+ parameters (loop path) and/or per-expert ``mlp.experts_{gate_up,down}.weight{i}``
+ parameters (Sequential-Op ``GroupedLinear``).
+
+ * :func:`replace_params_mxfp8` — used by ``te_mixtral_mxfp8.py`` (Improvement 3,
+ MXFP8). Expert gate (``w1``) and up (``w3``) rows are row-interleaved in
+ blocks of 32 to match the GLU layout the fused MXFP8 grouped-MLP kernel
+ expects.
+"""
+
+from __future__ import annotations
+
+import re
+
+import torch
+import torch.distributed as dist
+from transformers import MixtralConfig
+
+
+# Block size for the gate/up interleaved layout. Must match the
+# ``glu_interleave_size`` configured on ``ScaledSwiGLU`` in the MXFP8 MoE
+# block (see ``te_mixtral_mxfp8.py``).
+GLU_INTERLEAVE_SIZE = 32
+
+
+# ---------------------------------------------------------------------------
+# Low-level copy helpers
+# ---------------------------------------------------------------------------
+
+
+def _copy_param(target: torch.Tensor, source: torch.Tensor) -> None:
+ """Copy ``source`` into ``target`` preserving the target's dtype/device."""
+ target.copy_(source.to(device=target.device, dtype=target.dtype))
+
+
+def _copy_qkv_proj_to_fused(
+ fused_qkv: torch.Tensor,
+ proj_weight: torch.Tensor,
+ proj_kind: str,
+ config: MixtralConfig,
+) -> None:
+ """Copy one HF Q/K/V projection into the TE fused QKV layout.
+
+ TE interleaves the heads as ``[Q_g_0, ..., Q_g_{h-1}, K_g, V_g]`` per
+ KV group ``g``; HF stores Q/K/V as separate projections.
+ """
+ head_num = config.num_attention_heads
+ num_query_groups = config.num_key_value_heads
+ heads_per_group = head_num // num_query_groups
+ hidden_size = config.hidden_size
+ head_size = hidden_size // head_num
+ qkv_total_dim = head_num + 2 * num_query_groups
+
+ fused_view = fused_qkv.view(qkv_total_dim, head_size, hidden_size)
+ proj_weight = proj_weight.to(device=fused_view.device, dtype=fused_view.dtype)
+
+ if proj_kind == "q":
+ q_view = proj_weight.view(head_num, head_size, hidden_size)
+ for i in range(num_query_groups):
+ start = (heads_per_group + 2) * i
+ end = start + heads_per_group
+ fused_view[start:end].copy_(q_view[i * heads_per_group : (i + 1) * heads_per_group])
+ elif proj_kind == "k":
+ k_view = proj_weight.view(num_query_groups, head_size, hidden_size)
+ for i in range(num_query_groups):
+ fused_view[(heads_per_group + 2) * i + heads_per_group].copy_(k_view[i])
+ elif proj_kind == "v":
+ v_view = proj_weight.view(num_query_groups, head_size, hidden_size)
+ for i in range(num_query_groups):
+ fused_view[(heads_per_group + 2) * i + heads_per_group + 1].copy_(v_view[i])
+ else:
+ raise ValueError(f"Unsupported proj_kind: {proj_kind}")
+
+
+def _interleave_gate_up(
+ gate: torch.Tensor,
+ up: torch.Tensor,
+ interleave: int = GLU_INTERLEAVE_SIZE,
+) -> torch.Tensor:
+ """Interleave HF gate (``w1``) and up (``w3``) rows in blocks of ``interleave``.
+
+ HF stacks gate-then-up along the output dim (``[I gate rows; I up rows]``);
+ the fused MXFP8 kernel reads gate_up's output in the GLU-interleaved layout
+ ``[B gate; B up; B gate; B up; ...]`` with ``B = interleave``.
+ """
+ intermediate_size, hidden = gate.shape
+ if up.shape != (intermediate_size, hidden):
+ raise ValueError(f"gate and up shape mismatch: {gate.shape} vs {up.shape}")
+ if intermediate_size % interleave != 0:
+ raise ValueError(f"intermediate_size {intermediate_size} must be divisible by {interleave}")
+ g = gate.reshape(intermediate_size // interleave, interleave, hidden)
+ u = up.reshape(intermediate_size // interleave, interleave, hidden)
+ stacked = torch.stack([g, u], dim=1) # [I/B, 2, B, H]
+ return stacked.reshape(2 * intermediate_size, hidden).contiguous()
+
+
+# ---------------------------------------------------------------------------
+# Shared per-layer plumbing (top-level, attention, router gate)
+# ---------------------------------------------------------------------------
+
+
+def _ep_rank_from_config(config: MixtralConfig) -> tuple[int, int]:
+ """Return (ep_size, ep_rank). EP rank is global_rank % ep_size."""
+ ep_size = int(getattr(config, "expert_parallel_size", 1))
+ world_rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0
+ ep_rank = world_rank % ep_size if ep_size > 1 else 0
+ return ep_size, ep_rank
+
+
+def _collect_layer_prefixes(hf_state_dict: dict) -> set[str]:
+ prefixes = set()
+ for key in hf_state_dict.keys():
+ m = re.match(r"model\.layers\.\d+\.", key)
+ if m is not None:
+ prefixes.add(m.group())
+ return prefixes
+
+
+def _copy_top_level(hf_state_dict: dict, te_state_dict: dict) -> None:
+ direct = {
+ "model.embed_tokens.weight": "model.embed_tokens.weight",
+ "model.norm.weight": "model.norm.weight",
+ "lm_head.weight": "lm_head.weight",
+ "model.rotary_emb.inv_freq": "model.rotary_emb.inv_freq",
+ }
+ for hf_key, te_key in direct.items():
+ if hf_key in hf_state_dict and te_key in te_state_dict:
+ _copy_param(te_state_dict[te_key], hf_state_dict[hf_key])
+
+
+def _copy_attention_and_layernorms(
+ hf_state_dict: dict, te_state_dict: dict, layer_prefix: str, config: MixtralConfig
+) -> None:
+ direct = {
+ layer_prefix
+ + "input_layernorm.weight": layer_prefix
+ + "self_attention.layernorm_qkv.layer_norm_weight",
+ layer_prefix + "self_attn.o_proj.weight": layer_prefix + "self_attention.proj.weight",
+ layer_prefix
+ + "post_attention_layernorm.weight": layer_prefix
+ + "post_attention_layernorm.weight",
+ }
+ for hf_key, te_key in direct.items():
+ if hf_key in hf_state_dict and te_key in te_state_dict:
+ _copy_param(te_state_dict[te_key], hf_state_dict[hf_key])
+
+ fused_qkv_key = layer_prefix + "self_attention.layernorm_qkv.weight"
+ if fused_qkv_key in te_state_dict:
+ qkv_sources = {
+ "q": layer_prefix + "self_attn.q_proj.weight",
+ "k": layer_prefix + "self_attn.k_proj.weight",
+ "v": layer_prefix + "self_attn.v_proj.weight",
+ }
+ for proj_kind, hf_key in qkv_sources.items():
+ if hf_key in hf_state_dict:
+ _copy_qkv_proj_to_fused(
+ te_state_dict[fused_qkv_key], hf_state_dict[hf_key], proj_kind, config
+ )
+
+
+def _copy_router_gate(hf_state_dict: dict, te_state_dict: dict, layer_prefix: str) -> None:
+ candidates = (
+ layer_prefix + "mlp.gate.weight",
+ layer_prefix + "block_sparse_moe.gate.weight",
+ )
+ te_gate_key = layer_prefix + "mlp.gate.weight"
+ for hf_key in candidates:
+ if hf_key in hf_state_dict and te_gate_key in te_state_dict:
+ _copy_param(te_state_dict[te_gate_key], hf_state_dict[hf_key])
+ return
+
+
+def _packed_expert_candidates(layer_prefix: str) -> tuple[tuple[str, ...], tuple[str, ...]]:
+ gate_up = (
+ layer_prefix + "mlp.experts.gate_up_proj",
+ layer_prefix + "block_sparse_moe.experts.gate_up_proj",
+ )
+ down = (
+ layer_prefix + "mlp.experts.down_proj",
+ layer_prefix + "block_sparse_moe.experts.down_proj",
+ )
+ return gate_up, down
+
+
+def _sequential_op_keys(te_state_dict: dict, layer_prefix: str) -> tuple[list[str], list[str]]:
+ """Return sorted lists of TE Sequential-Op ``weight{i}`` keys, if present."""
+ gate_up_prefix = layer_prefix + "mlp.experts_gate_up."
+ down_prefix = layer_prefix + "mlp.experts_down."
+
+ def _weight_index(key: str) -> int:
+ match = re.search(r"weight(\d+)$", key)
+ assert match is not None
+ return int(match.group(1))
+
+ gate_up_keys = sorted(
+ (k for k in te_state_dict if k.startswith(gate_up_prefix) and re.search(r"weight\d+$", k)),
+ key=_weight_index,
+ )
+ down_keys = sorted(
+ (k for k in te_state_dict if k.startswith(down_prefix) and re.search(r"weight\d+$", k)),
+ key=_weight_index,
+ )
+ return gate_up_keys, down_keys
+
+
+# ---------------------------------------------------------------------------
+# Public entry points
+# ---------------------------------------------------------------------------
+
+
+def replace_params_bf16(
+ hf_state_dict: dict, te_state_dict: dict, config: MixtralConfig
+) -> set[str]:
+ """Map HF Mixtral weights into the BF16 TE state dict.
+
+ Expert weights are placed in stacked ``mlp.experts_{gate_up,down}_weight``
+ parameters (loop path) and/or per-expert Sequential-Op ``weight{i}``
+ parameters (grouped_op path). Both formats are written so a single
+ checkpoint loader supports either ``expert_ffn_mode``.
+
+ Supports both packed HF MoE tensors (``mlp.experts.gate_up_proj``) and
+ older per-expert tensors (``experts.{i}.w{1,2,3}.weight``).
+ """
+ ep_size, ep_rank = _ep_rank_from_config(config)
+ layer_prefixes = _collect_layer_prefixes(hf_state_dict)
+ _copy_top_level(hf_state_dict, te_state_dict)
+
+ for layer_prefix in layer_prefixes:
+ _copy_attention_and_layernorms(hf_state_dict, te_state_dict, layer_prefix, config)
+ _copy_router_gate(hf_state_dict, te_state_dict, layer_prefix)
+
+ packed_gate_up_candidates, packed_down_candidates = _packed_expert_candidates(layer_prefix)
+ te_gate_up_key = layer_prefix + "mlp.experts_gate_up_weight"
+ te_down_key = layer_prefix + "mlp.experts_down_weight"
+
+ # Path A: stacked param (loop path) <- packed HF tensor.
+ for hf_key in packed_gate_up_candidates:
+ if hf_key in hf_state_dict and te_gate_up_key in te_state_dict:
+ te_gate_up = te_state_dict[te_gate_up_key]
+ local_experts = te_gate_up.shape[0]
+ expert_start = ep_rank * local_experts if ep_size > 1 else 0
+ expert_end = expert_start + local_experts
+ _copy_param(
+ te_state_dict[te_gate_up_key],
+ hf_state_dict[hf_key][expert_start:expert_end],
+ )
+ break
+ for hf_key in packed_down_candidates:
+ if hf_key in hf_state_dict and te_down_key in te_state_dict:
+ te_down = te_state_dict[te_down_key]
+ local_experts = te_down.shape[0]
+ expert_start = ep_rank * local_experts if ep_size > 1 else 0
+ expert_end = expert_start + local_experts
+ _copy_param(
+ te_state_dict[te_down_key],
+ hf_state_dict[hf_key][expert_start:expert_end],
+ )
+ break
+
+ # Path B: Sequential-Op per-expert params <- packed HF tensor.
+ te_gate_up_op_keys, te_down_op_keys = _sequential_op_keys(te_state_dict, layer_prefix)
+ if te_gate_up_op_keys and te_down_op_keys:
+ num_local_experts = len(te_gate_up_op_keys)
+ expert_start = ep_rank * num_local_experts if ep_size > 1 else 0
+ expert_end = expert_start + num_local_experts
+ for hf_key in packed_gate_up_candidates:
+ if hf_key in hf_state_dict:
+ hf_gate_up = hf_state_dict[hf_key][expert_start:expert_end]
+ for expert_idx, te_key in enumerate(te_gate_up_op_keys):
+ _copy_param(te_state_dict[te_key], hf_gate_up[expert_idx])
+ break
+ for hf_key in packed_down_candidates:
+ if hf_key in hf_state_dict:
+ hf_down = hf_state_dict[hf_key][expert_start:expert_end]
+ for expert_idx, te_key in enumerate(te_down_op_keys):
+ _copy_param(te_state_dict[te_key], hf_down[expert_idx])
+ break
+
+ # Path C: older HF format with per-expert w1/w2/w3 -> stacked params.
+ if te_gate_up_key in te_state_dict and te_down_key in te_state_dict:
+ te_gate_up = te_state_dict[te_gate_up_key]
+ te_down = te_state_dict[te_down_key]
+ num_local_experts = te_gate_up.shape[0]
+ expert_start = ep_rank * num_local_experts if ep_size > 1 else 0
+ for expert_idx in range(num_local_experts):
+ global_expert_idx = expert_start + expert_idx
+ for expert_prefix in (
+ layer_prefix + f"mlp.experts.{global_expert_idx}.",
+ layer_prefix + f"block_sparse_moe.experts.{global_expert_idx}.",
+ ):
+ w1_key = expert_prefix + "w1.weight"
+ w3_key = expert_prefix + "w3.weight"
+ w2_key = expert_prefix + "w2.weight"
+ if w1_key in hf_state_dict:
+ te_gate_up[expert_idx, : config.intermediate_size].copy_(
+ hf_state_dict[w1_key].to(
+ device=te_gate_up.device, dtype=te_gate_up.dtype
+ )
+ )
+ if w3_key in hf_state_dict:
+ te_gate_up[expert_idx, config.intermediate_size :].copy_(
+ hf_state_dict[w3_key].to(
+ device=te_gate_up.device, dtype=te_gate_up.dtype
+ )
+ )
+ if w2_key in hf_state_dict:
+ te_down[expert_idx].copy_(
+ hf_state_dict[w2_key].to(device=te_down.device, dtype=te_down.dtype)
+ )
+
+ # Path D: older HF format -> Sequential-Op per-expert params.
+ if te_gate_up_op_keys and te_down_op_keys:
+ num_local_experts = len(te_gate_up_op_keys)
+ expert_start = ep_rank * num_local_experts if ep_size > 1 else 0
+ for expert_idx in range(num_local_experts):
+ global_expert_idx = expert_start + expert_idx
+ gate_up = te_state_dict[te_gate_up_op_keys[expert_idx]]
+ down = te_state_dict[te_down_op_keys[expert_idx]]
+ for expert_prefix in (
+ layer_prefix + f"mlp.experts.{global_expert_idx}.",
+ layer_prefix + f"block_sparse_moe.experts.{global_expert_idx}.",
+ ):
+ w1_key = expert_prefix + "w1.weight"
+ w3_key = expert_prefix + "w3.weight"
+ w2_key = expert_prefix + "w2.weight"
+ if w1_key in hf_state_dict:
+ gate_up[: config.intermediate_size].copy_(
+ hf_state_dict[w1_key].to(device=gate_up.device, dtype=gate_up.dtype)
+ )
+ if w3_key in hf_state_dict:
+ gate_up[config.intermediate_size :].copy_(
+ hf_state_dict[w3_key].to(device=gate_up.device, dtype=gate_up.dtype)
+ )
+ if w2_key in hf_state_dict:
+ down.copy_(hf_state_dict[w2_key].to(device=down.device, dtype=down.dtype))
+
+ return layer_prefixes
+
+
+def replace_params_mxfp8(
+ hf_state_dict: dict, te_state_dict: dict, config: MixtralConfig
+) -> set[str]:
+ """Map HF Mixtral weights into the MXFP8 TE state dict.
+
+ Per-expert gate/up rows are interleaved in blocks of 32 to match the GLU
+ layout that the fused MXFP8 grouped-MLP kernel reads.
+
+ Supports both packed HF tensors (``mlp.experts.gate_up_proj`` of shape
+ ``[E, 2I, H]``) and older per-expert tensors
+ (``mlp.experts.{i}.w{1,2,3}.weight``).
+ """
+ ep_size, ep_rank = _ep_rank_from_config(config)
+ layer_prefixes = _collect_layer_prefixes(hf_state_dict)
+ _copy_top_level(hf_state_dict, te_state_dict)
+
+ for layer_prefix in layer_prefixes:
+ _copy_attention_and_layernorms(hf_state_dict, te_state_dict, layer_prefix, config)
+ _copy_router_gate(hf_state_dict, te_state_dict, layer_prefix)
+
+ te_gate_up_op_keys, te_down_op_keys = _sequential_op_keys(te_state_dict, layer_prefix)
+ if not (te_gate_up_op_keys and te_down_op_keys):
+ continue
+
+ num_local_experts = len(te_gate_up_op_keys)
+ expert_start = ep_rank * num_local_experts if ep_size > 1 else 0
+ expert_end = expert_start + num_local_experts
+ intermediate_size = config.intermediate_size
+
+ packed_gate_up_candidates, packed_down_candidates = _packed_expert_candidates(layer_prefix)
+
+ # Path A: newer HF format with packed gate_up tensor [E, 2I, H].
+ packed_gate_up_done = False
+ for hf_key in packed_gate_up_candidates:
+ if hf_key not in hf_state_dict:
+ continue
+ hf_gate_up = hf_state_dict[hf_key][expert_start:expert_end]
+ for expert_idx, te_key in enumerate(te_gate_up_op_keys):
+ gate = hf_gate_up[expert_idx, :intermediate_size]
+ up = hf_gate_up[expert_idx, intermediate_size:]
+ interleaved = _interleave_gate_up(gate, up, GLU_INTERLEAVE_SIZE)
+ _copy_param(te_state_dict[te_key], interleaved)
+ packed_gate_up_done = True
+ break
+
+ packed_down_done = False
+ for hf_key in packed_down_candidates:
+ if hf_key not in hf_state_dict:
+ continue
+ hf_down = hf_state_dict[hf_key][expert_start:expert_end]
+ for expert_idx, te_key in enumerate(te_down_op_keys):
+ _copy_param(te_state_dict[te_key], hf_down[expert_idx])
+ packed_down_done = True
+ break
+
+ if packed_gate_up_done and packed_down_done:
+ continue
+
+ # Path B: older HF format with per-expert w1/w2/w3 weights.
+ for expert_idx in range(num_local_experts):
+ global_expert_idx = expert_start + expert_idx
+ for expert_prefix in (
+ layer_prefix + f"mlp.experts.{global_expert_idx}.",
+ layer_prefix + f"block_sparse_moe.experts.{global_expert_idx}.",
+ ):
+ w1_key = expert_prefix + "w1.weight"
+ w3_key = expert_prefix + "w3.weight"
+ w2_key = expert_prefix + "w2.weight"
+ if w1_key in hf_state_dict and w3_key in hf_state_dict and w2_key in hf_state_dict:
+ gate = hf_state_dict[w1_key]
+ up = hf_state_dict[w3_key]
+ interleaved = _interleave_gate_up(gate, up, GLU_INTERLEAVE_SIZE)
+ _copy_param(te_state_dict[te_gate_up_op_keys[expert_idx]], interleaved)
+ _copy_param(te_state_dict[te_down_op_keys[expert_idx]], hf_state_dict[w2_key])
+ break
+
+ return layer_prefixes
diff --git a/docs/examples/te_mixtral/media/dense_to_sparse.drawio b/docs/examples/te_mixtral/media/dense_to_sparse.drawio
new file mode 100644
index 0000000000..f222bc07e6
--- /dev/null
+++ b/docs/examples/te_mixtral/media/dense_to_sparse.drawio
@@ -0,0 +1,108 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/docs/examples/te_mixtral/media/dense_to_sparse.svg b/docs/examples/te_mixtral/media/dense_to_sparse.svg
new file mode 100644
index 0000000000..971a94a9fb
--- /dev/null
+++ b/docs/examples/te_mixtral/media/dense_to_sparse.svg
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/docs/examples/te_mixtral/media/fused_mlp_path.drawio b/docs/examples/te_mixtral/media/fused_mlp_path.drawio
new file mode 100644
index 0000000000..a70123ab71
--- /dev/null
+++ b/docs/examples/te_mixtral/media/fused_mlp_path.drawio
@@ -0,0 +1,82 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/docs/examples/te_mixtral/media/fused_mlp_path.svg b/docs/examples/te_mixtral/media/fused_mlp_path.svg
new file mode 100644
index 0000000000..cbd7779621
--- /dev/null
+++ b/docs/examples/te_mixtral/media/fused_mlp_path.svg
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/docs/examples/te_mixtral/media/mixtral_decoder_swap.drawio b/docs/examples/te_mixtral/media/mixtral_decoder_swap.drawio
new file mode 100644
index 0000000000..093d2dfa90
--- /dev/null
+++ b/docs/examples/te_mixtral/media/mixtral_decoder_swap.drawio
@@ -0,0 +1,100 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/docs/examples/te_mixtral/media/mixtral_decoder_swap.svg b/docs/examples/te_mixtral/media/mixtral_decoder_swap.svg
new file mode 100644
index 0000000000..c35a47c302
--- /dev/null
+++ b/docs/examples/te_mixtral/media/mixtral_decoder_swap.svg
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/docs/examples/te_mixtral/media/moe_loop_vs_grouped.drawio b/docs/examples/te_mixtral/media/moe_loop_vs_grouped.drawio
new file mode 100644
index 0000000000..72fb7ba9b1
--- /dev/null
+++ b/docs/examples/te_mixtral/media/moe_loop_vs_grouped.drawio
@@ -0,0 +1,85 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/docs/examples/te_mixtral/media/moe_loop_vs_grouped.svg b/docs/examples/te_mixtral/media/moe_loop_vs_grouped.svg
new file mode 100644
index 0000000000..7907a3238d
--- /dev/null
+++ b/docs/examples/te_mixtral/media/moe_loop_vs_grouped.svg
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/docs/examples/te_mixtral/requirements.txt b/docs/examples/te_mixtral/requirements.txt
new file mode 100644
index 0000000000..5ad5e71db2
--- /dev/null
+++ b/docs/examples/te_mixtral/requirements.txt
@@ -0,0 +1,10 @@
+torchao!=0.14.0
+transformer_engine[pytorch]
+transformers==5.8.0
+accelerate==1.13.0
+datasets==4.8.4
+safetensors==0.7.0
+huggingface_hub==1.10.1
+tokenizers==0.22.2
+flash-attn
+nvidia-cudnn-frontend>=1.23.0
diff --git a/docs/examples/te_mixtral/run_finetune_ep.py b/docs/examples/te_mixtral/run_finetune_ep.py
new file mode 100644
index 0000000000..e389b7986b
--- /dev/null
+++ b/docs/examples/te_mixtral/run_finetune_ep.py
@@ -0,0 +1,127 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""EP fine-tune launcher for TE Mixtral.
+
+ python3 run_finetune_ep.py --improvement 0
+ torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 1
+ torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 2
+ torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 3
+
+Improvements (``--ep-size 2`` => 4 experts/rank on 8 GPUs, DP=4):
+
+ 0 = HF baseline BF16 (single process, ``device_map="auto"``).
+ 1 = TE EP BF16, Python loop over experts.
+ 2 = TE EP BF16, GroupedLinear.
+ 3 = TE EP MXFP8 + fused MXFP8 grouped-MLP kernel.
+"""
+
+import argparse
+import os
+import sys
+
+# Improvement 3 needs ``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1`` set before TE is imported,
+# because the fused-grouped-MLP fusion is registered at module-import time
+# inside ``if ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported(): ...``.
+for _i, _arg in enumerate(sys.argv[1:]):
+ if _arg == "--improvement" and _i + 2 < len(sys.argv) and sys.argv[_i + 2] == "3":
+ os.environ["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] = "1"
+ break
+ if _arg == "--improvement=3":
+ os.environ["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] = "1"
+ break
+
+from utils import HyperParameters, run_hf_baseline_finetune, run_te_mixtral_finetune
+
+
+IMPROVEMENT_LABELS = {
+ 0: "HF baseline BF16",
+ 1: "TE EP BF16, Python expert loop",
+ 2: "TE EP BF16, GroupedLinear",
+ 3: "TE EP MXFP8 fused MLP",
+}
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ description="Run TE Mixtral fine-tuning with Expert Parallelism."
+ )
+ parser.add_argument("--hf-token", type=str, default=os.environ.get("HF_TOKEN", ""))
+ parser.add_argument(
+ "--ep-size",
+ type=int,
+ default=2,
+ help="Expert-parallel group size. Default 2 -> 4 experts/rank with 8 GPUs (DP=4).",
+ )
+ parser.add_argument(
+ "--improvement",
+ type=int,
+ choices=(0, 1, 2, 3),
+ default=1,
+ help=(
+ "Improvement: "
+ "0=HF baseline BF16, "
+ "1=TE EP BF16 Python loop, "
+ "2=TE EP BF16 GroupedLinear, "
+ "3=TE EP MXFP8 fused MLP."
+ ),
+ )
+ parser.add_argument("--batch-size", type=int, default=1)
+ parser.add_argument("--max-seq-length", type=int, default=256)
+ parser.add_argument("--warmup-steps", type=int, default=1)
+ parser.add_argument("--train-steps", type=int, default=2)
+ return parser.parse_args()
+
+
+def main() -> None:
+ args = parse_args()
+
+ hp = HyperParameters()
+ hp.model_name = "mistralai/Mixtral-8x7B-v0.1"
+ hp.hf_access_token = args.hf_token
+ hp.batch_size = args.batch_size
+ hp.max_seq_length = args.max_seq_length
+ hp.num_warmup_steps = args.warmup_steps
+ hp.num_training_steps = args.train_steps
+
+ if args.improvement == 0:
+ hp.expert_parallel_size = 1
+ hp.mixed_precision = "bf16"
+ hp.expert_ffn_mode = "loop" # unused: HF baseline doesn't use TE MoE
+ elif args.improvement == 1:
+ hp.expert_parallel_size = args.ep_size
+ hp.mixed_precision = "bf16"
+ hp.expert_ffn_mode = "loop"
+ elif args.improvement == 2:
+ hp.expert_parallel_size = args.ep_size
+ hp.mixed_precision = "bf16"
+ hp.expert_ffn_mode = "grouped_op"
+ elif args.improvement == 3:
+ hp.expert_parallel_size = args.ep_size
+ hp.mixed_precision = "mxfp8"
+ hp.expert_ffn_mode = "grouped_op"
+ hp.model_impl = "te_mixtral_mxfp8"
+
+ print(
+ f"[Improvement {args.improvement}] {IMPROVEMENT_LABELS[args.improvement]}\n"
+ f" mixed_precision={hp.mixed_precision}, ep_size={hp.expert_parallel_size}, "
+ f"expert_ffn_mode={hp.expert_ffn_mode}\n"
+ f" batch_size={hp.batch_size}, max_seq_length={hp.max_seq_length}"
+ )
+
+ if args.improvement == 0:
+ world_size = int(os.environ.get("WORLD_SIZE", "1"))
+ if world_size != 1:
+ raise ValueError(
+ "HF baseline must run as a single process (device_map='auto'). "
+ "Use plain python or torchrun --nproc_per_node=1."
+ )
+ run_hf_baseline_finetune(hp)
+ return
+
+ run_te_mixtral_finetune(hp)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/docs/examples/te_mixtral/te_mixtral.py b/docs/examples/te_mixtral/te_mixtral.py
new file mode 100644
index 0000000000..20c6cd6d61
--- /dev/null
+++ b/docs/examples/te_mixtral/te_mixtral.py
@@ -0,0 +1,1274 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""TransformerEngine-optimized Mixtral model with Mixture of Experts."""
+
+import logging
+import warnings
+from collections import OrderedDict
+from contextlib import nullcontext
+from dataclasses import dataclass
+
+from typing import Any, ClassVar, ContextManager, Protocol
+from typing_extensions import Unpack
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import transformer_engine.common.recipe
+import transformer_engine.pytorch
+import transformers
+
+from transformer_engine.pytorch.ops import GroupedLinear as TEOpsGroupedLinear
+from transformer_engine.pytorch.ops import Sequential as TEOpsSequential
+from transformer_engine.pytorch.ops import SwiGLU as TEOpsSwiGLU
+from transformer_engine.pytorch.attention import InferenceParams
+from transformer_engine.pytorch.attention.inference import PagedKVCacheManager
+from transformer_engine.pytorch.attention.rope import RotaryPositionEmbedding
+from transformer_engine.pytorch.quantization import FP8GlobalStateManager
+from transformer_engine.pytorch.router import fused_moe_aux_loss
+from transformers import MixtralConfig, PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding
+from transformers.utils.generic import TransformersKwargs
+
+logger = logging.getLogger(__name__)
+
+
+AUTO_MAP = {
+ "AutoConfig": "modeling_mixtral_te.TEMixtralConfig",
+ "AutoModel": "modeling_mixtral_te.TEMixtralModel",
+ "AutoModelForCausalLM": "modeling_mixtral_te.TEMixtralForCausalLM",
+}
+
+
+# HF->TE checkpoint mapping lives in ``hf_to_te_weights.py`` so both
+# ``te_mixtral.py`` (BF16) and ``te_mixtral_mxfp8.py`` (MXFP8) can share it.
+from hf_to_te_weights import ( # noqa: E402
+ replace_params_bf16 as replace_params,
+)
+
+
+class TEMixtralConfig(MixtralConfig):
+ """TEMixtral configuration."""
+
+ # Attention input format:
+ # "bshd" = Batch, Sequence, Head, Dimension (standard padded format)
+ # "thd" = Total tokens (packed/unpadded), Head, Dimension (sequence packing format)
+ attn_input_format: str = "thd"
+ self_attn_mask_type: str = "padding_causal"
+ layer_precision: list[str | None] | None = None
+ use_quantized_model_init: bool = False
+ expert_parallel_size: int = 1
+ moe_aux_loss_coeff: float = 0.0
+ # Expert FFN execution mode:
+ # "grouped_op" — fuse all per-rank experts via the TE Sequential-Op
+ # ``transformer_engine.pytorch.ops.GroupedLinear``. On
+ # Blackwell (SM100+) this automatically dispatches to the
+ # graph-safe ``general_grouped_gemm_for_grouped_tensor``
+ # path added in https://github.com/NVIDIA/TransformerEngine/pull/2923 .
+ # "loop" — naive Python loop, one F.linear per expert. Pedagogical
+ # baseline so the tutorial can isolate the GroupedLinear win.
+ expert_ffn_mode: str = "grouped_op"
+
+ def __init__(self, **kwargs):
+ """Initialize the TEMixtralConfig with additional TE-related config options."""
+ super().__init__(**kwargs)
+
+ if self.layer_precision is not None:
+ if len(self.layer_precision) != self.num_hidden_layers:
+ raise ValueError(
+ f"layer_precision must be a list of length {self.num_hidden_layers}"
+ )
+ for precision in self.layer_precision:
+ if precision not in {"fp8", "fp4", None}:
+ raise ValueError(
+ f'layer_precision element must be "fp8", "fp4", or None, got {precision!r}'
+ )
+
+ if self.expert_ffn_mode not in ("grouped_op", "loop"):
+ raise ValueError(
+ f'expert_ffn_mode must be "grouped_op" or "loop", got {self.expert_ffn_mode!r}'
+ )
+
+ if self.num_local_experts % self.expert_parallel_size != 0:
+ raise ValueError(
+ f"num_local_experts ({self.num_local_experts}) must be divisible by "
+ f"expert_parallel_size ({self.expert_parallel_size})"
+ )
+
+
+@dataclass
+class DispatchOutput:
+ """Output of TokenDispatcher.dispatch().
+
+ Attributes:
+ expert_input: Tokens sorted by local expert, shape ``[total_recv_tokens, H]``.
+ tokens_per_expert: Token count per local expert.
+ handle: Opaque state needed by ``combine()`` to reverse the dispatch.
+ """
+
+ expert_input: torch.Tensor
+ tokens_per_expert: list[int]
+ handle: Any
+
+
+class TokenDispatcher(Protocol):
+ """Protocol for MoE token dispatch/combine strategies.
+
+ Encapsulates the full dispatch cycle (permute -> communicate -> sort) and
+ combine cycle (unsort -> communicate -> unpermute) so that the MoE block
+ is agnostic to the communication backend (NCCL all-to-all, HybridEP, etc.).
+ """
+
+ def dispatch(
+ self,
+ hidden_states: torch.Tensor,
+ selected_experts: torch.Tensor,
+ routing_weights: torch.Tensor,
+ ) -> DispatchOutput:
+ """Dispatch tokens to their assigned experts.
+
+ Args:
+ hidden_states: Flattened input tensor of shape ``[N, H]``.
+ selected_experts: Expert assignments, shape ``[N, top_k]``, int.
+ routing_weights: Normalized routing probabilities, shape ``[N, top_k]``, float32.
+
+ Returns:
+ DispatchOutput with expert-sorted tokens, per-expert counts, and an opaque handle.
+ """
+ ...
+
+ def combine(
+ self,
+ expert_output: torch.Tensor,
+ handle: Any,
+ ) -> torch.Tensor:
+ """Combine expert outputs back to the original token order.
+
+ Args:
+ expert_output: Expert output tensor of shape ``[total_recv_tokens, H]``.
+ handle: Opaque state from ``dispatch()``.
+
+ Returns:
+ Combined output tensor of shape ``[N, H]`` with routing weights applied.
+ """
+ ...
+
+ def set_ep_group(self, ep_group: dist.ProcessGroup) -> None:
+ """Set the expert-parallel process group for communication."""
+ ...
+
+
+class TEMixtralPreTrainedModel(PreTrainedModel):
+ """Base class for TEMixtral models."""
+
+ config_class = TEMixtralConfig
+ base_model_prefix = "model"
+ _no_split_modules = ("TEMixtralDecoderLayer",)
+ _skip_keys_device_placement = ("past_key_values",)
+ _do_not_quantize = (
+ "lm_head",
+ "model.layers.*.mlp.gate",
+ ) # Flag for testing that these layers are not quantized.
+
+ def init_empty_weights(self):
+ """Handles moving the model from the meta device to the cuda device and initializing the weights."""
+ for module in self.modules():
+ if hasattr(module, "reset_parameters"):
+ module.reset_parameters()
+
+ # After reset_parameters materializes GroupedLinear views on CUDA,
+ # re-stack them into the authoritative stacked parameters.
+ for module in self.modules():
+ if isinstance(module, TEMixtralSparseMoeBlock):
+ module._restack_from_views()
+
+ self.model.embed_tokens.to_empty(device="cuda")
+ self.model.embed_tokens.apply(self._init_weights)
+
+ self.model.rotary_emb.inv_freq = LlamaRotaryEmbedding(config=self.model.config).inv_freq.to(
+ "cuda"
+ )
+
+ self.tie_weights()
+
+ def _init_weights(self, module):
+ """Initialize module weights.
+
+ We only use this method for standard pytorch modules, TE modules handle their own weight initialization through
+ `init_method` parameters and the `reset_parameters` method.
+ """
+ if module.__module__.startswith("transformer_engine.pytorch"):
+ return
+
+ super()._init_weights(module)
+
+ def state_dict(self, *args, **kwargs):
+ """Override state_dict to filter out TransformerEngine's _extra_state keys."""
+ state_dict = super().state_dict(*args, **kwargs)
+ return {k: v for k, v in state_dict.items() if not k.endswith("_extra_state")}
+
+
+class TEMixtralSparseMoeBlock(nn.Module):
+ """Mixture of Experts block using TransformerEngine GroupedLinear."""
+
+ def __init__(self, config: MixtralConfig, dispatcher: TokenDispatcher | None = None):
+ """Initialize the sparse MoE block."""
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.num_experts = config.num_local_experts
+ self.top_k = config.num_experts_per_tok
+ self.jitter_noise = config.router_jitter_noise
+
+ self.ep_size = getattr(config, "expert_parallel_size", 1)
+ self.num_local_experts = self.num_experts // self.ep_size
+ self.expert_ffn_mode = getattr(config, "expert_ffn_mode", "grouped_op")
+ self._uses_stacked_expert_weights = self.expert_ffn_mode != "grouped_op"
+ self.moe_aux_loss_coeff = getattr(config, "moe_aux_loss_coeff", 0.0)
+ self._aux_loss: torch.Tensor = torch.tensor(0.0)
+ self.initializer_range = config.initializer_range
+
+ self.dispatcher: TokenDispatcher = dispatcher or AllToAllTokenDispatcher(
+ self.num_experts,
+ self.num_local_experts,
+ self.hidden_size,
+ self.ep_size,
+ )
+
+ device = "meta" if torch.get_default_device() == torch.device("meta") else "cuda"
+
+ def _init_method(x):
+ torch.nn.init.normal_(x, mean=0.0, std=config.initializer_range)
+
+ # Router always outputs num_experts logits (replicated across EP ranks)
+ with transformer_engine.pytorch.quantized_model_init(enabled=False):
+ self.gate = transformer_engine.pytorch.Linear(
+ self.hidden_size,
+ self.num_experts,
+ bias=False,
+ device=device,
+ params_dtype=config.dtype,
+ init_method=_init_method,
+ )
+
+ # Expert FFNs — only num_local_experts per rank when EP > 1.
+ # Both ``grouped_op`` (improvement 2) and ``loop`` (improvement 1) allocate the same
+ # pair of GroupedLinear ops; ``loop`` just routes its tokens through
+ # them one-expert-at-a-time in ``_expert_ffn``.
+ self.experts_gate_up = TEOpsGroupedLinear(
+ num_groups=self.num_local_experts,
+ in_features=self.hidden_size,
+ out_features=2 * self.intermediate_size,
+ bias=False,
+ dtype=config.dtype,
+ device=device,
+ )
+ self.experts_down = TEOpsGroupedLinear(
+ num_groups=self.num_local_experts,
+ in_features=self.intermediate_size,
+ out_features=self.hidden_size,
+ bias=False,
+ dtype=config.dtype,
+ device=device,
+ )
+ # ``grouped_op`` runs the two GroupedLinears + SwiGLU through TE's
+ # fusible Sequential wrapper so the OperationFuser can collapse them.
+ if self.expert_ffn_mode == "grouped_op":
+ object.__setattr__(
+ self,
+ "_experts_ffn_op",
+ TEOpsSequential(self.experts_gate_up, TEOpsSwiGLU(), self.experts_down),
+ )
+
+ if self._uses_stacked_expert_weights:
+ # Stack per-expert weights into single parameters (authoritative weight store).
+ # GroupedLinear's _parameters dict is emptied; weight attributes are set as views
+ # so that reset_parameters() / _get_weight_tensors() can still find them.
+ self.experts_gate_up_weight = nn.Parameter(
+ torch.stack(
+ [
+ self.experts_gate_up._parameters.pop(f"weight{i}").data
+ for i in range(self.num_local_experts)
+ ]
+ )
+ ) # [num_local_experts, 2*intermediate_size, hidden_size]
+
+ self.experts_down_weight = nn.Parameter(
+ torch.stack(
+ [
+ self.experts_down._parameters.pop(f"weight{i}").data
+ for i in range(self.num_local_experts)
+ ]
+ )
+ ) # [num_local_experts, hidden_size, intermediate_size]
+
+ # Set views back on GroupedLinear so getattr(self, "weight{i}") still works
+ # (needed by GroupedLinear.reset_parameters and _get_weight_tensors).
+ self._sync_expert_views()
+
+ def _restack_from_views(self) -> None:
+ """Re-create stacked parameters on CUDA after meta init.
+
+ Called by ``init_empty_weights()`` after ``reset_parameters()`` has been called
+ on all TE modules. Since GroupedLinear has no registered parameters (we popped them),
+ its ``reset_parameters()`` cannot move them from meta to CUDA. This method explicitly
+ creates the stacked parameters on CUDA and reinitializes them.
+ """
+ if not self._uses_stacked_expert_weights:
+ return
+
+ device = torch.cuda.current_device()
+ for attr_name in ("experts_gate_up_weight", "experts_down_weight"):
+ old_param = getattr(self, attr_name)
+ new_data = torch.empty_like(old_param, device=device)
+ torch.nn.init.normal_(new_data, mean=0.0, std=self.initializer_range)
+ setattr(self, attr_name, nn.Parameter(new_data))
+
+ # Re-sync views to point to the new stacked parameter
+ self._sync_expert_views()
+
+ def _sync_expert_views(self) -> None:
+ """Set GroupedLinear weight attributes as views of the stacked parameters.
+
+ GroupedLinear internally uses ``getattr(self, f"weight{i}")`` in methods like
+ ``reset_parameters()`` and ``_get_weight_tensors()``. After popping the original
+ parameters, we set views of the stacked tensor so these methods keep working.
+ Uses ``object.__setattr__`` to bypass ``nn.Module.__setattr__`` and avoid
+ re-registering them as parameters.
+ """
+ if not self._uses_stacked_expert_weights:
+ return
+ gate_up_w = self.experts_gate_up_weight
+ 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
+ for i in range(self.num_local_experts):
+ object.__setattr__(self.experts_down, f"weight{i}", down_w[i])
+
+ def set_ep_group(self, ep_group: dist.ProcessGroup) -> None:
+ """Set the expert-parallel process group for token dispatch.
+
+ Must be called before the first forward pass when ``ep_size > 1``.
+ """
+ self.dispatcher.set_ep_group(ep_group)
+
+ def _expert_ffn(self, tokens: torch.Tensor, m_splits: list[int]) -> torch.Tensor:
+ """Run the expert SwiGLU FFN (gate_up -> silu -> down) per local expert."""
+ if self.expert_ffn_mode == "grouped_op":
+ # Run gate_up -> SwiGLU -> down as one TE fusible op group.
+ # Each GroupedLinear consumes the same per-expert split sizes.
+ split_sizes = torch.tensor(m_splits, dtype=torch.int32, device=tokens.device)
+ return self._experts_ffn_op(tokens, split_sizes, split_sizes)
+ elif self.expert_ffn_mode == "loop":
+ # Naive HF-style loop: one F.linear per expert against a slice of the
+ # stacked weight. Same checkpoint as the grouped path; only kernel
+ # dispatch differs.
+ # IMPORTANT: do NOT go through ``.data`` here — that detaches
+ # the tensor from autograd, so backward never reaches the
+ # expert ``nn.Parameter`` and the optimizer silently skips
+ # ~95% of the model.
+ gate_up_w = self.experts_gate_up_weight
+ down_w = self.experts_down_weight
+ outputs = []
+ for i, chunk in enumerate(torch.split(tokens, m_splits, dim=0)):
+ if chunk.shape[0] == 0:
+ outputs.append(chunk)
+ continue
+ gate, up = torch.nn.functional.linear(chunk, gate_up_w[i]).chunk(2, dim=-1)
+ outputs.append(
+ torch.nn.functional.linear(torch.nn.functional.silu(gate) * up, down_w[i])
+ )
+ return torch.cat(outputs, dim=0)
+
+ raise RuntimeError(f"Unknown expert_ffn_mode: {self.expert_ffn_mode!r}")
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ """Forward pass for the MoE block.
+
+ Args:
+ hidden_states: Input tensor of shape [B, S, H] (bshd) or [T, H] (thd).
+
+ Returns:
+ Output tensor of the same shape as the input.
+ """
+ original_shape = hidden_states.shape
+
+ # Apply multiplicative jitter noise to hidden states during training to encourage load balancing
+ if self.training and self.jitter_noise > 0:
+ hidden_states = hidden_states * torch.empty_like(hidden_states).uniform_(
+ 1.0 - self.jitter_noise, 1.0 + self.jitter_noise
+ )
+
+ # Flatten to [N, H] for routing
+ if hidden_states.dim() == 3:
+ hidden_states = hidden_states.reshape(-1, self.hidden_size)
+
+ # Router: compute expert assignments
+ with transformer_engine.pytorch.autocast(enabled=False):
+ # Keep the router logits in bf16 during FP8 training
+ router_logits = self.gate(hidden_states) # [N, num_experts]
+
+ # Compute the full (N, E) softmax probs once and reuse them for both
+ # top-k and the fused aux loss kernel.
+ softmax_probs = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float32)
+ routing_weights, selected_experts = torch.topk(
+ softmax_probs, self.top_k, dim=-1
+ ) # [N, top_k]
+ # Normalize routing weights
+ routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
+
+ # Auxiliary load-balancing loss (switch transformer style). Use TE's
+ # fused router kernel — fold bincount + softmax-mean + sum into one
+ # CUDA launch.
+ if self.moe_aux_loss_coeff > 0:
+ num_tokens = hidden_states.shape[0]
+ tokens_per_expert = torch.bincount(
+ selected_experts.reshape(-1), minlength=self.num_experts
+ ).to(torch.int32)
+ self._aux_loss = fused_moe_aux_loss(
+ probs=softmax_probs,
+ tokens_per_expert=tokens_per_expert,
+ total_num_tokens=num_tokens,
+ num_experts=self.num_experts,
+ topk=self.top_k,
+ coeff=self.moe_aux_loss_coeff,
+ )
+ else:
+ self._aux_loss = torch.tensor(0.0, device=hidden_states.device)
+
+ # Populate GroupedLinear weight attributes from stacked parameters.
+ self._sync_expert_views()
+
+ if isinstance(self.dispatcher, AllToAllTokenDispatcher):
+ pad_to_multiple = None
+ if (
+ self.expert_ffn_mode == "grouped_op"
+ and FP8GlobalStateManager.is_fp8_enabled()
+ and FP8GlobalStateManager.get_fp8_recipe().mxfp8()
+ ):
+ pad_to_multiple = 128
+ self.dispatcher.pad_to_multiple = pad_to_multiple
+
+ dispatch_output = self.dispatcher.dispatch(hidden_states, selected_experts, routing_weights)
+
+ expert_input = dispatch_output.expert_input
+ tokens_per_expert = dispatch_output.tokens_per_expert
+
+ expert_output = self._expert_ffn(expert_input, tokens_per_expert)
+
+ output = self.dispatcher.combine(expert_output, dispatch_output.handle)
+
+ return output.reshape(original_shape)
+
+
+class TEMixtralDecoderLayer(nn.Module):
+ """Mixtral decoder layer using TE attention and MoE MLP."""
+
+ def __init__(
+ self, config: MixtralConfig, layer_idx: int, dispatcher: TokenDispatcher | None = None
+ ):
+ """Initialize the decoder layer."""
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ device = "meta" if torch.get_default_device() == torch.device("meta") else "cuda"
+
+ def _init_method(x):
+ torch.nn.init.normal_(x, mean=0.0, std=config.initializer_range)
+
+ self.self_attention = transformer_engine.pytorch.MultiheadAttention(
+ hidden_size=config.hidden_size,
+ num_attention_heads=config.num_attention_heads,
+ num_gqa_groups=config.num_key_value_heads,
+ bias=False,
+ layernorm_epsilon=config.rms_norm_eps,
+ attention_dropout=0,
+ fuse_qkv_params=True,
+ qkv_weight_interleaved=True,
+ normalization="RMSNorm",
+ input_layernorm=True,
+ qkv_format=config.attn_input_format,
+ attn_mask_type=config.self_attn_mask_type,
+ layer_number=layer_idx + 1,
+ params_dtype=config.dtype,
+ device=device,
+ init_method=_init_method,
+ output_layer_init_method=_init_method,
+ )
+
+ self.post_attention_layernorm = transformer_engine.pytorch.RMSNorm(
+ config.hidden_size,
+ eps=config.rms_norm_eps,
+ dtype=config.dtype,
+ device=device,
+ )
+
+ self.mlp = TEMixtralSparseMoeBlock(config, dispatcher)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ rotary_pos_emb: torch.Tensor | None = None,
+ inference_params: InferenceParams | None = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ """Forward pass for the decoder layer."""
+ # Self attention with fused input layernorm
+ attn_output = self.self_attention(
+ hidden_states,
+ attention_mask=attention_mask,
+ rotary_pos_emb=rotary_pos_emb,
+ inference_params=inference_params,
+ cu_seqlens_q=kwargs.get("cu_seqlens_q", None),
+ cu_seqlens_kv=kwargs.get("cu_seqlens_kv", None),
+ cu_seqlens_q_padded=kwargs.get("cu_seqlens_q_padded", None),
+ cu_seqlens_kv_padded=kwargs.get("cu_seqlens_kv_padded", None),
+ max_seqlen_q=kwargs.get("max_seqlen_q", None),
+ max_seqlen_kv=kwargs.get("max_seqlen_kv", None),
+ pad_between_seqs=kwargs.get("pad_between_seqs", None),
+ )
+
+ # Residual connection
+ hidden_states = hidden_states + attn_output
+
+ # Post-attention layernorm + MoE MLP + residual
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+class TEMixtralModel(TEMixtralPreTrainedModel):
+ """Mixtral model implemented in Transformer Engine."""
+
+ def __init__(
+ self,
+ config: MixtralConfig,
+ fp8_recipe: transformer_engine.common.recipe.Recipe | None = None,
+ fp4_recipe: transformer_engine.common.recipe.Recipe | None = None,
+ dispatcher: TokenDispatcher | None = None,
+ ):
+ """Initialize the TEMixtral model.
+
+ Args:
+ config: The configuration of the model.
+ fp8_recipe: The FP8 recipe for the model.
+ fp4_recipe: The FP4 recipe for the model.
+ dispatcher: The token dispatcher for the model. If None, the default AllToAllTokenDispatcher will be used.
+ """
+ super().__init__(config)
+ self.config = config
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+ self._fp8_recipe: transformer_engine.common.recipe.Recipe | None = fp8_recipe
+ self._fp4_recipe: transformer_engine.common.recipe.Recipe | None = fp4_recipe
+
+ if fp8_recipe is not None and self.config.layer_precision is None:
+ if fp4_recipe is not None:
+ raise RuntimeError(
+ "Both FP8 and FP4 recipes provided, but no layer precision provided."
+ )
+
+ warnings.warn(
+ "No layer precision provided, using FP8 recipe for all layers.", UserWarning
+ )
+ self.config.layer_precision = ["fp8"] * self.config.num_hidden_layers
+
+ self.embed_tokens = nn.Embedding(
+ config.vocab_size, config.hidden_size, self.padding_idx, dtype=config.dtype
+ )
+
+ layers: list[TEMixtralDecoderLayer] = []
+ for layer_idx in range(config.num_hidden_layers):
+ with self.get_autocast_context(layer_idx, init=True):
+ layers += [TEMixtralDecoderLayer(config, layer_idx, dispatcher)]
+
+ self.layers = nn.ModuleList(layers)
+
+ self.norm = transformer_engine.pytorch.RMSNorm(
+ config.hidden_size,
+ eps=config.rms_norm_eps,
+ dtype=config.dtype,
+ device="meta" if torch.get_default_device() == torch.device("meta") else "cuda",
+ )
+
+ self.rotary_emb = RotaryPositionEmbedding(config.hidden_size // config.num_attention_heads)
+ self.rotary_emb.inv_freq = LlamaRotaryEmbedding(config=config).inv_freq
+
+ self.gradient_checkpointing = False
+
+ self.post_init()
+
+ def set_ep_groups(self, ep_group: dist.ProcessGroup) -> None:
+ """Propagate an expert-parallel process group to every MoE block.
+
+ Args:
+ ep_group: The EP process group to set on each ``TEMixtralSparseMoeBlock``.
+ """
+ for layer in self.layers:
+ layer.mlp.set_ep_group(ep_group)
+
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ past_key_values: InferenceParams | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPast:
+ """Forward pass for the TEMixtral model."""
+ all_hidden_states = []
+ output_hidden_states = kwargs.get("output_hidden_states", False)
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
+
+ hidden_states = inputs_embeds
+
+ # TE-specific input handling
+ has_thd_input = [
+ x in kwargs for x in ["cu_seq_lens_q", "cu_seq_lens_k", "max_length_q", "max_length_k"]
+ ]
+ decode_without_mask = (
+ isinstance(past_key_values, InferenceParams)
+ and hidden_states.dim() == 3
+ and hidden_states.size(1) == 1
+ )
+ should_pack_inputs = (
+ not any(has_thd_input)
+ and self.config.attn_input_format == "thd"
+ and not decode_without_mask
+ )
+
+ if should_pack_inputs:
+ assert (
+ attention_mask is not None
+ ), "Attention mask is required when packing BSHD inputs."
+ batch_size = hidden_states.size(0)
+ padded_seq_len = hidden_states.size(1)
+ hidden_states, indices, cu_seqlens, max_seqlen, _ = _unpad_input(
+ hidden_states, attention_mask
+ )
+
+ # MXFP8 block scaling requires the token dim divisible by 32.
+ # After THD unpadding the total token count is data-dependent.
+ thd_orig_tokens = hidden_states.shape[0]
+ thd_remainder = thd_orig_tokens % 32
+ if thd_remainder != 0:
+ thd_pad = 32 - thd_remainder
+ hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, thd_pad))
+ # Extend cu_seqlens: add padding tokens to the last sequence
+ cu_seqlens = cu_seqlens.clone()
+ cu_seqlens[-1] = cu_seqlens[-1] + thd_pad
+ max_seqlen = max_seqlen + thd_pad
+
+ kwargs["cu_seq_lens_q"] = kwargs["cu_seq_lens_k"] = cu_seqlens
+ kwargs["max_length_q"] = kwargs["max_length_k"] = max_seqlen
+
+ if (
+ self.config.attn_input_format == "thd"
+ and hidden_states.dim() == 3
+ and hidden_states.size(0) == 1
+ ):
+ hidden_states = hidden_states.squeeze(0)
+
+ if (
+ self.config.attn_input_format == "bshd"
+ and attention_mask is not None
+ and attention_mask.dim() == 2
+ ):
+ # Convert HF mask (1=attend, 0=pad) to TE boolean mask (True=masked, False=attend)
+ attention_mask = ~attention_mask[:, None, None, :].bool()
+
+ 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)))
+
+ with torch.autocast(device_type="cuda", enabled=False):
+ te_rope_emb = self.rotary_emb(max_seq_len=self.config.max_position_embeddings)
+
+ with self.get_autocast_context(None, outer=True):
+ for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
+ if output_hidden_states:
+ all_hidden_states = (*all_hidden_states, hidden_states)
+
+ with self.get_autocast_context(layer_idx):
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=(
+ None if self.config.attn_input_format == "thd" else attention_mask
+ ),
+ rotary_pos_emb=te_rope_emb,
+ inference_params=past_key_values,
+ cu_seqlens_q=kwargs.get("cu_seq_lens_q", None),
+ cu_seqlens_kv=kwargs.get("cu_seq_lens_k", None),
+ cu_seqlens_q_padded=kwargs.get("cu_seq_lens_q_padded", None),
+ cu_seqlens_kv_padded=kwargs.get("cu_seq_lens_k_padded", None),
+ max_seqlen_q=kwargs.get("max_length_q", None),
+ max_seqlen_kv=kwargs.get("max_length_k", None),
+ pad_between_seqs=kwargs.get("pad_between_seqs", None),
+ )
+
+ hidden_states = self.norm(hidden_states)
+
+ if output_hidden_states:
+ all_hidden_states = (*all_hidden_states, hidden_states)
+
+ if should_pack_inputs:
+ if thd_remainder != 0:
+ hidden_states = hidden_states[:thd_orig_tokens]
+ hidden_states = _pad_input(hidden_states, indices, batch_size, padded_seq_len)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ hidden_states=all_hidden_states if output_hidden_states else None,
+ )
+
+ def get_autocast_context(
+ self, layer_number: int | None, init: bool = False, outer: bool = False
+ ) -> ContextManager:
+ """Return the appropriate TE autocast context manager for a given layer.
+
+ This function handles both the quantized_model_init during layer creation and the te.autocast() during layer
+ forward pass.
+
+ Args:
+ layer_number: The 0-indexed layer number.
+ init: Whether to return a `quantized_model_init` context for layer initialization.
+ outer: Whether to return a global te.autocast() context to wrap the entire model stack.
+ """
+ if self.config.layer_precision is None:
+ return nullcontext()
+
+ if outer:
+ if "fp8" not in self.config.layer_precision:
+ return nullcontext()
+ if self._fp8_recipe is None:
+ warnings.warn("No FP8 recipe provided, using default recipe.", UserWarning)
+ return transformer_engine.pytorch.autocast(enabled=True, recipe=self._fp8_recipe)
+
+ precision = self.config.layer_precision[layer_number]
+ recipe = {"fp8": self._fp8_recipe, "fp4": self._fp4_recipe}.get(precision)
+
+ if init and self.config.use_quantized_model_init:
+ if precision in ("fp8", "fp4"):
+ return transformer_engine.pytorch.quantized_model_init(recipe=recipe)
+ return nullcontext()
+
+ if precision == "fp8":
+ if recipe is None:
+ warnings.warn("No FP8 recipe provided, using default recipe.", UserWarning)
+ return transformer_engine.pytorch.autocast(enabled=True, recipe=recipe)
+ if precision == "fp4":
+ if recipe is None:
+ raise RuntimeError("No FP4 recipe provided, but layer precision is set to FP4.")
+ return transformer_engine.pytorch.autocast(enabled=True, recipe=recipe)
+ return transformer_engine.pytorch.autocast(enabled=False)
+
+
+class TEMixtralForCausalLM(TEMixtralPreTrainedModel, transformers.GenerationMixin):
+ """Mixtral model with causal language head."""
+
+ _tied_weights_keys: ClassVar[list[str]] = []
+
+ def __init__(
+ self,
+ config,
+ fp8_recipe: transformer_engine.common.recipe.Recipe | None = None,
+ fp4_recipe: transformer_engine.common.recipe.Recipe | None = None,
+ dispatcher: TokenDispatcher | None = None,
+ ):
+ """Initialize the TEMixtralForCausalLM model.
+
+ Args:
+ config: The configuration of the model.
+ fp8_recipe: The FP8 recipe for the model.
+ fp4_recipe: The FP4 recipe for the model.
+ dispatcher: The token dispatcher for expert parallelism. If None, the default
+ AllToAllTokenDispatcher will be used.
+ """
+ super().__init__(config)
+ self.model = TEMixtralModel(
+ config, fp8_recipe=fp8_recipe, fp4_recipe=fp4_recipe, dispatcher=dispatcher
+ )
+ self.vocab_size = config.vocab_size
+
+ with transformer_engine.pytorch.quantized_model_init(enabled=False):
+ self.lm_head = transformer_engine.pytorch.Linear(
+ config.hidden_size,
+ config.vocab_size,
+ bias=False,
+ params_dtype=config.dtype,
+ device="meta" if torch.get_default_device() == torch.device("meta") else "cuda",
+ init_method=lambda x: torch.nn.init.normal_(
+ x, mean=0.0, std=config.initializer_range
+ ),
+ )
+
+ self.post_init()
+
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ past_key_values: tuple[tuple[torch.Tensor, ...], ...] | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ shift_labels: torch.Tensor | None = None,
+ use_cache: bool | None = None,
+ cache_position: torch.Tensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> CausalLMOutputWithPast:
+ """Forward pass for the TEMixtralForCausalLM model."""
+ outputs: BaseModelOutputWithPast = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ slice_indices = (
+ slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ )
+
+ with transformer_engine.pytorch.autocast(enabled=False):
+ if hidden_states.ndim == 3:
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+ else:
+ logits = self.lm_head(hidden_states[slice_indices, :])
+
+ loss = None
+ if labels is not None or shift_labels is not None:
+ loss = self.loss_function(
+ logits=logits,
+ labels=labels,
+ shift_labels=shift_labels,
+ vocab_size=self.config.vocab_size,
+ **kwargs,
+ )
+
+ # Collect auxiliary load-balancing loss from all MoE layers
+ if self.config.moe_aux_loss_coeff > 0 and loss is not None:
+ aux_loss = sum(layer.mlp._aux_loss for layer in self.model.layers)
+ loss = loss + aux_loss
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+# Required for torch.compile'd functions below (_pad_input, _unpad_input, _build_expert_sort_indices)
+# that use data-dependent scalar values (e.g., max_seqlen_in_batch.item()) or produce tensors
+# whose shape depends on input data (e.g., repeat_interleave with tensor counts).
+# These must be set at module level because torch.compile traces lazily on first call,
+# so a scoped setting would not be active at trace time.
+torch._dynamo.config.capture_scalar_outputs = True
+torch._dynamo.config.capture_dynamic_output_shape_ops = True
+
+
+@torch.compile
+def _pad_input(hidden_states, indices, batch, seqlen):
+ """Convert a THD tensor to a BSHD equivalent tensor.
+
+ Adapted from huggingface/transformers/modeling_flash_attention_utils.py
+ """
+ dim = hidden_states.shape[1:]
+ output = torch.zeros(
+ (batch * seqlen), *dim, device=hidden_states.device, dtype=hidden_states.dtype
+ )
+ output[indices] = hidden_states
+ return output.view(batch, seqlen, *dim)
+
+
+@torch.compile
+def _unpad_input(hidden_states, attention_mask, unused_mask=None):
+ """Convert a BSHD tensor to a THD equivalent tensor.
+
+ Adapted from huggingface/transformers/modeling_flash_attention_utils.py
+ """
+ batch_size = hidden_states.size(0)
+ seq_length = hidden_states.size(1)
+
+ if attention_mask.shape[1] != seq_length:
+ return (
+ hidden_states.squeeze(1),
+ torch.arange(batch_size, dtype=torch.int64, device=hidden_states.device),
+ torch.arange(batch_size + 1, dtype=torch.int32, device=hidden_states.device),
+ 1,
+ 1,
+ )
+
+ all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask
+ seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32)
+ used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+ indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten()
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
+ cu_seqlens = torch.nn.functional.pad(
+ torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)
+ )
+
+ return (
+ hidden_states.reshape(-1, *hidden_states.shape[2:])[indices],
+ indices,
+ cu_seqlens,
+ max_seqlen_in_batch,
+ used_seqlens_in_batch,
+ )
+
+
+class HFInferenceParams(InferenceParams):
+ """Extension of the InferenceParams class to support HF generate() and beam search."""
+
+ def get_seq_length(self, layer_idx: int = 0) -> int:
+ """Return the current cached sequence length.
+
+ Required by HuggingFace transformers generate() to determine how many
+ tokens have already been cached.
+ """
+ if not self.sequences:
+ return 0
+ return max(self.sequences.values())
+
+ def reorder_cache(self, beam_idx: torch.LongTensor):
+ """Reorder the cache based on the beam indices."""
+ if isinstance(self.cache_manager, PagedKVCacheManager):
+ raise NotImplementedError("Beam search is not supported for paged cache manager.")
+ for layer_number, (key_cache, value_cache) in self.cache_manager.cache.items():
+ updated_key_cache = key_cache.index_select(0, beam_idx)
+ updated_value_cache = value_cache.index_select(0, beam_idx)
+ self.cache_manager.cache[layer_number] = (updated_key_cache, updated_value_cache)
+
+
+@torch.compile(fullgraph=True)
+def _build_expert_sort_indices(recv_counts: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ """Build sort and unsort index tensors for reordering received tokens by local expert.
+
+ After all-to-all, tokens arrive grouped by source rank:
+ ``[src0_exp0..src0_expL, src1_exp0..src1_expL, ...]``. ``GroupedLinear`` expects them
+ grouped by expert: ``[all_exp0, all_exp1, ...]``.
+
+ Uses only vectorized tensor operations (no ``.item()`` calls or Python-level loops)
+ so that it is compatible with ``torch.compile(fullgraph=True)``.
+
+ Args:
+ recv_counts: Integer tensor of shape ``[ep_size, num_local_experts]`` giving the
+ number of tokens received from each source rank for each local expert.
+
+ Returns:
+ A ``(sort_indices, unsort_indices)`` pair of 1-D ``int64`` tensors that can be
+ used to reorder and restore the token dimension.
+ """
+ ep_size, num_local_experts = recv_counts.shape
+ device = recv_counts.device
+ num_blocks = ep_size * num_local_experts
+
+ # Source-grouped (row-major) block offsets: [s0e0, s0e1, ..., s1e0, s1e1, ...]
+ counts_src = recv_counts.reshape(-1).long()
+ offsets_src = torch.zeros(num_blocks, dtype=torch.long, device=device)
+ offsets_src[1:] = counts_src[:-1].cumsum(0)
+
+ # Expert-grouped (column-major) block offsets: [e0s0, e0s1, ..., e1s0, e1s1, ...]
+ counts_exp = recv_counts.t().contiguous().reshape(-1).long()
+ offsets_exp = torch.zeros(num_blocks, dtype=torch.long, device=device)
+ offsets_exp[1:] = counts_exp[:-1].cumsum(0)
+
+ total = counts_src.sum()
+
+ # Mapping from source block index (s * L + e) to expert block index (e * S + s)
+ s_idx = torch.arange(ep_size, device=device).unsqueeze(1).expand(ep_size, num_local_experts)
+ e_idx = (
+ torch.arange(num_local_experts, device=device)
+ .unsqueeze(0)
+ .expand(ep_size, num_local_experts)
+ )
+ src_to_exp = (e_idx * ep_size + s_idx).reshape(-1)
+
+ # Per-block positional shift from source layout to expert layout
+ shifts = offsets_exp[src_to_exp] - offsets_src
+
+ # Expand per-block shifts to per-token
+ token_shifts = shifts.repeat_interleave(counts_src)
+
+ # Map each source-grouped position to its expert-grouped destination
+ src_positions = torch.arange(total, device=device)
+ dst_positions = src_positions + token_shifts
+
+ # sort_indices[exp_pos] = src_pos (gathers source tokens into expert order)
+ sort_indices = torch.empty(total, dtype=torch.long, device=device)
+ sort_indices[dst_positions] = src_positions
+
+ # unsort_indices: inverse permutation (restores expert-ordered output to source order)
+ unsort_indices = torch.empty_like(sort_indices)
+ unsort_indices[sort_indices] = torch.arange(total, device=device)
+
+ return sort_indices, unsort_indices
+
+
+@dataclass
+class _AllToAllHandle:
+ """Opaque handle for AllToAllTokenDispatcher, storing state between dispatch and combine."""
+
+ row_id_map: torch.Tensor
+ routing_weights: torch.Tensor
+ restore_shape: torch.Size
+ map_type: str = "index"
+ pad_offsets: torch.Tensor | None = None
+ unsort_indices: torch.Tensor | None = None
+ input_split_sizes: list[int] | None = None
+ output_split_sizes: list[int] | None = None
+
+
+class _DifferentiableAllToAll(torch.autograd.Function):
+ """Differentiable wrapper around dist.all_to_all_single.
+
+ The forward pass performs the standard all-to-all communication.
+ The backward pass reverses the communication direction (swapping
+ input/output split sizes) so that gradients flow correctly.
+ """
+
+ @staticmethod
+ def forward(
+ ctx,
+ input: torch.Tensor,
+ output_split_sizes: list[int],
+ input_split_sizes: list[int],
+ group: dist.ProcessGroup,
+ ) -> torch.Tensor:
+ """Perform all-to-all forward and save sizes for backward."""
+ ctx.input_split_sizes = input_split_sizes
+ ctx.output_split_sizes = output_split_sizes
+ ctx.group = group
+ output = torch.empty(
+ sum(output_split_sizes),
+ input.shape[1],
+ device=input.device,
+ dtype=input.dtype,
+ )
+ dist.all_to_all_single(
+ output, input.contiguous(), output_split_sizes, input_split_sizes, group=group
+ )
+ return output
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None, None]:
+ """Reverse all-to-all: swap input and output split sizes."""
+ grad_input = torch.empty(
+ sum(ctx.input_split_sizes),
+ grad_output.shape[1],
+ device=grad_output.device,
+ dtype=grad_output.dtype,
+ )
+ dist.all_to_all_single(
+ grad_input,
+ grad_output.contiguous(),
+ ctx.input_split_sizes,
+ ctx.output_split_sizes,
+ group=ctx.group,
+ )
+ return grad_input, None, None, None
+
+
+class AllToAllTokenDispatcher:
+ """TokenDispatcher using NCCL all-to-all for expert-parallel communication.
+
+ Handles both EP=1 (no communication, just permute/unpermute) and EP>1
+ (all-to-all token exchange between ranks) cases transparently.
+
+ Args:
+ num_experts: Total number of experts (global).
+ num_local_experts: Number of experts on this rank.
+ hidden_size: Hidden dimension size.
+ ep_size: Expert parallel world size.
+ """
+
+ def __init__(self, num_experts: int, num_local_experts: int, hidden_size: int, ep_size: int):
+ """Initialize the AllToAllTokenDispatcher."""
+ self.num_experts = num_experts
+ self.num_local_experts = num_local_experts
+ self.hidden_size = hidden_size
+ self.ep_size = ep_size
+ self._ep_group: dist.ProcessGroup | None = None
+ self.pad_to_multiple: int | None = None
+
+ def set_ep_group(self, ep_group: dist.ProcessGroup) -> None:
+ """Set the expert-parallel process group for all-to-all communication."""
+ self._ep_group = ep_group
+
+ def dispatch(
+ self,
+ hidden_states: torch.Tensor,
+ selected_experts: torch.Tensor,
+ routing_weights: torch.Tensor,
+ ) -> DispatchOutput:
+ """Dispatch tokens to their assigned experts via permute and optional all-to-all.
+
+ Args:
+ hidden_states: Flattened input tensor of shape ``[N, H]``.
+ selected_experts: Expert assignments, shape ``[N, top_k]``, int.
+ routing_weights: Normalized routing probabilities, shape ``[N, top_k]``, float32.
+
+ Returns:
+ DispatchOutput with expert-sorted tokens, per-expert counts, and an opaque handle.
+ """
+ # Compute m_splits: number of tokens per expert
+ m_splits_tensor = torch.bincount(
+ selected_experts.reshape(-1), minlength=self.num_experts
+ ).int()
+
+ pad_offsets = None
+ if self.pad_to_multiple is not None:
+ routing_map = torch.zeros(
+ hidden_states.shape[0],
+ self.num_experts,
+ dtype=torch.bool,
+ device=hidden_states.device,
+ )
+ routing_map.scatter_(1, selected_experts, True)
+ routing_probs = torch.zeros(
+ hidden_states.shape[0],
+ self.num_experts,
+ dtype=routing_weights.dtype,
+ device=hidden_states.device,
+ )
+ routing_probs.scatter_(1, selected_experts, routing_weights)
+ (
+ permuted_hidden,
+ _,
+ row_id_map,
+ pad_offsets,
+ m_splits_tensor,
+ ) = transformer_engine.pytorch.moe_permute_and_pad_with_probs(
+ hidden_states,
+ routing_probs,
+ routing_map,
+ m_splits_tensor,
+ self.pad_to_multiple,
+ )
+ m_splits_tensor = m_splits_tensor.int()
+ routing_weights_for_unpermute = routing_probs
+ map_type = "mask"
+ else:
+ # Permute tokens by expert using TE moe_permute.
+ permuted_hidden, row_id_map = transformer_engine.pytorch.moe_permute(
+ hidden_states,
+ selected_experts.to(torch.int32),
+ num_out_tokens=selected_experts.numel(),
+ map_type="index",
+ )
+ routing_weights_for_unpermute = routing_weights
+ map_type = "index"
+
+ if self._ep_group is not None:
+ ep_group = self._ep_group
+
+ # Token counts per expert, reshaped to [ep_size, num_local_experts]
+ send_counts = m_splits_tensor.reshape(self.ep_size, self.num_local_experts)
+
+ # Exchange per-expert token counts between EP ranks
+ recv_counts = torch.empty_like(send_counts)
+ dist.all_to_all_single(recv_counts.flatten(), send_counts.flatten(), group=ep_group)
+
+ # Derive split sizes for the token all-to-all
+ input_split_sizes = send_counts.sum(dim=1).tolist()
+ output_split_sizes = recv_counts.sum(dim=1).tolist()
+ local_m_splits = recv_counts.sum(dim=0).int().tolist()
+
+ # Dispatch tokens to expert-owning ranks (differentiable)
+ recv_tokens = _DifferentiableAllToAll.apply(
+ permuted_hidden, output_split_sizes, input_split_sizes, ep_group
+ )
+
+ # Sort received tokens by local expert index.
+ # After all_to_all layout is [src0_exp0..src0_expL, src1_exp0..src1_expL, ...].
+ # GroupedLinear needs [all_exp0, all_exp1, ...].
+ sort_indices, unsort_indices = _build_expert_sort_indices(recv_counts)
+
+ handle = _AllToAllHandle(
+ row_id_map=row_id_map,
+ routing_weights=routing_weights_for_unpermute,
+ restore_shape=hidden_states.shape,
+ map_type=map_type,
+ pad_offsets=pad_offsets,
+ unsort_indices=unsort_indices,
+ input_split_sizes=input_split_sizes,
+ output_split_sizes=output_split_sizes,
+ )
+ return DispatchOutput(
+ expert_input=recv_tokens[sort_indices],
+ tokens_per_expert=local_m_splits,
+ handle=handle,
+ )
+
+ handle = _AllToAllHandle(
+ row_id_map=row_id_map,
+ routing_weights=routing_weights_for_unpermute,
+ restore_shape=hidden_states.shape,
+ map_type=map_type,
+ pad_offsets=pad_offsets,
+ )
+ return DispatchOutput(
+ expert_input=permuted_hidden,
+ tokens_per_expert=m_splits_tensor.tolist(),
+ handle=handle,
+ )
+
+ def combine(self, expert_output: torch.Tensor, handle: _AllToAllHandle) -> torch.Tensor:
+ """Combine expert outputs back to the original token order.
+
+ Args:
+ expert_output: Expert output tensor of shape ``[total_recv_tokens, H]``.
+ handle: Handle from ``dispatch()`` containing state for the reverse operation.
+
+ Returns:
+ Combined output tensor of shape ``[N, H]`` with routing weights applied.
+ """
+ if self._ep_group is not None:
+ assert handle.unsort_indices is not None
+ # Unsort back to source-rank-grouped order and reverse all_to_all (differentiable)
+ combined = _DifferentiableAllToAll.apply(
+ expert_output[handle.unsort_indices],
+ handle.input_split_sizes,
+ handle.output_split_sizes,
+ self._ep_group,
+ )
+ else:
+ combined = expert_output
+
+ # Unpermute and combine with routing weights (keep probs in float32 for numerical stability)
+ return transformer_engine.pytorch.moe_unpermute(
+ combined,
+ handle.row_id_map,
+ merging_probs=handle.routing_weights,
+ restore_shape=handle.restore_shape,
+ map_type=handle.map_type,
+ pad_offsets=handle.pad_offsets,
+ )
diff --git a/docs/examples/te_mixtral/te_mixtral_mxfp8.py b/docs/examples/te_mixtral/te_mixtral_mxfp8.py
new file mode 100644
index 0000000000..cca7d23085
--- /dev/null
+++ b/docs/examples/te_mixtral/te_mixtral_mxfp8.py
@@ -0,0 +1,568 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""TE-native MXFP8 Mixtral model (improvement 3).
+
+MoE FFN is a TE ``Sequential`` of three fusible ops — ``GroupedLinear``
+(gate_up), ``ScaledSwiGLU(glu_interleave_size=32)``, ``GroupedLinear``
+(down) — that the OperationFuser collapses into the fused
+``ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8`` and backward kernels under
+MXFP8. HF gate (``w1``) and up (``w3``) weights are row-interleaved in
+blocks of 32 to match the GLU interleaved layout that fused kernel reads.
+
+The fused kernel is enabled by ``utils._enable_fused_mxfp8_grouped_mlp()``
+(sets ``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1`` and patches the SM-version /
+cudnn-frontend signature checks). Requires
+``nvidia-cudnn-frontend >= 1.23.0`` and SM>=10 (Blackwell B100/B200/B300+).
+"""
+
+from __future__ import annotations
+
+import logging
+from collections import OrderedDict
+from contextlib import nullcontext
+from typing import Any, ClassVar, ContextManager
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+
+import transformer_engine.common.recipe as te_recipe
+import transformer_engine.pytorch as te
+from transformer_engine.pytorch.attention.inference import InferenceParams
+from transformer_engine.pytorch.attention.rope import RotaryPositionEmbedding
+from transformer_engine.pytorch.ops import (
+ GroupedLinear as TEOpsGroupedLinear,
+ ScaledSwiGLU,
+ Sequential as TEOpsSequential,
+)
+from transformer_engine.pytorch.router import fused_moe_aux_loss
+from transformers import MixtralConfig, PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding
+
+from te_moe_dispatch import AllToAllTokenDispatcher
+from te_mixtral import (
+ _pad_input,
+ _unpad_input,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# HF->TE checkpoint mapping is shared with the BF16 path in te_mixtral.py.
+# ``GLU_INTERLEAVE_SIZE`` is the gate/up interleave block (32); the fused
+# MXFP8 forward op only fires when ``ScaledSwiGLU`` is configured with it.
+from hf_to_te_weights import (
+ GLU_INTERLEAVE_SIZE,
+ replace_params_mxfp8 as replace_params,
+)
+
+
+class TEMixtralMXFP8Config(MixtralConfig):
+ """Improvement-3 config. Same surface as :class:`te_mixtral.TEMixtralConfig` but
+ with the FFN mode locked to ``grouped_op`` + MXFP8."""
+
+ attn_input_format: str = "thd"
+ self_attn_mask_type: str = "padding_causal"
+ expert_parallel_size: int = 1
+ moe_aux_loss_coeff: float = 0.0
+
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ if self.num_local_experts % self.expert_parallel_size != 0:
+ raise ValueError(
+ f"num_local_experts ({self.num_local_experts}) must be divisible by "
+ f"expert_parallel_size ({self.expert_parallel_size})"
+ )
+
+
+class TEMixtralMXFP8PreTrainedModel(PreTrainedModel):
+ """HF integration boilerplate for the improvement-3 model."""
+
+ config_class = TEMixtralMXFP8Config
+ base_model_prefix = "model"
+ _no_split_modules = ("TEMixtralMXFP8DecoderLayer",)
+ _skip_keys_device_placement = ("past_key_values",)
+ _do_not_quantize = ("lm_head", "model.layers.*.mlp.gate")
+
+ def _init_weights(self, module):
+ if module.__module__.startswith("transformer_engine.pytorch"):
+ return
+ super()._init_weights(module)
+
+ def state_dict(self, *args, **kwargs):
+ sd = super().state_dict(*args, **kwargs)
+ return {k: v for k, v in sd.items() if not k.endswith("_extra_state")}
+
+
+class TEMixtralMXFP8SparseMoeBlock(nn.Module):
+ """MoE block: router + EP dispatcher + fused MXFP8 grouped MLP."""
+
+ def __init__(
+ self,
+ config: TEMixtralMXFP8Config,
+ dispatcher: AllToAllTokenDispatcher | None = None,
+ ) -> None:
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.num_experts = config.num_local_experts
+ self.top_k = config.num_experts_per_tok
+ self.jitter_noise = config.router_jitter_noise
+
+ self.ep_size = getattr(config, "expert_parallel_size", 1)
+ self.num_local_experts = self.num_experts // self.ep_size
+ self.moe_aux_loss_coeff = getattr(config, "moe_aux_loss_coeff", 0.0)
+ self._aux_loss: torch.Tensor = torch.tensor(0.0)
+
+ if self.intermediate_size % GLU_INTERLEAVE_SIZE != 0:
+ raise ValueError(
+ f"intermediate_size ({self.intermediate_size}) must be divisible by "
+ f"GLU_INTERLEAVE_SIZE ({GLU_INTERLEAVE_SIZE})"
+ )
+
+ self.dispatcher = dispatcher or AllToAllTokenDispatcher(
+ num_experts=self.num_experts,
+ num_local_experts=self.num_local_experts,
+ hidden_size=self.hidden_size,
+ ep_size=self.ep_size,
+ )
+
+ device = "meta" if torch.get_default_device() == torch.device("meta") else "cuda"
+
+ def _init_method(x: torch.Tensor) -> None:
+ torch.nn.init.normal_(x, mean=0.0, std=config.initializer_range)
+
+ with te.quantized_model_init(enabled=False):
+ self.gate = te.Linear(
+ self.hidden_size,
+ self.num_experts,
+ bias=False,
+ device=device,
+ params_dtype=config.dtype,
+ init_method=_init_method,
+ )
+
+ self.experts_gate_up = TEOpsGroupedLinear(
+ num_groups=self.num_local_experts,
+ in_features=self.hidden_size,
+ out_features=2 * self.intermediate_size,
+ bias=False,
+ dtype=config.dtype,
+ device=device,
+ )
+ self.experts_swiglu = ScaledSwiGLU(glu_interleave_size=GLU_INTERLEAVE_SIZE)
+ self.experts_down = TEOpsGroupedLinear(
+ num_groups=self.num_local_experts,
+ in_features=self.intermediate_size,
+ out_features=self.hidden_size,
+ bias=False,
+ dtype=config.dtype,
+ device=device,
+ )
+ # Wrap as TE Sequential to enable forward/backward op fusion
+ # (ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 / dswiglu).
+ object.__setattr__(
+ self,
+ "_experts_ffn_op",
+ TEOpsSequential(self.experts_gate_up, self.experts_swiglu, self.experts_down),
+ )
+
+ def set_ep_group(self, ep_group: dist.ProcessGroup) -> None:
+ """Set the EP communication group on the dispatcher.
+
+ Each EP rank owns its local slice of expert weights as ordinary
+ Parameters (``weight0..weight{N-1}``) because per-expert parameters
+ are never replicated across the EP group.
+ """
+ self.dispatcher.set_ep_group(ep_group)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ original_shape = hidden_states.shape
+
+ if self.training and self.jitter_noise > 0:
+ hidden_states = hidden_states * torch.empty_like(hidden_states).uniform_(
+ 1.0 - self.jitter_noise, 1.0 + self.jitter_noise
+ )
+
+ if hidden_states.dim() == 3:
+ hidden_states = hidden_states.reshape(-1, self.hidden_size)
+
+ with te.autocast(enabled=False):
+ router_logits = self.gate(hidden_states) # [N, E]
+
+ # Top-k routing weights, two algebraically equivalent forms.
+ # Old::
+ #
+ # probs = softmax(logits) # (N, E)
+ # weights, idx = topk(probs, k)
+ # weights = weights / weights.sum(-1, keepdim=True)
+ #
+ # New (used here)::
+ #
+ # topk_logits, idx = topk(logits, k)
+ # weights = softmax(topk_logits) # softmax over (N, k)
+ topk_logits, selected_experts = torch.topk(router_logits, self.top_k, dim=-1)
+ routing_weights = torch.nn.functional.softmax(topk_logits, dim=-1, dtype=torch.float32)
+
+ # Bincount once, in the MoE block. ``AllToAllTokenDispatcher``
+ # takes this as a required argument, so the dispatcher never
+ # bincounts again.
+ tokens_per_expert = torch.bincount(
+ selected_experts.reshape(-1), minlength=self.num_experts
+ ).to(torch.int32)
+
+ if self.moe_aux_loss_coeff > 0:
+ num_tokens = hidden_states.shape[0]
+ softmax_probs = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float32)
+ self._aux_loss = fused_moe_aux_loss(
+ probs=softmax_probs,
+ tokens_per_expert=tokens_per_expert,
+ total_num_tokens=num_tokens,
+ num_experts=self.num_experts,
+ topk=self.top_k,
+ coeff=self.moe_aux_loss_coeff,
+ )
+ else:
+ self._aux_loss = torch.tensor(0.0, device=hidden_states.device)
+
+ dispatch_out = self.dispatcher.dispatch(
+ hidden_states,
+ selected_experts,
+ routing_weights,
+ tokens_per_expert,
+ )
+ expert_input = dispatch_out.expert_input
+ expert_probs = dispatch_out.expert_probs
+ split_sizes = torch.tensor(
+ dispatch_out.tokens_per_expert, dtype=torch.int32, device=expert_input.device
+ )
+
+ # Fused gate_up -> ScaledSwiGLU(probs) -> down.
+ expert_output = self._experts_ffn_op(expert_input, split_sizes, expert_probs, split_sizes)
+
+ output = self.dispatcher.combine(expert_output, dispatch_out.handle)
+ return output.reshape(original_shape)
+
+
+class TEMixtralMXFP8DecoderLayer(nn.Module):
+ """Self-attention + improvement-3 MoE block."""
+
+ def __init__(
+ self,
+ config: TEMixtralMXFP8Config,
+ layer_idx: int,
+ dispatcher: AllToAllTokenDispatcher | None = None,
+ ) -> None:
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ device = "meta" if torch.get_default_device() == torch.device("meta") else "cuda"
+
+ def _init_method(x: torch.Tensor) -> None:
+ torch.nn.init.normal_(x, mean=0.0, std=config.initializer_range)
+
+ self.self_attention = te.MultiheadAttention(
+ hidden_size=config.hidden_size,
+ num_attention_heads=config.num_attention_heads,
+ num_gqa_groups=config.num_key_value_heads,
+ bias=False,
+ layernorm_epsilon=config.rms_norm_eps,
+ attention_dropout=0,
+ fuse_qkv_params=True,
+ qkv_weight_interleaved=True,
+ normalization="RMSNorm",
+ input_layernorm=True,
+ qkv_format=config.attn_input_format,
+ attn_mask_type=config.self_attn_mask_type,
+ layer_number=layer_idx + 1,
+ params_dtype=config.dtype,
+ device=device,
+ init_method=_init_method,
+ output_layer_init_method=_init_method,
+ )
+ self.post_attention_layernorm = te.RMSNorm(
+ config.hidden_size,
+ eps=config.rms_norm_eps,
+ dtype=config.dtype,
+ device=device,
+ )
+ self.mlp = TEMixtralMXFP8SparseMoeBlock(config, dispatcher)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ rotary_pos_emb: torch.Tensor | None = None,
+ inference_params: InferenceParams | None = None,
+ **kwargs: Any,
+ ) -> torch.Tensor:
+ attn_output = self.self_attention(
+ hidden_states,
+ attention_mask=attention_mask,
+ rotary_pos_emb=rotary_pos_emb,
+ inference_params=inference_params,
+ cu_seqlens_q=kwargs.get("cu_seqlens_q", None),
+ cu_seqlens_kv=kwargs.get("cu_seqlens_kv", None),
+ cu_seqlens_q_padded=kwargs.get("cu_seqlens_q_padded", None),
+ cu_seqlens_kv_padded=kwargs.get("cu_seqlens_kv_padded", None),
+ max_seqlen_q=kwargs.get("max_seqlen_q", None),
+ max_seqlen_kv=kwargs.get("max_seqlen_kv", None),
+ pad_between_seqs=kwargs.get("pad_between_seqs", None),
+ )
+ hidden_states = hidden_states + attn_output
+
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ return residual + hidden_states
+
+
+class TEMixtralMXFP8Model(TEMixtralMXFP8PreTrainedModel):
+ """Embedding + N decoder layers + RMSNorm. THD-packed under MXFP8."""
+
+ def __init__(
+ self,
+ config: TEMixtralMXFP8Config,
+ fp8_recipe: te_recipe.Recipe | None = None,
+ dispatcher: AllToAllTokenDispatcher | None = None,
+ ) -> None:
+ super().__init__(config)
+ self.config = config
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+ self._fp8_recipe = fp8_recipe
+
+ self.embed_tokens = nn.Embedding(
+ config.vocab_size, config.hidden_size, self.padding_idx, dtype=config.dtype
+ )
+
+ layers: list[TEMixtralMXFP8DecoderLayer] = [
+ TEMixtralMXFP8DecoderLayer(config, i, dispatcher)
+ for i in range(config.num_hidden_layers)
+ ]
+ self.layers = nn.ModuleList(layers)
+
+ self.norm = te.RMSNorm(
+ config.hidden_size,
+ eps=config.rms_norm_eps,
+ dtype=config.dtype,
+ device="meta" if torch.get_default_device() == torch.device("meta") else "cuda",
+ )
+
+ self.rotary_emb = RotaryPositionEmbedding(config.hidden_size // config.num_attention_heads)
+ self.rotary_emb.inv_freq = LlamaRotaryEmbedding(config=config).inv_freq
+
+ self.gradient_checkpointing = False
+ self.post_init()
+
+ def set_ep_groups(self, ep_group: dist.ProcessGroup) -> None:
+ for layer in self.layers:
+ layer.mlp.set_ep_group(ep_group)
+
+ def _outer_autocast(self) -> ContextManager:
+ if self._fp8_recipe is None:
+ return nullcontext()
+ return te.autocast(enabled=True, recipe=self._fp8_recipe)
+
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ past_key_values: InferenceParams | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Any,
+ ) -> BaseModelOutputWithPast:
+ del position_ids, use_cache # not used in this minimal forward
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("Specify exactly one of input_ids or inputs_embeds")
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ hidden_states = inputs_embeds
+
+ has_thd_input = [
+ x in kwargs for x in ("cu_seq_lens_q", "cu_seq_lens_k", "max_length_q", "max_length_k")
+ ]
+ decode_without_mask = (
+ isinstance(past_key_values, InferenceParams)
+ and hidden_states.dim() == 3
+ and hidden_states.size(1) == 1
+ )
+ should_pack_inputs = (
+ not any(has_thd_input)
+ and self.config.attn_input_format == "thd"
+ and not decode_without_mask
+ )
+
+ thd_remainder = 0
+ thd_orig_tokens = 0
+ indices = None
+ batch_size = 0
+ padded_seq_len = 0
+ if should_pack_inputs:
+ assert attention_mask is not None, "attention_mask required when packing BSHD."
+ batch_size = hidden_states.size(0)
+ padded_seq_len = hidden_states.size(1)
+ hidden_states, indices, cu_seqlens, max_seqlen, _ = _unpad_input(
+ hidden_states, attention_mask
+ )
+
+ # MXFP8 requires total tokens divisible by 32; pad the last seq.
+ thd_orig_tokens = hidden_states.shape[0]
+ thd_remainder = thd_orig_tokens % 32
+ if thd_remainder != 0:
+ thd_pad = 32 - thd_remainder
+ hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, thd_pad))
+ cu_seqlens = cu_seqlens.clone()
+ cu_seqlens[-1] = cu_seqlens[-1] + thd_pad
+ max_seqlen = max_seqlen + thd_pad
+
+ kwargs["cu_seq_lens_q"] = kwargs["cu_seq_lens_k"] = cu_seqlens
+ kwargs["max_length_q"] = kwargs["max_length_k"] = max_seqlen
+
+ if (
+ self.config.attn_input_format == "thd"
+ and hidden_states.dim() == 3
+ and hidden_states.size(0) == 1
+ ):
+ hidden_states = hidden_states.squeeze(0)
+
+ if (
+ self.config.attn_input_format == "bshd"
+ and attention_mask is not None
+ and attention_mask.dim() == 2
+ ):
+ attention_mask = ~attention_mask[:, None, None, :].bool()
+
+ 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)))
+
+ with torch.autocast(device_type="cuda", enabled=False):
+ te_rope_emb = self.rotary_emb(max_seq_len=self.config.max_position_embeddings)
+
+ with self._outer_autocast():
+ for layer_idx, decoder_layer in enumerate(self.layers):
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=(
+ None if self.config.attn_input_format == "thd" else attention_mask
+ ),
+ rotary_pos_emb=te_rope_emb,
+ inference_params=past_key_values,
+ cu_seqlens_q=kwargs.get("cu_seq_lens_q", None),
+ cu_seqlens_kv=kwargs.get("cu_seq_lens_k", None),
+ cu_seqlens_q_padded=kwargs.get("cu_seq_lens_q_padded", None),
+ cu_seqlens_kv_padded=kwargs.get("cu_seq_lens_k_padded", None),
+ max_seqlen_q=kwargs.get("max_length_q", None),
+ max_seqlen_kv=kwargs.get("max_length_k", None),
+ pad_between_seqs=kwargs.get("pad_between_seqs", None),
+ )
+
+ hidden_states = self.norm(hidden_states)
+
+ if should_pack_inputs:
+ if thd_remainder != 0:
+ hidden_states = hidden_states[:thd_orig_tokens]
+ hidden_states = _pad_input(hidden_states, indices, batch_size, padded_seq_len)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ hidden_states=None,
+ )
+
+
+class TEMixtralMXFP8ForCausalLM(TEMixtralMXFP8PreTrainedModel):
+ """Causal LM wrapper with MXFP8 autocast."""
+
+ _tied_weights_keys: ClassVar[list[str]] = []
+
+ def __init__(
+ self,
+ config: TEMixtralMXFP8Config,
+ fp8_recipe: te_recipe.Recipe | None = None,
+ dispatcher: AllToAllTokenDispatcher | None = None,
+ ) -> None:
+ super().__init__(config)
+ self.model = TEMixtralMXFP8Model(config, fp8_recipe=fp8_recipe, dispatcher=dispatcher)
+ self.vocab_size = config.vocab_size
+ with te.quantized_model_init(enabled=False):
+ self.lm_head = te.Linear(
+ config.hidden_size,
+ config.vocab_size,
+ bias=False,
+ params_dtype=config.dtype,
+ device="meta" if torch.get_default_device() == torch.device("meta") else "cuda",
+ init_method=lambda x: torch.nn.init.normal_(
+ x, mean=0.0, std=config.initializer_range
+ ),
+ )
+ self.post_init()
+
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ past_key_values: Any = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ shift_labels: torch.Tensor | None = None,
+ use_cache: bool | None = None,
+ cache_position: torch.Tensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Any,
+ ) -> CausalLMOutputWithPast:
+ outputs = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ slice_indices = (
+ slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ )
+ with te.autocast(enabled=False):
+ if hidden_states.ndim == 3:
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+ else:
+ logits = self.lm_head(hidden_states[slice_indices, :])
+
+ loss = None
+ if labels is not None or shift_labels is not None:
+ loss = self.loss_function(
+ logits=logits,
+ labels=labels,
+ shift_labels=shift_labels,
+ vocab_size=self.config.vocab_size,
+ **kwargs,
+ )
+
+ if self.config.moe_aux_loss_coeff > 0 and loss is not None:
+ aux_loss = sum(layer.mlp._aux_loss for layer in self.model.layers)
+ loss = loss + aux_loss
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
diff --git a/docs/examples/te_mixtral/te_moe_dispatch.py b/docs/examples/te_mixtral/te_moe_dispatch.py
new file mode 100644
index 0000000000..0fc06a6391
--- /dev/null
+++ b/docs/examples/te_mixtral/te_moe_dispatch.py
@@ -0,0 +1,298 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Token dispatch / combine for the MXFP8 TE-native MoE block.
+
+Wraps the permute/pad/all-to-all/sort-by-expert plumbing that moves tokens
+from data-parallel ranks to their owning expert ranks, and reverses the
+operation on the way back. The transport is NCCL ``all_to_all_single``;
+the all-to-all is just the mechanism — the public API is ``dispatch()``
+/ ``combine()``.
+
+Per-expert MoE permute pads to 128 (grouped MXFP8 GEMM M-tile). Both
+hidden states *and* per-token routing probabilities are transmitted so
+the destination-side ``ScaledSwiGLU(glu_interleave_size=32)`` has its
+scales locally — that's what trips the fused
+``ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8`` kernel. ``combine()`` does not
+re-apply routing weights (already applied inside ScaledSwiGLU).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.distributed as dist
+
+import transformer_engine.pytorch as te
+
+# Required so the ``@torch.compile`` helpers below can capture data-dependent
+# tensor shapes (e.g. ``repeat_interleave`` with tensor counts) without
+# bailing out to Python. Must be set at module level — torch.compile traces
+# lazily on the first call, so a scoped setting wouldn't be active.
+torch._dynamo.config.capture_scalar_outputs = True
+torch._dynamo.config.capture_dynamic_output_shape_ops = True
+
+
+# Per-expert token-count alignment required by the fused MXFP8 grouped-MLP
+# CuTe-DSL kernel (ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8). The kernel
+# rejects an input whose per-group token count is not a multiple of 256
+# ("Invalid a.shape[0] ... expected to be divisible by 256"). 128 was the
+# old value that worked when the fused kernel wasn't firing on B300; with
+# the SM10.x gate now passing on B300 we must pad to 256.
+_MXFP8_GROUP_ALIGN = 256
+
+
+@dataclass
+class DispatchOutput:
+ """Tokens, per-token probs, and split sizes routed to the local experts."""
+
+ expert_input: torch.Tensor
+ expert_probs: torch.Tensor
+ tokens_per_expert: list[int]
+ handle: Any
+
+
+@dataclass
+class _Handle:
+ row_id_map: torch.Tensor
+ restore_shape: torch.Size
+ pad_offsets: torch.Tensor | None
+ unsort_indices: torch.Tensor | None = None
+ input_split_sizes: list[int] | None = None
+ output_split_sizes: list[int] | None = None
+
+
+class _DifferentiableAllToAll(torch.autograd.Function):
+ """``dist.all_to_all_single`` wrapped in autograd (works for 1-D and 2-D)."""
+
+ @staticmethod
+ def forward(
+ ctx,
+ input: torch.Tensor,
+ output_split_sizes: list[int],
+ input_split_sizes: list[int],
+ group: dist.ProcessGroup,
+ ) -> torch.Tensor:
+ ctx.input_split_sizes = input_split_sizes
+ ctx.output_split_sizes = output_split_sizes
+ ctx.group = group
+ total_out = sum(output_split_sizes)
+ if input.dim() == 1:
+ output = torch.empty(total_out, device=input.device, dtype=input.dtype)
+ else:
+ output = torch.empty(
+ total_out, *input.shape[1:], device=input.device, dtype=input.dtype
+ )
+ dist.all_to_all_single(
+ output, input.contiguous(), output_split_sizes, input_split_sizes, group=group
+ )
+ return output
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor):
+ total_in = sum(ctx.input_split_sizes)
+ if grad_output.dim() == 1:
+ grad_input = torch.empty(total_in, device=grad_output.device, dtype=grad_output.dtype)
+ else:
+ grad_input = torch.empty(
+ total_in, *grad_output.shape[1:], device=grad_output.device, dtype=grad_output.dtype
+ )
+ dist.all_to_all_single(
+ grad_input,
+ grad_output.contiguous(),
+ ctx.input_split_sizes,
+ ctx.output_split_sizes,
+ group=ctx.group,
+ )
+ return grad_input, None, None, None
+
+
+@torch.compile(fullgraph=True)
+def _build_expert_sort_indices(recv_counts: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ """Build sort/unsort indices that regroup ``[src*expert]`` tokens by expert.
+
+ After ``all_to_all`` tokens arrive grouped by source rank
+ (``[src0_exp0..src0_expL, src1_exp0..src1_expL, ...]``). ``GroupedLinear``
+ needs them grouped by expert (``[all_exp0, all_exp1, ...]``).
+
+ Uses only vectorized tensor ops (no ``.item()`` calls or Python-level
+ loops) so this is ``torch.compile(fullgraph=True)``-safe.
+ """
+ ep_size, num_local_experts = recv_counts.shape
+ device = recv_counts.device
+ num_blocks = ep_size * num_local_experts
+
+ counts_src = recv_counts.reshape(-1).long()
+ offsets_src = torch.zeros(num_blocks, dtype=torch.long, device=device)
+ offsets_src[1:] = counts_src[:-1].cumsum(0)
+
+ counts_exp = recv_counts.t().contiguous().reshape(-1).long()
+ offsets_exp = torch.zeros(num_blocks, dtype=torch.long, device=device)
+ offsets_exp[1:] = counts_exp[:-1].cumsum(0)
+
+ total = counts_src.sum()
+
+ s_idx = torch.arange(ep_size, device=device).unsqueeze(1).expand(ep_size, num_local_experts)
+ e_idx = (
+ torch.arange(num_local_experts, device=device)
+ .unsqueeze(0)
+ .expand(ep_size, num_local_experts)
+ )
+ src_to_exp = (e_idx * ep_size + s_idx).reshape(-1)
+ shifts = offsets_exp[src_to_exp] - offsets_src
+ token_shifts = shifts.repeat_interleave(counts_src)
+ src_positions = torch.arange(total, device=device)
+ dst_positions = src_positions + token_shifts
+ sort_indices = torch.empty(total, dtype=torch.long, device=device)
+ sort_indices[dst_positions] = src_positions
+ unsort_indices = torch.empty_like(sort_indices)
+ unsort_indices[sort_indices] = torch.arange(total, device=device)
+ return sort_indices, unsort_indices
+
+
+class AllToAllTokenDispatcher:
+ """NCCL all-to-all dispatcher for the TE-native MXFP8 MoE block.
+
+ Args:
+ num_experts: Total global experts.
+ num_local_experts: Experts owned by this rank.
+ hidden_size: Hidden feature dim.
+ ep_size: Expert-parallel world size (1 = single-process).
+ pad_align: Per-expert split alignment. Must be a multiple of 128 for
+ the grouped MXFP8 GEMM.
+ """
+
+ def __init__(
+ self,
+ num_experts: int,
+ num_local_experts: int,
+ hidden_size: int,
+ ep_size: int,
+ pad_align: int = _MXFP8_GROUP_ALIGN,
+ ) -> None:
+ self.num_experts = num_experts
+ self.num_local_experts = num_local_experts
+ self.hidden_size = hidden_size
+ self.ep_size = ep_size
+ self.pad_align = pad_align
+ self._ep_group: dist.ProcessGroup | None = None
+
+ def set_ep_group(self, ep_group: dist.ProcessGroup) -> None:
+ self._ep_group = ep_group
+
+ def dispatch(
+ self,
+ hidden_states: torch.Tensor,
+ selected_experts: torch.Tensor,
+ routing_weights: torch.Tensor,
+ tokens_per_expert: torch.Tensor,
+ ) -> DispatchOutput:
+ """Permute -> pad -> (all-to-all) -> sort-by-expert.
+
+ ``tokens_per_expert`` is required: the MoE block already computes the
+ per-expert token count (for the fused aux loss + the routing tables),
+ so the dispatcher takes it as input rather than launching another
+ ``torch.bincount`` kernel.
+ """
+ num_tokens = hidden_states.shape[0]
+
+ # Dense per-expert routing tables required by ``moe_permute_and_pad_with_probs``.
+ routing_map = torch.zeros(
+ num_tokens, self.num_experts, dtype=torch.bool, device=hidden_states.device
+ )
+ routing_map.scatter_(1, selected_experts, True)
+ routing_probs = torch.zeros(
+ num_tokens, self.num_experts, dtype=routing_weights.dtype, device=hidden_states.device
+ )
+ routing_probs.scatter_(1, selected_experts, routing_weights)
+
+ (
+ permuted_hidden,
+ permuted_probs,
+ row_id_map,
+ pad_offsets,
+ padded_tokens_per_expert,
+ ) = te.moe_permute_and_pad_with_probs(
+ hidden_states,
+ routing_probs,
+ routing_map,
+ tokens_per_expert,
+ self.pad_align,
+ )
+ padded_tokens_per_expert = padded_tokens_per_expert.int()
+
+ if self._ep_group is None or self.ep_size == 1:
+ handle = _Handle(
+ row_id_map=row_id_map,
+ restore_shape=hidden_states.shape,
+ pad_offsets=pad_offsets,
+ )
+ return DispatchOutput(
+ expert_input=permuted_hidden,
+ expert_probs=permuted_probs,
+ tokens_per_expert=padded_tokens_per_expert.tolist(),
+ handle=handle,
+ )
+
+ # EP > 1: ship both tokens and probs across ranks. A single packed
+ # all_to_all was tried; the extra ``.contiguous()`` slicing on the
+ # receive side cost more than the saved NCCL collective at Mixtral
+ # batch=8 / seq=8192 (the probs comm is ~1/2048 of the token comm,
+ # so the all_to_all is bandwidth-bound, not latency-bound).
+ ep_group = self._ep_group
+ send_counts = padded_tokens_per_expert.reshape(self.ep_size, self.num_local_experts)
+ recv_counts = torch.empty_like(send_counts)
+ dist.all_to_all_single(recv_counts.flatten(), send_counts.flatten(), group=ep_group)
+
+ input_split_sizes = send_counts.sum(dim=1).tolist()
+ output_split_sizes = recv_counts.sum(dim=1).tolist()
+ local_m_splits = recv_counts.sum(dim=0).int().tolist()
+
+ recv_tokens = _DifferentiableAllToAll.apply(
+ permuted_hidden, output_split_sizes, input_split_sizes, ep_group
+ )
+ recv_probs = _DifferentiableAllToAll.apply(
+ permuted_probs, output_split_sizes, input_split_sizes, ep_group
+ )
+
+ sort_indices, unsort_indices = _build_expert_sort_indices(recv_counts)
+
+ handle = _Handle(
+ row_id_map=row_id_map,
+ restore_shape=hidden_states.shape,
+ pad_offsets=pad_offsets,
+ unsort_indices=unsort_indices,
+ input_split_sizes=input_split_sizes,
+ output_split_sizes=output_split_sizes,
+ )
+ return DispatchOutput(
+ expert_input=recv_tokens[sort_indices],
+ expert_probs=recv_probs[sort_indices],
+ tokens_per_expert=local_m_splits,
+ handle=handle,
+ )
+
+ def combine(self, expert_output: torch.Tensor, handle: _Handle) -> torch.Tensor:
+ """Reverse the dispatch. ``ScaledSwiGLU`` already applied per-token probs,
+ so ``moe_unpermute`` is called without ``merging_probs``."""
+ if handle.unsort_indices is not None:
+ combined = _DifferentiableAllToAll.apply(
+ expert_output[handle.unsort_indices],
+ handle.input_split_sizes,
+ handle.output_split_sizes,
+ self._ep_group,
+ )
+ else:
+ combined = expert_output
+
+ return te.moe_unpermute(
+ combined,
+ handle.row_id_map,
+ merging_probs=None,
+ restore_shape=handle.restore_shape,
+ map_type="mask",
+ pad_offsets=handle.pad_offsets,
+ )
diff --git a/docs/examples/te_mixtral/test_accuracy.py b/docs/examples/te_mixtral/test_accuracy.py
new file mode 100644
index 0000000000..46ed2a5bf2
--- /dev/null
+++ b/docs/examples/te_mixtral/test_accuracy.py
@@ -0,0 +1,188 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+"""Forward + backward parity check for te_mixtral (BF16 and MXFP8) vs HF.
+
+Compares logits, loss, and weight gradients (``model.embed_tokens.weight``
+and ``lm_head.weight``) between the HuggingFace reference and the
+TransformerEngine port in both BF16 and MXFP8 modes.
+"""
+
+import torch
+from transformers import MixtralConfig, MixtralForCausalLM
+
+import transformer_engine.pytorch as te
+from transformer_engine.common import recipe as te_recipe
+
+from te_mixtral import TEMixtralForCausalLM, replace_params as replace_params_bf16
+from te_mixtral_mxfp8 import (
+ TEMixtralMXFP8ForCausalLM,
+ replace_params as replace_params_mxfp8,
+)
+
+
+# BF16 should match HF very closely.
+BF16_TOL = 0.01
+
+# MXFP8 quantizes activations to FP8 with per-tile bf16 scales, so we expect
+# larger logit / gradient drift than the BF16 case.
+MXFP8_LOGITS_ATOL = 1.5
+MXFP8_LOGITS_RTOL = 0.05
+MXFP8_LOSS_ATOL = 0.05
+MXFP8_LOSS_RTOL = 0.05
+MXFP8_GRAD_ATOL = 1.0
+MXFP8_GRAD_RTOL = 0.1
+
+
+def _build_config():
+ return MixtralConfig(
+ hidden_size=256,
+ intermediate_size=512,
+ num_local_experts=4,
+ num_experts_per_tok=2,
+ num_hidden_layers=2,
+ num_attention_heads=8,
+ num_key_value_heads=8,
+ vocab_size=1024,
+ max_position_embeddings=128,
+ router_jitter_noise=0.0,
+ rms_norm_eps=1e-5,
+ )
+
+
+def _load_te_weights(model_te, model_hf, replace_params_fn):
+ te_state_dict = model_te.state_dict()
+ replace_params_fn(model_hf.state_dict(), te_state_dict, model_te.config)
+ missing, unexpected = model_te.load_state_dict(te_state_dict, strict=False)
+ if unexpected:
+ raise RuntimeError(f"Unexpected TE keys during load: {unexpected}")
+ allowed_missing = [k for k in missing if k.endswith("_extra_state")]
+ if len(allowed_missing) != len(missing):
+ raise RuntimeError(f"Unexpected missing TE keys during load: {missing}")
+
+
+def _zero_grads(model):
+ for p in model.parameters():
+ if p.grad is not None:
+ p.grad = None
+
+
+def _forward_backward(model, input_ids, attention_mask, labels, *, fp8_recipe=None):
+ """Return (logits, loss, embed_grad, lm_head_grad), all detached as float32."""
+ _zero_grads(model)
+ if fp8_recipe is not None:
+ with te.autocast(enabled=True, recipe=fp8_recipe):
+ out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
+ out.loss.backward()
+ else:
+ out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
+ out.loss.backward()
+ return (
+ out.logits.detach().float(),
+ out.loss.detach().float(),
+ model.model.embed_tokens.weight.grad.detach().float().clone(),
+ model.lm_head.weight.grad.detach().float().clone(),
+ )
+
+
+def _compare(label, hf, te_, *, atol, rtol):
+ diff = (hf - te_).abs()
+ max_diff = diff.max().item()
+ mean_diff = diff.mean().item()
+ print(f" {label:<14s} max={max_diff:.6f} mean={mean_diff:.6f}")
+ torch.testing.assert_close(te_, hf, atol=atol, rtol=rtol)
+
+
+def _make_inputs(cfg, device):
+ torch.manual_seed(1)
+ # seq divisible by 32 so the MXFP8 path is happy.
+ input_ids = torch.randint(0, cfg.vocab_size, (2, 64), device=device)
+ attention_mask = torch.ones_like(input_ids, device=device)
+ labels = input_ids.clone()
+ return input_ids, attention_mask, labels
+
+
+def _build_hf(cfg, device, dtype):
+ torch.manual_seed(0)
+ model = MixtralForCausalLM(cfg).to(device=device, dtype=dtype)
+ model.eval()
+ return model
+
+
+def _run_bf16(cfg, model_hf, inputs, device, dtype):
+ print("=" * 64)
+ print("BF16 parity check (forward + backward)")
+ print("=" * 64)
+
+ te_cfg = TEMixtralForCausalLM.config_class(**cfg.to_dict())
+ model_te = TEMixtralForCausalLM(te_cfg).to(device=device, dtype=dtype)
+ _load_te_weights(model_te, model_hf, replace_params_bf16)
+ model_te.eval()
+
+ input_ids, attention_mask, labels = inputs
+ hf_logits, hf_loss, hf_embed_g, hf_lm_g = _forward_backward(
+ model_hf, input_ids, attention_mask, labels
+ )
+ te_logits, te_loss, te_embed_g, te_lm_g = _forward_backward(
+ model_te, input_ids, attention_mask, labels
+ )
+
+ print(f" logits shape {tuple(hf_logits.shape)}")
+ print(f" HF loss {hf_loss.item():.6f}")
+ print(f" TE loss {te_loss.item():.6f}")
+ _compare("logits", hf_logits, te_logits, atol=BF16_TOL, rtol=0.0)
+ _compare("loss", hf_loss, te_loss, atol=BF16_TOL, rtol=0.0)
+ _compare("embed.grad", hf_embed_g, te_embed_g, atol=BF16_TOL, rtol=0.0)
+ _compare("lm_head.grad", hf_lm_g, te_lm_g, atol=BF16_TOL, rtol=0.0)
+ print("BF16 parity OK.\n")
+
+
+def _run_mxfp8(cfg, model_hf, inputs, device, dtype):
+ print("=" * 64)
+ print("MXFP8 parity check (forward + backward)")
+ print("=" * 64)
+
+ te_cfg = TEMixtralMXFP8ForCausalLM.config_class(**cfg.to_dict())
+ te_cfg.attn_input_format = "bshd"
+ te_cfg.self_attn_mask_type = "causal"
+ te_cfg.expert_parallel_size = 1
+ te_cfg.dtype = dtype
+ recipe = te_recipe.MXFP8BlockScaling(fp8_format=te_recipe.Format.E4M3)
+ model_te = TEMixtralMXFP8ForCausalLM(te_cfg, fp8_recipe=recipe).to(device=device, dtype=dtype)
+ _load_te_weights(model_te, model_hf, replace_params_mxfp8)
+ model_te.eval()
+
+ input_ids, attention_mask, labels = inputs
+ hf_logits, hf_loss, hf_embed_g, hf_lm_g = _forward_backward(
+ model_hf, input_ids, attention_mask, labels
+ )
+ te_logits, te_loss, te_embed_g, te_lm_g = _forward_backward(
+ model_te, input_ids, attention_mask, labels, fp8_recipe=recipe
+ )
+
+ print(f" logits shape {tuple(hf_logits.shape)}")
+ print(f" HF loss {hf_loss.item():.6f}")
+ print(f" TE loss {te_loss.item():.6f}")
+ _compare("logits", hf_logits, te_logits, atol=MXFP8_LOGITS_ATOL, rtol=MXFP8_LOGITS_RTOL)
+ _compare("loss", hf_loss, te_loss, atol=MXFP8_LOSS_ATOL, rtol=MXFP8_LOSS_RTOL)
+ _compare("embed.grad", hf_embed_g, te_embed_g, atol=MXFP8_GRAD_ATOL, rtol=MXFP8_GRAD_RTOL)
+ _compare("lm_head.grad", hf_lm_g, te_lm_g, atol=MXFP8_GRAD_ATOL, rtol=MXFP8_GRAD_RTOL)
+ print("MXFP8 parity OK.\n")
+
+
+def main() -> None:
+ assert torch.cuda.is_available(), "CUDA required."
+
+ cfg = _build_config()
+ device = "cuda"
+ dtype = torch.bfloat16
+
+ model_hf = _build_hf(cfg, device, dtype)
+ inputs = _make_inputs(cfg, device)
+
+ _run_bf16(cfg, model_hf, inputs, device, dtype)
+ _run_mxfp8(cfg, model_hf, inputs, device, dtype)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/docs/examples/te_mixtral/tutorial_accelerate_hf_mixtral_with_te.ipynb b/docs/examples/te_mixtral/tutorial_accelerate_hf_mixtral_with_te.ipynb
new file mode 100644
index 0000000000..decdcff31e
--- /dev/null
+++ b/docs/examples/te_mixtral/tutorial_accelerate_hf_mixtral_with_te.ipynb
@@ -0,0 +1,462 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Accelerating Hugging Face Mixtral MoE Fine-Tuning with Transformer Engine\n",
+ "\n",
+ "
\n",
+ "\n",
+ "Goal\n",
+ "\n",
+ "This tutorial showcases how to accelerate fine-tuning a mixture-of-experts model, [Mixtral-8x7B](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1), with Transformer Engine (TE) in `BF16` and `MXFP8` precision.\n",
+ "
"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "**Setup**\n",
+ "\n",
+ "Mixtral-8x7B has 8 experts and roughly 47B total parameters. In `BF16` the model weights alone consume ~93 GB, and full `AdamW` fine-tuning needs ~370 GB. This tutorial is tested on 8x B300 GPUs with `Expert Parallelism (EP) = 2` and `Data Parallelism (DP) = 4`, so the experts are divided across 2 GPUs and there are 4 replicas. The container used is [pytorch-26.04-py3](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch?version=26.04-py3). A sequence length of 8192 and a global batch size of 48 are used across the experiments.\n",
+ "\n",
+ "Install the required Python packages using the following command in a terminal:"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "```bash\n",
+ "pip install -r requirements.txt \n",
+ "```"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Table of Contents\n",
+ "\n",
+ "1. [Baseline] Running HF Mixtral -- Without Expert Parallelism (Precision: `BF16`)\n",
+ "2. [Improvement 1] Transformer Engine with Expert Parallelism (Precision: `BF16`)\n",
+ "3. [Improvement 2] Batched Expert Execution with `GroupedLinear` (Precision: `BF16`)\n",
+ "4. [Improvement 3] Precision Optimization and Fused MLP (Precision: `MXFP8`)\n",
+ "5. Conclusion\n",
+ "6. Appendix: Dependencies"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## [Baseline] Running HF Mixtral -- Without Expert Parallelism (Precision: `BF16`)\n",
+ "\n",
+ "Before applying any Transformer Engine optimizations, we establish a Hugging Face (HF) baseline. Mixtral replaces the standard Transformer feed-forward network (FFN) with a sparse **Mixture of Experts** (MoE): a learned router selects the top-2 experts out of 8 per token, as shown in Fig 1. \n",
+ "\n",
+ "\n",
+ "\n",
+ " Fig 1: Dense Transformer block (left) vs Sparse MoE Transformer block (right).\n",
+ "\n",
+ "\n",
+ "\n",
+ "The current HF implementation has two limitations.\n",
+ "\n",
+ "1. **Pipeline parallelism**. Because the full model does not fit on one GPU, the baseline uses pipeline parallelism to split the model across GPUs. This is the simplest way to partition a model, but GPU utilization is limited by pipeline bubbles and sequential layer dependencies.\n",
+ "\n",
+ "\n",
+ "2. **Excessive kernel launches.** [HF's MixtralSparseMoeBlock](https://github.com/huggingface/transformers/blob/3ef278124e47832f34406ca3ca85bc50ad8b79bb/src/transformers/models/mixtral/modeling_mixtral.py) iterates over all 8 experts in a Python loop. Each expert triggers individual kernel launches. \n",
+ "\n",
+ "```python\n",
+ "for expert_idx, expert_layer in enumerate(self.experts):\n",
+ " idx, top_x = torch.where(expert_mask[expert_idx])\n",
+ " current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)\n",
+ " current_hidden = expert_layer(current_state) * routing_weights[top_x, idx, None]\n",
+ " final_hidden_states.index_add_(0, top_x, current_hidden)\n",
+ "```\n",
+ "\n",
+ "For each layer, HF loops through the experts sequentially. Each expert is much smaller than the dense FFN, so each expert GEMM is small and cannot saturate the GPU's tensor cores. Looping over many experts therefore launches many small GEMMs, leaving the FFN dominated by orchestration overhead and memory movement.\n",
+ "\n",
+ "\n",
+ "The script [run_finetune_ep.py](run_finetune_ep.py) initializes Hugging Face and then runs fine-tuning. For the full implementation, refer to [utils.py](utils.py). Now, let's execute the following command in the terminal."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "```bash\n",
+ "python3 run_finetune_ep.py --improvement 0 --batch-size 48 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n",
+ "```"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Here is the expected output:\n",
+ "\n",
+ "```\n",
+ "30 fine-tuning steps complete!\n",
+ "Median time per step: 2472 ms\n",
+ "```\n",
+ "\n",
+ "Let's add this information in a table and keep comparing it with a few possible improvements in future sections:\n",
+ "\n",
+ "| Models | Precision | Step Time | Speedup (over baseline) |\n",
+ "|---|---|---:|---:|\n",
+ "| HF baseline | BF16 | 2472 ms | 1 |"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## [Improvement 1] Transformer Engine with Expert Parallelism (Precision: `BF16`)\n",
+ "\n",
+ "Now that we have a baseline, let's bring in Transformer Engine. This section replaces the HF Transformer block with TE modules and introduces expert parallelism (EP). \n",
+ "\n",
+ "\n",
+ "\n",
+ " Fig 2: HF MixtralDecoderLayer (left) wrapped by TE modules (right).\n",
+ "\n",
+ "\n",
+ "**Fused Building Blocks**\n",
+ "\n",
+ "- **Attention block.** Instead of using one module for layer norm and another module for attention, TE combines them (`RMSNorm` and attention) with `te.MultiheadAttention`, where the input `RMSNorm` is bundled with the fused QKV projection. The layer norm weights and QKV weights are stored in the same building block: `self_attention.layernorm_qkv.weight`. Here is how you can use TE's attention block:\n",
+ "\n",
+ " ```python\n",
+ " self.self_attention = transformer_engine.pytorch.MultiheadAttention(\n",
+ " hidden_size=config.hidden_size,\n",
+ " fuse_qkv_params=True,\n",
+ " qkv_weight_interleaved=True,\n",
+ " normalization=\"RMSNorm\",\n",
+ " input_layernorm=True,\n",
+ " ...\n",
+ " )\n",
+ " ```\n",
+ "\n",
+ "- **MoE block.** TE provides the building blocks for the MoE layer. First, the gate computes router probabilities, then `softmax` and `top-k` select the top two experts out of eight for each token. The selected experts are then passed to the dispatcher. The dispatcher determines which EP ranks host the selected MoE experts, and NCCL handles the all-to-all communication that moves tokens to those ranks. The dispatcher also uses `transformer_engine.pytorch.moe_permute_and_pad_with_probs` to handle padding requirements, such as padding to multiples of 32 required by `MXFP8`.\n",
+ "\n",
+ " Below is the overview pseudocode for the MoE block:\n",
+ "\n",
+ " ```python\n",
+ " router_logits = self.gate(hidden_states) \n",
+ "\n",
+ " softmax_probs = torch.nn.functional.softmax(router_logits, dim=-1)\n",
+ "\n",
+ " routing_weights, selected_experts = torch.topk(softmax_probs, self.top_k, dim=-1)\n",
+ "\n",
+ " dispatch_output = self.dispatcher.dispatch(hidden_states, selected_experts, routing_weights)\n",
+ " ```\n",
+ "\n",
+ "**Parallelism layout**\n",
+ "\n",
+ "In this tutorial, EP=2 is used to split the model between 2 GPUs, each hosting 4 experts. In this 8-GPU setup, the model is replicated 4 times, creating 4 data-parallel groups.\n",
+ "\n",
+ "Here is how to set up EP:\n",
+ "\n",
+ "```python\n",
+ "config.expert_parallel_size = 2\n",
+ "ep_size = config.expert_parallel_size\n",
+ "dp_size = world_size // ep_size\n",
+ "ep_group = None\n",
+ "for dp_rank in range(dp_size):\n",
+ " ranks = list(range(dp_rank * ep_size, (dp_rank + 1) * ep_size))\n",
+ " group = dist.new_group(ranks=ranks)\n",
+ " if dist.get_rank() in ranks:\n",
+ " ep_group = group\n",
+ "model.model.set_ep_groups(ep_group=ep_group)\n",
+ "```\n",
+ "\n",
+ "**Mapping the HF checkpoint to TE**\n",
+ "\n",
+ "Some weights/parameters need to be reshaped and also remapped to corresponding weight names in TE modules. The `replace_params` helper in [te_mixtral.py](te_mixtral.py) performs the mapping (also illustrated in Fig 2 above). The two non-trivial groups are:\n",
+ "\n",
+ "- **Attention.** HF stores Q, K, V as separate projections; TE fuses them into a single QKV weight that lives under the `layernorm_qkv` submodule:\n",
+ "\n",
+ "| HF key | TE key |\n",
+ "|---|---|\n",
+ "| `self_attn.q_proj.weight` | `self_attention.layernorm_qkv.weight` (Q slice) |\n",
+ "| `self_attn.k_proj.weight` | `self_attention.layernorm_qkv.weight` (K slice) |\n",
+ "| `self_attn.v_proj.weight` | `self_attention.layernorm_qkv.weight` (V slice) |\n",
+ "| `input_layernorm.weight` | `self_attention.layernorm_qkv.layer_norm_weight` |\n",
+ "\n",
+ "- **MoE experts.** HF packs all experts' `SwiGLU` projections into two tensors per layer; TE keeps the same packing under different attribute names so `replace_params` is essentially a copy:\n",
+ "\n",
+ "| HF key | TE key |\n",
+ "|---|---|\n",
+ "| `mlp.experts.gate_up_proj` `[num_experts, 2*ffn, h]` | `mlp.experts_gate_up_weight` |\n",
+ "| `mlp.experts.down_proj` `[num_experts, h, ffn]` | `mlp.experts_down_weight` |\n",
+ "| `mlp.gate.weight` | `mlp.gate.weight` |\n",
+ "\n",
+ "All other weights (embeddings, norms, LM head) are direct copies. See `replace_params` in `te_mixtral.py` for the full mapping.\n",
+ "\n",
+ "Let's launch the same fine-tuning loop -- this time across 8 GPUs via `torchrun`. See `run_finetune_ep.py` and `utils.py` for the full implementation.\n",
+ "\n",
+ "Now, let's execute the following command."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "```bash\n",
+ "torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 1 --ep-size 2 --batch-size 12 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n",
+ "```"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Here is the expected output:\n",
+ "\n",
+ "```\n",
+ "30 fine-tuning steps complete!\n",
+ "Median time per step: 747 ms\n",
+ "```\n",
+ "\n",
+ "Compared to the baseline implementation, we see the following result:\n",
+ "\n",
+ "| Models | Precision | Step Time | Speedup (over baseline) |\n",
+ "|---|---|---:|---:|\n",
+ "| HF baseline | BF16 | 2472 ms | 1 |\n",
+ "| TE decoder, TE building blocks, and MoE layer | BF16 | 747 ms | 3.31 |\n",
+ "\n",
+ "Improvement 1 is 3.31x faster than the baseline, a **231%** speedup."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## [Improvement 2] Batched Expert Execution with `GroupedLinear` (Precision: `BF16`)\n",
+ "\n",
+ "Improvement 1 kept the per-expert FFN as a Python loop. If each rank owns 4 local experts, the loop launches the per-expert GEMMs one by one. Another limitation is that the expert GEMMs are usually small: each one sees only the tokens routed to it, which is too small to feed the tensor cores efficiently. The following section shows how to execute the GEMMs in a batch. \n",
+ "\n",
+ "\n",
+ "\n",
+ " Fig 3: Left: looping through experts one-by-one. Right: one grouped-GEMM over all experts.\n",
+ ""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "`GroupedLinear` applies multiple linear transformations in one call. It gathers the experts' weights and input tokens. Although each expert receives a different number of tokens, `GroupedLinear` supports this by accepting per-expert token counts (`split_sizes`). `GroupedLinear` submits the local experts through TE’s grouped GEMM path instead of launching one PyTorch Linear operation per expert. This reduces per-expert launch and scheduling overhead. \n",
+ "\n",
+ "Here are the steps to use `GroupedLinear`. Each expert keeps its own weight tensor (`weight0`, `weight1`, ...), and the call takes the per-expert token counts as an extra positional argument:\n",
+ "\n",
+ "```python\n",
+ "from transformer_engine.pytorch.ops import GroupedLinear\n",
+ "\n",
+ "experts_gate_up = GroupedLinear(\n",
+ " num_groups=num_local_experts,\n",
+ " in_features=hidden_size,\n",
+ " out_features=2 * intermediate_size,\n",
+ " bias=False,\n",
+ " dtype=torch.bfloat16,\n",
+ " device=\"cuda\",\n",
+ ")\n",
+ "\n",
+ "gate_up_output = experts_gate_up(tokens, split_sizes)\n",
+ "```\n",
+ "\n",
+ "Compared with the Python loop in Improvement 1, this becomes one gate-up projection per layer instead of 4 separate calls (4 is the number of experts on a GPU). The expert weights can be imported from HF. In `te_mixtral.py`, the grouped-op path keeps each local expert as a normal per-expert `weight{i}` parameter and loads the owning expert slice directly.\n",
+ "\n",
+ "To see the effect of `GroupedLinear`, we keep everything else unchanged. Execute the following command in the terminal. "
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "```bash\n",
+ "torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 2 --ep-size 2 --batch-size 12 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n",
+ "```"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Here is the expected output:\n",
+ "\n",
+ "```\n",
+ "30 fine-tuning steps complete!\n",
+ "Median time per step: 635 ms\n",
+ "```\n",
+ "\n",
+ "Adding the `GroupedLinear` result gives us:\n",
+ "\n",
+ "| Models | Precision | Step Time | Speedup (over baseline) |\n",
+ "|---|---|---:|---:|\n",
+ "| HF baseline | BF16 | 2472 ms | 1 |\n",
+ "| TE decoder, TE building blocks, and MoE layer | BF16 | 747 ms | 3.31 |\n",
+ "| TE with `GroupedLinear` | BF16 | 635 ms | 3.89 |\n",
+ "\n",
+ "`GroupedLinear` reaches a 3.89x speedup over the baseline, or **289%**."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## [Improvement 3] Precision Optimization and Fused MLP (Precision: `MXFP8`)\n",
+ "\n",
+ "With EP and grouped expert GEMMs in place, the next improvement lowers precision from `BF16` to `MXFP8`. `MXFP8` converts the weight and activation values to 8 bits instead of 16 bits. To preserve dynamic range, it keeps one `E8M0` scale factor for every 32 values; applying that scale recovers a wider numerical range. On Blackwell GPUs, `MXFP8` is native and hardware accelerated, so `MXFP8` GEMMs can run through specialized Tensor Core instructions. Read more about `MXFP8` and block scaling in the [Transformer Engine FP8 primer](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html#MXFP8-and-block-scaling).\n",
+ "\n",
+ "The model still keeps its master weights in `BF16`, so using `MXFP8` adds `Quantization` and `De-Quantization` work around the GEMMs. `Quantization` converts `BF16` weights and activations to `MXFP8` before the low-precision GEMM; `De-Quantization` converts the result back to the higher-precision format. The naive path performs these as separate operations. This motivates the fused MLP path shown below.\n",
+ "\n",
+ "\n",
+ "\n",
+ " Fig 4: The MXFP8 path fuses multiple operations into one kernel before the down projection.\n",
+ "\n",
+ "\n",
+ "To use `MXFP8`, we simply define a recipe and pass it to the model. \n",
+ "\n",
+ "```python\n",
+ "fp8_recipe = te_recipe.MXFP8BlockScaling()\n",
+ "model = TEMixtralMXFP8ForCausalLM(config, fp8_recipe=fp8_recipe, dispatcher=dispatcher)\n",
+ "```\n",
+ "\n",
+ "Now, the model's forward and backward passes run under `MXFP8` precision which is enabled through TE's `autocast` API:\n",
+ "\n",
+ "```python\n",
+ "with te.autocast(enabled=True, recipe=self._fp8_recipe):\n",
+ " for decoder_layer in self.layers:\n",
+ " hidden_states = decoder_layer(hidden_states)\n",
+ "```"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "To use fused MLP, we import TE's `Sequential` API to chain together `gate_up`, `ScaledSwiGLU`, and `down`. It also folds the `De-Quantization` step into the fused path. `ScaledSwiGLU` is chosen to combine the routing probabilities (\"scales\") with the expert FFN computations. \n",
+ "\n",
+ "```python\n",
+ "from transformer_engine.pytorch.ops import GroupedLinear, ScaledSwiGLU, Sequential\n",
+ "\n",
+ "experts_ffn = Sequential(GroupedLinear(gate_up), ScaledSwiGLU(), GroupedLinear(down))\n",
+ "```\n",
+ "\n",
+ "TE's `Sequential` scans the ops and, if the pattern matches, replaces the `GroupedLinear -> ScaledSwiGLU -> GroupedLinear` pattern with a fused operation object: `ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8` for forward and a matching fused backward op. It reduces framework overhead, fuses the SwiGLU/probability-scaling work into the grouped MLP path, and avoids some intermediate materialization.\n",
+ "\n",
+ "
\n",
+ "\n",
+ "Note\n",
+ "\n",
+ "`NVTE_CUTEDSL_FUSED_GROUPED_MLP=1` must be set before TE imports the fused op registration. In this tutorial, `run_finetune_ep.py` already does that automatically for improvement 3.\n",
+ "\n",
+ "
"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Execute the following in the terminal:\n",
+ "\n",
+ "```bash\n",
+ "torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 3 --ep-size 2 --batch-size 12 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n",
+ "```\n",
+ "\n",
+ "Here is the expected result:\n",
+ "```\n",
+ "30 fine-tuning steps complete!\n",
+ "Median time per step: 542 ms\n",
+ "```"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "With MXFP8 fused MLP included, the final comparison is:\n",
+ "\n",
+ "| Models | Precision | Step Time | Speedup (over baseline) |\n",
+ "|---|---|---:|---:|\n",
+ "| HF baseline | BF16 | 2472 ms | 1 |\n",
+ "| TE EP Python loop | BF16 | 747 ms | 3.31 |\n",
+ "| TE with `GroupedLinear` | BF16 | 635 ms | 3.89 |\n",
+ "| TE with MXFP8 fused MLP | MXFP8 | 542 ms | 4.56 |\n",
+ "\n",
+ "For Mixtral-8x7B, we get the largest speedup with MXFP8 fused MLP: 4.56x faster than the baseline, or **356%**."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Conclusion\n",
+ "\n",
+ "This tutorial walks through three progressive optimization improvements that speed up the fine-tuning of the Mixtral-8x7B model by replacing building blocks in a Hugging Face baseline with TE-native blocks like MXFP8 and fused grouped MLP. The tutorial uses **global batch 48, seq 8192** on 8x B300 to demonstrate the speedups.\n",
+ "\n",
+ "To run all improvements together, execute the following in a terminal. All four runs use the same global batch size of 48. The TE runs use DP=4, so the per-rank batch size is 12.\n",
+ "\n",
+ "```bash\n",
+ "python3 run_finetune_ep.py --improvement 0 --batch-size 48 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n",
+ "\n",
+ "torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 1 --ep-size 2 --batch-size 12 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n",
+ "\n",
+ "torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 2 --ep-size 2 --batch-size 12 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n",
+ "\n",
+ "torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 3 --ep-size 2 --batch-size 12 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n",
+ "```"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "
![](\"media/dense_to_sparse.svg\")
\n",
+ " ![](\"media/mixtral_decoder_swap.svg\")
\n",
+ " ![](\"media/moe_loop_vs_grouped.svg\")
\n",
+ " ![](\"media/fused_mlp_path.svg\")
\n",
+ "