From 45ed767a7af61a64dc41e29be1b6c0cc14c21d9c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 18 Aug 2026 12:37:21 +0200 Subject: [PATCH 01/41] [PyTorch] Add DeepSeekV3Layer skeleton (MLA + MoE) Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- .../pytorch/deepseek/__init__.py | 11 ++++++++ transformer_engine/pytorch/deepseek/moe.py | 27 +++++++++++++++++++ .../deepseek/multi_latent_attention.py | 24 +++++++++++++++++ .../pytorch/deepseek/transformer_layer.py | 25 +++++++++++++++++ 4 files changed, 87 insertions(+) create mode 100644 transformer_engine/pytorch/deepseek/__init__.py create mode 100644 transformer_engine/pytorch/deepseek/moe.py create mode 100644 transformer_engine/pytorch/deepseek/multi_latent_attention.py create mode 100644 transformer_engine/pytorch/deepseek/transformer_layer.py diff --git a/transformer_engine/pytorch/deepseek/__init__.py b/transformer_engine/pytorch/deepseek/__init__.py new file mode 100644 index 0000000000..5dafdf1fff --- /dev/null +++ b/transformer_engine/pytorch/deepseek/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DeepSeekV3 transformer layer built from Transformer Engine MoE building blocks.""" + +from transformer_engine.pytorch.deepseek.multi_latent_attention import MultiLatentAttention +from transformer_engine.pytorch.deepseek.moe import DeepSeekV3MoE +from transformer_engine.pytorch.deepseek.transformer_layer import DeepSeekV3Layer + +__all__ = ["DeepSeekV3Layer", "DeepSeekV3MoE", "MultiLatentAttention"] diff --git a/transformer_engine/pytorch/deepseek/moe.py b/transformer_engine/pytorch/deepseek/moe.py new file mode 100644 index 0000000000..f4f787743b --- /dev/null +++ b/transformer_engine/pytorch/deepseek/moe.py @@ -0,0 +1,27 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DeepSeekV3 MoE block: sigmoid router with aux-loss-free bias, shared + +routed experts.""" + +import torch + +__all__ = ["DeepSeekV3MoE"] + + +class DeepSeekV3MoE(torch.nn.Module): + """ + DeepSeekV3-style Mixture of Experts block composed from TE MoE + primitives: ``fused_topk_with_score_function`` (sigmoid score function, + expert bias, grouped top-k), ``moe_permute_with_probs``/``moe_unpermute``, + :class:`GroupedLinear` routed experts, a shared expert + (:class:`LayerNormMLP`), ``Fp8Padding``/``Fp8Unpadding`` and optional + expert parallelism via ``ep_dispatch``/``ep_combine``. + + .. warning:: Work in progress, not functional yet. + """ + + def __init__(self, *args, **kwargs): + super().__init__() + raise NotImplementedError("DeepSeekV3MoE is under development") diff --git a/transformer_engine/pytorch/deepseek/multi_latent_attention.py b/transformer_engine/pytorch/deepseek/multi_latent_attention.py new file mode 100644 index 0000000000..6c2bb7420b --- /dev/null +++ b/transformer_engine/pytorch/deepseek/multi_latent_attention.py @@ -0,0 +1,24 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Multi-Latent Attention (MLA) block as used in DeepSeekV3.""" + +import torch + +__all__ = ["MultiLatentAttention"] + + +class MultiLatentAttention(torch.nn.Module): + """ + Multi-Latent Attention with low-rank Q/KV down-projections and a + decoupled RoPE/NoPE head split, composed from :class:`Linear`, + :class:`LayerNormLinear` and :class:`DotProductAttention` + (``kv_channels=(head_dim_qk, head_dim_v)``). + + .. warning:: Work in progress, not functional yet. + """ + + def __init__(self, *args, **kwargs): + super().__init__() + raise NotImplementedError("MultiLatentAttention is under development") diff --git a/transformer_engine/pytorch/deepseek/transformer_layer.py b/transformer_engine/pytorch/deepseek/transformer_layer.py new file mode 100644 index 0000000000..2a28a6ceb3 --- /dev/null +++ b/transformer_engine/pytorch/deepseek/transformer_layer.py @@ -0,0 +1,25 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DeepSeekV3 transformer layer.""" + +import torch + +__all__ = ["DeepSeekV3Layer"] + + +class DeepSeekV3Layer(torch.nn.Module): + """ + A full DeepSeekV3 transformer layer, analogous to + :class:`TransformerLayer`: :class:`MultiLatentAttention` followed by + either a dense :class:`LayerNormMLP` (first layers) or + :class:`DeepSeekV3MoE`, with the same residual and fused + bias-dropout-add plumbing as :class:`TransformerLayer`. + + .. warning:: Work in progress, not functional yet. + """ + + def __init__(self, *args, **kwargs): + super().__init__() + raise NotImplementedError("DeepSeekV3Layer is under development") From c306c6f840bfe187cbdc3ecad39ad5aa517b669e Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 18 Aug 2026 13:56:21 +0200 Subject: [PATCH 02/41] Move DeepSeekV3 skeleton to models/deepseek_v3 subpackage Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/models/__init__.py | 13 +++++++++++++ .../{deepseek => models/deepseek_v3}/__init__.py | 8 +++++--- .../pytorch/{deepseek => models/deepseek_v3}/moe.py | 0 .../deepseek_v3}/multi_latent_attention.py | 0 .../deepseek_v3}/transformer_layer.py | 0 5 files changed, 18 insertions(+), 3 deletions(-) create mode 100644 transformer_engine/pytorch/models/__init__.py rename transformer_engine/pytorch/{deepseek => models/deepseek_v3}/__init__.py (50%) rename transformer_engine/pytorch/{deepseek => models/deepseek_v3}/moe.py (100%) rename transformer_engine/pytorch/{deepseek => models/deepseek_v3}/multi_latent_attention.py (100%) rename transformer_engine/pytorch/{deepseek => models/deepseek_v3}/transformer_layer.py (100%) diff --git a/transformer_engine/pytorch/models/__init__.py b/transformer_engine/pytorch/models/__init__.py new file mode 100644 index 0000000000..bee5474c81 --- /dev/null +++ b/transformer_engine/pytorch/models/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Model-specific transformer layers composed from Transformer Engine modules.""" + +from transformer_engine.pytorch.models.deepseek_v3 import ( + DeepSeekV3Layer, + DeepSeekV3MoE, + MultiLatentAttention, +) + +__all__ = ["DeepSeekV3Layer", "DeepSeekV3MoE", "MultiLatentAttention"] diff --git a/transformer_engine/pytorch/deepseek/__init__.py b/transformer_engine/pytorch/models/deepseek_v3/__init__.py similarity index 50% rename from transformer_engine/pytorch/deepseek/__init__.py rename to transformer_engine/pytorch/models/deepseek_v3/__init__.py index 5dafdf1fff..a7cbb50ae2 100644 --- a/transformer_engine/pytorch/deepseek/__init__.py +++ b/transformer_engine/pytorch/models/deepseek_v3/__init__.py @@ -4,8 +4,10 @@ """DeepSeekV3 transformer layer built from Transformer Engine MoE building blocks.""" -from transformer_engine.pytorch.deepseek.multi_latent_attention import MultiLatentAttention -from transformer_engine.pytorch.deepseek.moe import DeepSeekV3MoE -from transformer_engine.pytorch.deepseek.transformer_layer import DeepSeekV3Layer +from transformer_engine.pytorch.models.deepseek_v3.multi_latent_attention import ( + MultiLatentAttention, +) +from transformer_engine.pytorch.models.deepseek_v3.moe import DeepSeekV3MoE +from transformer_engine.pytorch.models.deepseek_v3.transformer_layer import DeepSeekV3Layer __all__ = ["DeepSeekV3Layer", "DeepSeekV3MoE", "MultiLatentAttention"] diff --git a/transformer_engine/pytorch/deepseek/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py similarity index 100% rename from transformer_engine/pytorch/deepseek/moe.py rename to transformer_engine/pytorch/models/deepseek_v3/moe.py diff --git a/transformer_engine/pytorch/deepseek/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py similarity index 100% rename from transformer_engine/pytorch/deepseek/multi_latent_attention.py rename to transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py diff --git a/transformer_engine/pytorch/deepseek/transformer_layer.py b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py similarity index 100% rename from transformer_engine/pytorch/deepseek/transformer_layer.py rename to transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py From f73d04edaf0e16740f428b537c78d70d0c9c1ec1 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 18 Aug 2026 14:00:07 +0200 Subject: [PATCH 03/41] Add DeepSeekV3 layer entries to PyTorch API docs Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- docs/api/pytorch.rst | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 5fac0a89a6..bd3099b590 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -59,6 +59,15 @@ PyTorch .. autoapifunction:: transformer_engine.pytorch.deinterleave_glu_tensor +Model-specific layers +--------------------- + +.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3Layer(**kwargs) + +.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3MoE(**kwargs) + +.. autoapiclass:: transformer_engine.pytorch.models.MultiLatentAttention(**kwargs) + Data types ---------- From 09f28a9a3903b85ea28acff8ef63149738f38ea8 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 18 Aug 2026 14:13:44 +0200 Subject: [PATCH 04/41] [PyTorch] Implement DeepSeekV3Layer: MLA + MoE from TE building blocks MultiLatentAttention: low-rank q/kv latents (RMSNorm fused into LayerNormLinear up-projections), decoupled RoPE/NoPE head split with a shared key rope head, DotProductAttention with kv_channels=(qk, v) for the cuDNN fused backend. DeepSeekV3MoE: fused sigmoid router with aux-loss-free expert bias and grouped top-k, routed experts as te.ops GroupedLinear+ScaledSwiGLU+ GroupedLinear (CuTe fused grouped MLP on supported HW), probs applied per-token in the activation, local permute/unpermute or NCCL expert parallelism via ep_dispatch/ep_combine, optional shared expert. DeepSeekV3Layer: pre-RMSNorm + MLA and dense LayerNormMLP (RMSNorm, swiglu) or MoE with residual connections. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_deepseek.py | 124 +++++++++ transformer_engine/pytorch/__init__.py | 1 + .../pytorch/models/deepseek_v3/moe.py | 245 +++++++++++++++++- .../deepseek_v3/multi_latent_attention.py | 193 +++++++++++++- .../models/deepseek_v3/transformer_layer.py | 163 +++++++++++- 5 files changed, 702 insertions(+), 24 deletions(-) create mode 100644 tests/pytorch/test_deepseek.py diff --git a/tests/pytorch/test_deepseek.py b/tests/pytorch/test_deepseek.py new file mode 100644 index 0000000000..7778d0448c --- /dev/null +++ b/tests/pytorch/test_deepseek.py @@ -0,0 +1,124 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import pytest +import torch + +from transformer_engine.pytorch.utils import deinterleave_glu_tensor +from transformer_engine.pytorch.models import ( + DeepSeekV3Layer, + DeepSeekV3MoE, + MultiLatentAttention, +) + +SEQ_LEN = 128 +BATCH = 2 +HIDDEN = 256 +HEADS = 4 +DTYPE = torch.bfloat16 + +MLA_KWARGS = dict( + q_lora_rank=96, + kv_lora_rank=64, + qk_nope_head_dim=64, + qk_rope_head_dim=32, + v_head_dim=64, +) + + +def _input(requires_grad=True): + torch.manual_seed(1234) + return torch.randn( + SEQ_LEN, BATCH, HIDDEN, dtype=DTYPE, device="cuda", requires_grad=requires_grad + ) + + +def test_mla_forward_backward(): + torch.manual_seed(0) + mla = MultiLatentAttention(HIDDEN, HEADS, params_dtype=DTYPE, **MLA_KWARGS) + x = _input() + out = mla(x) + assert out.shape == x.shape + out.sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() + + +@pytest.mark.parametrize("shared", [False, True], ids=["no_shared", "shared"]) +@pytest.mark.parametrize("grouped", [False, True], ids=["ungrouped", "grouped"]) +def test_moe_forward_backward(shared, grouped): + torch.manual_seed(0) + moe = DeepSeekV3MoE( + HIDDEN, + moe_ffn_hidden_size=128, + num_experts=8, + topk=2, + num_groups=4 if grouped else None, + group_topk=2 if grouped else None, + shared_expert_ffn_hidden_size=128 if shared else None, + params_dtype=DTYPE, + ) + x = _input() + out = moe(x) + assert out.shape == x.shape + out.sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() + + counts = moe._last_tokens_per_expert + assert counts.sum().item() == SEQ_LEN * BATCH * 2 + bias_before = moe.expert_bias.clone() + moe.update_expert_bias() + assert not torch.equal(bias_before, moe.expert_bias) + + +def test_moe_matches_dense_reference(): + """topk == num_experts with uniform probs must reduce to a sum of expert MLPs.""" + torch.manual_seed(0) + num_experts = 4 + moe = DeepSeekV3MoE( + HIDDEN, + moe_ffn_hidden_size=128, + num_experts=num_experts, + topk=num_experts, + routed_scaling_factor=1.0, + params_dtype=DTYPE, + ) + x = _input(requires_grad=False) + out = moe(x) + + tokens = x.reshape(-1, HIDDEN) + probs, _ = moe._route(moe.gate(tokens).float()) + fc1, _, fc2 = moe.experts + ref = torch.zeros_like(tokens) + for e in range(num_experts): + w1 = deinterleave_glu_tensor(getattr(fc1, f"weight{e}"), 32) + w2 = getattr(fc2, f"weight{e}") + gate_part, lin_part = (tokens @ w1.t()).chunk(2, dim=-1) + act = torch.nn.functional.silu(gate_part.float()) * lin_part.float() + ref += (act.to(DTYPE) * probs[:, e : e + 1].to(DTYPE)) @ w2.t() + torch.testing.assert_close(out.reshape(-1, HIDDEN), ref, rtol=0.05, atol=0.05) + + +@pytest.mark.parametrize("num_experts", [None, 8], ids=["dense", "moe"]) +def test_layer_forward_backward(num_experts): + torch.manual_seed(0) + layer = ( + DeepSeekV3Layer( + HIDDEN, + HEADS, + ffn_hidden_size=512, + num_experts=num_experts, + moe_ffn_hidden_size=128 if num_experts else None, + topk=2 if num_experts else None, + shared_expert_ffn_hidden_size=128 if num_experts else None, + params_dtype=DTYPE, + **MLA_KWARGS, + ) + if num_experts + else DeepSeekV3Layer(HIDDEN, HEADS, ffn_hidden_size=512, params_dtype=DTYPE, **MLA_KWARGS) + ) + x = _input() + out = layer(x) + assert out.shape == x.shape + out.sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 2b1803bfb2..fae4d973e5 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -34,6 +34,7 @@ from transformer_engine.pytorch.attention import InferenceParams from transformer_engine.pytorch.attention import RotaryPositionEmbedding from transformer_engine.pytorch.transformer import TransformerLayer +from transformer_engine.pytorch import models from transformer_engine.pytorch.permutation import ( moe_permute, moe_permute_with_probs, diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index f4f787743b..5a1c8d650c 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -5,23 +5,248 @@ """DeepSeekV3 MoE block: sigmoid router with aux-loss-free bias, shared + routed experts.""" +from typing import Optional, Union + import torch +import transformer_engine.pytorch.ops as te_ops +from transformer_engine.pytorch.router import fused_topk_with_score_function +from transformer_engine.pytorch.permutation import moe_permute_with_probs, moe_unpermute + __all__ = ["DeepSeekV3MoE"] +def _make_expert_mlp(num_experts, hidden_size, ffn_hidden_size, dtype, device): + # GroupedLinear + ScaledSwiGLU + GroupedLinear fuses into a single CuTe + # grouped MLP on supported hardware; elsewhere it runs as three ops with + # the same API and checkpoint layout. + return te_ops.Sequential( + te_ops.GroupedLinear( + num_experts, hidden_size, 2 * ffn_hidden_size, bias=False, dtype=dtype, device=device + ), + te_ops.ScaledSwiGLU(glu_interleave_size=32), + te_ops.GroupedLinear( + num_experts, ffn_hidden_size, hidden_size, bias=False, dtype=dtype, device=device + ), + ) + + class DeepSeekV3MoE(torch.nn.Module): """ - DeepSeekV3-style Mixture of Experts block composed from TE MoE - primitives: ``fused_topk_with_score_function`` (sigmoid score function, - expert bias, grouped top-k), ``moe_permute_with_probs``/``moe_unpermute``, - :class:`GroupedLinear` routed experts, a shared expert - (:class:`LayerNormMLP`), ``Fp8Padding``/``Fp8Unpadding`` and optional - expert parallelism via ``ep_dispatch``/``ep_combine``. - - .. warning:: Work in progress, not functional yet. + DeepSeekV3-style Mixture of Experts block. + + Routing uses the fused sigmoid router with aux-loss-free expert bias and + node-limited (grouped) top-k (``fused_topk_with_score_function``). Routed + experts run as a grouped SwiGLU MLP built from ``te.ops`` (fusable into a + single CuTe grouped-GEMM kernel); routing probabilities are applied + per-token inside the expert MLP, so unpermute/combine is a plain + accumulation. Token routing is either local + (``moe_permute_with_probs``/``moe_unpermute``) or, when ``ep_group`` is + given, expert-parallel over NCCL (``ep_dispatch``/``ep_combine``). + + When expert parallelism is used, ``transformer_engine.pytorch.ep.ep_bootstrap`` + must be called once per process before the first forward, and inputs must + be bfloat16. + + Parameters + ---------- + hidden_size : int + size of each input sample. + moe_ffn_hidden_size : int + ffn size of each routed expert. + num_experts : int + total number of routed experts. + topk : int, default = 8 + number of experts per token. + num_groups : int, optional + number of expert groups for node-limited routing. + group_topk : int, optional + number of groups each token is limited to. + routed_scaling_factor : float, default = 2.5 + scaling applied to the routing probabilities. + shared_expert_ffn_hidden_size : int, optional + ffn size of the shared expert; ``None`` + disables the shared expert. + expert_bias_update_rate : float, default = 1e-3 + step size of the aux-loss-free bias update + (see :meth:`update_expert_bias`). + params_dtype : torch.dtype, optional + dtype of module parameters. + ep_group : ProcessGroup, optional + expert-parallel process group; enables the NCCL EP path. + ep_max_tokens_per_rank : int, optional + max local tokens per forward (required with EP). + ep_recv_capacity_per_rank : int, optional + receive-buffer capacity; defaults to + ``ep_size * ep_max_tokens_per_rank * topk``. + ep_alignment : int, default = 128 + per-expert row alignment of the EP receive buffer. """ - def __init__(self, *args, **kwargs): + def __init__( + self, + hidden_size: int, + moe_ffn_hidden_size: int, + num_experts: int, + topk: int = 8, + num_groups: Optional[int] = None, + group_topk: Optional[int] = None, + routed_scaling_factor: float = 2.5, + shared_expert_ffn_hidden_size: Optional[int] = None, + expert_bias_update_rate: float = 1e-3, + params_dtype: Optional[torch.dtype] = None, + device: Union[torch.device, str] = "cuda", + ep_group: Optional[torch.distributed.ProcessGroup] = None, + ep_max_tokens_per_rank: Optional[int] = None, + ep_recv_capacity_per_rank: Optional[int] = None, + ep_alignment: int = 128, + ) -> None: super().__init__() - raise NotImplementedError("DeepSeekV3MoE is under development") + + dtype = params_dtype if params_dtype is not None else torch.get_default_dtype() + self.hidden_size = hidden_size + self.num_experts = num_experts + self.topk = topk + self.num_groups = num_groups + self.group_topk = group_topk + self.routed_scaling_factor = routed_scaling_factor + self.expert_bias_update_rate = expert_bias_update_rate + + self.gate = torch.nn.Linear( + hidden_size, num_experts, bias=False, dtype=dtype, device=device + ) + self.register_buffer( + "expert_bias", torch.zeros(num_experts, dtype=torch.float32, device=device) + ) + self._last_tokens_per_expert: Optional[torch.Tensor] = None + + self.ep_group = ep_group + self.ep_size = 1 if ep_group is None else torch.distributed.get_world_size(ep_group) + assert num_experts % self.ep_size == 0 + num_local_experts = num_experts // self.ep_size + + self.experts = _make_expert_mlp( + num_local_experts, hidden_size, moe_ffn_hidden_size, dtype, device + ) + + self.shared_expert = None + if shared_expert_ffn_hidden_size is not None: + self.shared_expert = te_ops.Sequential( + te_ops.Linear( + hidden_size, + 2 * shared_expert_ffn_hidden_size, + bias=False, + dtype=dtype, + device=device, + ), + te_ops.SwiGLU(), + te_ops.Linear( + shared_expert_ffn_hidden_size, + hidden_size, + bias=False, + dtype=dtype, + device=device, + ), + ) + + self.ep_buffer = None + if ep_group is not None: + from transformer_engine.pytorch.ep import EpBuffer + + assert ep_max_tokens_per_rank is not None, "EP requires ep_max_tokens_per_rank." + if ep_recv_capacity_per_rank is None: + ep_recv_capacity_per_rank = self.ep_size * ep_max_tokens_per_rank * topk + self.ep_buffer = EpBuffer( + top_k=topk, + max_tokens_per_rank=ep_max_tokens_per_rank, + hidden_dim=hidden_size, + num_local_experts=num_local_experts, + recv_capacity_per_rank=ep_recv_capacity_per_rank, + alignment=ep_alignment, + device=device, + ) + + def _route(self, logits: torch.Tensor, topk_indices: Optional[torch.Tensor] = None): + return fused_topk_with_score_function( + logits=logits, + topk=self.topk, + use_pre_softmax=False, + num_groups=self.num_groups, + group_topk=self.group_topk, + scaling_factor=self.routed_scaling_factor, + score_function="sigmoid", + expert_bias=self.expert_bias, + topk_indices=topk_indices, + ) + + def _forward_local(self, tokens: torch.Tensor) -> torch.Tensor: + probs, routing_map = self._route(self.gate(tokens).float()) + tokens_per_expert = routing_map.sum(dim=0) + self._last_tokens_per_expert = tokens_per_expert.detach() + + num_out = tokens.shape[0] * self.topk + permuted, permuted_probs, row_id_map = moe_permute_with_probs( + tokens, probs, routing_map, num_out_tokens=num_out + ) + + # The fused grouped MLP requires the total row count to be a multiple + # of 128; rows beyond sum(tokens_per_expert) fall outside every group. + pad = (-num_out) % 128 + if pad: + permuted = torch.nn.functional.pad(permuted, (0, 0, 0, pad)) + permuted_probs = torch.nn.functional.pad(permuted_probs, (0, pad)) + + out = self.experts( + permuted, tokens_per_expert, permuted_probs.to(tokens.dtype), tokens_per_expert + ) + return moe_unpermute(out[:num_out], row_id_map, restore_shape=tokens.shape) + + def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: + from transformer_engine.pytorch.ep import ep_dispatch, ep_combine + + assert tokens.dtype == torch.bfloat16, "The EP path requires bfloat16 inputs." + topk_idx = torch.empty( + (tokens.shape[0], self.topk), dtype=torch.int64, device=tokens.device + ) + probs, topk_idx = self._route(self.gate(tokens).float(), topk_indices=topk_idx) + self._last_tokens_per_expert = torch.bincount( + topk_idx.flatten(), minlength=self.num_experts + ) + topk_weights = probs.gather(1, topk_idx).float() + + recv_tokens, recv_weights, tokens_per_expert = ep_dispatch( + self.ep_buffer, tokens, topk_idx, topk_weights + ) + expert_out = self.experts( + recv_tokens, tokens_per_expert, recv_weights.to(tokens.dtype), tokens_per_expert + ) + return ep_combine(self.ep_buffer, expert_out, num_local_tokens=tokens.shape[0]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """ + Parameters + ---------- + hidden_states : torch.Tensor + input of shape ``[..., hidden_size]``. + """ + tokens = hidden_states.reshape(-1, self.hidden_size) + if self.ep_group is not None: + out = self._forward_ep(tokens) + else: + out = self._forward_local(tokens) + if self.shared_expert is not None: + out = out + self.shared_expert(tokens) + return out.view_as(hidden_states) + + @torch.no_grad() + def update_expert_bias(self) -> None: + """Aux-loss-free bias update from the last forward's routing counts. + + With data/expert parallelism, all-reduce ``_last_tokens_per_expert`` + across ranks before calling (or call on identically-routed ranks). + """ + counts = self._last_tokens_per_expert + if counts is None: + return + err = counts.float().mean() - counts.float() + self.expert_bias += self.expert_bias_update_rate * torch.sign(err) diff --git a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py index 6c2bb7420b..a36075f2c5 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py +++ b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py @@ -4,21 +4,200 @@ """Multi-Latent Attention (MLA) block as used in DeepSeekV3.""" +from typing import Optional, Union + import torch +from transformer_engine.pytorch.module import Linear, LayerNormLinear +from transformer_engine.pytorch.attention import DotProductAttention, RotaryPositionEmbedding +from transformer_engine.pytorch.attention.rope import apply_rotary_pos_emb + __all__ = ["MultiLatentAttention"] class MultiLatentAttention(torch.nn.Module): """ - Multi-Latent Attention with low-rank Q/KV down-projections and a - decoupled RoPE/NoPE head split, composed from :class:`Linear`, - :class:`LayerNormLinear` and :class:`DotProductAttention` - (``kv_channels=(head_dim_qk, head_dim_v)``). + Multi-Latent Attention as used in DeepSeekV3. - .. warning:: Work in progress, not functional yet. + Queries and key-values are projected through low-rank latents + (``q_lora_rank``, ``kv_lora_rank``); RMSNorm on each latent is fused into + the up-projection (:class:`LayerNormLinear` with RMSNorm). Each query/key + head is split into a ``qk_nope_head_dim`` part and a ``qk_rope_head_dim`` + part; RoPE is applied only to the rope part, and the key rope part comes + from a single shared head broadcast to all heads. Attention runs through + :class:`DotProductAttention` with asymmetric head dims + ``kv_channels=(qk_nope_head_dim + qk_rope_head_dim, v_head_dim)``, which + supports the cuDNN fused attention backend. + + Parameters + ---------- + hidden_size : int + size of each input sample. + num_attention_heads : int + number of attention heads. + q_lora_rank : int, default = 1536 + rank of the query latent. + kv_lora_rank : int, default = 512 + rank of the key-value latent. + qk_nope_head_dim : int, default = 128 + per-head dim of the non-rotary query/key part. + qk_rope_head_dim : int, default = 64 + per-head dim of the rotary query/key part. + v_head_dim : int, default = 128 + per-head dim of the values. + attention_dropout : float, default = 0.0 + dropout probability on attention scores. + attn_mask_type : str, default = "causal" + attention mask type passed to :class:`DotProductAttention`. + rotary_base : float, default = 10000.0 + RoPE base. + softmax_scale : float, optional + softmax scale; defaults to ``1/sqrt(qk head dim)`` inside + :class:`DotProductAttention`. + qkv_format : str, default = "sbhd" + layout of the input/output tensors. + params_dtype : torch.dtype, optional + dtype of module parameters. + tp_group : ProcessGroup, optional + tensor-parallel process group for the up/output projections. + tp_size : int, default = 1 + tensor-parallel world size. """ - def __init__(self, *args, **kwargs): + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + q_lora_rank: int = 1536, + kv_lora_rank: int = 512, + qk_nope_head_dim: int = 128, + qk_rope_head_dim: int = 64, + v_head_dim: int = 128, + attention_dropout: float = 0.0, + attn_mask_type: str = "causal", + rotary_base: float = 10000.0, + softmax_scale: Optional[float] = None, + qkv_format: str = "sbhd", + params_dtype: Optional[torch.dtype] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_size: int = 1, + device: Union[torch.device, str] = "cuda", + ) -> None: super().__init__() - raise NotImplementedError("MultiLatentAttention is under development") + + assert qkv_format in ("sbhd", "bshd"), "MultiLatentAttention supports sbhd/bshd formats." + assert num_attention_heads % tp_size == 0 + + self.qkv_format = qkv_format + self.num_attention_heads = num_attention_heads + self.num_attention_heads_per_partition = num_attention_heads // tp_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.kv_lora_rank = kv_lora_rank + + common = {"bias": False, "params_dtype": params_dtype, "device": device} + tp = {"tp_group": tp_group, "tp_size": tp_size} + + self.q_down_proj = Linear(hidden_size, q_lora_rank, **common) + self.q_up_proj = LayerNormLinear( + q_lora_rank, + num_attention_heads * self.qk_head_dim, + normalization="RMSNorm", + parallel_mode="column" if tp_size > 1 else None, + **tp, + **common, + ) + self.kv_down_proj = Linear(hidden_size, kv_lora_rank + qk_rope_head_dim, **common) + self.kv_up_proj = LayerNormLinear( + kv_lora_rank, + num_attention_heads * (qk_nope_head_dim + v_head_dim), + normalization="RMSNorm", + parallel_mode="column" if tp_size > 1 else None, + **tp, + **common, + ) + self.out_proj = Linear( + num_attention_heads * v_head_dim, + hidden_size, + parallel_mode="row" if tp_size > 1 else None, + **tp, + **common, + ) + + self.rope = RotaryPositionEmbedding(qk_rope_head_dim, rotary_base=rotary_base) + self._rope_freqs: Optional[torch.Tensor] = None + + self.core_attention = DotProductAttention( + num_attention_heads, + kv_channels=(self.qk_head_dim, v_head_dim), + attention_dropout=attention_dropout, + qkv_format=qkv_format, + attn_mask_type=attn_mask_type, + softmax_scale=softmax_scale, + tp_group=tp_group, + tp_size=tp_size, + ) + + def _rope_freqs_for(self, seq_len: int, device: torch.device) -> torch.Tensor: + if self._rope_freqs is None or self._rope_freqs.shape[0] < seq_len: + self._rope_freqs = self.rope(seq_len).to(device) + return self._rope_freqs[:seq_len] + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + attn_mask_type: Optional[str] = None, + checkpoint_core_attention: bool = False, + ) -> torch.Tensor: + """ + Parameters + ---------- + hidden_states : torch.Tensor + input of shape ``[sq, b, h]`` (sbhd) or ``[b, sq, h]`` (bshd). + attention_mask : torch.Tensor, optional + boolean mask passed to :class:`DotProductAttention`. + attn_mask_type : str, optional + override of the constructor's mask type. + checkpoint_core_attention : bool, default = False + checkpoint the core attention computation. + """ + seq_dim = 0 if self.qkv_format == "sbhd" else 1 + seq_len = hidden_states.shape[seq_dim] + heads = self.num_attention_heads_per_partition + + q = self.q_up_proj(self.q_down_proj(hidden_states)) + q = q.view(*q.shape[:-1], heads, self.qk_head_dim) + + kv_down = self.kv_down_proj(hidden_states) + kv_latent, k_pos = torch.split(kv_down, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + kv = self.kv_up_proj(kv_latent) + kv = kv.view(*kv.shape[:-1], heads, self.qk_nope_head_dim + self.v_head_dim) + k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + + freqs = self._rope_freqs_for(seq_len, hidden_states.device) + q_rope = apply_rotary_pos_emb( + q[..., self.qk_nope_head_dim :].contiguous(), + freqs, + tensor_format=self.qkv_format, + fused=True, + ) + k_rope = apply_rotary_pos_emb( + k_pos.unsqueeze(-2), freqs, tensor_format=self.qkv_format, fused=True + ) + + q = torch.cat([q[..., : self.qk_nope_head_dim], q_rope], dim=-1) + k = torch.cat([k_nope, k_rope.expand(*k_nope.shape[:-1], -1)], dim=-1) + + context = self.core_attention( + q, + k, + v.contiguous(), + attention_mask=attention_mask, + qkv_format=self.qkv_format, + attn_mask_type=attn_mask_type, + checkpoint_core_attention=checkpoint_core_attention, + ) + return self.out_proj(context) diff --git a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py index 2a28a6ceb3..af1eeb1a95 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py +++ b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py @@ -4,22 +4,171 @@ """DeepSeekV3 transformer layer.""" +from typing import Optional, Union + import torch +from transformer_engine.pytorch.module import LayerNormMLP, RMSNorm +from transformer_engine.pytorch.models.deepseek_v3.multi_latent_attention import ( + MultiLatentAttention, +) +from transformer_engine.pytorch.models.deepseek_v3.moe import DeepSeekV3MoE + __all__ = ["DeepSeekV3Layer"] class DeepSeekV3Layer(torch.nn.Module): """ A full DeepSeekV3 transformer layer, analogous to - :class:`TransformerLayer`: :class:`MultiLatentAttention` followed by - either a dense :class:`LayerNormMLP` (first layers) or - :class:`DeepSeekV3MoE`, with the same residual and fused - bias-dropout-add plumbing as :class:`TransformerLayer`. + :class:`TransformerLayer`: pre-RMSNorm + :class:`MultiLatentAttention`, + then either a dense SwiGLU MLP (:class:`LayerNormMLP` with RMSNorm, used + for the first dense layers of DeepSeekV3) or :class:`DeepSeekV3MoE`, each + with a residual connection. - .. warning:: Work in progress, not functional yet. + Parameters + ---------- + hidden_size : int + size of each input sample. + num_attention_heads : int + number of attention heads. + ffn_hidden_size : int + ffn size of the dense MLP (used when ``num_experts`` is + ``None``). + num_experts : int, optional + number of routed experts; ``None`` makes this a dense layer. + moe_ffn_hidden_size : int, optional + ffn size of each routed expert (required with MoE). + hidden_dropout : float, default = 0.0 + dropout probability on the residual branches. + kwargs common to the submodules (``q_lora_rank``, ``kv_lora_rank``, + ``qk_nope_head_dim``, ``qk_rope_head_dim``, ``v_head_dim``, + ``attention_dropout``, ``attn_mask_type``, ``qkv_format``, ``topk``, + ``num_groups``, ``group_topk``, ``routed_scaling_factor``, + ``shared_expert_ffn_hidden_size``, EP options, ...) are forwarded to + :class:`MultiLatentAttention` and :class:`DeepSeekV3MoE`. """ - def __init__(self, *args, **kwargs): + _MLA_KWARGS = frozenset( + { + "q_lora_rank", + "kv_lora_rank", + "qk_nope_head_dim", + "qk_rope_head_dim", + "v_head_dim", + "attention_dropout", + "attn_mask_type", + "rotary_base", + "softmax_scale", + "qkv_format", + "tp_group", + "tp_size", + } + ) + _MOE_KWARGS = frozenset( + { + "topk", + "num_groups", + "group_topk", + "routed_scaling_factor", + "shared_expert_ffn_hidden_size", + "expert_bias_update_rate", + "ep_group", + "ep_max_tokens_per_rank", + "ep_recv_capacity_per_rank", + "ep_alignment", + } + ) + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + ffn_hidden_size: Optional[int] = None, + num_experts: Optional[int] = None, + moe_ffn_hidden_size: Optional[int] = None, + hidden_dropout: float = 0.0, + layernorm_epsilon: float = 1e-5, + params_dtype: Optional[torch.dtype] = None, + device: Union[torch.device, str] = "cuda", + **kwargs, + ) -> None: super().__init__() - raise NotImplementedError("DeepSeekV3Layer is under development") + + unknown = set(kwargs) - self._MLA_KWARGS - self._MOE_KWARGS + if unknown: + raise TypeError(f"Unexpected keyword arguments: {sorted(unknown)}") + mla_kwargs = {k: v for k, v in kwargs.items() if k in self._MLA_KWARGS} + moe_kwargs = {k: v for k, v in kwargs.items() if k in self._MOE_KWARGS} + + self.hidden_dropout = hidden_dropout + + self.input_layernorm = RMSNorm( + hidden_size, eps=layernorm_epsilon, device=device, dtype=params_dtype + ) + self.self_attention = MultiLatentAttention( + hidden_size, + num_attention_heads, + params_dtype=params_dtype, + device=device, + **mla_kwargs, + ) + + if num_experts is None: + assert ffn_hidden_size is not None, "Dense layers require ffn_hidden_size." + self.pre_mlp_layernorm = None + self.mlp = LayerNormMLP( + hidden_size, + ffn_hidden_size, + eps=layernorm_epsilon, + normalization="RMSNorm", + activation="swiglu", + bias=False, + params_dtype=params_dtype, + device=device, + ) + else: + assert moe_ffn_hidden_size is not None, "MoE layers require moe_ffn_hidden_size." + self.pre_mlp_layernorm = RMSNorm( + hidden_size, eps=layernorm_epsilon, device=device, dtype=params_dtype + ) + self.mlp = DeepSeekV3MoE( + hidden_size, + moe_ffn_hidden_size, + num_experts, + params_dtype=params_dtype, + device=device, + **moe_kwargs, + ) + + def _residual_add(self, out: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + out = torch.nn.functional.dropout(out, p=self.hidden_dropout, training=self.training) + return residual + out + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + checkpoint_core_attention: bool = False, + ) -> torch.Tensor: + """ + Parameters + ---------- + hidden_states : torch.Tensor + input of shape ``[sq, b, h]`` (sbhd) or ``[b, sq, h]`` (bshd). + attention_mask : torch.Tensor, optional + boolean attention mask. + checkpoint_core_attention : bool, default = False + checkpoint the core attention computation. + """ + attention_out = self.self_attention( + self.input_layernorm(hidden_states), + attention_mask=attention_mask, + checkpoint_core_attention=checkpoint_core_attention, + ) + hidden_states = self._residual_add(attention_out, hidden_states) + + if self.pre_mlp_layernorm is not None: + mlp_out = self.mlp(self.pre_mlp_layernorm(hidden_states)) + else: + mlp_out = self.mlp(hidden_states) + return self._residual_add(mlp_out, hidden_states) From e23100b73cef8e7f7c955af8f626f83177aef06e Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 18 Aug 2026 14:28:46 +0200 Subject: [PATCH 05/41] Add distributed EP test for DeepSeekV3 MoE/layer run_deepseek_ep.py checks the EP path against the all-experts-local path numerically (forward, input/gate grads, all-reduced expert wgrads) and smoke-tests the full layer with EP. Also size the default EP recv capacity for per-expert alignment padding and the fused grouped MLP's row-count requirement. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_deepseek_ep.py | 185 ++++++++++++++++++ .../distributed/run_test_deepseek_ep.sh | 52 +++++ tests/pytorch/distributed/test_deepseek_ep.py | 26 +++ .../pytorch/models/deepseek_v3/moe.py | 6 +- 4 files changed, 268 insertions(+), 1 deletion(-) create mode 100644 tests/pytorch/distributed/run_deepseek_ep.py create mode 100644 tests/pytorch/distributed/run_test_deepseek_ep.sh create mode 100644 tests/pytorch/distributed/test_deepseek_ep.py diff --git a/tests/pytorch/distributed/run_deepseek_ep.py b/tests/pytorch/distributed/run_deepseek_ep.py new file mode 100644 index 0000000000..0edae05961 --- /dev/null +++ b/tests/pytorch/distributed/run_deepseek_ep.py @@ -0,0 +1,185 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Multi-process DeepSeekV3 MoE/layer EP tests, launched via torchrun.""" + +import os +import sys +import unittest + +import torch +import torch.distributed as dist + +from transformer_engine.pytorch.ep import ep_bootstrap, ep_finalize, release_symm_mem_pool +from transformer_engine.pytorch.models import DeepSeekV3Layer, DeepSeekV3MoE + +HIDDEN = 256 +MOE_FFN = 128 +SHARED_FFN = 128 +NUM_LOCAL_EXPERTS = 2 +TOP_K = 2 +TOKENS_PER_RANK = 64 +HEADS = 4 +DTYPE = torch.bfloat16 + +MLA_KWARGS = dict( + q_lora_rank=96, + kv_lora_rank=64, + qk_nope_head_dim=64, + qk_rope_head_dim=32, + v_head_dim=64, +) + + +def _device_sm() -> int: + major, minor = torch.cuda.get_device_capability() + return major * 10 + minor + + +def _recv_capacity(ep_size: int) -> int: + cap = ep_size * TOKENS_PER_RANK * TOP_K + NUM_LOCAL_EXPERTS * 128 + return -(-cap // 128) * 128 + + +def _broadcast_params(module: torch.nn.Module) -> None: + for t in list(module.parameters()) + list(module.buffers()): + dist.broadcast(t.detach(), src=0) + + +class TestDeepSeekEP(unittest.TestCase): + @classmethod + def setUpClass(cls): + if _device_sm() < 90: + raise unittest.SkipTest(f"NCCL EP requires SM>=90 (got SM{_device_sm()})") + cls.rank = dist.get_rank() + cls.ep_size = dist.get_world_size() + cls.num_experts = NUM_LOCAL_EXPERTS * cls.ep_size + world_pg = dist.distributed_c10d._get_default_group() + cls.ep_group = dist.new_group(ranks=list(range(world_pg.size())), backend="nccl") + ep_bootstrap( + cls.ep_group, + num_experts=cls.num_experts, + max_tokens_per_rank=TOKENS_PER_RANK, + hidden_dim=HIDDEN, + num_topk=TOP_K, + recv_capacity_per_rank=_recv_capacity(cls.ep_size), + ) + + def _make_moe(self, ep: bool, shared: bool = True) -> DeepSeekV3MoE: + return DeepSeekV3MoE( + HIDDEN, + moe_ffn_hidden_size=MOE_FFN, + num_experts=self.num_experts, + topk=TOP_K, + shared_expert_ffn_hidden_size=SHARED_FFN if shared else None, + params_dtype=DTYPE, + ep_group=self.ep_group if ep else None, + ep_max_tokens_per_rank=TOKENS_PER_RANK if ep else None, + ep_recv_capacity_per_rank=_recv_capacity(self.ep_size) if ep else None, + ) + + def _copy_local_expert_weights(self, ep_moe: DeepSeekV3MoE, ref: DeepSeekV3MoE) -> None: + with torch.no_grad(): + ep_moe.gate.weight.copy_(ref.gate.weight) + if ref.shared_expert is not None: + for dst, src in zip( + ep_moe.shared_expert.parameters(), ref.shared_expert.parameters() + ): + dst.copy_(src) + ep_fc1, _, ep_fc2 = ep_moe.experts + ref_fc1, _, ref_fc2 = ref.experts + for local_e in range(NUM_LOCAL_EXPERTS): + global_e = self.rank * NUM_LOCAL_EXPERTS + local_e + getattr(ep_fc1, f"weight{local_e}").copy_(getattr(ref_fc1, f"weight{global_e}")) + getattr(ep_fc2, f"weight{local_e}").copy_(getattr(ref_fc2, f"weight{global_e}")) + + def test_moe_ep_matches_local(self): + """EP MoE must match the single-GPU (all-experts-local) path numerically.""" + torch.manual_seed(0) + ref = self._make_moe(ep=False) + _broadcast_params(ref) + ep_moe = self._make_moe(ep=True) + self._copy_local_expert_weights(ep_moe, ref) + + torch.manual_seed(1234 + self.rank) + x = torch.randn(TOKENS_PER_RANK, HIDDEN, dtype=DTYPE, device="cuda") + x_ep = x.clone().requires_grad_(True) + x_ref = x.clone().requires_grad_(True) + + out_ep = ep_moe(x_ep) + out_ref = ref(x_ref) + torch.testing.assert_close(out_ep, out_ref, rtol=0.05, atol=0.05) + + grad_out = torch.randn_like(out_ep) + out_ep.backward(grad_out) + out_ref.backward(grad_out) + torch.testing.assert_close(x_ep.grad, x_ref.grad, rtol=0.05, atol=0.05) + torch.testing.assert_close( + ep_moe.gate.weight.grad, ref.gate.weight.grad, rtol=0.1, atol=0.1 + ) + + # A local expert's wgrad on its owner rank equals the sum of the + # reference wgrads over all ranks. + ep_fc1, _, ep_fc2 = ep_moe.experts + ref_fc1, _, ref_fc2 = ref.experts + for ep_fc, ref_fc in ((ep_fc1, ref_fc1), (ep_fc2, ref_fc2)): + for local_e in range(NUM_LOCAL_EXPERTS): + global_e = self.rank * NUM_LOCAL_EXPERTS + local_e + ref_grad = getattr(ref_fc, f"weight{global_e}").grad.float() + dist.all_reduce(ref_grad) + ep_grad = getattr(ep_fc, f"weight{local_e}").grad.float() + torch.testing.assert_close(ep_grad, ref_grad, rtol=0.1, atol=0.1) + + counts = ep_moe._last_tokens_per_expert.clone() + dist.all_reduce(counts) + self.assertEqual(counts.sum().item(), self.ep_size * TOKENS_PER_RANK * TOP_K) + + def test_layer_ep_forward_backward(self): + """Full DeepSeekV3Layer smoke test with an EP MoE block.""" + torch.manual_seed(10 + self.rank) + layer = DeepSeekV3Layer( + HIDDEN, + HEADS, + num_experts=self.num_experts, + moe_ffn_hidden_size=MOE_FFN, + topk=TOP_K, + shared_expert_ffn_hidden_size=SHARED_FFN, + params_dtype=DTYPE, + ep_group=self.ep_group, + ep_max_tokens_per_rank=TOKENS_PER_RANK, + ep_recv_capacity_per_rank=_recv_capacity(self.ep_size), + **MLA_KWARGS, + ) + x = torch.randn( + TOKENS_PER_RANK // 2, 2, HIDDEN, dtype=DTYPE, device="cuda", requires_grad=True + ) + out = layer(x) + self.assertEqual(out.shape, x.shape) + out.sum().backward() + self.assertIsNotNone(x.grad) + self.assertTrue(torch.isfinite(x.grad).all()) + + layer.mlp.update_expert_bias() + self.assertTrue(torch.isfinite(layer.mlp.expert_bias).all()) + + +def _init_distributed(): + dist.init_process_group(backend="nccl") + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + try: + from torch.distributed import _symmetric_memory as _symm_mem + + _symm_mem.set_backend("NCCL") + except (ImportError, RuntimeError): + pass + + +if __name__ == "__main__": + _init_distributed() + suite = unittest.TestLoader().loadTestsFromTestCase(TestDeepSeekEP) + result = unittest.TextTestRunner(stream=sys.stdout, verbosity=2).run(suite) + dist.barrier() + ep_finalize() + release_symm_mem_pool() + dist.destroy_process_group() + sys.exit(0 if result.wasSuccessful() else 1) diff --git a/tests/pytorch/distributed/run_test_deepseek_ep.sh b/tests/pytorch/distributed/run_test_deepseek_ep.sh new file mode 100644 index 0000000000..8c0bbbc5b9 --- /dev/null +++ b/tests/pytorch/distributed/run_test_deepseek_ep.sh @@ -0,0 +1,52 @@ +#!/bin/bash +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +# +# Launcher for tests/pytorch/distributed/run_deepseek_ep.py. Auto-detects GPU count. + +set -uo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +if [ "${DETECTED_GPUS}" -lt 2 ]; then + echo "DeepSeek EP test requires >= 2 GPUs (found ${DETECTED_GPUS}); SKIPPING." + exit 0 +fi + +# NCCL EP requires active NVLink P2P among ranks on the node. +if ! nvidia-smi nvlink --status 2>/dev/null | grep -qE 'Link [0-9]+:.*GB/s'; then + echo "No NVLink between GPUs (PCIe-only fabric); NCCL EP is unsupported here. SKIPPING." + exit 0 +fi + +NUM_RANKS="${NVTE_TEST_EP_NUM_RANKS:-${DETECTED_GPUS}}" +if [ "${NUM_RANKS}" -gt 8 ]; then NUM_RANKS=8; fi + +TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-180}" + +: ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"} +export NCCL_EP_JIT_CACHE_DIR +mkdir -p "$NCCL_EP_JIT_CACHE_DIR" + +SCRIPT="${SCRIPT_DIR}/run_deepseek_ep.py" +LOG="stdout_deepseek_ep.txt" + +echo "=== Running ${SCRIPT} on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ===" +setsid timeout --foreground --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \ + torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" \ + "${SCRIPT}" 2>&1 | tee "${LOG}" +RC=${PIPESTATUS[0]} +pkill -9 -f "tests/pytorch/distributed/run_deepseek_ep.py" 2>/dev/null || true + +RET=0 +if [ "${RC}" -ne 0 ]; then echo "torchrun exited with ${RC}"; RET=1; fi +if grep -qE "(^|]:)FAILED|(^|]:)Traceback" "${LOG}"; then RET=1; fi +if ! grep -qE "Ran [0-9]+ test|^OK$" "${LOG}"; then + echo "ERROR: no test summary — likely hang or early crash" + RET=1 +fi +if [ -z "${KEEP_EP_LOGS:-}" ]; then rm -f "${LOG}"; fi + +exit $RET diff --git a/tests/pytorch/distributed/test_deepseek_ep.py b/tests/pytorch/distributed/test_deepseek_ep.py new file mode 100644 index 0000000000..4a4d9a8dea --- /dev/null +++ b/tests/pytorch/distributed/test_deepseek_ep.py @@ -0,0 +1,26 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Pytest driver — spawns run_deepseek_ep.py under torchrun and asserts it passed.""" + +import os +import subprocess +from pathlib import Path + +import pytest +import torch + +TEST_ROOT = Path(__file__).parent.resolve() +LAUNCHER = TEST_ROOT / "run_test_deepseek_ep.sh" + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="DeepSeek EP requires >= 2 GPUs") +def test_multi_process_deepseek_ep(): + timeout_s = int(os.environ.get("NVTE_TEST_EP_TIMEOUT_S", "180")) + proc = subprocess.run( + ["bash", str(LAUNCHER)], + env={**os.environ, "KEEP_EP_LOGS": "1", "TEST_TIMEOUT_S": str(timeout_s)}, + timeout=timeout_s + 30, + check=False, + ) + assert proc.returncode == 0, f"DeepSeek EP test suite failed (rc={proc.returncode})" diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index 5a1c8d650c..f413221bd1 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -155,7 +155,11 @@ def __init__( assert ep_max_tokens_per_rank is not None, "EP requires ep_max_tokens_per_rank." if ep_recv_capacity_per_rank is None: - ep_recv_capacity_per_rank = self.ep_size * ep_max_tokens_per_rank * topk + # Worst case plus per-expert alignment padding, rounded up to + # the multiple of 128 required by the fused grouped MLP. + cap = self.ep_size * ep_max_tokens_per_rank * topk + cap += num_local_experts * max(ep_alignment, 1) + ep_recv_capacity_per_rank = -(-cap // 128) * 128 self.ep_buffer = EpBuffer( top_k=topk, max_tokens_per_rank=ep_max_tokens_per_rank, From 4c6e1e8aff62def1cd0bd72ce9bfb960a4aab035 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 18 Aug 2026 15:49:08 +0200 Subject: [PATCH 06/41] Fix EP wgrad test collective + zero EP recv/grad buffers The per-expert wgrad check called all_reduce on different tensors per rank (rank-local experts), corrupting the reference grads; reduce every expert's grad on every rank instead. Also pass zero-filled recv/grad buffers to ep_dispatch/ep_combine so alignment-padding rows inside the grouped-GEMM m_splits can never poison expert wgrads. Verified on lyris (4x GB300, arm64): run_test_deepseek_ep.sh passes on all ranks (EP forward/dgrad/gate-grad/expert-wgrad match the all-local reference; full-layer EP smoke passes). Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_deepseek_ep.py | 12 +++++++---- .../pytorch/models/deepseek_v3/moe.py | 20 +++++++++++++++++-- 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/tests/pytorch/distributed/run_deepseek_ep.py b/tests/pytorch/distributed/run_deepseek_ep.py index 0edae05961..bf756b69ad 100644 --- a/tests/pytorch/distributed/run_deepseek_ep.py +++ b/tests/pytorch/distributed/run_deepseek_ep.py @@ -119,16 +119,20 @@ def test_moe_ep_matches_local(self): ) # A local expert's wgrad on its owner rank equals the sum of the - # reference wgrads over all ranks. + # reference wgrads over all ranks. all_reduce is collective, so every + # rank must reduce every expert's grad (in the same order). ep_fc1, _, ep_fc2 = ep_moe.experts ref_fc1, _, ref_fc2 = ref.experts for ep_fc, ref_fc in ((ep_fc1, ref_fc1), (ep_fc2, ref_fc2)): + ref_grads = [ + getattr(ref_fc, f"weight{e}").grad.float().clone() for e in range(self.num_experts) + ] + for g in ref_grads: + dist.all_reduce(g) for local_e in range(NUM_LOCAL_EXPERTS): global_e = self.rank * NUM_LOCAL_EXPERTS + local_e - ref_grad = getattr(ref_fc, f"weight{global_e}").grad.float() - dist.all_reduce(ref_grad) ep_grad = getattr(ep_fc, f"weight{local_e}").grad.float() - torch.testing.assert_close(ep_grad, ref_grad, rtol=0.1, atol=0.1) + torch.testing.assert_close(ep_grad, ref_grads[global_e], rtol=0.1, atol=0.1) counts = ep_moe._last_tokens_per_expert.clone() dist.all_reduce(counts) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index f413221bd1..3c182b4405 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -218,13 +218,29 @@ def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: ) topk_weights = probs.gather(1, topk_idx).float() + # Zero-filled recv/grad buffers: per-expert alignment padding lands + # inside the grouped-GEMM m_splits, so uninitialized rows would poison + # the expert wgrads. + cap = self.ep_buffer.recv_capacity_per_rank recv_tokens, recv_weights, tokens_per_expert = ep_dispatch( - self.ep_buffer, tokens, topk_idx, topk_weights + self.ep_buffer, + tokens, + topk_idx, + topk_weights, + recv_tokens=torch.zeros( + (cap, self.hidden_size), dtype=tokens.dtype, device=tokens.device + ), + recv_topk_weights=torch.zeros((cap,), dtype=torch.float32, device=tokens.device), ) expert_out = self.experts( recv_tokens, tokens_per_expert, recv_weights.to(tokens.dtype), tokens_per_expert ) - return ep_combine(self.ep_buffer, expert_out, num_local_tokens=tokens.shape[0]) + return ep_combine( + self.ep_buffer, + expert_out, + num_local_tokens=tokens.shape[0], + grad_out=torch.zeros_like(expert_out), + ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """ From aa17c37fb9a0e8cd74c3b5d67a5d34365f4133d7 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 18 Aug 2026 17:23:10 +0200 Subject: [PATCH 07/41] Use fused MLA RoPE kernels in MultiLatentAttention Move the Triton MLA RoPE kernels (Megatron-LM fused_mla_yarn_rope_apply port) from tests/pytorch/attention/ mla_rope_utils.py into models/deepseek_v3/mla_rope.py and use them in MultiLatentAttention: the q kernel rotates the rope slice in place and the kv kernel assembles key/value in a single pass, removing the torch.cat/expand/contiguous copies (~10% of layer GPU time). PyTorch fallback (same convention) covers missing Triton and bshd. Fix a latent bug from the test util: the q backward kernel assumed a contiguous incoming gradient, but cuDNN attention backward can hand over a strided one (allocator-state dependent IMA). The old test file stays as a compat shim. Add a Triton-vs-PyTorch parity test. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/mla_rope_utils.py | 652 +----------------- tests/pytorch/test_deepseek.py | 60 ++ .../pytorch/models/deepseek_v3/mla_rope.py | 495 +++++++++++++ .../deepseek_v3/multi_latent_attention.py | 55 +- 4 files changed, 601 insertions(+), 661 deletions(-) create mode 100644 transformer_engine/pytorch/models/deepseek_v3/mla_rope.py diff --git a/tests/pytorch/attention/mla_rope_utils.py b/tests/pytorch/attention/mla_rope_utils.py index 90eebfc66a..d022757886 100644 --- a/tests/pytorch/attention/mla_rope_utils.py +++ b/tests/pytorch/attention/mla_rope_utils.py @@ -2,26 +2,17 @@ # # See LICENSE for license information. -"""MLA RoPE for DSv3 671B - Triton forward and backward kernels. - -Source: Megatron-LM megatron/core/fusions/fused_mla_yarn_rope_apply.py -Falls back to pure PyTorch when Triton is unavailable. - -Note: DSv3 uses YaRN-scaled RoPE for long-context extrapolation. This test -intentionally uses plain RoPE (base=10000) because it only validates MXFP8 -attention path wiring, tensor shapes, forward/backward flow, and relative BF16 -vs MXFP8 behavior. Both reference and MXFP8 paths use the same RoPE tables. -""" +"""Compat shim: the MLA RoPE kernels moved to +``transformer_engine.pytorch.models.deepseek_v3.mla_rope``.""" import torch -try: - import triton - import triton.language as tl - - HAVE_TRITON = True -except ImportError: - HAVE_TRITON = False +from transformer_engine.pytorch.models.deepseek_v3.mla_rope import ( # noqa: F401 + HAVE_TRITON, + apply_mla_rope_kv, + apply_mla_rope_q, + build_rope_tables, +) HEAD_DIM_ROPE = 64 HEAD_DIM_NOPE = 128 @@ -29,576 +20,6 @@ ROTARY_BASE = 10000 -def build_rope_tables( - seq_len: int, - emb_dim: int = HEAD_DIM_ROPE, - base: int = ROTARY_BASE, - device: torch.device = None, -) -> tuple[torch.Tensor, torch.Tensor]: - inv_freq = 1.0 / ( - base ** (torch.arange(0, emb_dim, 2, dtype=torch.float32, device=device) / emb_dim) - ) - t = torch.arange(seq_len, device=device, dtype=torch.float32) - freqs = torch.outer(t, inv_freq) - freqs = torch.cat([freqs, freqs], dim=-1) - return torch.cos(freqs).contiguous(), torch.sin(freqs).contiguous() - - -if HAVE_TRITON: - - # Not used for non-packed batches; kept for THD compatibility. - @triton.jit - def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): - token_idx = -1 - this_seq_len = 0 - seq_idx = 0 - last_cum_seqlen = tl.load(cu_seqlens) // cp_size - while seq_idx < seq_num: - cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1) // cp_size - if token_idx == -1 and cur_cum_seqlen > pid_m: - token_idx = pid_m - last_cum_seqlen - this_seq_len = cur_cum_seqlen - last_cum_seqlen - last_cum_seqlen = cur_cum_seqlen - seq_idx += 1 - if cp_size > 1: - if token_idx < this_seq_len // 2: - token_idx = token_idx + cp_rank * this_seq_len // 2 - else: - token_idx = (token_idx - this_seq_len // 2) + ( - 2 * cp_size - cp_rank - 1 - ) * this_seq_len // 2 - return token_idx - - @triton.autotune( - configs=[ - triton.Config({"BLOCK_H": 1}), - triton.Config({"BLOCK_H": 2}), - triton.Config({"BLOCK_H": 4}), - triton.Config({"BLOCK_H": 8}), - triton.Config({"BLOCK_H": 16}), - triton.Config({"BLOCK_H": 32}), - triton.Config({"BLOCK_H": 64}), - triton.Config({"BLOCK_H": 128}), - ], - key=["emb_dim", "head_num"], - restore_value=["Q"], - ) - @triton.jit - def rotary_fwd_q_kernel( - Q, - COS, - SIN, - qk_head_dim, - emb_dim: tl.constexpr, - head_num: tl.constexpr, - batch_size, - seq_num, - cu_seqlens_q, - stride_x_seq, - stride_x_nheads, - cp_rank, - cp_size, - BLOCK_H: tl.constexpr, - ): - pid_m = tl.program_id(axis=0) - pid_head = tl.program_id(axis=1) - if cu_seqlens_q is None: - token_idx = pid_m // batch_size - else: - token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - sin_right = sin_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - Q = Q + pid_m * stride_x_seq - x_off = head_offsets[:, None] * stride_x_nheads + qk_head_dim - mask = head_offsets[:, None] < head_num - x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 - x_2_off = x_1_off + 1 - x_1 = tl.load(Q + x_1_off, mask=mask) - x_2 = tl.load(Q + x_2_off, mask=mask) - x_left = x_1 * cos_left - x_2 * sin_left - x_right = x_2 * cos_right + x_1 * sin_right - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - tl.store(Q + x_left_off, x_left, mask=mask) - tl.store(Q + x_right_off, x_right, mask=mask) - - @triton.autotune( - configs=[ - triton.Config({"BLOCK_H": 1}), - triton.Config({"BLOCK_H": 2}), - triton.Config({"BLOCK_H": 4}), - triton.Config({"BLOCK_H": 8}), - triton.Config({"BLOCK_H": 16}), - triton.Config({"BLOCK_H": 32}), - triton.Config({"BLOCK_H": 64}), - triton.Config({"BLOCK_H": 128}), - ], - key=["emb_dim", "head_num"], - restore_value=["DO"], - ) - @triton.jit - def rotary_bwd_q_kernel( - DO, - COS, - SIN, - qk_head_dim, - emb_dim: tl.constexpr, - head_num: tl.constexpr, - batch_size, - seq_num, - cu_seqlens_q, - stride_x_seq, - stride_x_nheads, - cp_rank, - cp_size, - BLOCK_H: tl.constexpr, - ): - pid_m = tl.program_id(axis=0) - pid_head = tl.program_id(axis=1) - if cu_seqlens_q is None: - token_idx = pid_m // batch_size - else: - token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - sin_right = sin_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - DO = DO + pid_m * stride_x_seq - x_off = head_offsets[:, None] * stride_x_nheads + qk_head_dim - mask = head_offsets[:, None] < head_num - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - x_left = tl.load(DO + x_left_off, mask=mask) - x_right = tl.load(DO + x_right_off, mask=mask) - x_1 = x_left * cos_left + x_right * sin_right - x_2 = -x_left * sin_left + x_right * cos_right - x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 - x_2_off = x_1_off + 1 - tl.store(DO + x_1_off, x_1, mask=mask) - tl.store(DO + x_2_off, x_2, mask=mask) - - @triton.autotune( - configs=[ - triton.Config({"BLOCK_H": 1}), - triton.Config({"BLOCK_H": 2}), - triton.Config({"BLOCK_H": 4}), - triton.Config({"BLOCK_H": 8}), - triton.Config({"BLOCK_H": 16}), - triton.Config({"BLOCK_H": 32}), - triton.Config({"BLOCK_H": 64}), - triton.Config({"BLOCK_H": 128}), - ], - key=["emb_dim", "k_dim", "v_dim", "head_num"], - ) - @triton.jit - def rotary_fwd_kv_kernel( - KV, - K_POS_EMB, - O_KEY, - O_VALUE, - COS, - SIN, - emb_dim: tl.constexpr, - k_dim: tl.constexpr, - v_dim: tl.constexpr, - head_num: tl.constexpr, - batch_size, - seq_num, - cu_seqlens_kv, - stride_kv_seq, - stride_kv_nheads, - stride_emb_seq, - stride_k_seq, - stride_k_nheads, - stride_v_seq, - stride_v_nheads, - cp_rank, - cp_size, - BLOCK_H: tl.constexpr, - ): - pid_m = tl.program_id(axis=0) - pid_head = tl.program_id(axis=1) - if cu_seqlens_kv is None: - token_idx = pid_m // batch_size - else: - token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - KV_ptr = KV + pid_m * stride_kv_seq - kv_off = head_offsets[:, None] * stride_kv_nheads - mask = head_offsets[:, None] < head_num - k_in_off = kv_off + tl.arange(0, k_dim)[None, :] - v_in_off = kv_off + k_dim + tl.arange(0, v_dim)[None, :] - k = tl.load(KV_ptr + k_in_off, mask=mask) - v = tl.load(KV_ptr + v_in_off, mask=mask) - K_ptr = O_KEY + pid_m * stride_k_seq + pid_head * BLOCK_H * stride_k_nheads - V_ptr = O_VALUE + pid_m * stride_v_seq + pid_head * BLOCK_H * stride_v_nheads - k_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + tl.arange(0, k_dim)[None, :] - v_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_v_nheads + tl.arange(0, v_dim)[None, :] - tl.store(K_ptr + k_out_off, k, mask=mask) - tl.store(V_ptr + v_out_off, v, mask=mask) - EMB = K_POS_EMB + pid_m * stride_emb_seq - x_1 = tl.load(EMB + tl.arange(0, emb_dim // 2) * 2) - x_2 = tl.load(EMB + tl.arange(0, emb_dim // 2) * 2 + 1) - x_left = x_1 * cos_left - x_2 * sin_left - x_right = x_2 * cos_right + x_1 * sin_right - x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - x_left_off = ( - tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads - + k_dim - + tl.arange(0, emb_dim // 2)[None, :] - ) - x_right_off = x_left_off + emb_dim // 2 - tl.store(K_ptr + x_left_off, x_left, mask=mask) - tl.store(K_ptr + x_right_off, x_right, mask=mask) - - @triton.autotune( - configs=[ - triton.Config({"BLOCK_H": 1}), - triton.Config({"BLOCK_H": 2}), - triton.Config({"BLOCK_H": 4}), - triton.Config({"BLOCK_H": 8}), - triton.Config({"BLOCK_H": 16}), - triton.Config({"BLOCK_H": 32}), - triton.Config({"BLOCK_H": 64}), - triton.Config({"BLOCK_H": 128}), - ], - key=["emb_dim", "k_dim", "v_dim", "head_num"], - ) - @triton.jit - def rotary_bwd_kv_kernel( - dK, - dV, - dKV, - dEMB, - COS, - SIN, - emb_dim: tl.constexpr, - k_dim: tl.constexpr, - v_dim: tl.constexpr, - head_num: tl.constexpr, - batch_size, - seq_num, - cu_seqlens_kv, - stride_dk_seq, - stride_dk_nheads, - stride_dv_seq, - stride_dv_nheads, - stride_dkv_seq, - stride_dkv_nheads, - stride_demb_seq, - cp_rank, - cp_size, - BLOCK_H: tl.constexpr, - ): - pid_m = tl.program_id(axis=0) - pid_head = tl.program_id(axis=1) - if cu_seqlens_kv is None: - token_idx = pid_m // batch_size - else: - token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size) - head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - dKV_ptr = dKV + pid_m * stride_dkv_seq - dkv_off = head_offsets[:, None] * stride_dkv_nheads - mask = head_offsets[:, None] < head_num - dk_out_off = dkv_off + tl.arange(0, k_dim)[None, :] - dv_out_off = dkv_off + k_dim + tl.arange(0, v_dim)[None, :] - dK_ptr = dK + pid_m * stride_dk_seq + pid_head * BLOCK_H * stride_dk_nheads - dV_ptr = dV + pid_m * stride_dv_seq + pid_head * BLOCK_H * stride_dv_nheads - dk_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + tl.arange(0, k_dim)[None, :] - dv_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dv_nheads + tl.arange(0, v_dim)[None, :] - dk = tl.load(dK_ptr + dk_in_off, mask=mask) - dv = tl.load(dV_ptr + dv_in_off, mask=mask) - tl.store(dKV_ptr + dk_out_off, dk, mask=mask) - tl.store(dKV_ptr + dv_out_off, dv, mask=mask) - if pid_head == 0: - x_left_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) - x_right_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) - for i in tl.static_range(triton.cdiv(head_num, BLOCK_H)): - head_offsets_i = i * BLOCK_H + tl.arange(0, BLOCK_H) - dK_ptr_i = dK + pid_m * stride_dk_seq - x_off = head_offsets_i[:, None] * stride_dk_nheads + k_dim - mask_i = head_offsets_i[:, None] < head_num - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - x_left_accum += tl.load(dK_ptr_i + x_left_off, mask=mask_i) - x_right_accum += tl.load(dK_ptr_i + x_right_off, mask=mask_i) - x_left_accum = tl.sum(x_left_accum, axis=0) - x_right_accum = tl.sum(x_right_accum, axis=0) - x_left_accum = x_left_accum.to(dEMB.dtype.element_ty) - x_right_accum = x_right_accum.to(dEMB.dtype.element_ty) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load( - COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2) - ) - sin_right = tl.load( - SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2) - ) - x_1 = x_left_accum * cos_left + x_right_accum * sin_right - x_2 = -x_left_accum * sin_left + x_right_accum * cos_right - dEMB_ptr = dEMB + pid_m * stride_demb_seq - tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2, x_1) - tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2 + 1, x_2) - - def _flattened_token_stride(tensor: torch.Tensor) -> int: - if tensor.dim() == 4: - return tensor.stride(1) - return tensor.stride(0) - - class _MLARoPEQTriton(torch.autograd.Function): - @staticmethod - def forward(ctx, q, cos, sin, head_dim_nope, head_dim_rope): - s, b, nheads, _ = q.shape - total = s * b - - grid_q = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_fwd_q_kernel[grid_q]( - q, - cos, - sin, - head_dim_nope, - head_dim_rope, - nheads, - b, - None, - None, - _flattened_token_stride(q), - q.stride(2), - 0, - 1, - ) - - ctx.save_for_backward(cos, sin) - ctx.head_dim_nope = head_dim_nope - ctx.head_dim_rope = head_dim_rope - ctx.nheads = nheads - ctx.s = s - ctx.b = b - return q - - @staticmethod - def backward(ctx, dq): - cos, sin = ctx.saved_tensors - s, b, nheads = ctx.s, ctx.b, ctx.nheads - total = s * b - - grid_q = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_bwd_q_kernel[grid_q]( - dq, - cos, - sin, - ctx.head_dim_nope, - ctx.head_dim_rope, - nheads, - b, - None, - None, - _flattened_token_stride(dq), - dq.stride(2), - 0, - 1, - ) - return dq, None, None, None, None - - class _MLARoPEKVTriton(torch.autograd.Function): - @staticmethod - def forward(ctx, kv, k_pos_emb, cos, sin, head_dim_nope, head_dim_rope, head_dim_v): - s, b, nheads, _ = kv.shape - total = s * b - - o_key = kv.new_empty(s, b, nheads, head_dim_nope + head_dim_rope) - o_value = kv.new_empty(s, b, nheads, head_dim_v) - grid_kv = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_fwd_kv_kernel[grid_kv]( - kv, - k_pos_emb, - o_key, - o_value, - cos, - sin, - head_dim_rope, - head_dim_nope, - head_dim_v, - nheads, - b, - None, - None, - _flattened_token_stride(kv), - kv.stride(2), - _flattened_token_stride(k_pos_emb), - _flattened_token_stride(o_key), - o_key.stride(2), - _flattened_token_stride(o_value), - o_value.stride(2), - 0, - 1, - ) - - ctx.save_for_backward(cos, sin) - ctx.head_dim_nope = head_dim_nope - ctx.head_dim_rope = head_dim_rope - ctx.head_dim_v = head_dim_v - ctx.nheads = nheads - ctx.s = s - ctx.b = b - return o_key, o_value - - @staticmethod - def backward(ctx, dk_out, dv_out): - cos, sin = ctx.saved_tensors - s, b, nheads = ctx.s, ctx.b, ctx.nheads - ndp, ndr, ndv = ctx.head_dim_nope, ctx.head_dim_rope, ctx.head_dim_v - total = s * b - - d_kv = dk_out.new_empty(s, b, nheads, ndp + ndv) - d_emb = dk_out.new_empty(s, b, 1, ndr) - grid_kv = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_bwd_kv_kernel[grid_kv]( - dk_out, - dv_out, - d_kv, - d_emb, - cos, - sin, - ndr, - ndp, - ndv, - nheads, - b, - None, - None, - _flattened_token_stride(dk_out), - dk_out.stride(2), - _flattened_token_stride(dv_out), - dv_out.stride(2), - _flattened_token_stride(d_kv), - d_kv.stride(2), - _flattened_token_stride(d_emb), - 0, - 1, - ) - return d_kv, d_emb, None, None, None, None, None - - -def _apply_mla_rope_q_with_tables( - q: torch.Tensor, - cos_table: torch.Tensor, - sin_table: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, -) -> torch.Tensor: - if HAVE_TRITON: - return _MLARoPEQTriton.apply( - q, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - ) - return _apply_pytorch_q(q, cos_table, sin_table, head_dim_nope, head_dim_rope) - - -def _apply_mla_rope_kv_with_tables( - kv: torch.Tensor, - k_pos_emb: torch.Tensor, - cos_table: torch.Tensor, - sin_table: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, - head_dim_v: int = HEAD_DIM_V, -) -> tuple[torch.Tensor, torch.Tensor]: - if HAVE_TRITON: - return _MLARoPEKVTriton.apply( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, - ) - return _apply_pytorch_kv( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, - ) - - -def apply_mla_rope_q( - q: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, - base: int = ROTARY_BASE, - cos_table: torch.Tensor | None = None, - sin_table: torch.Tensor | None = None, -) -> torch.Tensor: - if cos_table is None or sin_table is None: - s = q.shape[0] - cos_table, sin_table = build_rope_tables( - s, - emb_dim=head_dim_rope, - base=base, - device=q.device, - ) - return _apply_mla_rope_q_with_tables( - q, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - ) - - -def apply_mla_rope_kv( - kv: torch.Tensor, - k_pos_emb: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, - head_dim_v: int = HEAD_DIM_V, - base: int = ROTARY_BASE, - cos_table: torch.Tensor | None = None, - sin_table: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - if cos_table is None or sin_table is None: - s = kv.shape[0] - cos_table, sin_table = build_rope_tables( - s, - emb_dim=head_dim_rope, - base=base, - device=kv.device, - ) - return _apply_mla_rope_kv_with_tables( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, - ) - - def apply_mla_rope( q: torch.Tensor, kv: torch.Tensor, @@ -609,60 +30,13 @@ def apply_mla_rope( base: int = ROTARY_BASE, cos_table: torch.Tensor | None = None, sin_table: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +): if cos_table is None or sin_table is None: - s = q.shape[0] cos_table, sin_table = build_rope_tables( - s, - emb_dim=head_dim_rope, - base=base, - device=q.device, + q.shape[0], head_dim_rope, base=base, device=q.device ) - q = _apply_mla_rope_q_with_tables(q, cos_table, sin_table, head_dim_nope, head_dim_rope) - k, v = _apply_mla_rope_kv_with_tables( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, + q = apply_mla_rope_q(q, cos_table, sin_table, head_dim_nope, head_dim_rope) + k, v = apply_mla_rope_kv( + kv, k_pos_emb, cos_table, sin_table, head_dim_nope, head_dim_rope, head_dim_v ) return q, k, v - - -def _rotate_interleaved_to_neox( - x: torch.Tensor, cos_table: torch.Tensor, sin_table: torch.Tensor -) -> torch.Tensor: - cos_ = cos_table[:, None, None, :].to(x.dtype) - sin_ = sin_table[:, None, None, :].to(x.dtype) - half_dim = x.shape[-1] // 2 - x_1 = x[..., 0::2] - x_2 = x[..., 1::2] - x_left = x_1 * cos_[..., :half_dim] - x_2 * sin_[..., :half_dim] - x_right = x_2 * cos_[..., half_dim:] + x_1 * sin_[..., half_dim:] - return torch.cat((x_left, x_right), dim=-1) - - -def _apply_pytorch_q(q, cos_table, sin_table, head_dim_nope, head_dim_rope): - q_nope = q[..., :head_dim_nope] - q_rope = q[..., head_dim_nope : head_dim_nope + head_dim_rope] - q_rope = _rotate_interleaved_to_neox(q_rope, cos_table, sin_table) - return torch.cat((q_nope, q_rope), dim=-1) - - -def _apply_pytorch_kv( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, -): - k_nope = kv[..., :head_dim_nope] - v = kv[..., head_dim_nope : head_dim_nope + head_dim_v] - k_rope = _rotate_interleaved_to_neox(k_pos_emb, cos_table, sin_table).expand( - -1, -1, kv.shape[2], -1 - ) - return torch.cat((k_nope, k_rope), dim=-1), v diff --git a/tests/pytorch/test_deepseek.py b/tests/pytorch/test_deepseek.py index 7778d0448c..4c4aea0a92 100644 --- a/tests/pytorch/test_deepseek.py +++ b/tests/pytorch/test_deepseek.py @@ -34,6 +34,66 @@ def _input(requires_grad=True): ) +def test_mla_rope_triton_matches_pytorch(): + from transformer_engine.pytorch.models.deepseek_v3 import mla_rope + + if not mla_rope.HAVE_TRITON: + pytest.skip("Triton unavailable") + s, b, h = 64, 2, 4 + nope, rope, vdim = 64, 32, 64 + cos, sin = mla_rope.build_rope_tables(s, rope, device="cuda") + + torch.manual_seed(0) + q_leaf = torch.randn(s, b, h, nope + rope, device="cuda", requires_grad=True) + kv_leaf = torch.randn(s, b, h, nope + vdim, device="cuda", requires_grad=True) + pos_leaf = torch.randn(s, b, 1, rope, device="cuda", requires_grad=True) + grad_q = torch.randn(s, b, h, nope + rope, device="cuda") + grad_k = torch.randn(s, b, h, nope + rope, device="cuda") + grad_v = torch.randn(s, b, h, vdim, device="cuda") + + def run(fmt): + # non-leaf copies: the Triton q kernel rotates in place + q, kv, pos = q_leaf * 1.0, kv_leaf * 1.0, pos_leaf * 1.0 + q_out = mla_rope.apply_mla_rope_q(q, cos, sin, nope, rope, fmt) + k_out, v_out = mla_rope.apply_mla_rope_kv(kv, pos, cos, sin, nope, rope, vdim, fmt) + # fresh grad clones: the Triton q backward modifies its input grad in place + torch.autograd.backward( + [q_out, k_out, v_out], [grad_q.clone(), grad_k.clone(), grad_v.clone()] + ) + grads = (q_leaf.grad.clone(), kv_leaf.grad.clone(), pos_leaf.grad.clone()) + q_leaf.grad = kv_leaf.grad = pos_leaf.grad = None + return (q_out.clone(), k_out, v_out), grads + + (q_t, k_t, v_t), grads_t = run("sbhd") + + seq_dim = 0 + q_ref = torch.cat( + ( + (q_leaf * 1.0)[..., :nope], + mla_rope._rotate_interleaved_to_neox((q_leaf * 1.0)[..., nope:], cos, sin, seq_dim), + ), + dim=-1, + ) + k_ref = torch.cat( + ( + (kv_leaf * 1.0)[..., :nope], + mla_rope._rotate_interleaved_to_neox(pos_leaf * 1.0, cos, sin, seq_dim).expand( + s, b, h, rope + ), + ), + dim=-1, + ) + v_ref = (kv_leaf * 1.0)[..., nope:] + torch.autograd.backward([q_ref, k_ref, v_ref], [grad_q.clone(), grad_k.clone(), grad_v.clone()]) + + torch.testing.assert_close(q_t, q_ref, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(k_t, k_ref, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(v_t, v_ref, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(grads_t[0], q_leaf.grad, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(grads_t[1], kv_leaf.grad, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(grads_t[2], pos_leaf.grad, rtol=1e-5, atol=1e-5) + + def test_mla_forward_backward(): torch.manual_seed(0) mla = MultiLatentAttention(HIDDEN, HEADS, params_dtype=DTYPE, **MLA_KWARGS) diff --git a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py new file mode 100644 index 0000000000..350bedb69b --- /dev/null +++ b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py @@ -0,0 +1,495 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused MLA RoPE kernels (DeepSeekV3-style decoupled RoPE/NoPE). + +Triton forward/backward kernels adapted from Megatron-LM +``megatron/core/fusions/fused_mla_yarn_rope_apply.py``. The query kernel +rotates the trailing ``head_dim_rope`` slice in place (no concat); the KV +kernel builds the final key (nope | broadcast-rotated shared rope head) and +value tensors in a single pass. Falls back to pure PyTorch when Triton is +unavailable or for the ``bshd`` layout (the Triton path is ``sbhd``-only). + +Rotation convention: the rope slice is read interleaved (as stored in +HF/Megatron DeepSeekV3 checkpoints) and written in NeoX half-split layout, +matching the Megatron fused kernel semantics. +""" + +from typing import Optional, Tuple + +import torch + +try: + import triton + import triton.language as tl + + HAVE_TRITON = True +except ImportError: + HAVE_TRITON = False + +__all__ = ["build_rope_tables", "apply_mla_rope_q", "apply_mla_rope_kv"] + + +def build_rope_tables( + seq_len: int, + emb_dim: int, + base: float = 10000.0, + device: Optional[torch.device] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """cos/sin tables of shape ``[seq_len, emb_dim]`` (fp32, NeoX duplicated halves).""" + inv_freq = 1.0 / ( + base ** (torch.arange(0, emb_dim, 2, dtype=torch.float32, device=device) / emb_dim) + ) + t = torch.arange(seq_len, device=device, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + freqs = torch.cat([freqs, freqs], dim=-1) + return torch.cos(freqs).contiguous(), torch.sin(freqs).contiguous() + + +if HAVE_TRITON: + + # Not used for non-packed batches; kept for THD compatibility. + @triton.jit + def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): + token_idx = -1 + this_seq_len = 0 + seq_idx = 0 + last_cum_seqlen = tl.load(cu_seqlens) // cp_size + while seq_idx < seq_num: + cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1) // cp_size + if token_idx == -1 and cur_cum_seqlen > pid_m: + token_idx = pid_m - last_cum_seqlen + this_seq_len = cur_cum_seqlen - last_cum_seqlen + last_cum_seqlen = cur_cum_seqlen + seq_idx += 1 + if cp_size > 1: + if token_idx < this_seq_len // 2: + token_idx = token_idx + cp_rank * this_seq_len // 2 + else: + token_idx = (token_idx - this_seq_len // 2) + ( + 2 * cp_size - cp_rank - 1 + ) * this_seq_len // 2 + return token_idx + + _AUTOTUNE_CONFIGS = [triton.Config({"BLOCK_H": h}) for h in (1, 2, 4, 8, 16, 32, 64, 128)] + + @triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["emb_dim", "head_num"], restore_value=["Q"]) + @triton.jit + def rotary_fwd_q_kernel( + Q, + COS, + SIN, + qk_head_dim, + emb_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_q, + stride_x_seq, + stride_x_nheads, + cp_rank, + cp_size, + BLOCK_H: tl.constexpr, + ): + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + if cu_seqlens_q is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) + cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_right = sin_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) + Q = Q + pid_m * stride_x_seq + x_off = head_offsets[:, None] * stride_x_nheads + qk_head_dim + mask = head_offsets[:, None] < head_num + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_1 = tl.load(Q + x_1_off, mask=mask) + x_2 = tl.load(Q + x_2_off, mask=mask) + x_left = x_1 * cos_left - x_2 * sin_left + x_right = x_2 * cos_right + x_1 * sin_right + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + tl.store(Q + x_left_off, x_left, mask=mask) + tl.store(Q + x_right_off, x_right, mask=mask) + + @triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["emb_dim", "head_num"], restore_value=["DO"]) + @triton.jit + def rotary_bwd_q_kernel( + DO, + COS, + SIN, + qk_head_dim, + emb_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_q, + stride_x_seq, + stride_x_nheads, + cp_rank, + cp_size, + BLOCK_H: tl.constexpr, + ): + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + if cu_seqlens_q is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) + cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_right = sin_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) + DO = DO + pid_m * stride_x_seq + x_off = head_offsets[:, None] * stride_x_nheads + qk_head_dim + mask = head_offsets[:, None] < head_num + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left = tl.load(DO + x_left_off, mask=mask) + x_right = tl.load(DO + x_right_off, mask=mask) + x_1 = x_left * cos_left + x_right * sin_right + x_2 = -x_left * sin_left + x_right * cos_right + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + tl.store(DO + x_1_off, x_1, mask=mask) + tl.store(DO + x_2_off, x_2, mask=mask) + + @triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["emb_dim", "k_dim", "v_dim", "head_num"]) + @triton.jit + def rotary_fwd_kv_kernel( + KV, + K_POS_EMB, + O_KEY, + O_VALUE, + COS, + SIN, + emb_dim: tl.constexpr, + k_dim: tl.constexpr, + v_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_kv, + stride_kv_seq, + stride_kv_nheads, + stride_emb_seq, + stride_k_seq, + stride_k_nheads, + stride_v_seq, + stride_v_nheads, + cp_rank, + cp_size, + BLOCK_H: tl.constexpr, + ): + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + if cu_seqlens_kv is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size) + cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) + KV_ptr = KV + pid_m * stride_kv_seq + kv_off = head_offsets[:, None] * stride_kv_nheads + mask = head_offsets[:, None] < head_num + k_in_off = kv_off + tl.arange(0, k_dim)[None, :] + v_in_off = kv_off + k_dim + tl.arange(0, v_dim)[None, :] + k = tl.load(KV_ptr + k_in_off, mask=mask) + v = tl.load(KV_ptr + v_in_off, mask=mask) + K_ptr = O_KEY + pid_m * stride_k_seq + pid_head * BLOCK_H * stride_k_nheads + V_ptr = O_VALUE + pid_m * stride_v_seq + pid_head * BLOCK_H * stride_v_nheads + k_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + tl.arange(0, k_dim)[None, :] + v_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_v_nheads + tl.arange(0, v_dim)[None, :] + tl.store(K_ptr + k_out_off, k, mask=mask) + tl.store(V_ptr + v_out_off, v, mask=mask) + EMB = K_POS_EMB + pid_m * stride_emb_seq + x_1 = tl.load(EMB + tl.arange(0, emb_dim // 2) * 2) + x_2 = tl.load(EMB + tl.arange(0, emb_dim // 2) * 2 + 1) + x_left = x_1 * cos_left - x_2 * sin_left + x_right = x_2 * cos_right + x_1 * sin_right + x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + x_left_off = ( + tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] + ) + x_right_off = x_left_off + emb_dim // 2 + tl.store(K_ptr + x_left_off, x_left, mask=mask) + tl.store(K_ptr + x_right_off, x_right, mask=mask) + + @triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["emb_dim", "k_dim", "v_dim", "head_num"]) + @triton.jit + def rotary_bwd_kv_kernel( + dK, + dV, + dKV, + dEMB, + COS, + SIN, + emb_dim: tl.constexpr, + k_dim: tl.constexpr, + v_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_kv, + stride_dk_seq, + stride_dk_nheads, + stride_dv_seq, + stride_dv_nheads, + stride_dkv_seq, + stride_dkv_nheads, + stride_demb_seq, + cp_rank, + cp_size, + BLOCK_H: tl.constexpr, + ): + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + if cu_seqlens_kv is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size) + head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) + dKV_ptr = dKV + pid_m * stride_dkv_seq + dkv_off = head_offsets[:, None] * stride_dkv_nheads + mask = head_offsets[:, None] < head_num + dk_out_off = dkv_off + tl.arange(0, k_dim)[None, :] + dv_out_off = dkv_off + k_dim + tl.arange(0, v_dim)[None, :] + dK_ptr = dK + pid_m * stride_dk_seq + pid_head * BLOCK_H * stride_dk_nheads + dV_ptr = dV + pid_m * stride_dv_seq + pid_head * BLOCK_H * stride_dv_nheads + dk_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + tl.arange(0, k_dim)[None, :] + dv_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dv_nheads + tl.arange(0, v_dim)[None, :] + dk = tl.load(dK_ptr + dk_in_off, mask=mask) + dv = tl.load(dV_ptr + dv_in_off, mask=mask) + tl.store(dKV_ptr + dk_out_off, dk, mask=mask) + tl.store(dKV_ptr + dv_out_off, dv, mask=mask) + if pid_head == 0: + x_left_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) + x_right_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) + for i in tl.static_range(triton.cdiv(head_num, BLOCK_H)): + head_offsets_i = i * BLOCK_H + tl.arange(0, BLOCK_H) + dK_ptr_i = dK + pid_m * stride_dk_seq + x_off = head_offsets_i[:, None] * stride_dk_nheads + k_dim + mask_i = head_offsets_i[:, None] < head_num + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left_accum += tl.load(dK_ptr_i + x_left_off, mask=mask_i) + x_right_accum += tl.load(dK_ptr_i + x_right_off, mask=mask_i) + x_left_accum = tl.sum(x_left_accum, axis=0) + x_right_accum = tl.sum(x_right_accum, axis=0) + x_left_accum = x_left_accum.to(dEMB.dtype.element_ty) + x_right_accum = x_right_accum.to(dEMB.dtype.element_ty) + cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + cos_right = tl.load( + COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + sin_right = tl.load( + SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + x_1 = x_left_accum * cos_left + x_right_accum * sin_right + x_2 = -x_left_accum * sin_left + x_right_accum * cos_right + dEMB_ptr = dEMB + pid_m * stride_demb_seq + tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2, x_1) + tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2 + 1, x_2) + + def _token_stride(tensor: torch.Tensor) -> int: + return tensor.stride(1) if tensor.dim() == 4 else tensor.stride(0) + + class _MLARoPEQTriton(torch.autograd.Function): + """In-place RoPE on the trailing rope slice of q [s, b, h, nope+rope].""" + + @staticmethod + def forward(ctx, q, cos, sin, head_dim_nope, head_dim_rope): + if not q.is_contiguous(): + q = q.contiguous() + s, b, nheads, _ = q.shape + grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + rotary_fwd_q_kernel[grid]( + q, + cos, + sin, + head_dim_nope, + head_dim_rope, + nheads, + b, + None, + None, + _token_stride(q), + q.stride(2), + 0, + 1, + ) + ctx.save_for_backward(cos, sin) + ctx.dims = (s, b, nheads, head_dim_nope, head_dim_rope) + return q + + @staticmethod + def backward(ctx, dq): + cos, sin = ctx.saved_tensors + # attention backward may hand over a strided grad; the kernel + # assumes a contiguous [s, b, h, d] layout + dq = dq.contiguous() + s, b, nheads, head_dim_nope, head_dim_rope = ctx.dims + grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + rotary_bwd_q_kernel[grid]( + dq, + cos, + sin, + head_dim_nope, + head_dim_rope, + nheads, + b, + None, + None, + _token_stride(dq), + dq.stride(2), + 0, + 1, + ) + return dq, None, None, None, None + + class _MLARoPEKVTriton(torch.autograd.Function): + """kv [s, b, h, nope+v] + shared rope head [s, b, 1, rope] -> (k, v).""" + + @staticmethod + def forward(ctx, kv, k_pos_emb, cos, sin, head_dim_nope, head_dim_rope, head_dim_v): + if not kv.is_contiguous(): + kv = kv.contiguous() + s, b, nheads, _ = kv.shape + o_key = kv.new_empty(s, b, nheads, head_dim_nope + head_dim_rope) + o_value = kv.new_empty(s, b, nheads, head_dim_v) + grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + rotary_fwd_kv_kernel[grid]( + kv, + k_pos_emb, + o_key, + o_value, + cos, + sin, + head_dim_rope, + head_dim_nope, + head_dim_v, + nheads, + b, + None, + None, + _token_stride(kv), + kv.stride(2), + _token_stride(k_pos_emb), + _token_stride(o_key), + o_key.stride(2), + _token_stride(o_value), + o_value.stride(2), + 0, + 1, + ) + ctx.save_for_backward(cos, sin) + ctx.dims = (s, b, nheads, head_dim_nope, head_dim_rope, head_dim_v) + return o_key, o_value + + @staticmethod + def backward(ctx, dk_out, dv_out): + cos, sin = ctx.saved_tensors + s, b, nheads, ndp, ndr, ndv = ctx.dims + dk_out = dk_out.contiguous() + dv_out = dv_out.contiguous() + d_kv = dk_out.new_empty(s, b, nheads, ndp + ndv) + d_emb = dk_out.new_empty(s, b, 1, ndr) + grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + rotary_bwd_kv_kernel[grid]( + dk_out, + dv_out, + d_kv, + d_emb, + cos, + sin, + ndr, + ndp, + ndv, + nheads, + b, + None, + None, + _token_stride(dk_out), + dk_out.stride(2), + _token_stride(dv_out), + dv_out.stride(2), + _token_stride(d_kv), + d_kv.stride(2), + _token_stride(d_emb), + 0, + 1, + ) + return d_kv, d_emb, None, None, None, None, None + + +def _rotate_interleaved_to_neox(x, cos_table, sin_table, seq_dim): + shape = [1, 1, 1, cos_table.shape[-1]] + shape[seq_dim] = cos_table.shape[0] + cos_ = cos_table.view(shape).to(x.dtype) + sin_ = sin_table.view(shape).to(x.dtype) + half = x.shape[-1] // 2 + x_1 = x[..., 0::2] + x_2 = x[..., 1::2] + x_left = x_1 * cos_[..., :half] - x_2 * sin_[..., :half] + x_right = x_2 * cos_[..., half:] + x_1 * sin_[..., half:] + return torch.cat((x_left, x_right), dim=-1) + + +def apply_mla_rope_q( + q: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + head_dim_nope: int, + head_dim_rope: int, + tensor_format: str = "sbhd", +) -> torch.Tensor: + """RoPE on the trailing ``head_dim_rope`` slice of q; in place on the Triton path.""" + if HAVE_TRITON and tensor_format == "sbhd": + return _MLARoPEQTriton.apply(q, cos_table, sin_table, head_dim_nope, head_dim_rope) + seq_dim = 0 if tensor_format == "sbhd" else 1 + q_rope = _rotate_interleaved_to_neox(q[..., head_dim_nope:], cos_table, sin_table, seq_dim) + return torch.cat((q[..., :head_dim_nope], q_rope), dim=-1) + + +def apply_mla_rope_kv( + kv: torch.Tensor, + k_pos_emb: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + head_dim_nope: int, + head_dim_rope: int, + head_dim_v: int, + tensor_format: str = "sbhd", +) -> Tuple[torch.Tensor, torch.Tensor]: + """Build (k, v) from kv ``[.., h, nope+v]`` and the shared rope head ``[.., 1, rope]``.""" + if HAVE_TRITON and tensor_format == "sbhd": + return _MLARoPEKVTriton.apply( + kv, k_pos_emb, cos_table, sin_table, head_dim_nope, head_dim_rope, head_dim_v + ) + seq_dim = 0 if tensor_format == "sbhd" else 1 + k_nope = kv[..., :head_dim_nope] + v = kv[..., head_dim_nope : head_dim_nope + head_dim_v] + k_rope = _rotate_interleaved_to_neox(k_pos_emb, cos_table, sin_table, seq_dim) + k_rope = k_rope.expand(*k_nope.shape[:-1], -1) + return torch.cat((k_nope, k_rope), dim=-1), v.contiguous() diff --git a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py index a36075f2c5..e5840a812f 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py +++ b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py @@ -9,8 +9,12 @@ import torch from transformer_engine.pytorch.module import Linear, LayerNormLinear -from transformer_engine.pytorch.attention import DotProductAttention, RotaryPositionEmbedding -from transformer_engine.pytorch.attention.rope import apply_rotary_pos_emb +from transformer_engine.pytorch.attention import DotProductAttention +from transformer_engine.pytorch.models.deepseek_v3.mla_rope import ( + apply_mla_rope_kv, + apply_mla_rope_q, + build_rope_tables, +) __all__ = ["MultiLatentAttention"] @@ -29,6 +33,10 @@ class MultiLatentAttention(torch.nn.Module): ``kv_channels=(qk_nope_head_dim + qk_rope_head_dim, v_head_dim)``, which supports the cuDNN fused attention backend. + RoPE uses the fused MLA kernels from :mod:`.mla_rope` (in-place on the + query rope slice, single-pass key/value assembly); the rope slice follows + the HF/Megatron DeepSeekV3 convention (interleaved weights, NeoX output). + Parameters ---------- hidden_size : int @@ -126,8 +134,8 @@ def __init__( **common, ) - self.rope = RotaryPositionEmbedding(qk_rope_head_dim, rotary_base=rotary_base) - self._rope_freqs: Optional[torch.Tensor] = None + self.rotary_base = rotary_base + self._rope_tables: Optional[tuple] = None self.core_attention = DotProductAttention( num_attention_heads, @@ -140,10 +148,13 @@ def __init__( tp_size=tp_size, ) - def _rope_freqs_for(self, seq_len: int, device: torch.device) -> torch.Tensor: - if self._rope_freqs is None or self._rope_freqs.shape[0] < seq_len: - self._rope_freqs = self.rope(seq_len).to(device) - return self._rope_freqs[:seq_len] + def _rope_tables_for(self, seq_len: int, device: torch.device): + if self._rope_tables is None or self._rope_tables[0].shape[0] < seq_len: + self._rope_tables = build_rope_tables( + seq_len, self.qk_rope_head_dim, base=self.rotary_base, device=device + ) + cos, sin = self._rope_tables + return cos[:seq_len], sin[:seq_len] def forward( self, @@ -175,26 +186,26 @@ def forward( kv_latent, k_pos = torch.split(kv_down, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) kv = self.kv_up_proj(kv_latent) kv = kv.view(*kv.shape[:-1], heads, self.qk_nope_head_dim + self.v_head_dim) - k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - - freqs = self._rope_freqs_for(seq_len, hidden_states.device) - q_rope = apply_rotary_pos_emb( - q[..., self.qk_nope_head_dim :].contiguous(), - freqs, - tensor_format=self.qkv_format, - fused=True, + + cos, sin = self._rope_tables_for(seq_len, hidden_states.device) + q = apply_mla_rope_q( + q, cos, sin, self.qk_nope_head_dim, self.qk_rope_head_dim, self.qkv_format ) - k_rope = apply_rotary_pos_emb( - k_pos.unsqueeze(-2), freqs, tensor_format=self.qkv_format, fused=True + k, v = apply_mla_rope_kv( + kv, + k_pos.unsqueeze(-2), + cos, + sin, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + self.qkv_format, ) - q = torch.cat([q[..., : self.qk_nope_head_dim], q_rope], dim=-1) - k = torch.cat([k_nope, k_rope.expand(*k_nope.shape[:-1], -1)], dim=-1) - context = self.core_attention( q, k, - v.contiguous(), + v, attention_mask=attention_mask, qkv_format=self.qkv_format, attn_mask_type=attn_mask_type, From c713af7bf06fca793e48c86be6fd203b4836bf14 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 21 Aug 2026 11:42:40 +0200 Subject: [PATCH 08/41] Add HF transformers numeric reference test for DeepSeekV3Layer Maps HF DeepseekV3DecoderLayer weights into DeepSeekV3Layer (GLU interleave for routed experts, fused latent norms) and checks forward and input grads match within bf16 tolerance. Expose layernorm_epsilon on MultiLatentAttention (HF latent RMSNorms use 1e-6). Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_deepseek_hf.py | 144 ++++++++++++++++++ .../deepseek_v3/multi_latent_attention.py | 5 + 2 files changed, 149 insertions(+) create mode 100644 tests/pytorch/test_deepseek_hf.py diff --git a/tests/pytorch/test_deepseek_hf.py b/tests/pytorch/test_deepseek_hf.py new file mode 100644 index 0000000000..08f1c06ccd --- /dev/null +++ b/tests/pytorch/test_deepseek_hf.py @@ -0,0 +1,144 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Numeric comparison of DeepSeekV3Layer against the HF transformers reference.""" + +import pytest +import torch + +transformers = pytest.importorskip("transformers") +from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config +from transformers.models.deepseek_v3.modeling_deepseek_v3 import ( + DeepseekV3DecoderLayer, + DeepseekV3RotaryEmbedding, +) + +from transformer_engine.pytorch.models import DeepSeekV3Layer +from transformer_engine.pytorch.utils import interleave_glu_tensor + +SEQ, BATCH = 64, 2 +HIDDEN, HEADS = 256, 4 +Q_LORA, KV_LORA = 96, 64 +NOPE, ROPE, VDIM = 64, 32, 64 +NUM_EXPERTS, TOPK, N_GROUP, TOPK_GROUP = 16, 4, 4, 2 +MOE_FFN, N_SHARED = 128, 1 +DTYPE = torch.bfloat16 + + +def _hf_config(): + return DeepseekV3Config( + hidden_size=HIDDEN, + intermediate_size=4 * HIDDEN, + moe_intermediate_size=MOE_FFN, + num_hidden_layers=1, + num_attention_heads=HEADS, + num_key_value_heads=HEADS, + n_shared_experts=N_SHARED, + n_routed_experts=NUM_EXPERTS, + routed_scaling_factor=2.5, + kv_lora_rank=KV_LORA, + q_lora_rank=Q_LORA, + qk_rope_head_dim=ROPE, + v_head_dim=VDIM, + qk_nope_head_dim=NOPE, + n_group=N_GROUP, + topk_group=TOPK_GROUP, + num_experts_per_tok=TOPK, + first_k_dense_replace=0, + norm_topk_prob=True, + rms_norm_eps=1e-5, + attention_bias=False, + attention_dropout=0.0, + rope_interleave=True, + _attn_implementation="eager", + ) + + +def _init_hf_layer(config): + torch.manual_seed(0) + layer = DeepseekV3DecoderLayer(config, layer_idx=0).to(device="cuda", dtype=DTYPE) + with torch.no_grad(): + for name, p in layer.named_parameters(): + if "layernorm" in name or "norm" in name: + p.copy_(1.0 + 0.1 * torch.randn_like(p)) + else: + p.normal_(0.0, 0.02) + bias = layer.mlp.gate.e_score_correction_bias + bias.copy_(0.1 * torch.randn_like(bias)) + return layer + + +def _build_te_layer(hf): + te_layer = DeepSeekV3Layer( + HIDDEN, + HEADS, + num_experts=NUM_EXPERTS, + moe_ffn_hidden_size=MOE_FFN, + topk=TOPK, + num_groups=N_GROUP, + group_topk=TOPK_GROUP, + routed_scaling_factor=2.5, + shared_expert_ffn_hidden_size=MOE_FFN * N_SHARED, + q_lora_rank=Q_LORA, + kv_lora_rank=KV_LORA, + qk_nope_head_dim=NOPE, + qk_rope_head_dim=ROPE, + v_head_dim=VDIM, + params_dtype=DTYPE, + ) + attn, mla = hf.self_attn, te_layer.self_attention + with torch.no_grad(): + te_layer.input_layernorm.weight.copy_(hf.input_layernorm.weight) + te_layer.pre_mlp_layernorm.weight.copy_(hf.post_attention_layernorm.weight) + + mla.q_down_proj.weight.copy_(attn.q_a_proj.weight) + mla.q_up_proj.layer_norm_weight.copy_(attn.q_a_layernorm.weight) + mla.q_up_proj.weight.copy_(attn.q_b_proj.weight) + mla.kv_down_proj.weight.copy_(attn.kv_a_proj_with_mqa.weight) + mla.kv_up_proj.layer_norm_weight.copy_(attn.kv_a_layernorm.weight) + mla.kv_up_proj.weight.copy_(attn.kv_b_proj.weight) + mla.out_proj.weight.copy_(attn.o_proj.weight) + + moe = te_layer.mlp + moe.gate.weight.copy_(hf.mlp.gate.weight) + moe.expert_bias.copy_(hf.mlp.gate.e_score_correction_bias) + fc1, _, fc2 = moe.experts + for e in range(NUM_EXPERTS): + getattr(fc1, f"weight{e}").copy_( + interleave_glu_tensor(hf.mlp.experts.gate_up_proj[e], 32) + ) + getattr(fc2, f"weight{e}").copy_(hf.mlp.experts.down_proj[e]) + shared = hf.mlp.shared_experts + moe.shared_expert[0].weight.copy_( + torch.cat([shared.gate_proj.weight, shared.up_proj.weight], dim=0) + ) + moe.shared_expert[2].weight.copy_(shared.down_proj.weight) + return te_layer + + +def test_layer_matches_hf(): + config = _hf_config() + hf = _init_hf_layer(config) + te_layer = _build_te_layer(hf) + + torch.manual_seed(1) + x = torch.randn(BATCH, SEQ, HIDDEN, dtype=DTYPE, device="cuda") + x_hf = x.clone().requires_grad_(True) + x_te = x.transpose(0, 1).contiguous().requires_grad_(True) # sbhd + + rotary = DeepseekV3RotaryEmbedding(config).to("cuda") + position_ids = torch.arange(SEQ, device="cuda").unsqueeze(0).expand(BATCH, -1) + cos, sin = rotary(x_hf, position_ids) + causal = torch.full((SEQ, SEQ), float("-inf"), device="cuda", dtype=DTYPE).triu(1) + causal = causal[None, None].expand(BATCH, 1, SEQ, SEQ) + + out_hf = hf(x_hf, attention_mask=causal, position_embeddings=(cos, sin)) + out_te = te_layer(x_te) + + torch.testing.assert_close(out_te.transpose(0, 1), out_hf, rtol=5e-2, atol=5e-2) + + grad = torch.randn_like(out_hf) + out_hf.backward(grad) + out_te.backward(grad.transpose(0, 1).contiguous()) + torch.testing.assert_close(x_te.grad.transpose(0, 1), x_hf.grad, rtol=5e-2, atol=5e-2) diff --git a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py index e5840a812f..0ddf2f75b2 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py +++ b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py @@ -57,6 +57,8 @@ class MultiLatentAttention(torch.nn.Module): dropout probability on attention scores. attn_mask_type : str, default = "causal" attention mask type passed to :class:`DotProductAttention`. + layernorm_epsilon : float, default = 1e-6 + epsilon of the latent RMSNorms (matches DeepSeekV3). rotary_base : float, default = 10000.0 RoPE base. softmax_scale : float, optional @@ -83,6 +85,7 @@ def __init__( v_head_dim: int = 128, attention_dropout: float = 0.0, attn_mask_type: str = "causal", + layernorm_epsilon: float = 1e-6, rotary_base: float = 10000.0, softmax_scale: Optional[float] = None, qkv_format: str = "sbhd", @@ -113,6 +116,7 @@ def __init__( q_lora_rank, num_attention_heads * self.qk_head_dim, normalization="RMSNorm", + eps=layernorm_epsilon, parallel_mode="column" if tp_size > 1 else None, **tp, **common, @@ -122,6 +126,7 @@ def __init__( kv_lora_rank, num_attention_heads * (qk_nope_head_dim + v_head_dim), normalization="RMSNorm", + eps=layernorm_epsilon, parallel_mode="column" if tp_size > 1 else None, **tp, **common, From 88028a94f061e0ce075cfd788cf01e580d571ad4 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 21 Aug 2026 12:08:12 +0200 Subject: [PATCH 09/41] Docstring cleanups for lint and docs build Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- .../pytorch/models/deepseek_v3/mla_rope.py | 28 ++++++++++++++++--- .../models/deepseek_v3/transformer_layer.py | 13 +++++---- 2 files changed, 31 insertions(+), 10 deletions(-) diff --git a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py index 350bedb69b..f8c01f75d6 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py +++ b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py @@ -92,6 +92,7 @@ def rotary_fwd_q_kernel( cp_size, BLOCK_H: tl.constexpr, ): + """In-place RoPE fwd on the trailing rope slice of q.""" pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) if cu_seqlens_q is None: @@ -139,6 +140,7 @@ def rotary_bwd_q_kernel( cp_size, BLOCK_H: tl.constexpr, ): + """In-place RoPE bwd on the trailing rope slice of dq.""" pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) if cu_seqlens_q is None: @@ -195,6 +197,7 @@ def rotary_fwd_kv_kernel( cp_size, BLOCK_H: tl.constexpr, ): + """Fwd: build (key, value) from kv and the shared rotated rope head.""" pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) if cu_seqlens_kv is None: @@ -262,6 +265,7 @@ def rotary_bwd_kv_kernel( cp_size, BLOCK_H: tl.constexpr, ): + """Bwd: scatter (dk, dv) into dkv and reduce rope-slice grads into demb.""" pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) if cu_seqlens_kv is None: @@ -320,10 +324,14 @@ class _MLARoPEQTriton(torch.autograd.Function): @staticmethod def forward(ctx, q, cos, sin, head_dim_nope, head_dim_rope): + """Rotate the rope slice of q in place.""" if not q.is_contiguous(): q = q.contiguous() s, b, nheads, _ = q.shape - grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + + def grid(meta): + return (s * b, triton.cdiv(nheads, meta["BLOCK_H"])) + rotary_fwd_q_kernel[grid]( q, cos, @@ -345,12 +353,16 @@ def forward(ctx, q, cos, sin, head_dim_nope, head_dim_rope): @staticmethod def backward(ctx, dq): + """Counter-rotate the rope slice of dq (in place on the copy).""" cos, sin = ctx.saved_tensors # attention backward may hand over a strided grad; the kernel # assumes a contiguous [s, b, h, d] layout dq = dq.contiguous() s, b, nheads, head_dim_nope, head_dim_rope = ctx.dims - grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + + def grid(meta): + return (s * b, triton.cdiv(nheads, meta["BLOCK_H"])) + rotary_bwd_q_kernel[grid]( dq, cos, @@ -373,12 +385,16 @@ class _MLARoPEKVTriton(torch.autograd.Function): @staticmethod def forward(ctx, kv, k_pos_emb, cos, sin, head_dim_nope, head_dim_rope, head_dim_v): + """Build (k, v) from kv and the shared rope head.""" if not kv.is_contiguous(): kv = kv.contiguous() s, b, nheads, _ = kv.shape o_key = kv.new_empty(s, b, nheads, head_dim_nope + head_dim_rope) o_value = kv.new_empty(s, b, nheads, head_dim_v) - grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + + def grid(meta): + return (s * b, triton.cdiv(nheads, meta["BLOCK_H"])) + rotary_fwd_kv_kernel[grid]( kv, k_pos_emb, @@ -409,13 +425,17 @@ def forward(ctx, kv, k_pos_emb, cos, sin, head_dim_nope, head_dim_rope, head_dim @staticmethod def backward(ctx, dk_out, dv_out): + """Gradients for (kv, k_pos_emb) from (dk, dv).""" cos, sin = ctx.saved_tensors s, b, nheads, ndp, ndr, ndv = ctx.dims dk_out = dk_out.contiguous() dv_out = dv_out.contiguous() d_kv = dk_out.new_empty(s, b, nheads, ndp + ndv) d_emb = dk_out.new_empty(s, b, 1, ndr) - grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + + def grid(meta): + return (s * b, triton.cdiv(nheads, meta["BLOCK_H"])) + rotary_bwd_kv_kernel[grid]( dk_out, dv_out, diff --git a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py index af1eeb1a95..f41ec0061c 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py +++ b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py @@ -40,12 +40,13 @@ class DeepSeekV3Layer(torch.nn.Module): ffn size of each routed expert (required with MoE). hidden_dropout : float, default = 0.0 dropout probability on the residual branches. - kwargs common to the submodules (``q_lora_rank``, ``kv_lora_rank``, - ``qk_nope_head_dim``, ``qk_rope_head_dim``, ``v_head_dim``, - ``attention_dropout``, ``attn_mask_type``, ``qkv_format``, ``topk``, - ``num_groups``, ``group_topk``, ``routed_scaling_factor``, - ``shared_expert_ffn_hidden_size``, EP options, ...) are forwarded to - :class:`MultiLatentAttention` and :class:`DeepSeekV3MoE`. + **kwargs + kwargs common to the submodules (``q_lora_rank``, ``kv_lora_rank``, + ``qk_nope_head_dim``, ``qk_rope_head_dim``, ``v_head_dim``, + ``attention_dropout``, ``attn_mask_type``, ``qkv_format``, ``topk``, + ``num_groups``, ``group_topk``, ``routed_scaling_factor``, + ``shared_expert_ffn_hidden_size``, EP options, ...), forwarded to + :class:`MultiLatentAttention` and :class:`DeepSeekV3MoE`. """ _MLA_KWARGS = frozenset( From 5f68c9bbdaa984abb35ba071f44c0271a8fce986 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 21 Aug 2026 14:09:24 +0200 Subject: [PATCH 10/41] Move model-specific layers to a dedicated docs page docs/api/pytorch_models.rst: usage (local and EP), fused-path notes, HF checkpoint weight mapping, and the class API; linked from the PyTorch API page via a toctree entry. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- docs/api/pytorch.rst | 8 ++- docs/api/pytorch_models.rst | 113 ++++++++++++++++++++++++++++++++++++ 2 files changed, 118 insertions(+), 3 deletions(-) create mode 100644 docs/api/pytorch_models.rst diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index bd3099b590..4fa279cfbc 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -62,11 +62,13 @@ PyTorch Model-specific layers --------------------- -.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3Layer(**kwargs) +Full transformer layers for specific model families live in +``transformer_engine.pytorch.models``: -.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3MoE(**kwargs) +.. toctree:: + :maxdepth: 1 -.. autoapiclass:: transformer_engine.pytorch.models.MultiLatentAttention(**kwargs) + pytorch_models Data types ---------- diff --git a/docs/api/pytorch_models.rst b/docs/api/pytorch_models.rst new file mode 100644 index 0000000000..0534acb282 --- /dev/null +++ b/docs/api/pytorch_models.rst @@ -0,0 +1,113 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Model-specific layers (te.models) +================================= + +The ``transformer_engine.pytorch.models`` namespace holds full transformer +layers for specific model families, composed from Transformer Engine modules +and fused kernels. Each family lives in its own subpackage. + +DeepSeek-V3 +----------- + +A DeepSeek-V3 transformer layer analogous to +:class:`transformer_engine.pytorch.TransformerLayer`: Multi-Latent Attention +(MLA) with low-rank q/kv latents and decoupled RoPE/NoPE heads, plus a +DeepSeek-style Mixture of Experts block (fused sigmoid router with +aux-loss-free expert bias and node-limited grouped top-k, grouped-GEMM SwiGLU +experts, optional shared expert). The same architecture is used by other +model families (e.g. GLM-5, Kimi K2), which can reuse these modules. + +Basic usage (single GPU, all experts local): + +.. code-block:: python + + import torch + import transformer_engine.pytorch as te + + layer = te.models.DeepSeekV3Layer( + hidden_size=7168, + num_attention_heads=128, + num_experts=64, + moe_ffn_hidden_size=2048, + topk=8, + shared_expert_ffn_hidden_size=2048, + params_dtype=torch.bfloat16, + ) + x = torch.randn(seq_len, batch, 7168, dtype=torch.bfloat16, device="cuda") + y = layer(x) # sbhd layout + +Expert parallelism routes tokens between GPUs with the NCCL EP backend +(``transformer_engine.pytorch.ep``). Call ``ep_bootstrap`` once per process +before the first forward; EP requires bfloat16 inputs and NCCL >= 2.30.4: + +.. code-block:: python + + from transformer_engine.pytorch.ep import ep_bootstrap + + ep_bootstrap(ep_group, num_experts=64, max_tokens_per_rank=tokens, + hidden_dim=7168, num_topk=8, recv_capacity_per_rank=capacity) + layer = te.models.DeepSeekV3Layer( + ..., + ep_group=ep_group, + ep_max_tokens_per_rank=tokens, + ) + +On SM100-class GPUs the routed experts fuse into a single CuTe grouped-GEMM +MLP when running under ``te.autocast`` with an MXFP8/NVFP4 recipe and +``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1``; elsewhere the same modules run unfused +with an identical checkpoint layout. + +Loading HuggingFace checkpoints +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The layer follows the HF/Megatron DeepSeek-V3 conventions (interleaved rope +weights, sigmoid router bias used for selection only). Weights map from +``transformers`` ``DeepseekV3DecoderLayer`` as follows (latent RMSNorms are +fused into the up-projections): + +.. list-table:: + :header-rows: 1 + + * - Transformer Engine + - HuggingFace + * - ``input_layernorm.weight`` + - ``input_layernorm.weight`` + * - ``pre_mlp_layernorm.weight`` + - ``post_attention_layernorm.weight`` + * - ``self_attention.q_down_proj.weight`` + - ``self_attn.q_a_proj.weight`` + * - ``self_attention.q_up_proj.{layer_norm_weight, weight}`` + - ``self_attn.{q_a_layernorm, q_b_proj}.weight`` + * - ``self_attention.kv_down_proj.weight`` + - ``self_attn.kv_a_proj_with_mqa.weight`` + * - ``self_attention.kv_up_proj.{layer_norm_weight, weight}`` + - ``self_attn.{kv_a_layernorm, kv_b_proj}.weight`` + * - ``self_attention.out_proj.weight`` + - ``self_attn.o_proj.weight`` + * - ``mlp.gate.weight`` / ``mlp.expert_bias`` + - ``mlp.gate.weight`` / ``mlp.gate.e_score_correction_bias`` + * - ``mlp.experts[0].weight{i}`` + - ``interleave_glu_tensor(cat([gate_proj, up_proj]), 32)`` of expert *i* + * - ``mlp.experts[2].weight{i}`` + - ``mlp.experts.down_proj[i]`` + * - ``mlp.shared_expert[0].weight`` / ``[2].weight`` + - ``cat([gate_proj, up_proj])`` / ``down_proj`` of ``shared_experts`` + +See ``tests/pytorch/test_deepseek_hf.py`` for a complete, numerically +verified mapping. + +API +^^^ + +.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3Layer(hidden_size, num_attention_heads, **kwargs) + :members: forward + +.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3MoE(hidden_size, moe_ffn_hidden_size, num_experts, **kwargs) + :members: forward, update_expert_bias + +.. autoapiclass:: transformer_engine.pytorch.models.MultiLatentAttention(hidden_size, num_attention_heads, **kwargs) + :members: forward From 9db495a6577397f898f8d7a812fff557e1102e6e Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 21 Aug 2026 14:53:21 +0200 Subject: [PATCH 11/41] Drop HF-transformers comparison test from the repo Keep the verified weight-mapping table in the docs; the comparison itself stays as an out-of-tree script. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- docs/api/pytorch_models.rst | 4 +- tests/pytorch/test_deepseek_hf.py | 144 ------------------------------ 2 files changed, 2 insertions(+), 146 deletions(-) delete mode 100644 tests/pytorch/test_deepseek_hf.py diff --git a/docs/api/pytorch_models.rst b/docs/api/pytorch_models.rst index 0534acb282..665d547cc8 100644 --- a/docs/api/pytorch_models.rst +++ b/docs/api/pytorch_models.rst @@ -97,8 +97,8 @@ fused into the up-projections): * - ``mlp.shared_expert[0].weight`` / ``[2].weight`` - ``cat([gate_proj, up_proj])`` / ``down_proj`` of ``shared_experts`` -See ``tests/pytorch/test_deepseek_hf.py`` for a complete, numerically -verified mapping. +The routed-expert fc1 layout can be produced with +:func:`transformer_engine.pytorch.interleave_glu_tensor`. API ^^^ diff --git a/tests/pytorch/test_deepseek_hf.py b/tests/pytorch/test_deepseek_hf.py deleted file mode 100644 index 08f1c06ccd..0000000000 --- a/tests/pytorch/test_deepseek_hf.py +++ /dev/null @@ -1,144 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Numeric comparison of DeepSeekV3Layer against the HF transformers reference.""" - -import pytest -import torch - -transformers = pytest.importorskip("transformers") -from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config -from transformers.models.deepseek_v3.modeling_deepseek_v3 import ( - DeepseekV3DecoderLayer, - DeepseekV3RotaryEmbedding, -) - -from transformer_engine.pytorch.models import DeepSeekV3Layer -from transformer_engine.pytorch.utils import interleave_glu_tensor - -SEQ, BATCH = 64, 2 -HIDDEN, HEADS = 256, 4 -Q_LORA, KV_LORA = 96, 64 -NOPE, ROPE, VDIM = 64, 32, 64 -NUM_EXPERTS, TOPK, N_GROUP, TOPK_GROUP = 16, 4, 4, 2 -MOE_FFN, N_SHARED = 128, 1 -DTYPE = torch.bfloat16 - - -def _hf_config(): - return DeepseekV3Config( - hidden_size=HIDDEN, - intermediate_size=4 * HIDDEN, - moe_intermediate_size=MOE_FFN, - num_hidden_layers=1, - num_attention_heads=HEADS, - num_key_value_heads=HEADS, - n_shared_experts=N_SHARED, - n_routed_experts=NUM_EXPERTS, - routed_scaling_factor=2.5, - kv_lora_rank=KV_LORA, - q_lora_rank=Q_LORA, - qk_rope_head_dim=ROPE, - v_head_dim=VDIM, - qk_nope_head_dim=NOPE, - n_group=N_GROUP, - topk_group=TOPK_GROUP, - num_experts_per_tok=TOPK, - first_k_dense_replace=0, - norm_topk_prob=True, - rms_norm_eps=1e-5, - attention_bias=False, - attention_dropout=0.0, - rope_interleave=True, - _attn_implementation="eager", - ) - - -def _init_hf_layer(config): - torch.manual_seed(0) - layer = DeepseekV3DecoderLayer(config, layer_idx=0).to(device="cuda", dtype=DTYPE) - with torch.no_grad(): - for name, p in layer.named_parameters(): - if "layernorm" in name or "norm" in name: - p.copy_(1.0 + 0.1 * torch.randn_like(p)) - else: - p.normal_(0.0, 0.02) - bias = layer.mlp.gate.e_score_correction_bias - bias.copy_(0.1 * torch.randn_like(bias)) - return layer - - -def _build_te_layer(hf): - te_layer = DeepSeekV3Layer( - HIDDEN, - HEADS, - num_experts=NUM_EXPERTS, - moe_ffn_hidden_size=MOE_FFN, - topk=TOPK, - num_groups=N_GROUP, - group_topk=TOPK_GROUP, - routed_scaling_factor=2.5, - shared_expert_ffn_hidden_size=MOE_FFN * N_SHARED, - q_lora_rank=Q_LORA, - kv_lora_rank=KV_LORA, - qk_nope_head_dim=NOPE, - qk_rope_head_dim=ROPE, - v_head_dim=VDIM, - params_dtype=DTYPE, - ) - attn, mla = hf.self_attn, te_layer.self_attention - with torch.no_grad(): - te_layer.input_layernorm.weight.copy_(hf.input_layernorm.weight) - te_layer.pre_mlp_layernorm.weight.copy_(hf.post_attention_layernorm.weight) - - mla.q_down_proj.weight.copy_(attn.q_a_proj.weight) - mla.q_up_proj.layer_norm_weight.copy_(attn.q_a_layernorm.weight) - mla.q_up_proj.weight.copy_(attn.q_b_proj.weight) - mla.kv_down_proj.weight.copy_(attn.kv_a_proj_with_mqa.weight) - mla.kv_up_proj.layer_norm_weight.copy_(attn.kv_a_layernorm.weight) - mla.kv_up_proj.weight.copy_(attn.kv_b_proj.weight) - mla.out_proj.weight.copy_(attn.o_proj.weight) - - moe = te_layer.mlp - moe.gate.weight.copy_(hf.mlp.gate.weight) - moe.expert_bias.copy_(hf.mlp.gate.e_score_correction_bias) - fc1, _, fc2 = moe.experts - for e in range(NUM_EXPERTS): - getattr(fc1, f"weight{e}").copy_( - interleave_glu_tensor(hf.mlp.experts.gate_up_proj[e], 32) - ) - getattr(fc2, f"weight{e}").copy_(hf.mlp.experts.down_proj[e]) - shared = hf.mlp.shared_experts - moe.shared_expert[0].weight.copy_( - torch.cat([shared.gate_proj.weight, shared.up_proj.weight], dim=0) - ) - moe.shared_expert[2].weight.copy_(shared.down_proj.weight) - return te_layer - - -def test_layer_matches_hf(): - config = _hf_config() - hf = _init_hf_layer(config) - te_layer = _build_te_layer(hf) - - torch.manual_seed(1) - x = torch.randn(BATCH, SEQ, HIDDEN, dtype=DTYPE, device="cuda") - x_hf = x.clone().requires_grad_(True) - x_te = x.transpose(0, 1).contiguous().requires_grad_(True) # sbhd - - rotary = DeepseekV3RotaryEmbedding(config).to("cuda") - position_ids = torch.arange(SEQ, device="cuda").unsqueeze(0).expand(BATCH, -1) - cos, sin = rotary(x_hf, position_ids) - causal = torch.full((SEQ, SEQ), float("-inf"), device="cuda", dtype=DTYPE).triu(1) - causal = causal[None, None].expand(BATCH, 1, SEQ, SEQ) - - out_hf = hf(x_hf, attention_mask=causal, position_embeddings=(cos, sin)) - out_te = te_layer(x_te) - - torch.testing.assert_close(out_te.transpose(0, 1), out_hf, rtol=5e-2, atol=5e-2) - - grad = torch.randn_like(out_hf) - out_hf.backward(grad) - out_te.backward(grad.transpose(0, 1).contiguous()) - torch.testing.assert_close(x_te.grad.transpose(0, 1), x_hf.grad, rtol=5e-2, atol=5e-2) From 4883b1728218c11cba4c14d77f46f98a66d57f5a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 21 Aug 2026 14:54:35 +0200 Subject: [PATCH 12/41] Docs: reduce models page to a plain API listing Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- docs/api/pytorch.rst | 7 +-- docs/api/pytorch_models.rst | 98 +------------------------------------ 2 files changed, 4 insertions(+), 101 deletions(-) diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 4fa279cfbc..1c39469f2c 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -59,11 +59,8 @@ PyTorch .. autoapifunction:: transformer_engine.pytorch.deinterleave_glu_tensor -Model-specific layers ---------------------- - -Full transformer layers for specific model families live in -``transformer_engine.pytorch.models``: +Models +------ .. toctree:: :maxdepth: 1 diff --git a/docs/api/pytorch_models.rst b/docs/api/pytorch_models.rst index 665d547cc8..2cde879ffb 100644 --- a/docs/api/pytorch_models.rst +++ b/docs/api/pytorch_models.rst @@ -3,106 +3,12 @@ See LICENSE for license information. -Model-specific layers (te.models) -================================= - -The ``transformer_engine.pytorch.models`` namespace holds full transformer -layers for specific model families, composed from Transformer Engine modules -and fused kernels. Each family lives in its own subpackage. +Models +====== DeepSeek-V3 ----------- -A DeepSeek-V3 transformer layer analogous to -:class:`transformer_engine.pytorch.TransformerLayer`: Multi-Latent Attention -(MLA) with low-rank q/kv latents and decoupled RoPE/NoPE heads, plus a -DeepSeek-style Mixture of Experts block (fused sigmoid router with -aux-loss-free expert bias and node-limited grouped top-k, grouped-GEMM SwiGLU -experts, optional shared expert). The same architecture is used by other -model families (e.g. GLM-5, Kimi K2), which can reuse these modules. - -Basic usage (single GPU, all experts local): - -.. code-block:: python - - import torch - import transformer_engine.pytorch as te - - layer = te.models.DeepSeekV3Layer( - hidden_size=7168, - num_attention_heads=128, - num_experts=64, - moe_ffn_hidden_size=2048, - topk=8, - shared_expert_ffn_hidden_size=2048, - params_dtype=torch.bfloat16, - ) - x = torch.randn(seq_len, batch, 7168, dtype=torch.bfloat16, device="cuda") - y = layer(x) # sbhd layout - -Expert parallelism routes tokens between GPUs with the NCCL EP backend -(``transformer_engine.pytorch.ep``). Call ``ep_bootstrap`` once per process -before the first forward; EP requires bfloat16 inputs and NCCL >= 2.30.4: - -.. code-block:: python - - from transformer_engine.pytorch.ep import ep_bootstrap - - ep_bootstrap(ep_group, num_experts=64, max_tokens_per_rank=tokens, - hidden_dim=7168, num_topk=8, recv_capacity_per_rank=capacity) - layer = te.models.DeepSeekV3Layer( - ..., - ep_group=ep_group, - ep_max_tokens_per_rank=tokens, - ) - -On SM100-class GPUs the routed experts fuse into a single CuTe grouped-GEMM -MLP when running under ``te.autocast`` with an MXFP8/NVFP4 recipe and -``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1``; elsewhere the same modules run unfused -with an identical checkpoint layout. - -Loading HuggingFace checkpoints -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -The layer follows the HF/Megatron DeepSeek-V3 conventions (interleaved rope -weights, sigmoid router bias used for selection only). Weights map from -``transformers`` ``DeepseekV3DecoderLayer`` as follows (latent RMSNorms are -fused into the up-projections): - -.. list-table:: - :header-rows: 1 - - * - Transformer Engine - - HuggingFace - * - ``input_layernorm.weight`` - - ``input_layernorm.weight`` - * - ``pre_mlp_layernorm.weight`` - - ``post_attention_layernorm.weight`` - * - ``self_attention.q_down_proj.weight`` - - ``self_attn.q_a_proj.weight`` - * - ``self_attention.q_up_proj.{layer_norm_weight, weight}`` - - ``self_attn.{q_a_layernorm, q_b_proj}.weight`` - * - ``self_attention.kv_down_proj.weight`` - - ``self_attn.kv_a_proj_with_mqa.weight`` - * - ``self_attention.kv_up_proj.{layer_norm_weight, weight}`` - - ``self_attn.{kv_a_layernorm, kv_b_proj}.weight`` - * - ``self_attention.out_proj.weight`` - - ``self_attn.o_proj.weight`` - * - ``mlp.gate.weight`` / ``mlp.expert_bias`` - - ``mlp.gate.weight`` / ``mlp.gate.e_score_correction_bias`` - * - ``mlp.experts[0].weight{i}`` - - ``interleave_glu_tensor(cat([gate_proj, up_proj]), 32)`` of expert *i* - * - ``mlp.experts[2].weight{i}`` - - ``mlp.experts.down_proj[i]`` - * - ``mlp.shared_expert[0].weight`` / ``[2].weight`` - - ``cat([gate_proj, up_proj])`` / ``down_proj`` of ``shared_experts`` - -The routed-expert fc1 layout can be produced with -:func:`transformer_engine.pytorch.interleave_glu_tensor`. - -API -^^^ - .. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3Layer(hidden_size, num_attention_heads, **kwargs) :members: forward From 2f52af137544575d6aedac704ae6cf49989a72bc Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 11:24:45 +0200 Subject: [PATCH 13/41] Rename distributed DeepSeek EP tests to generic test_models Signed-off-by: Pawel Gadzinski --- .../distributed/{run_deepseek_ep.py => run_models.py} | 2 +- .../{run_test_deepseek_ep.sh => run_test_models.sh} | 8 ++++---- .../distributed/{test_deepseek_ep.py => test_models.py} | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) rename tests/pytorch/distributed/{run_deepseek_ep.py => run_models.py} (98%) rename tests/pytorch/distributed/{run_test_deepseek_ep.sh => run_test_models.sh} (86%) rename tests/pytorch/distributed/{test_deepseek_ep.py => test_models.py} (84%) diff --git a/tests/pytorch/distributed/run_deepseek_ep.py b/tests/pytorch/distributed/run_models.py similarity index 98% rename from tests/pytorch/distributed/run_deepseek_ep.py rename to tests/pytorch/distributed/run_models.py index bf756b69ad..ff5e233503 100644 --- a/tests/pytorch/distributed/run_deepseek_ep.py +++ b/tests/pytorch/distributed/run_models.py @@ -1,7 +1,7 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. -"""Multi-process DeepSeekV3 MoE/layer EP tests, launched via torchrun.""" +"""Multi-process tests for model-specific layers (te.models), launched via torchrun.""" import os import sys diff --git a/tests/pytorch/distributed/run_test_deepseek_ep.sh b/tests/pytorch/distributed/run_test_models.sh similarity index 86% rename from tests/pytorch/distributed/run_test_deepseek_ep.sh rename to tests/pytorch/distributed/run_test_models.sh index 8c0bbbc5b9..e3411ee7d9 100644 --- a/tests/pytorch/distributed/run_test_deepseek_ep.sh +++ b/tests/pytorch/distributed/run_test_models.sh @@ -3,7 +3,7 @@ # # See LICENSE for license information. # -# Launcher for tests/pytorch/distributed/run_deepseek_ep.py. Auto-detects GPU count. +# Launcher for tests/pytorch/distributed/run_models.py (model-specific layers). Auto-detects GPU count. set -uo pipefail @@ -30,15 +30,15 @@ TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-180}" export NCCL_EP_JIT_CACHE_DIR mkdir -p "$NCCL_EP_JIT_CACHE_DIR" -SCRIPT="${SCRIPT_DIR}/run_deepseek_ep.py" -LOG="stdout_deepseek_ep.txt" +SCRIPT="${SCRIPT_DIR}/run_models.py" +LOG="stdout_models.txt" echo "=== Running ${SCRIPT} on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ===" setsid timeout --foreground --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \ torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" \ "${SCRIPT}" 2>&1 | tee "${LOG}" RC=${PIPESTATUS[0]} -pkill -9 -f "tests/pytorch/distributed/run_deepseek_ep.py" 2>/dev/null || true +pkill -9 -f "tests/pytorch/distributed/run_models.py" 2>/dev/null || true RET=0 if [ "${RC}" -ne 0 ]; then echo "torchrun exited with ${RC}"; RET=1; fi diff --git a/tests/pytorch/distributed/test_deepseek_ep.py b/tests/pytorch/distributed/test_models.py similarity index 84% rename from tests/pytorch/distributed/test_deepseek_ep.py rename to tests/pytorch/distributed/test_models.py index 4a4d9a8dea..83213bd6a8 100644 --- a/tests/pytorch/distributed/test_deepseek_ep.py +++ b/tests/pytorch/distributed/test_models.py @@ -1,7 +1,7 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. -"""Pytest driver — spawns run_deepseek_ep.py under torchrun and asserts it passed.""" +"""Pytest driver — spawns run_models.py (model-specific layers, multi-GPU) under torchrun.""" import os import subprocess @@ -11,7 +11,7 @@ import torch TEST_ROOT = Path(__file__).parent.resolve() -LAUNCHER = TEST_ROOT / "run_test_deepseek_ep.sh" +LAUNCHER = TEST_ROOT / "run_test_models.sh" @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="DeepSeek EP requires >= 2 GPUs") From e841f966ea345ac11c5c612d3f311dfc0b6bf7f7 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 11:26:27 +0200 Subject: [PATCH 14/41] Rename test_deepseek.py to test_models.py and add models tests to QA scripts Signed-off-by: Pawel Gadzinski --- qa/L0_pytorch_unittest/test.sh | 1 + qa/L1_pytorch_distributed_unittest/test.sh | 1 + tests/pytorch/{test_deepseek.py => test_models.py} | 0 3 files changed, 2 insertions(+) rename tests/pytorch/{test_deepseek.py => test_models.py} (100%) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 14a5f4fe3d..ae4e183cda 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -58,6 +58,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_overrid python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml $TE_PATH/tests/pytorch/test_permutation.py || test_fail "test_permutation.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_entropy.xml $TE_PATH/tests/pytorch/test_parallel_cross_entropy.py || test_fail "test_parallel_cross_entropy.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_models.xml $TE_PATH/tests/pytorch/test_models.py || test_fail "test_models.py" NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading_v1.xml $TE_PATH/tests/pytorch/test_cpu_offloading_v1.py || test_fail "test_cpu_offloading_v1.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hybrid_quantization.xml $TE_PATH/tests/pytorch/test_hybrid_quantization.py || test_fail "test_hybrid_quantization.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_identity_quantizer.xml $TE_PATH/tests/pytorch/test_identity_quantizer.py || test_fail "test_identity_quantizer.py" diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index ec19492ee7..6773055b19 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -54,6 +54,7 @@ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cast_master_weights_to_fp8.xml $TE_PATH/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py || test_fail "test_cast_master_weights_to_fp8.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_newton_schulz.xml $TE_PATH/tests/pytorch/distributed/test_newton_schulz.py || test_fail "test_newton_schulz.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_ep.xml $TE_PATH/tests/pytorch/distributed/test_ep.py || test_fail "test_ep.py" +python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_models.xml $TE_PATH/tests/pytorch/distributed/test_models.py || test_fail "distributed/test_models.py" # debug tests diff --git a/tests/pytorch/test_deepseek.py b/tests/pytorch/test_models.py similarity index 100% rename from tests/pytorch/test_deepseek.py rename to tests/pytorch/test_models.py From 935e475c1e4c267f92974111715036b96c949443 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 12:08:14 +0200 Subject: [PATCH 15/41] Add YaRN RoPE scaling to DeepSeek V3 MLA Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_models.py | 47 +++++++++++++ .../pytorch/models/deepseek_v3/mla_rope.py | 68 +++++++++++++++++-- .../deepseek_v3/multi_latent_attention.py | 43 +++++++++++- .../models/deepseek_v3/transformer_layer.py | 6 ++ 4 files changed, 155 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/test_models.py b/tests/pytorch/test_models.py index 4c4aea0a92..678c104bb8 100644 --- a/tests/pytorch/test_models.py +++ b/tests/pytorch/test_models.py @@ -2,6 +2,8 @@ # # See LICENSE for license information. +import math + import pytest import torch @@ -104,6 +106,51 @@ def test_mla_forward_backward(): assert x.grad is not None and torch.isfinite(x.grad).all() +def test_rope_tables_yarn(): + from transformer_engine.pytorch.models.deepseek_v3 import mla_rope + + s, rope = 8192, 64 + cos, sin = mla_rope.build_rope_tables(s, rope, device="cuda") + cos_none, sin_none = mla_rope.build_rope_tables(s, rope, device="cuda", scaling_factor=None) + assert torch.equal(cos, cos_none) and torch.equal(sin, sin_none) + + yarn = dict(scaling_factor=40.0, original_max_position_embeddings=4096) + cos_y, sin_y = mla_rope.build_rope_tables(s, rope, device="cuda", **yarn) + factor = mla_rope.yarn_concentration_factor(40.0, 1.0, 0.0) + assert factor == pytest.approx(0.1 * math.log(40.0) + 1.0) + # amplitude scaled by the concentration factor + torch.testing.assert_close(cos_y**2 + sin_y**2, torch.full_like(cos_y, factor**2)) + # high-frequency dims untouched, low-frequency dims interpolated by 1/scaling_factor + torch.testing.assert_close(cos_y[:, 0] / factor, cos[:, 0]) + angle_y = torch.atan2(sin_y[:, rope // 2 - 1], cos_y[:, rope // 2 - 1]) + angle = torch.atan2(sin[:, rope // 2 - 1], cos[:, rope // 2 - 1]) + torch.testing.assert_close(angle_y[:64], angle[:64] / 40.0, atol=1e-4, rtol=0) + + +@pytest.mark.parametrize("mscale_all_dim", [0.0, 1.0]) +def test_mla_yarn_forward_backward(mscale_all_dim): + torch.manual_seed(0) + mla = MultiLatentAttention( + HIDDEN, + HEADS, + params_dtype=DTYPE, + rope_scaling_factor=40.0, + original_max_position_embeddings=64, + mscale_all_dim=mscale_all_dim, + **MLA_KWARGS, + ) + m = 0.1 * mscale_all_dim * math.log(40.0) + 1.0 + qk_head_dim = MLA_KWARGS["qk_nope_head_dim"] + MLA_KWARGS["qk_rope_head_dim"] + assert mla.core_attention.unfused_attention.softmax_scale == pytest.approx( + m * m / math.sqrt(qk_head_dim) + ) + x = _input() + out = mla(x) + assert out.shape == x.shape + out.sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() + + @pytest.mark.parametrize("shared", [False, True], ids=["no_shared", "shared"]) @pytest.mark.parametrize("grouped", [False, True], ids=["ungrouped", "grouped"]) def test_moe_forward_backward(shared, grouped): diff --git a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py index f8c01f75d6..341850fddb 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py +++ b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py @@ -16,6 +16,7 @@ matching the Megatron fused kernel semantics. """ +import math from typing import Optional, Tuple import torch @@ -28,7 +29,44 @@ except ImportError: HAVE_TRITON = False -__all__ = ["build_rope_tables", "apply_mla_rope_q", "apply_mla_rope_kv"] +__all__ = [ + "build_rope_tables", + "apply_mla_rope_q", + "apply_mla_rope_kv", + "yarn_mscale", + "yarn_concentration_factor", +] + + +def _yarn_correction_dim(num_rotations, dim, base, max_pos): + return (dim * math.log(max_pos / (num_rotations * 2 * math.pi))) / (2 * math.log(base)) + + +def _yarn_correction_range(beta_fast, beta_slow, dim, base, max_pos, round_to_int=True): + low = _yarn_correction_dim(beta_fast, dim, base, max_pos) + high = _yarn_correction_dim(beta_slow, dim, base, max_pos) + if round_to_int: + low, high = math.floor(low), math.ceil(high) + return max(low, 0), min(high, dim - 1) + + +def _yarn_linear_ramp(low, high, dim, device): + if low == high: + high += 0.001 + ramp = (torch.arange(dim, dtype=torch.float32, device=device) - low) / (high - low) + return torch.clamp(ramp, 0, 1) + + +def yarn_mscale(scale: float, mscale: float = 1.0) -> float: + """YaRN attention temperature factor ``0.1 * mscale * ln(scale) + 1`` (1 for scale <= 1).""" + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +def yarn_concentration_factor(scaling_factor: float, mscale: float, mscale_all_dim: float) -> float: + """Factor multiplied into cos/sin tables (as in Megatron-Core).""" + return yarn_mscale(scaling_factor, mscale) / yarn_mscale(scaling_factor, mscale_all_dim) def build_rope_tables( @@ -36,15 +74,33 @@ def build_rope_tables( emb_dim: int, base: float = 10000.0, device: Optional[torch.device] = None, + scaling_factor: Optional[float] = None, + original_max_position_embeddings: int = 4096, + beta_fast: float = 32.0, + beta_slow: float = 1.0, + mscale: float = 1.0, + mscale_all_dim: float = 0.0, ) -> Tuple[torch.Tensor, torch.Tensor]: - """cos/sin tables of shape ``[seq_len, emb_dim]`` (fp32, NeoX duplicated halves).""" - inv_freq = 1.0 / ( - base ** (torch.arange(0, emb_dim, 2, dtype=torch.float32, device=device) / emb_dim) - ) + """cos/sin tables of shape ``[seq_len, emb_dim]`` (fp32, NeoX duplicated halves). + + With ``scaling_factor`` set, frequencies follow YaRN (NTK-by-parts ramp between + ``beta_fast``/``beta_slow`` rotations over ``original_max_position_embeddings``) and the + tables are scaled by the YaRN concentration factor. + """ + exponent = torch.arange(0, emb_dim, 2, dtype=torch.float32, device=device) / emb_dim + inv_freq = 1.0 / (base**exponent) + factor = 1.0 + if scaling_factor is not None: + low, high = _yarn_correction_range( + beta_fast, beta_slow, emb_dim, base, original_max_position_embeddings + ) + extra_mask = 1.0 - _yarn_linear_ramp(low, high, emb_dim // 2, device) + inv_freq = (inv_freq / scaling_factor) * (1 - extra_mask) + inv_freq * extra_mask + factor = yarn_concentration_factor(scaling_factor, mscale, mscale_all_dim) t = torch.arange(seq_len, device=device, dtype=torch.float32) freqs = torch.outer(t, inv_freq) freqs = torch.cat([freqs, freqs], dim=-1) - return torch.cos(freqs).contiguous(), torch.sin(freqs).contiguous() + return (torch.cos(freqs) * factor).contiguous(), (torch.sin(freqs) * factor).contiguous() if HAVE_TRITON: diff --git a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py index 0ddf2f75b2..ca9c5ed3a7 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py +++ b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py @@ -4,6 +4,7 @@ """Multi-Latent Attention (MLA) block as used in DeepSeekV3.""" +import math from typing import Optional, Union import torch @@ -14,6 +15,7 @@ apply_mla_rope_kv, apply_mla_rope_q, build_rope_tables, + yarn_mscale, ) __all__ = ["MultiLatentAttention"] @@ -61,9 +63,22 @@ class MultiLatentAttention(torch.nn.Module): epsilon of the latent RMSNorms (matches DeepSeekV3). rotary_base : float, default = 10000.0 RoPE base. + rope_scaling_factor : float, optional + YaRN context-extension factor; ``None`` disables YaRN. + original_max_position_embeddings : int, default = 4096 + pre-extension context length (YaRN). + beta_fast : float, default = 32.0 + YaRN high-frequency rotation bound. + beta_slow : float, default = 1.0 + YaRN low-frequency rotation bound. + mscale : float, default = 1.0 + YaRN mscale of the rope part. + mscale_all_dim : float, default = 0.0 + YaRN mscale of all dims; sets the default softmax scale to + ``m**2 / sqrt(qk head dim)`` with ``m = 0.1 * mscale_all_dim * ln(factor) + 1``. softmax_scale : float, optional - softmax scale; defaults to ``1/sqrt(qk head dim)`` inside - :class:`DotProductAttention`. + softmax scale; defaults to ``1/sqrt(qk head dim)`` (times the YaRN + ``m**2`` when YaRN is enabled). qkv_format : str, default = "sbhd" layout of the input/output tensors. params_dtype : torch.dtype, optional @@ -87,6 +102,12 @@ def __init__( attn_mask_type: str = "causal", layernorm_epsilon: float = 1e-6, rotary_base: float = 10000.0, + rope_scaling_factor: Optional[float] = None, + original_max_position_embeddings: int = 4096, + beta_fast: float = 32.0, + beta_slow: float = 1.0, + mscale: float = 1.0, + mscale_all_dim: float = 0.0, softmax_scale: Optional[float] = None, qkv_format: str = "sbhd", params_dtype: Optional[torch.dtype] = None, @@ -140,8 +161,20 @@ def __init__( ) self.rotary_base = rotary_base + self._yarn_kwargs = dict( + scaling_factor=rope_scaling_factor, + original_max_position_embeddings=original_max_position_embeddings, + beta_fast=beta_fast, + beta_slow=beta_slow, + mscale=mscale, + mscale_all_dim=mscale_all_dim, + ) self._rope_tables: Optional[tuple] = None + if softmax_scale is None and rope_scaling_factor is not None: + m = yarn_mscale(rope_scaling_factor, mscale_all_dim) + softmax_scale = m * m / math.sqrt(self.qk_head_dim) + self.core_attention = DotProductAttention( num_attention_heads, kv_channels=(self.qk_head_dim, v_head_dim), @@ -156,7 +189,11 @@ def __init__( def _rope_tables_for(self, seq_len: int, device: torch.device): if self._rope_tables is None or self._rope_tables[0].shape[0] < seq_len: self._rope_tables = build_rope_tables( - seq_len, self.qk_rope_head_dim, base=self.rotary_base, device=device + seq_len, + self.qk_rope_head_dim, + base=self.rotary_base, + device=device, + **self._yarn_kwargs, ) cos, sin = self._rope_tables return cos[:seq_len], sin[:seq_len] diff --git a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py index f41ec0061c..aa2fab232d 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py +++ b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py @@ -59,6 +59,12 @@ class DeepSeekV3Layer(torch.nn.Module): "attention_dropout", "attn_mask_type", "rotary_base", + "rope_scaling_factor", + "original_max_position_embeddings", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", "softmax_scale", "qkv_format", "tp_group", From 7eaefd97fc0790f0a1805b0d3f098fe911c45dc5 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 12:10:54 +0200 Subject: [PATCH 16/41] Drop tests/pytorch/attention/mla_rope_utils.py shim; use models.deepseek_v3.mla_rope directly Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/mla_rope_utils.py | 42 ------------------- .../attention/test_linear_mxfp8_attention.py | 29 +++++++------ 2 files changed, 17 insertions(+), 54 deletions(-) delete mode 100644 tests/pytorch/attention/mla_rope_utils.py diff --git a/tests/pytorch/attention/mla_rope_utils.py b/tests/pytorch/attention/mla_rope_utils.py deleted file mode 100644 index d022757886..0000000000 --- a/tests/pytorch/attention/mla_rope_utils.py +++ /dev/null @@ -1,42 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Compat shim: the MLA RoPE kernels moved to -``transformer_engine.pytorch.models.deepseek_v3.mla_rope``.""" - -import torch - -from transformer_engine.pytorch.models.deepseek_v3.mla_rope import ( # noqa: F401 - HAVE_TRITON, - apply_mla_rope_kv, - apply_mla_rope_q, - build_rope_tables, -) - -HEAD_DIM_ROPE = 64 -HEAD_DIM_NOPE = 128 -HEAD_DIM_V = 128 -ROTARY_BASE = 10000 - - -def apply_mla_rope( - q: torch.Tensor, - kv: torch.Tensor, - k_pos_emb: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, - head_dim_v: int = HEAD_DIM_V, - base: int = ROTARY_BASE, - cos_table: torch.Tensor | None = None, - sin_table: torch.Tensor | None = None, -): - if cos_table is None or sin_table is None: - cos_table, sin_table = build_rope_tables( - q.shape[0], head_dim_rope, base=base, device=q.device - ) - q = apply_mla_rope_q(q, cos_table, sin_table, head_dim_nope, head_dim_rope) - k, v = apply_mla_rope_kv( - kv, k_pos_emb, cos_table, sin_table, head_dim_nope, head_dim_rope, head_dim_v - ) - return q, k, v diff --git a/tests/pytorch/attention/test_linear_mxfp8_attention.py b/tests/pytorch/attention/test_linear_mxfp8_attention.py index f1bba7bc9a..95770a6bb8 100644 --- a/tests/pytorch/attention/test_linear_mxfp8_attention.py +++ b/tests/pytorch/attention/test_linear_mxfp8_attention.py @@ -36,7 +36,11 @@ _current_file = pathlib.Path(__file__).resolve() sys.path = [str(_current_file.parent.parent)] + sys.path from utils import ModelConfig, compare_and_assert, get_available_attention_backends -from mla_rope_utils import apply_mla_rope, build_rope_tables +from transformer_engine.pytorch.models.deepseek_v3.mla_rope import ( + apply_mla_rope_kv, + apply_mla_rope_q, + build_rope_tables, +) try: @@ -183,6 +187,13 @@ def _run_projections( return q_flat, kv_flat, q, kv, k_pos_emb +def _apply_rope(q, kv, k_pos_emb, rope_tables): + cos, sin = rope_tables + q = apply_mla_rope_q(q, cos, sin, HEAD_DIM_NOPE, HEAD_DIM_ROPE) + k, v = apply_mla_rope_kv(kv, k_pos_emb, cos, sin, HEAD_DIM_NOPE, HEAD_DIM_ROPE, HEAD_DIM_V) + return q, k, v + + def _run_forward_bf16( modules: tuple, x: torch.Tensor, @@ -190,7 +201,7 @@ def _run_forward_bf16( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: q_proj, kv_proj, dpa, out_linear = modules _, _, q, kv, k_pos_emb = _run_projections(q_proj, kv_proj, x) - q, k, v = apply_mla_rope(q, kv, k_pos_emb, cos_table=rope_tables[0], sin_table=rope_tables[1]) + q, k, v = _apply_rope(q, kv, k_pos_emb, rope_tables) attn_out = dpa(q, k, v, qkv_format="sbhd") return q, k, v, out_linear(attn_out.view(x.shape[0], x.shape[1], HIDDEN_SIZE)) @@ -212,13 +223,7 @@ def _run_forward_mxfp8( x, is_first_microbatch, ) - q, k, v = apply_mla_rope( - q, - kv, - k_pos_emb, - cos_table=rope_tables[0], - sin_table=rope_tables[1], - ) + q, k, v = _apply_rope(q, kv, k_pos_emb, rope_tables) attn_out = dpa(q, k, v, qkv_format="sbhd") out = out_linear( attn_out.view(x.shape[0], x.shape[1], HIDDEN_SIZE), @@ -292,7 +297,7 @@ def test_accuracy(self, batch_size: int, seq_len: int) -> None: _set_seed() baseline_modules, mxfp8_modules = _build_modules() x = torch.randn(seq_len, batch_size, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") - rope_tables = build_rope_tables(seq_len, device=x.device) + rope_tables = build_rope_tables(seq_len, HEAD_DIM_ROPE, device=x.device) q_bf16, k_bf16, v_bf16, out_bf16 = _run_forward_bf16(baseline_modules, x, rope_tables) q_mxfp8, k_mxfp8, v_mxfp8, out_mxfp8 = _run_forward_mxfp8( @@ -378,7 +383,7 @@ def test_backward(self, batch_size: int, seq_len: int) -> None: device="cuda", requires_grad=True, ) - rope_tables = build_rope_tables(seq_len, device=x.device) + rope_tables = build_rope_tables(seq_len, HEAD_DIM_ROPE, device=x.device) *_, out_mxfp8 = _run_forward_mxfp8(mxfp8_modules, x, fp8_recipe, rope_tables) out_mxfp8.sum().backward() @@ -412,7 +417,7 @@ def test_performance(self, batch_size: int, seq_len: int) -> None: device="cuda", requires_grad=True, ) - rope_tables = build_rope_tables(seq_len, device=x.device) + rope_tables = build_rope_tables(seq_len, HEAD_DIM_ROPE, device=x.device) mxfp8_fprop_ms, mxfp8_bprop_ms = _benchmark_training_step( _run_forward_mxfp8, mxfp8_modules, x, fp8_recipe, rope_tables From cef2e39607767a03954824de5f257ee02b7bcb0b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 12:26:37 +0200 Subject: [PATCH 17/41] Distributed models test: single full DeepSeekV3Layer EP-vs-local numerical comparison Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_models.py | 91 ++++++++++--------------- 1 file changed, 37 insertions(+), 54 deletions(-) diff --git a/tests/pytorch/distributed/run_models.py b/tests/pytorch/distributed/run_models.py index ff5e233503..8bfdd7bfb3 100644 --- a/tests/pytorch/distributed/run_models.py +++ b/tests/pytorch/distributed/run_models.py @@ -11,7 +11,7 @@ import torch.distributed as dist from transformer_engine.pytorch.ep import ep_bootstrap, ep_finalize, release_symm_mem_pool -from transformer_engine.pytorch.models import DeepSeekV3Layer, DeepSeekV3MoE +from transformer_engine.pytorch.models import DeepSeekV3Layer HIDDEN = 256 MOE_FFN = 128 @@ -65,64 +65,72 @@ def setUpClass(cls): recv_capacity_per_rank=_recv_capacity(cls.ep_size), ) - def _make_moe(self, ep: bool, shared: bool = True) -> DeepSeekV3MoE: - return DeepSeekV3MoE( + def _make_layer(self, ep: bool) -> DeepSeekV3Layer: + return DeepSeekV3Layer( HIDDEN, - moe_ffn_hidden_size=MOE_FFN, + HEADS, num_experts=self.num_experts, + moe_ffn_hidden_size=MOE_FFN, topk=TOP_K, - shared_expert_ffn_hidden_size=SHARED_FFN if shared else None, + shared_expert_ffn_hidden_size=SHARED_FFN, params_dtype=DTYPE, ep_group=self.ep_group if ep else None, ep_max_tokens_per_rank=TOKENS_PER_RANK if ep else None, ep_recv_capacity_per_rank=_recv_capacity(self.ep_size) if ep else None, + **MLA_KWARGS, ) - def _copy_local_expert_weights(self, ep_moe: DeepSeekV3MoE, ref: DeepSeekV3MoE) -> None: + def _copy_weights(self, ep_layer: DeepSeekV3Layer, ref: DeepSeekV3Layer) -> None: + ref_params = dict(ref.named_parameters()) + ref_bufs = dict(ref.named_buffers()) with torch.no_grad(): - ep_moe.gate.weight.copy_(ref.gate.weight) - if ref.shared_expert is not None: - for dst, src in zip( - ep_moe.shared_expert.parameters(), ref.shared_expert.parameters() - ): - dst.copy_(src) - ep_fc1, _, ep_fc2 = ep_moe.experts - ref_fc1, _, ref_fc2 = ref.experts + for name, p in ep_layer.named_parameters(): + if not name.startswith("mlp.experts."): + p.copy_(ref_params[name]) + for name, b in ep_layer.named_buffers(): + if name in ref_bufs and b.shape == ref_bufs[name].shape: + b.copy_(ref_bufs[name]) + ep_fc1, _, ep_fc2 = ep_layer.mlp.experts + ref_fc1, _, ref_fc2 = ref.mlp.experts for local_e in range(NUM_LOCAL_EXPERTS): global_e = self.rank * NUM_LOCAL_EXPERTS + local_e getattr(ep_fc1, f"weight{local_e}").copy_(getattr(ref_fc1, f"weight{global_e}")) getattr(ep_fc2, f"weight{local_e}").copy_(getattr(ref_fc2, f"weight{global_e}")) - def test_moe_ep_matches_local(self): - """EP MoE must match the single-GPU (all-experts-local) path numerically.""" + def test_layer_ep_matches_local(self): + """Full DeepSeekV3Layer with EP must match the all-experts-local layer numerically.""" torch.manual_seed(0) - ref = self._make_moe(ep=False) + ref = self._make_layer(ep=False) _broadcast_params(ref) - ep_moe = self._make_moe(ep=True) - self._copy_local_expert_weights(ep_moe, ref) + ep_layer = self._make_layer(ep=True) + self._copy_weights(ep_layer, ref) torch.manual_seed(1234 + self.rank) - x = torch.randn(TOKENS_PER_RANK, HIDDEN, dtype=DTYPE, device="cuda") + x = torch.randn(TOKENS_PER_RANK // 2, 2, HIDDEN, dtype=DTYPE, device="cuda") x_ep = x.clone().requires_grad_(True) x_ref = x.clone().requires_grad_(True) - out_ep = ep_moe(x_ep) + out_ep = ep_layer(x_ep) out_ref = ref(x_ref) + self.assertEqual(out_ep.shape, x.shape) torch.testing.assert_close(out_ep, out_ref, rtol=0.05, atol=0.05) grad_out = torch.randn_like(out_ep) out_ep.backward(grad_out) out_ref.backward(grad_out) torch.testing.assert_close(x_ep.grad, x_ref.grad, rtol=0.05, atol=0.05) - torch.testing.assert_close( - ep_moe.gate.weight.grad, ref.gate.weight.grad, rtol=0.1, atol=0.1 - ) + + ref_params = dict(ref.named_parameters()) + for name, p in ep_layer.named_parameters(): + if name.startswith("mlp.experts.") or p.grad is None: + continue + torch.testing.assert_close(p.grad, ref_params[name].grad, rtol=0.1, atol=0.1, msg=name) # A local expert's wgrad on its owner rank equals the sum of the # reference wgrads over all ranks. all_reduce is collective, so every # rank must reduce every expert's grad (in the same order). - ep_fc1, _, ep_fc2 = ep_moe.experts - ref_fc1, _, ref_fc2 = ref.experts + ep_fc1, _, ep_fc2 = ep_layer.mlp.experts + ref_fc1, _, ref_fc2 = ref.mlp.experts for ep_fc, ref_fc in ((ep_fc1, ref_fc1), (ep_fc2, ref_fc2)): ref_grads = [ getattr(ref_fc, f"weight{e}").grad.float().clone() for e in range(self.num_experts) @@ -134,37 +142,12 @@ def test_moe_ep_matches_local(self): ep_grad = getattr(ep_fc, f"weight{local_e}").grad.float() torch.testing.assert_close(ep_grad, ref_grads[global_e], rtol=0.1, atol=0.1) - counts = ep_moe._last_tokens_per_expert.clone() + counts = ep_layer.mlp._last_tokens_per_expert.clone() dist.all_reduce(counts) self.assertEqual(counts.sum().item(), self.ep_size * TOKENS_PER_RANK * TOP_K) - def test_layer_ep_forward_backward(self): - """Full DeepSeekV3Layer smoke test with an EP MoE block.""" - torch.manual_seed(10 + self.rank) - layer = DeepSeekV3Layer( - HIDDEN, - HEADS, - num_experts=self.num_experts, - moe_ffn_hidden_size=MOE_FFN, - topk=TOP_K, - shared_expert_ffn_hidden_size=SHARED_FFN, - params_dtype=DTYPE, - ep_group=self.ep_group, - ep_max_tokens_per_rank=TOKENS_PER_RANK, - ep_recv_capacity_per_rank=_recv_capacity(self.ep_size), - **MLA_KWARGS, - ) - x = torch.randn( - TOKENS_PER_RANK // 2, 2, HIDDEN, dtype=DTYPE, device="cuda", requires_grad=True - ) - out = layer(x) - self.assertEqual(out.shape, x.shape) - out.sum().backward() - self.assertIsNotNone(x.grad) - self.assertTrue(torch.isfinite(x.grad).all()) - - layer.mlp.update_expert_bias() - self.assertTrue(torch.isfinite(layer.mlp.expert_bias).all()) + ep_layer.mlp.update_expert_bias() + self.assertTrue(torch.isfinite(ep_layer.mlp.expert_bias).all()) def _init_distributed(): From 5453fea6f0f760d6e3c05e4cbe653a412889ae96 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 12:32:05 +0200 Subject: [PATCH 18/41] run_models.py: plain main() instead of unittest, simplify launcher Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_models.py | 224 ++++++++++--------- tests/pytorch/distributed/run_test_models.sh | 19 +- tests/pytorch/distributed/test_models.py | 2 +- 3 files changed, 118 insertions(+), 127 deletions(-) diff --git a/tests/pytorch/distributed/run_models.py b/tests/pytorch/distributed/run_models.py index 8bfdd7bfb3..ed41f020f9 100644 --- a/tests/pytorch/distributed/run_models.py +++ b/tests/pytorch/distributed/run_models.py @@ -5,7 +5,6 @@ import os import sys -import unittest import torch import torch.distributed as dist @@ -46,111 +45,94 @@ def _broadcast_params(module: torch.nn.Module) -> None: dist.broadcast(t.detach(), src=0) -class TestDeepSeekEP(unittest.TestCase): - @classmethod - def setUpClass(cls): - if _device_sm() < 90: - raise unittest.SkipTest(f"NCCL EP requires SM>=90 (got SM{_device_sm()})") - cls.rank = dist.get_rank() - cls.ep_size = dist.get_world_size() - cls.num_experts = NUM_LOCAL_EXPERTS * cls.ep_size - world_pg = dist.distributed_c10d._get_default_group() - cls.ep_group = dist.new_group(ranks=list(range(world_pg.size())), backend="nccl") - ep_bootstrap( - cls.ep_group, - num_experts=cls.num_experts, - max_tokens_per_rank=TOKENS_PER_RANK, - hidden_dim=HIDDEN, - num_topk=TOP_K, - recv_capacity_per_rank=_recv_capacity(cls.ep_size), - ) - - def _make_layer(self, ep: bool) -> DeepSeekV3Layer: - return DeepSeekV3Layer( - HIDDEN, - HEADS, - num_experts=self.num_experts, - moe_ffn_hidden_size=MOE_FFN, - topk=TOP_K, - shared_expert_ffn_hidden_size=SHARED_FFN, - params_dtype=DTYPE, - ep_group=self.ep_group if ep else None, - ep_max_tokens_per_rank=TOKENS_PER_RANK if ep else None, - ep_recv_capacity_per_rank=_recv_capacity(self.ep_size) if ep else None, - **MLA_KWARGS, - ) - - def _copy_weights(self, ep_layer: DeepSeekV3Layer, ref: DeepSeekV3Layer) -> None: - ref_params = dict(ref.named_parameters()) - ref_bufs = dict(ref.named_buffers()) - with torch.no_grad(): - for name, p in ep_layer.named_parameters(): - if not name.startswith("mlp.experts."): - p.copy_(ref_params[name]) - for name, b in ep_layer.named_buffers(): - if name in ref_bufs and b.shape == ref_bufs[name].shape: - b.copy_(ref_bufs[name]) - ep_fc1, _, ep_fc2 = ep_layer.mlp.experts - ref_fc1, _, ref_fc2 = ref.mlp.experts - for local_e in range(NUM_LOCAL_EXPERTS): - global_e = self.rank * NUM_LOCAL_EXPERTS + local_e - getattr(ep_fc1, f"weight{local_e}").copy_(getattr(ref_fc1, f"weight{global_e}")) - getattr(ep_fc2, f"weight{local_e}").copy_(getattr(ref_fc2, f"weight{global_e}")) - - def test_layer_ep_matches_local(self): - """Full DeepSeekV3Layer with EP must match the all-experts-local layer numerically.""" - torch.manual_seed(0) - ref = self._make_layer(ep=False) - _broadcast_params(ref) - ep_layer = self._make_layer(ep=True) - self._copy_weights(ep_layer, ref) - - torch.manual_seed(1234 + self.rank) - x = torch.randn(TOKENS_PER_RANK // 2, 2, HIDDEN, dtype=DTYPE, device="cuda") - x_ep = x.clone().requires_grad_(True) - x_ref = x.clone().requires_grad_(True) - - out_ep = ep_layer(x_ep) - out_ref = ref(x_ref) - self.assertEqual(out_ep.shape, x.shape) - torch.testing.assert_close(out_ep, out_ref, rtol=0.05, atol=0.05) - - grad_out = torch.randn_like(out_ep) - out_ep.backward(grad_out) - out_ref.backward(grad_out) - torch.testing.assert_close(x_ep.grad, x_ref.grad, rtol=0.05, atol=0.05) - - ref_params = dict(ref.named_parameters()) +def _make_layer(ep_group, ep_size: int, num_experts: int) -> DeepSeekV3Layer: + ep = ep_group is not None + return DeepSeekV3Layer( + HIDDEN, + HEADS, + num_experts=num_experts, + moe_ffn_hidden_size=MOE_FFN, + topk=TOP_K, + shared_expert_ffn_hidden_size=SHARED_FFN, + params_dtype=DTYPE, + ep_group=ep_group, + ep_max_tokens_per_rank=TOKENS_PER_RANK if ep else None, + ep_recv_capacity_per_rank=_recv_capacity(ep_size) if ep else None, + **MLA_KWARGS, + ) + + +def _copy_weights(ep_layer: DeepSeekV3Layer, ref: DeepSeekV3Layer, rank: int) -> None: + ref_params = dict(ref.named_parameters()) + ref_bufs = dict(ref.named_buffers()) + with torch.no_grad(): for name, p in ep_layer.named_parameters(): - if name.startswith("mlp.experts.") or p.grad is None: - continue - torch.testing.assert_close(p.grad, ref_params[name].grad, rtol=0.1, atol=0.1, msg=name) - - # A local expert's wgrad on its owner rank equals the sum of the - # reference wgrads over all ranks. all_reduce is collective, so every - # rank must reduce every expert's grad (in the same order). + if not name.startswith("mlp.experts."): + p.copy_(ref_params[name]) + for name, b in ep_layer.named_buffers(): + if name in ref_bufs and b.shape == ref_bufs[name].shape: + b.copy_(ref_bufs[name]) ep_fc1, _, ep_fc2 = ep_layer.mlp.experts ref_fc1, _, ref_fc2 = ref.mlp.experts - for ep_fc, ref_fc in ((ep_fc1, ref_fc1), (ep_fc2, ref_fc2)): - ref_grads = [ - getattr(ref_fc, f"weight{e}").grad.float().clone() for e in range(self.num_experts) - ] - for g in ref_grads: - dist.all_reduce(g) - for local_e in range(NUM_LOCAL_EXPERTS): - global_e = self.rank * NUM_LOCAL_EXPERTS + local_e - ep_grad = getattr(ep_fc, f"weight{local_e}").grad.float() - torch.testing.assert_close(ep_grad, ref_grads[global_e], rtol=0.1, atol=0.1) - - counts = ep_layer.mlp._last_tokens_per_expert.clone() - dist.all_reduce(counts) - self.assertEqual(counts.sum().item(), self.ep_size * TOKENS_PER_RANK * TOP_K) - - ep_layer.mlp.update_expert_bias() - self.assertTrue(torch.isfinite(ep_layer.mlp.expert_bias).all()) - - -def _init_distributed(): + for local_e in range(NUM_LOCAL_EXPERTS): + global_e = rank * NUM_LOCAL_EXPERTS + local_e + getattr(ep_fc1, f"weight{local_e}").copy_(getattr(ref_fc1, f"weight{global_e}")) + getattr(ep_fc2, f"weight{local_e}").copy_(getattr(ref_fc2, f"weight{global_e}")) + + +def test_layer_ep_matches_local(rank: int, ep_size: int, ep_group) -> None: + """Full DeepSeekV3Layer with EP must match the all-experts-local layer numerically.""" + num_experts = NUM_LOCAL_EXPERTS * ep_size + torch.manual_seed(0) + ref = _make_layer(None, ep_size, num_experts) + _broadcast_params(ref) + ep_layer = _make_layer(ep_group, ep_size, num_experts) + _copy_weights(ep_layer, ref, rank) + + torch.manual_seed(1234 + rank) + x = torch.randn(TOKENS_PER_RANK // 2, 2, HIDDEN, dtype=DTYPE, device="cuda") + x_ep = x.clone().requires_grad_(True) + x_ref = x.clone().requires_grad_(True) + + out_ep = ep_layer(x_ep) + out_ref = ref(x_ref) + assert out_ep.shape == x.shape + torch.testing.assert_close(out_ep, out_ref, rtol=0.05, atol=0.05) + + grad_out = torch.randn_like(out_ep) + out_ep.backward(grad_out) + out_ref.backward(grad_out) + torch.testing.assert_close(x_ep.grad, x_ref.grad, rtol=0.05, atol=0.05) + + ref_params = dict(ref.named_parameters()) + for name, p in ep_layer.named_parameters(): + if name.startswith("mlp.experts.") or p.grad is None: + continue + torch.testing.assert_close(p.grad, ref_params[name].grad, rtol=0.1, atol=0.1, msg=name) + + # A local expert's wgrad on its owner rank equals the sum of the + # reference wgrads over all ranks. all_reduce is collective, so every + # rank must reduce every expert's grad (in the same order). + ep_fc1, _, ep_fc2 = ep_layer.mlp.experts + ref_fc1, _, ref_fc2 = ref.mlp.experts + for ep_fc, ref_fc in ((ep_fc1, ref_fc1), (ep_fc2, ref_fc2)): + ref_grads = [getattr(ref_fc, f"weight{e}").grad.float().clone() for e in range(num_experts)] + for g in ref_grads: + dist.all_reduce(g) + for local_e in range(NUM_LOCAL_EXPERTS): + global_e = rank * NUM_LOCAL_EXPERTS + local_e + ep_grad = getattr(ep_fc, f"weight{local_e}").grad.float() + torch.testing.assert_close(ep_grad, ref_grads[global_e], rtol=0.1, atol=0.1) + + counts = ep_layer.mlp._last_tokens_per_expert.clone() + dist.all_reduce(counts) + assert counts.sum().item() == ep_size * TOKENS_PER_RANK * TOP_K + + ep_layer.mlp.update_expert_bias() + assert torch.isfinite(ep_layer.mlp.expert_bias).all() + + +def main() -> int: dist.init_process_group(backend="nccl") torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) try: @@ -160,13 +142,33 @@ def _init_distributed(): except (ImportError, RuntimeError): pass + rank = dist.get_rank() + ep_size = dist.get_world_size() + if _device_sm() < 90: + if rank == 0: + print(f"NCCL EP requires SM>=90 (got SM{_device_sm()}); skipping.") + dist.destroy_process_group() + return 0 + + ep_group = dist.new_group(ranks=list(range(ep_size)), backend="nccl") + ep_bootstrap( + ep_group, + num_experts=NUM_LOCAL_EXPERTS * ep_size, + max_tokens_per_rank=TOKENS_PER_RANK, + hidden_dim=HIDDEN, + num_topk=TOP_K, + recv_capacity_per_rank=_recv_capacity(ep_size), + ) + try: + test_layer_ep_matches_local(rank, ep_size, ep_group) + print(f"[rank {rank}] PASSED") + finally: + dist.barrier() + ep_finalize() + release_symm_mem_pool() + dist.destroy_process_group() + return 0 + if __name__ == "__main__": - _init_distributed() - suite = unittest.TestLoader().loadTestsFromTestCase(TestDeepSeekEP) - result = unittest.TextTestRunner(stream=sys.stdout, verbosity=2).run(suite) - dist.barrier() - ep_finalize() - release_symm_mem_pool() - dist.destroy_process_group() - sys.exit(0 if result.wasSuccessful() else 1) + sys.exit(main()) diff --git a/tests/pytorch/distributed/run_test_models.sh b/tests/pytorch/distributed/run_test_models.sh index e3411ee7d9..b0230b10ac 100644 --- a/tests/pytorch/distributed/run_test_models.sh +++ b/tests/pytorch/distributed/run_test_models.sh @@ -31,22 +31,11 @@ export NCCL_EP_JIT_CACHE_DIR mkdir -p "$NCCL_EP_JIT_CACHE_DIR" SCRIPT="${SCRIPT_DIR}/run_models.py" -LOG="stdout_models.txt" echo "=== Running ${SCRIPT} on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ===" setsid timeout --foreground --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \ - torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" \ - "${SCRIPT}" 2>&1 | tee "${LOG}" -RC=${PIPESTATUS[0]} + torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" "${SCRIPT}" +RC=$? pkill -9 -f "tests/pytorch/distributed/run_models.py" 2>/dev/null || true - -RET=0 -if [ "${RC}" -ne 0 ]; then echo "torchrun exited with ${RC}"; RET=1; fi -if grep -qE "(^|]:)FAILED|(^|]:)Traceback" "${LOG}"; then RET=1; fi -if ! grep -qE "Ran [0-9]+ test|^OK$" "${LOG}"; then - echo "ERROR: no test summary — likely hang or early crash" - RET=1 -fi -if [ -z "${KEEP_EP_LOGS:-}" ]; then rm -f "${LOG}"; fi - -exit $RET +if [ "${RC}" -ne 0 ]; then echo "torchrun exited with ${RC}"; fi +exit $RC diff --git a/tests/pytorch/distributed/test_models.py b/tests/pytorch/distributed/test_models.py index 83213bd6a8..d7aa3a478e 100644 --- a/tests/pytorch/distributed/test_models.py +++ b/tests/pytorch/distributed/test_models.py @@ -19,7 +19,7 @@ def test_multi_process_deepseek_ep(): timeout_s = int(os.environ.get("NVTE_TEST_EP_TIMEOUT_S", "180")) proc = subprocess.run( ["bash", str(LAUNCHER)], - env={**os.environ, "KEEP_EP_LOGS": "1", "TEST_TIMEOUT_S": str(timeout_s)}, + env={**os.environ, "TEST_TIMEOUT_S": str(timeout_s)}, timeout=timeout_s + 30, check=False, ) From f24835b5ab7072665fc45e8c2261db2b6fe3469d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 12:37:56 +0200 Subject: [PATCH 19/41] run_models.py: fail hard instead of swallowing symm-mem/cleanup errors Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_models.py | 22 +++++++++------------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/distributed/run_models.py b/tests/pytorch/distributed/run_models.py index ed41f020f9..a77b18e205 100644 --- a/tests/pytorch/distributed/run_models.py +++ b/tests/pytorch/distributed/run_models.py @@ -135,12 +135,9 @@ def test_layer_ep_matches_local(rank: int, ep_size: int, ep_group) -> None: def main() -> int: dist.init_process_group(backend="nccl") torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) - try: - from torch.distributed import _symmetric_memory as _symm_mem + from torch.distributed import _symmetric_memory as _symm_mem - _symm_mem.set_backend("NCCL") - except (ImportError, RuntimeError): - pass + _symm_mem.set_backend("NCCL") rank = dist.get_rank() ep_size = dist.get_world_size() @@ -159,14 +156,13 @@ def main() -> int: num_topk=TOP_K, recv_capacity_per_rank=_recv_capacity(ep_size), ) - try: - test_layer_ep_matches_local(rank, ep_size, ep_group) - print(f"[rank {rank}] PASSED") - finally: - dist.barrier() - ep_finalize() - release_symm_mem_pool() - dist.destroy_process_group() + test_layer_ep_matches_local(rank, ep_size, ep_group) + print(f"[rank {rank}] PASSED") + + dist.barrier() + ep_finalize() + release_symm_mem_pool() + dist.destroy_process_group() return 0 From d66fc1c872d50d31bc82cf9f2a7d3e27fe41ef6d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 12:39:30 +0200 Subject: [PATCH 20/41] DeepSeekV3MoE docstring: ep_bootstrap must precede construction Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/models/deepseek_v3/moe.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index 3c182b4405..caa2359e5a 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -45,7 +45,8 @@ class DeepSeekV3MoE(torch.nn.Module): given, expert-parallel over NCCL (``ep_dispatch``/``ep_combine``). When expert parallelism is used, ``transformer_engine.pytorch.ep.ep_bootstrap`` - must be called once per process before the first forward, and inputs must + must be called once per process before constructing the module (it allocates + the ``EpBuffer`` in ``__init__``), and inputs must be bfloat16. Parameters From 80041fadbef3c9d39a673f59a0c972bab672b413 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 12:43:28 +0200 Subject: [PATCH 21/41] Distributed models test: launch torchrun directly from pytest, drop shell launcher Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_test_models.sh | 41 -------------------- tests/pytorch/distributed/test_models.py | 27 +++++++------ 2 files changed, 16 insertions(+), 52 deletions(-) delete mode 100644 tests/pytorch/distributed/run_test_models.sh diff --git a/tests/pytorch/distributed/run_test_models.sh b/tests/pytorch/distributed/run_test_models.sh deleted file mode 100644 index b0230b10ac..0000000000 --- a/tests/pytorch/distributed/run_test_models.sh +++ /dev/null @@ -1,41 +0,0 @@ -#!/bin/bash -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. -# -# Launcher for tests/pytorch/distributed/run_models.py (model-specific layers). Auto-detects GPU count. - -set -uo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" - -DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) -if [ "${DETECTED_GPUS}" -lt 2 ]; then - echo "DeepSeek EP test requires >= 2 GPUs (found ${DETECTED_GPUS}); SKIPPING." - exit 0 -fi - -# NCCL EP requires active NVLink P2P among ranks on the node. -if ! nvidia-smi nvlink --status 2>/dev/null | grep -qE 'Link [0-9]+:.*GB/s'; then - echo "No NVLink between GPUs (PCIe-only fabric); NCCL EP is unsupported here. SKIPPING." - exit 0 -fi - -NUM_RANKS="${NVTE_TEST_EP_NUM_RANKS:-${DETECTED_GPUS}}" -if [ "${NUM_RANKS}" -gt 8 ]; then NUM_RANKS=8; fi - -TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-180}" - -: ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"} -export NCCL_EP_JIT_CACHE_DIR -mkdir -p "$NCCL_EP_JIT_CACHE_DIR" - -SCRIPT="${SCRIPT_DIR}/run_models.py" - -echo "=== Running ${SCRIPT} on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ===" -setsid timeout --foreground --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \ - torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" "${SCRIPT}" -RC=$? -pkill -9 -f "tests/pytorch/distributed/run_models.py" 2>/dev/null || true -if [ "${RC}" -ne 0 ]; then echo "torchrun exited with ${RC}"; fi -exit $RC diff --git a/tests/pytorch/distributed/test_models.py b/tests/pytorch/distributed/test_models.py index d7aa3a478e..1b96eae2aa 100644 --- a/tests/pytorch/distributed/test_models.py +++ b/tests/pytorch/distributed/test_models.py @@ -1,7 +1,6 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. -"""Pytest driver — spawns run_models.py (model-specific layers, multi-GPU) under torchrun.""" import os import subprocess @@ -11,16 +10,22 @@ import torch TEST_ROOT = Path(__file__).parent.resolve() -LAUNCHER = TEST_ROOT / "run_test_models.sh" +NUM_PROCS = min(8, torch.cuda.device_count()) +LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"] -@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="DeepSeek EP requires >= 2 GPUs") -def test_multi_process_deepseek_ep(): - timeout_s = int(os.environ.get("NVTE_TEST_EP_TIMEOUT_S", "180")) - proc = subprocess.run( - ["bash", str(LAUNCHER)], - env={**os.environ, "TEST_TIMEOUT_S": str(timeout_s)}, - timeout=timeout_s + 30, - check=False, +def _has_nvlink() -> bool: + # NCCL EP falls back to the network transport and deadlocks on PCIe-only nodes. + out = subprocess.run( + ["nvidia-smi", "nvlink", "--status"], capture_output=True, text=True, check=False + ).stdout + return "GB/s" in out + + +@pytest.mark.skipif(NUM_PROCS < 2, reason="EP requires >= 2 GPUs") +@pytest.mark.skipif(not _has_nvlink(), reason="NCCL EP requires NVLink") +def test_deepseek_layer_ep(): + result = subprocess.run( + LAUNCH_CMD + [str(TEST_ROOT / "run_models.py")], env=os.environ, check=False, timeout=300 ) - assert proc.returncode == 0, f"DeepSeek EP test suite failed (rc={proc.returncode})" + assert result.returncode == 0 From 2475a9113d988f784ec3b649799431224a899223 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 13:04:23 +0200 Subject: [PATCH 22/41] Add DeepSeekV3Layer to test_sanity; pad per-expert rows for quantized grouped GEMM in local MoE path Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_sanity.py | 36 +++++++++++++++++++ .../pytorch/models/deepseek_v3/moe.py | 34 ++++++++++++++---- 2 files changed, 63 insertions(+), 7 deletions(-) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index c9b620fa1e..835813d5c6 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -35,6 +35,7 @@ is_bf16_available, ) from transformer_engine.common import recipe +from transformer_engine.pytorch.models import DeepSeekV3Layer from transformer_engine.pytorch.cpp_extensions import general_gemm from transformer_engine.pytorch.tensor.utils import replace_raw_data from transformer_engine.pytorch.module import is_module_grouped_tensor_path_supported @@ -736,6 +737,41 @@ def test_sanity_layernorm_mlp( _test_sanity_common(block, dtype, config, fp8_recipe, skip_wgrad, skip_dgrad, microbatching) +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) +@pytest.mark.parametrize("model", ["small"]) +@pytest.mark.parametrize("skip_wgrad", all_boolean) +@pytest.mark.parametrize("moe", all_boolean) +def test_sanity_deepseek_v3_layer(dtype, fp8_recipe, model, skip_wgrad, moe): + config = model_configs[model] + + if fp8_recipe is not None: + if not is_fp8_supported(config): + pytest.skip("Model config does not support FP8") + if fp8_recipe.nvfp4() and dtype == torch.float16: + pytest.skip("FP16 output for NVFP4 not supported") + + mlp_kwargs = ( + dict(num_experts=4, topk=2, moe_ffn_hidden_size=32, shared_expert_ffn_hidden_size=32) + if moe + else dict(ffn_hidden_size=4 * config.hidden_size) + ) + block = DeepSeekV3Layer( + config.hidden_size, + config.num_heads, + q_lora_rank=16, + kv_lora_rank=16, + qk_nope_head_dim=16, + qk_rope_head_dim=16, + v_head_dim=16, + params_dtype=dtype, + device="cuda", + **mlp_kwargs, + ) + + _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad) + + @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("model", ["small"]) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index caa2359e5a..60080d8816 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -11,7 +11,15 @@ import transformer_engine.pytorch.ops as te_ops from transformer_engine.pytorch.router import fused_topk_with_score_function -from transformer_engine.pytorch.permutation import moe_permute_with_probs, moe_unpermute +from transformer_engine.pytorch.permutation import ( + moe_permute_and_pad_with_probs, + moe_permute_with_probs, + moe_unpermute, +) +from transformer_engine.pytorch.quantization import ( + FP8GlobalStateManager, + get_align_size_for_quantization, +) __all__ = ["DeepSeekV3MoE"] @@ -189,14 +197,24 @@ def _forward_local(self, tokens: torch.Tensor) -> torch.Tensor: tokens_per_expert = routing_map.sum(dim=0) self._last_tokens_per_expert = tokens_per_expert.detach() - num_out = tokens.shape[0] * self.topk - permuted, permuted_probs, row_id_map = moe_permute_with_probs( - tokens, probs, routing_map, num_out_tokens=num_out - ) + # Quantized grouped GEMMs need every expert's row count aligned. + align = 1 + if FP8GlobalStateManager.is_fp8_enabled(): + align = get_align_size_for_quantization(FP8GlobalStateManager.get_fp8_recipe()) + if align > 1: + permuted, permuted_probs, row_id_map, pad_offsets, tokens_per_expert = ( + moe_permute_and_pad_with_probs(tokens, probs, routing_map, tokens_per_expert, align) + ) + else: + permuted, permuted_probs, row_id_map = moe_permute_with_probs( + tokens, probs, routing_map, num_out_tokens=tokens.shape[0] * self.topk + ) + pad_offsets = None # The fused grouped MLP requires the total row count to be a multiple # of 128; rows beyond sum(tokens_per_expert) fall outside every group. - pad = (-num_out) % 128 + num_rows = permuted.shape[0] + pad = (-num_rows) % 128 if pad: permuted = torch.nn.functional.pad(permuted, (0, 0, 0, pad)) permuted_probs = torch.nn.functional.pad(permuted_probs, (0, pad)) @@ -204,7 +222,9 @@ def _forward_local(self, tokens: torch.Tensor) -> torch.Tensor: out = self.experts( permuted, tokens_per_expert, permuted_probs.to(tokens.dtype), tokens_per_expert ) - return moe_unpermute(out[:num_out], row_id_map, restore_shape=tokens.shape) + return moe_unpermute( + out[:num_rows], row_id_map, restore_shape=tokens.shape, pad_offsets=pad_offsets + ) def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: from transformer_engine.pytorch.ep import ep_dispatch, ep_combine From babc5e7102b5251f978a83a312a31ee2c0a9a060 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 13:21:08 +0200 Subject: [PATCH 23/41] Tests: drop fwd/bwd smoke tests covered by sanity, trim sanity combos; expose MLA softmax_scale; clean docstrings Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_models.py | 101 ++++-------------- tests/pytorch/test_sanity.py | 10 +- .../pytorch/models/deepseek_v3/mla_rope.py | 20 ++-- .../deepseek_v3/multi_latent_attention.py | 3 +- 4 files changed, 37 insertions(+), 97 deletions(-) diff --git a/tests/pytorch/test_models.py b/tests/pytorch/test_models.py index 678c104bb8..cd1903e79f 100644 --- a/tests/pytorch/test_models.py +++ b/tests/pytorch/test_models.py @@ -8,11 +8,7 @@ import torch from transformer_engine.pytorch.utils import deinterleave_glu_tensor -from transformer_engine.pytorch.models import ( - DeepSeekV3Layer, - DeepSeekV3MoE, - MultiLatentAttention, -) +from transformer_engine.pytorch.models import DeepSeekV3MoE, MultiLatentAttention SEQ_LEN = 128 BATCH = 2 @@ -96,16 +92,6 @@ def run(fmt): torch.testing.assert_close(grads_t[2], pos_leaf.grad, rtol=1e-5, atol=1e-5) -def test_mla_forward_backward(): - torch.manual_seed(0) - mla = MultiLatentAttention(HIDDEN, HEADS, params_dtype=DTYPE, **MLA_KWARGS) - x = _input() - out = mla(x) - assert out.shape == x.shape - out.sum().backward() - assert x.grad is not None and torch.isfinite(x.grad).all() - - def test_rope_tables_yarn(): from transformer_engine.pytorch.models.deepseek_v3 import mla_rope @@ -128,8 +114,7 @@ def test_rope_tables_yarn(): @pytest.mark.parametrize("mscale_all_dim", [0.0, 1.0]) -def test_mla_yarn_forward_backward(mscale_all_dim): - torch.manual_seed(0) +def test_mla_yarn_softmax_scale(mscale_all_dim): mla = MultiLatentAttention( HIDDEN, HEADS, @@ -141,27 +126,23 @@ def test_mla_yarn_forward_backward(mscale_all_dim): ) m = 0.1 * mscale_all_dim * math.log(40.0) + 1.0 qk_head_dim = MLA_KWARGS["qk_nope_head_dim"] + MLA_KWARGS["qk_rope_head_dim"] - assert mla.core_attention.unfused_attention.softmax_scale == pytest.approx( - m * m / math.sqrt(qk_head_dim) - ) - x = _input() - out = mla(x) - assert out.shape == x.shape - out.sum().backward() - assert x.grad is not None and torch.isfinite(x.grad).all() + assert mla.softmax_scale == pytest.approx(m * m / math.sqrt(qk_head_dim)) @pytest.mark.parametrize("shared", [False, True], ids=["no_shared", "shared"]) @pytest.mark.parametrize("grouped", [False, True], ids=["ungrouped", "grouped"]) -def test_moe_forward_backward(shared, grouped): +@pytest.mark.parametrize("topk", [2, 4]) +def test_moe_matches_dense_reference(shared, grouped, topk): + """Routed output must equal the prob-weighted sum of the selected expert MLPs.""" torch.manual_seed(0) + num_experts = 4 moe = DeepSeekV3MoE( HIDDEN, moe_ffn_hidden_size=128, - num_experts=8, - topk=2, - num_groups=4 if grouped else None, - group_topk=2 if grouped else None, + num_experts=num_experts, + topk=topk, + num_groups=2 if grouped else None, + group_topk=topk // 2 if grouped else None, shared_expert_ffn_hidden_size=128 if shared else None, params_dtype=DTYPE, ) @@ -169,32 +150,13 @@ def test_moe_forward_backward(shared, grouped): out = moe(x) assert out.shape == x.shape out.sum().backward() - assert x.grad is not None and torch.isfinite(x.grad).all() - - counts = moe._last_tokens_per_expert - assert counts.sum().item() == SEQ_LEN * BATCH * 2 - bias_before = moe.expert_bias.clone() - moe.update_expert_bias() - assert not torch.equal(bias_before, moe.expert_bias) - - -def test_moe_matches_dense_reference(): - """topk == num_experts with uniform probs must reduce to a sum of expert MLPs.""" - torch.manual_seed(0) - num_experts = 4 - moe = DeepSeekV3MoE( - HIDDEN, - moe_ffn_hidden_size=128, - num_experts=num_experts, - topk=num_experts, - routed_scaling_factor=1.0, - params_dtype=DTYPE, - ) - x = _input(requires_grad=False) - out = moe(x) + assert torch.isfinite(x.grad).all() - tokens = x.reshape(-1, HIDDEN) + tokens = x.detach().reshape(-1, HIDDEN) probs, _ = moe._route(moe.gate(tokens).float()) + assert (probs > 0).sum(dim=1).eq(topk).all() + assert moe._last_tokens_per_expert.sum().item() == tokens.shape[0] * topk + fc1, _, fc2 = moe.experts ref = torch.zeros_like(tokens) for e in range(num_experts): @@ -203,29 +165,12 @@ def test_moe_matches_dense_reference(): gate_part, lin_part = (tokens @ w1.t()).chunk(2, dim=-1) act = torch.nn.functional.silu(gate_part.float()) * lin_part.float() ref += (act.to(DTYPE) * probs[:, e : e + 1].to(DTYPE)) @ w2.t() + if shared: + ref += moe.shared_expert(tokens) torch.testing.assert_close(out.reshape(-1, HIDDEN), ref, rtol=0.05, atol=0.05) - -@pytest.mark.parametrize("num_experts", [None, 8], ids=["dense", "moe"]) -def test_layer_forward_backward(num_experts): - torch.manual_seed(0) - layer = ( - DeepSeekV3Layer( - HIDDEN, - HEADS, - ffn_hidden_size=512, - num_experts=num_experts, - moe_ffn_hidden_size=128 if num_experts else None, - topk=2 if num_experts else None, - shared_expert_ffn_hidden_size=128 if num_experts else None, - params_dtype=DTYPE, - **MLA_KWARGS, - ) - if num_experts - else DeepSeekV3Layer(HIDDEN, HEADS, ffn_hidden_size=512, params_dtype=DTYPE, **MLA_KWARGS) - ) - x = _input() - out = layer(x) - assert out.shape == x.shape - out.sum().backward() - assert x.grad is not None and torch.isfinite(x.grad).all() + bias_before = moe.expert_bias.clone() + moe.update_expert_bias() + assert torch.isfinite(moe.expert_bias).all() + if topk < num_experts: + assert not torch.equal(bias_before, moe.expert_bias) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 835813d5c6..60d17e39af 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -739,11 +739,9 @@ def test_sanity_layernorm_mlp( @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) -@pytest.mark.parametrize("model", ["small"]) -@pytest.mark.parametrize("skip_wgrad", all_boolean) -@pytest.mark.parametrize("moe", all_boolean) -def test_sanity_deepseek_v3_layer(dtype, fp8_recipe, model, skip_wgrad, moe): - config = model_configs[model] +@pytest.mark.parametrize("moe", all_boolean, ids=["dense", "moe"]) +def test_sanity_deepseek_v3_layer(dtype, fp8_recipe, moe): + config = model_configs["small"] if fp8_recipe is not None: if not is_fp8_supported(config): @@ -769,7 +767,7 @@ def test_sanity_deepseek_v3_layer(dtype, fp8_recipe, model, skip_wgrad, moe): **mlp_kwargs, ) - _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad) + _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad=False) @pytest.mark.parametrize("dtype", param_types) diff --git a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py index 341850fddb..0aea284afd 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py +++ b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py @@ -4,17 +4,13 @@ """Fused MLA RoPE kernels (DeepSeekV3-style decoupled RoPE/NoPE). -Triton forward/backward kernels adapted from Megatron-LM -``megatron/core/fusions/fused_mla_yarn_rope_apply.py``. The query kernel -rotates the trailing ``head_dim_rope`` slice in place (no concat); the KV -kernel builds the final key (nope | broadcast-rotated shared rope head) and -value tensors in a single pass. Falls back to pure PyTorch when Triton is -unavailable or for the ``bshd`` layout (the Triton path is ``sbhd``-only). - -Rotation convention: the rope slice is read interleaved (as stored in -HF/Megatron DeepSeekV3 checkpoints) and written in NeoX half-split layout, -matching the Megatron fused kernel semantics. -""" +The query kernel rotates the trailing ``head_dim_rope`` slice in place; the KV +kernel builds the key (nope | broadcast-rotated shared rope head) and value +tensors in a single pass. Falls back to pure PyTorch when Triton is unavailable +or for the ``bshd`` layout. + +The rope slice is read interleaved (checkpoint layout) and written in NeoX +half-split layout.""" import math from typing import Optional, Tuple @@ -65,7 +61,7 @@ def yarn_mscale(scale: float, mscale: float = 1.0) -> float: def yarn_concentration_factor(scaling_factor: float, mscale: float, mscale_all_dim: float) -> float: - """Factor multiplied into cos/sin tables (as in Megatron-Core).""" + """Factor multiplied into cos/sin tables.""" return yarn_mscale(scaling_factor, mscale) / yarn_mscale(scaling_factor, mscale_all_dim) diff --git a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py index ca9c5ed3a7..4b4b1c9926 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py +++ b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py @@ -37,7 +37,7 @@ class MultiLatentAttention(torch.nn.Module): RoPE uses the fused MLA kernels from :mod:`.mla_rope` (in-place on the query rope slice, single-pass key/value assembly); the rope slice follows - the HF/Megatron DeepSeekV3 convention (interleaved weights, NeoX output). + the DeepSeekV3 checkpoint convention (interleaved weights, NeoX output). Parameters ---------- @@ -174,6 +174,7 @@ def __init__( if softmax_scale is None and rope_scaling_factor is not None: m = yarn_mscale(rope_scaling_factor, mscale_all_dim) softmax_scale = m * m / math.sqrt(self.qk_head_dim) + self.softmax_scale = softmax_scale self.core_attention = DotProductAttention( num_attention_heads, From 33549066ccd65e623e103975bd44664aae7c3a02 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 13:55:17 +0200 Subject: [PATCH 24/41] Rewrite DeepSeekV3MoE class docstring Signed-off-by: Pawel Gadzinski --- .../pytorch/models/deepseek_v3/moe.py | 28 ++++++++++--------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index 60080d8816..0b26933c4c 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -41,21 +41,23 @@ def _make_expert_mlp(num_experts, hidden_size, ffn_hidden_size, dtype, device): class DeepSeekV3MoE(torch.nn.Module): """ - DeepSeekV3-style Mixture of Experts block. + DeepSeekV3 Mixture-of-Experts block. - Routing uses the fused sigmoid router with aux-loss-free expert bias and - node-limited (grouped) top-k (``fused_topk_with_score_function``). Routed - experts run as a grouped SwiGLU MLP built from ``te.ops`` (fusable into a - single CuTe grouped-GEMM kernel); routing probabilities are applied - per-token inside the expert MLP, so unpermute/combine is a plain - accumulation. Token routing is either local - (``moe_permute_with_probs``/``moe_unpermute``) or, when ``ep_group`` is - given, expert-parallel over NCCL (``ep_dispatch``/``ep_combine``). + Each token is scored by a sigmoid router with a non-trainable expert bias + updated by ``update_expert_bias()`` (aux-loss-free load balancing) and, + optionally, group-limited routing: experts are split into ``num_groups`` + groups, the top ``group_topk`` groups are selected by their summed scores, + and the final ``topk`` experts are chosen only from those groups. Selected + tokens run through the routed experts, a SwiGLU MLP shared across experts + as a grouped GEMM, with the routing probability applied inside the MLP. An + optional shared expert (dense SwiGLU MLP) is added to every token. On + hardware that supports it the expert MLP runs as a single fused + grouped-GEMM kernel. - When expert parallelism is used, ``transformer_engine.pytorch.ep.ep_bootstrap`` - must be called once per process before constructing the module (it allocates - the ``EpBuffer`` in ``__init__``), and inputs must - be bfloat16. + Without ``ep_group`` all experts live on the local device. With + ``ep_group`` the experts are split across the group and tokens are + exchanged over NCCL; this requires ``ep_bootstrap`` to be called once per + process before constructing the module, and bfloat16 inputs. Parameters ---------- From 397733e75f8acb6f72bc6509b1baa5b9341aed9b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 13:59:02 +0200 Subject: [PATCH 25/41] DeepSeekV3MoE: drop ep_recv_capacity_per_rank and ep_alignment parameters Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_models.py | 1 - .../pytorch/models/deepseek_v3/moe.py | 25 ++++++++----------- .../models/deepseek_v3/transformer_layer.py | 2 -- 3 files changed, 10 insertions(+), 18 deletions(-) diff --git a/tests/pytorch/distributed/run_models.py b/tests/pytorch/distributed/run_models.py index a77b18e205..9561d20117 100644 --- a/tests/pytorch/distributed/run_models.py +++ b/tests/pytorch/distributed/run_models.py @@ -57,7 +57,6 @@ def _make_layer(ep_group, ep_size: int, num_experts: int) -> DeepSeekV3Layer: params_dtype=DTYPE, ep_group=ep_group, ep_max_tokens_per_rank=TOKENS_PER_RANK if ep else None, - ep_recv_capacity_per_rank=_recv_capacity(ep_size) if ep else None, **MLA_KWARGS, ) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index 0b26933c4c..b35cedab0f 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -24,6 +24,9 @@ __all__ = ["DeepSeekV3MoE"] +_EP_ALIGNMENT = 128 + + def _make_expert_mlp(num_experts, hidden_size, ffn_hidden_size, dtype, device): # GroupedLinear + ScaledSwiGLU + GroupedLinear fuses into a single CuTe # grouped MLP on supported hardware; elsewhere it runs as three ops with @@ -87,11 +90,6 @@ class DeepSeekV3MoE(torch.nn.Module): expert-parallel process group; enables the NCCL EP path. ep_max_tokens_per_rank : int, optional max local tokens per forward (required with EP). - ep_recv_capacity_per_rank : int, optional - receive-buffer capacity; defaults to - ``ep_size * ep_max_tokens_per_rank * topk``. - ep_alignment : int, default = 128 - per-expert row alignment of the EP receive buffer. """ def __init__( @@ -109,8 +107,6 @@ def __init__( device: Union[torch.device, str] = "cuda", ep_group: Optional[torch.distributed.ProcessGroup] = None, ep_max_tokens_per_rank: Optional[int] = None, - ep_recv_capacity_per_rank: Optional[int] = None, - ep_alignment: int = 128, ) -> None: super().__init__() @@ -165,19 +161,18 @@ def __init__( from transformer_engine.pytorch.ep import EpBuffer assert ep_max_tokens_per_rank is not None, "EP requires ep_max_tokens_per_rank." - if ep_recv_capacity_per_rank is None: - # Worst case plus per-expert alignment padding, rounded up to - # the multiple of 128 required by the fused grouped MLP. - cap = self.ep_size * ep_max_tokens_per_rank * topk - cap += num_local_experts * max(ep_alignment, 1) - ep_recv_capacity_per_rank = -(-cap // 128) * 128 + # Worst case plus per-expert alignment padding, rounded up to + # the multiple of 128 required by the fused grouped MLP. + cap = self.ep_size * ep_max_tokens_per_rank * topk + cap += num_local_experts * _EP_ALIGNMENT + cap = -(-cap // _EP_ALIGNMENT) * _EP_ALIGNMENT self.ep_buffer = EpBuffer( top_k=topk, max_tokens_per_rank=ep_max_tokens_per_rank, hidden_dim=hidden_size, num_local_experts=num_local_experts, - recv_capacity_per_rank=ep_recv_capacity_per_rank, - alignment=ep_alignment, + recv_capacity_per_rank=cap, + alignment=_EP_ALIGNMENT, device=device, ) diff --git a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py index aa2fab232d..ab11fe6394 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py +++ b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py @@ -81,8 +81,6 @@ class DeepSeekV3Layer(torch.nn.Module): "expert_bias_update_rate", "ep_group", "ep_max_tokens_per_rank", - "ep_recv_capacity_per_rank", - "ep_alignment", } ) From cc10354bf59d6342e411aadf38c74a8452a68ff1 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 14:01:30 +0200 Subject: [PATCH 26/41] DeepSeekV3MoE: build shared expert with the same SwiGLU MLP helper as routed experts Signed-off-by: Pawel Gadzinski --- .../pytorch/models/deepseek_v3/moe.py | 46 ++++++++----------- 1 file changed, 18 insertions(+), 28 deletions(-) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index b35cedab0f..e78e94fc2a 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -27,18 +27,22 @@ _EP_ALIGNMENT = 128 -def _make_expert_mlp(num_experts, hidden_size, ffn_hidden_size, dtype, device): - # GroupedLinear + ScaledSwiGLU + GroupedLinear fuses into a single CuTe - # grouped MLP on supported hardware; elsewhere it runs as three ops with - # the same API and checkpoint layout. +def _make_swiglu_mlp(hidden_size, ffn_hidden_size, dtype, device, num_experts=None): + """Dense SwiGLU MLP, or a grouped one (probs applied inside the activation) per expert. + + The grouped variant fuses into a single CuTe grouped MLP on supported hardware. + """ + common = dict(bias=False, dtype=dtype, device=device) + if num_experts is None: + return te_ops.Sequential( + te_ops.Linear(hidden_size, 2 * ffn_hidden_size, **common), + te_ops.SwiGLU(), + te_ops.Linear(ffn_hidden_size, hidden_size, **common), + ) return te_ops.Sequential( - te_ops.GroupedLinear( - num_experts, hidden_size, 2 * ffn_hidden_size, bias=False, dtype=dtype, device=device - ), + te_ops.GroupedLinear(num_experts, hidden_size, 2 * ffn_hidden_size, **common), te_ops.ScaledSwiGLU(glu_interleave_size=32), - te_ops.GroupedLinear( - num_experts, ffn_hidden_size, hidden_size, bias=False, dtype=dtype, device=device - ), + te_ops.GroupedLinear(num_experts, ffn_hidden_size, hidden_size, **common), ) @@ -132,28 +136,14 @@ def __init__( assert num_experts % self.ep_size == 0 num_local_experts = num_experts // self.ep_size - self.experts = _make_expert_mlp( - num_local_experts, hidden_size, moe_ffn_hidden_size, dtype, device + self.experts = _make_swiglu_mlp( + hidden_size, moe_ffn_hidden_size, dtype, device, num_experts=num_local_experts ) self.shared_expert = None if shared_expert_ffn_hidden_size is not None: - self.shared_expert = te_ops.Sequential( - te_ops.Linear( - hidden_size, - 2 * shared_expert_ffn_hidden_size, - bias=False, - dtype=dtype, - device=device, - ), - te_ops.SwiGLU(), - te_ops.Linear( - shared_expert_ffn_hidden_size, - hidden_size, - bias=False, - dtype=dtype, - device=device, - ), + self.shared_expert = _make_swiglu_mlp( + hidden_size, shared_expert_ffn_hidden_size, dtype, device ) self.ep_buffer = None From 2c1cd4c5fb8df0bc33e4efd59582a2f08245a7af Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 14:15:35 +0200 Subject: [PATCH 27/41] DeepSeekV3MoE EP path: count tokens per expert with scatter_add instead of syncing bincount Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/models/deepseek_v3/moe.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index e78e94fc2a..2bb47d49ce 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -221,10 +221,11 @@ def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: (tokens.shape[0], self.topk), dtype=torch.int64, device=tokens.device ) probs, topk_idx = self._route(self.gate(tokens).float(), topk_indices=topk_idx) - self._last_tokens_per_expert = torch.bincount( - topk_idx.flatten(), minlength=self.num_experts - ) - topk_weights = probs.gather(1, topk_idx).float() + flat_idx = topk_idx.flatten() + self._last_tokens_per_expert = torch.zeros( + self.num_experts, dtype=torch.long, device=tokens.device + ).scatter_add_(0, flat_idx, torch.ones_like(flat_idx)) + topk_weights = probs.gather(1, topk_idx) # Zero-filled recv/grad buffers: per-expert alignment padding lands # inside the grouped-GEMM m_splits, so uninitialized rows would poison From 6bae1ba180081748fbc54bcd1d69387519135dfa Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 15:18:05 +0200 Subject: [PATCH 28/41] Docs: list model-specific layers inline on the PyTorch API page; group standard layers, autocast and other utilities Signed-off-by: Pawel Gadzinski --- docs/api/pytorch.rst | 35 +++++++++++++++++++++++------------ docs/api/pytorch_models.rst | 19 ------------------- 2 files changed, 23 insertions(+), 31 deletions(-) delete mode 100644 docs/api/pytorch_models.rst diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 8b2b742372..497f414aee 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -6,6 +6,11 @@ PyTorch ======= +.. autoapiclass:: transformer_engine.pytorch.autocast(enabled=True, calibrating=False, recipe=None, amax_reduction_group=None) + +Standard layers +--------------- + .. autoapiclass:: transformer_engine.pytorch.Linear(in_features, out_features, bias=True, **kwargs) :members: forward, set_tensor_parallel_group @@ -34,20 +39,34 @@ PyTorch .. autoapiclass:: transformer_engine.pytorch.TransformerLayer(hidden_size, ffn_hidden_size, num_attention_heads, **kwargs) :members: forward, set_context_parallel_group, set_tensor_parallel_group +Model-specific layers +--------------------- + +DeepSeek-V3 +^^^^^^^^^^^ + +.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3Layer(hidden_size, num_attention_heads, **kwargs) + :members: forward + +.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3MoE(hidden_size, moe_ffn_hidden_size, num_experts, **kwargs) + :members: forward, update_expert_bias + +.. autoapiclass:: transformer_engine.pytorch.models.MultiLatentAttention(hidden_size, num_attention_heads, **kwargs) + :members: forward + +Other +----- + .. autoapiclass:: transformer_engine.pytorch.dot_product_attention.inference.InferenceParams(max_batch_size, max_sequence_length) :members: reset, allocate_memory, pre_step, get_seqlens_pre_step, convert_paged_to_nonpaged, step .. autoapiclass:: transformer_engine.pytorch.CudaRNGStatesTracker() :members: reset, get_states, set_states, add, fork - -.. autoapiclass:: transformer_engine.pytorch.autocast(enabled=True, calibrating=False, recipe=None, amax_reduction_group=None) - .. autoapifunction:: transformer_engine.pytorch.quantized_model_init .. autoapifunction:: transformer_engine.pytorch.checkpoint - .. autoapifunction:: transformer_engine.pytorch.make_graphed_callables .. autoapifunction:: transformer_engine.pytorch.get_cpu_offload_context @@ -62,14 +81,6 @@ PyTorch .. autoapifunction:: transformer_engine.pytorch.deinterleave_glu_tensor -Models ------- - -.. toctree:: - :maxdepth: 1 - - pytorch_models - Data types ---------- diff --git a/docs/api/pytorch_models.rst b/docs/api/pytorch_models.rst deleted file mode 100644 index 2cde879ffb..0000000000 --- a/docs/api/pytorch_models.rst +++ /dev/null @@ -1,19 +0,0 @@ -.. - Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - - See LICENSE for license information. - -Models -====== - -DeepSeek-V3 ------------ - -.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3Layer(hidden_size, num_attention_heads, **kwargs) - :members: forward - -.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3MoE(hidden_size, moe_ffn_hidden_size, num_experts, **kwargs) - :members: forward, update_expert_bias - -.. autoapiclass:: transformer_engine.pytorch.models.MultiLatentAttention(hidden_size, num_attention_heads, **kwargs) - :members: forward From 610a1e236b10c0f275dc18f6e81d19b96be1958d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 15:56:42 +0200 Subject: [PATCH 29/41] Lint: use dict literals in models.deepseek_v3 Signed-off-by: Pawel Gadzinski --- .../pytorch/models/deepseek_v3/moe.py | 2 +- .../models/deepseek_v3/multi_latent_attention.py | 16 ++++++++-------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index 2bb47d49ce..42f5048e6e 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -32,7 +32,7 @@ def _make_swiglu_mlp(hidden_size, ffn_hidden_size, dtype, device, num_experts=No The grouped variant fuses into a single CuTe grouped MLP on supported hardware. """ - common = dict(bias=False, dtype=dtype, device=device) + common = {"bias": False, "dtype": dtype, "device": device} if num_experts is None: return te_ops.Sequential( te_ops.Linear(hidden_size, 2 * ffn_hidden_size, **common), diff --git a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py index 4b4b1c9926..56b4d3d0d1 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py +++ b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py @@ -161,14 +161,14 @@ def __init__( ) self.rotary_base = rotary_base - self._yarn_kwargs = dict( - scaling_factor=rope_scaling_factor, - original_max_position_embeddings=original_max_position_embeddings, - beta_fast=beta_fast, - beta_slow=beta_slow, - mscale=mscale, - mscale_all_dim=mscale_all_dim, - ) + self._yarn_kwargs = { + "scaling_factor": rope_scaling_factor, + "original_max_position_embeddings": original_max_position_embeddings, + "beta_fast": beta_fast, + "beta_slow": beta_slow, + "mscale": mscale, + "mscale_all_dim": mscale_all_dim, + } self._rope_tables: Optional[tuple] = None if softmax_scale is None and rope_scaling_factor is not None: From b86d8f9c7ab1506c185f802a88bfd16f409addf3 Mon Sep 17 00:00:00 2001 From: William Yang <77467499+wilyan09007@users.noreply.github.com> Date: Fri, 4 Sep 2026 11:35:54 -0400 Subject: [PATCH 30/41] [PyTorch] Decline fused grouped MLP when the backward format is not E4M3 (#3352) * [PyTorch] Decline fused grouped MLP when the backward format is not E4M3 The fused grouped MLP packs the incoming activation gradient by reinterpreting its storage as E4M3, conditioned only on NVFP4 and never on the FP8 format. Under MXFP8BlockScaling(fp8_format=Format.HYBRID) the backward quantizers emit E5M2, so those bytes are read as the wrong format rather than converted, and every gradient out of the fusion is wrong. The forward pass is unaffected, so this shows up as a model that trains too slowly instead of one that fails. Fall back to the unfused ops when the recipe's backward format is not E4M3, and raise instead of reinterpreting if such a gradient reaches the kernel path. Signed-off-by: William * [PyTorch] Gate the grouped MLP backward format check on MXFP8 fp8_format describes the FP8 formats of an MXFP8 recipe. NVFP4BlockScaling carries one too, pinned to E4M3, but its gradients are quantized to FP4 and the value says nothing about them, so testing it for an NVFP4 recipe reached the right answer for the wrong reason. Restrict the check to recipes where it means something. The runtime check at the pack site is unchanged: it reads the grad output quantizer's own dtype on the non-NVFP4 branch, not the recipe. Signed-off-by: William * Test grouped MLP format fusion with real ops Signed-off-by: Przemek Tredak --------- Signed-off-by: William Signed-off-by: Przemek Tredak Co-authored-by: Przemek Tredak --- tests/pytorch/test_grouped_mlp.py | 40 +++++++++++++++++++ .../pytorch/ops/fused/grouped_mlp.py | 17 +++++++- 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index d48f7afae6..1d173c76d8 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -1067,6 +1067,46 @@ def train_step( class TestGroupedMLPFusedOp: """Tests for grouped MLP fused op""" + def test_fusion_requires_supported_grad_output_format(self, monkeypatch) -> None: + """Fuse E4M3 MXFP8 and NVFP4, but decline MXFP8 with an E5M2 backward.""" + from transformer_engine.common.recipe import Format, MXFP8BlockScaling, NVFP4BlockScaling + + fused_op_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMGLU + monkeypatch.setattr(fused_op_cls, "is_supported", classmethod(lambda cls: True)) + + fc1 = te.ops.GroupedLinear(1, 64, 128, bias=False, device="cuda") + activation = te.ops.ScaledSwiGLU(glu_interleave_size=32) + fc2 = te.ops.GroupedLinear(1, 64, 64, bias=False, device="cuda") + ops = [fc1, activation, fc2] + + def fuse(recipe): + return grouped_mlp_module.fuse_grouped_mlp_ops( + ops, + recipe=recipe, + fused_op_cls=fused_op_cls, + ) + + def assert_fused(recipe): + fused_ops = fuse(recipe) + assert len(fused_ops) == 1 + fused_op = fused_ops[0] + assert isinstance(fused_op, fused_op_cls) + assert list(fused_op.basic_ops) == ops + + hybrid = MXFP8BlockScaling(fp8_format=Format.HYBRID) + assert fuse(hybrid) is ops + + e4m3 = MXFP8BlockScaling(fp8_format=Format.E4M3) + assert_fused(e4m3) + + # NVFP4 quantizes gradients to FP4, so the FP8 format must not gate it. Forcing the + # lookup to E5M2 is what an NVFP4 recipe would hit if the check were not MXFP8-only. + monkeypatch.setattr( + grouped_mlp_module, "get_fp8_torch_dtype", lambda *_, **__: torch.float8_e5m2 + ) + nvfp4 = NVFP4BlockScaling(disable_rht=False) + assert_fused(nvfp4) + @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list) @pytest.mark.parametrize("single_grouped_weight", (False, True)) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 61f80b9d9f..2128bae1ec 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -17,7 +17,7 @@ from packaging.version import Version as PkgVersion import transformer_engine_torch as tex -from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType +from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed_weight import ( @@ -27,7 +27,7 @@ finalize_weight_grads, ) from ...module.base import _2X_ACC_WGRAD -from ...quantization import Recipe +from ...quantization import Recipe, get_fp8_torch_dtype from ...tensor import NVFP4Quantizer, NVFP4Tensor, NVFP4TensorStorage, Quantizer from ...tensor.grouped_tensor import GroupedTensor from ...tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor @@ -877,6 +877,12 @@ def fuse_grouped_mlp_ops( # NVFP4 fused grouped MLP uses graph-safe grouped quantize, which currently requires RHT. if recipe.nvfp4() and recipe.disable_rht: return ops + # The fused MXFP8 backward reinterprets the grad output's storage as E4M3, so an E5M2 + # backward format would have its gradients misread rather than converted. This declines + # MXFP8 with Format.HYBRID. fp8_format does not describe NVFP4 gradients, so NVFP4 is + # excluded from the check rather than relying on its value. + if recipe.mxfp8() and get_fp8_torch_dtype(recipe, fprop_tensor=False) != torch.float8_e4m3fn: + return ops if activation_op_types is None: activation_op_types = [ScaledSwiGLU, ScaledClampedQGeGLU] if _cudnn_frontend_supports_grouped_gemm_situglu(): @@ -2026,6 +2032,13 @@ def fuser_backward( or isinstance(fc1_weight_param, NVFP4Tensor) or isinstance(fc2_weight_param, NVFP4Tensor) ) + if not use_nvfp4 and fc2_grad_output_quantizer.dtype != DType.kFloat8E4M3: + # The pack below reinterprets the grad output's storage as E4M3 rather than + # converting it, so anything else would be read as the wrong format. + raise RuntimeError( + "Fused grouped MLP backward requires an E4M3 grad output, but the recipe " + f"produced {fc2_grad_output_quantizer.dtype}." + ) data_dtype = torch.float4_e2m1fn_x2 if use_nvfp4 else torch.float8_e4m3fn scale_view_dtype = torch.float8_e4m3fn if use_nvfp4 else torch.float8_e8m0fnu sf_vec_size = NVFP4_BLOCK_SCALING_SIZE if use_nvfp4 else MXFP8_BLOCK_SCALING_SIZE From e61b66d6efd6c0f34c517cc333f57ce853bf4633 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Fri, 4 Sep 2026 09:21:15 -0700 Subject: [PATCH 31/41] [PyTorch] Fix: Resolve EP symm-mem window offset for both old and new torch version (#3466) [PyTorch] Resolve EP symm-mem window offset against both torch symm-mem layouts Signed-off-by: Phuong Nguyen --- .../pytorch/csrc/extensions/ep.cpp | 23 +++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp index c74d4ddb6d..97bc70bddd 100644 --- a/transformer_engine/pytorch/csrc/extensions/ep.cpp +++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp @@ -59,6 +59,21 @@ std::atomic g_zero_copy_enabled{false}; // is not symm-mem-backed; the backend treats it as "no window, use staged copy". constexpr NVTECommWindow kNoWindow = {nullptr, 0}; +#ifdef NCCL_HAS_SYMMEM_SUPPORT +// Offset of a symm-mem allocation relative to the start of its NCCL window. +// Newer torch places the signal pad at the front of the allocation and exposes +// get_window_offset() for it; on older torch the window starts at the buffer +// base, where get_offset() is already window-relative. +template +auto symm_mem_window_offset(T* sm, int) -> decltype(sm->get_window_offset()) { + return sm->get_window_offset(); +} +template +size_t symm_mem_window_offset(T* sm, ...) { + return sm->get_offset(); +} +#endif + // Resolve ``t`` to an NCCL symm-mem window for the zero-copy one-sided path. // Returns ``kNoWindow`` when symm-mem support isn't compiled in, zero-copy is // disabled, no group is set, or ``t`` isn't symm-mem-backed; callers pass the @@ -78,11 +93,11 @@ NVTECommWindow maybe_make_window(const at::Tensor& t) { NVTE_CHECK(nccl_sm != nullptr, "Symm-mem backend mismatch: expected NCCLSymmetricMemory. Set the backend to " "\"NCCL\" before allocating EP payload buffers."); - // NCCL EP consumes window-relative offsets (the NCCL window starts at the signal pad, - // not at the buffer base). get_window_offset() = buffer_offset + get_offset(); add - // ``t``'s own storage offset for slice/view positioning. + // NCCL EP consumes window-relative offsets. Add ``t``'s own storage offset so a + // slice/view of a symm-mem allocation (e.g. the scale region carved from a shared + // recv buffer) resolves to its true position in the window. const uint64_t offset = - static_cast(nccl_sm->get_window_offset()) + + static_cast(symm_mem_window_offset(nccl_sm, 0)) + static_cast(t.storage_offset()) * static_cast(t.element_size()); return NVTECommWindow{static_cast(nccl_sm->get_window()), offset}; #else From 1f25d6ab1f30938d143bbdaf88a27c9fc045159e Mon Sep 17 00:00:00 2001 From: Wei Wang <143543872+nWEIdia@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:31:48 -0700 Subject: [PATCH 32/41] [NCCL][PyTorch] Fix test-fusible-ops-file-rendezvous-bug (#3478) * world_group() used a single hardcoded init_method="file:///tmp/rdzv", shared across every world_size this test parametrizes over ([device_count(), 1, 2]), and neither world_group() nor test_distributed_fuser_ops ever calls destroy_process_group() or removes the file afterwards. A world_size=N run's leftover FileStore content can then corrupt a differently-shaped world_size=M run that reuses the same path in the same pytest session, surfacing as a confusing NCCL bootstrap failure: torch.distributed.DistBackendError: NCCL error ... ncclOsSocketPollConnect: connect to ... Connection refused, exceeded error retry count after 35 attempts instead of a clear rendezvous error. Confirmed 100% deterministic: running the full file fresh (no pre-existing /tmp/rdzv) still fails test_distributed_fuser_ops[2] every time, because the [4] parametrization (which runs first) leaves /tmp/rdzv behind for [2] to trip over. This was previously masked by older bundled NCCL (2.30.7), which apparently tolerated the stale/mismatched FileStore well enough to still succeed; a newer NCCL (2.31.2) surfaces it as a hard failure. See the investigation writeup for the full comparison: Fix: key the rendezvous path by world_size (file:///tmp/rdzv_test_fusible_ops_{world_size}), and defensively remove any pre-existing file at that path in test_distributed_fuser_ops before launching each subprocess job, to also cover a leftover file from an earlier crashed run of the same world_size. Verified on a GB200 node NCCL 2.31.2, the container that previously failed 2 of 3 parametrizations): all 3 world_size parametrizations now pass, including two runs of the full file back to back with no manual cleanup in between. Co-Authored-By: Claude Sonnet 5 Signed-off-by: Wei Wang * Use unique rendezvous files in fusible ops tests Signed-off-by: Przemek Tredak --------- Signed-off-by: Wei Wang Signed-off-by: Przemek Tredak Co-authored-by: Claude Sonnet 5 Co-authored-by: Przemek Tredak --- tests/pytorch/distributed/test_fusible_ops.py | 12 +++++++----- .../distributed/test_fusible_ops_with_userbuffers.py | 8 ++++++-- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/pytorch/distributed/test_fusible_ops.py b/tests/pytorch/distributed/test_fusible_ops.py index d733286093..c314bf5bb6 100644 --- a/tests/pytorch/distributed/test_fusible_ops.py +++ b/tests/pytorch/distributed/test_fusible_ops.py @@ -12,6 +12,7 @@ import pathlib import subprocess import sys +import tempfile from typing import Optional import pytest @@ -58,7 +59,8 @@ def world_group() -> torch.distributed.ProcessGroup: torch.cuda.set_device(rank) group = torch.distributed.init_process_group( "nccl", - init_method="file:///tmp/rdzv", + # Each parallel job must use a fresh FileStore shared by only its ranks. + init_method=f"file://{os.environ['NVTE_TEST_RDZV_PATH']}", world_size=world_size, rank=rank, ) @@ -1053,10 +1055,10 @@ def test_distributed_fuser_ops(world_size: int) -> None: current_file, "--parallel", ] - result = subprocess.run( - command, - check=True, - ) + with tempfile.TemporaryDirectory(prefix="te-test-fusible-ops-") as temp_dir: + env = dict(os.environ) + env["NVTE_TEST_RDZV_PATH"] = str(pathlib.Path(temp_dir) / "rdzv") + subprocess.run(command, check=True, env=env) def main() -> None: diff --git a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py index 07dffebf5f..38f49a96cb 100644 --- a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py +++ b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py @@ -12,6 +12,7 @@ import pathlib import subprocess import sys +import tempfile import pytest import torch @@ -107,7 +108,8 @@ def world_group() -> torch.distributed.ProcessGroup: torch.cuda.set_device(local_rank) group = torch.distributed.init_process_group( "nccl", - init_method="file:///tmp/rdzv", + # Each parallel job must use a fresh FileStore shared by only its ranks. + init_method=f"file://{os.environ['NVTE_TEST_RDZV_PATH']}", world_size=world_size, rank=rank, device_id=torch.device(f"cuda:{local_rank}"), @@ -471,7 +473,9 @@ def test_fuser_ops_with_userbuffers( env["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" # Launch parallel job - run_distributed(command, env=env) + with tempfile.TemporaryDirectory(prefix="te-test-fusible-ops-userbuffers-") as temp_dir: + env["NVTE_TEST_RDZV_PATH"] = str(pathlib.Path(temp_dir) / "rdzv") + run_distributed(command, env=env) def main() -> None: From 80a89adc272887fba3999a6a4b434347327d2c74 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:47:15 -0700 Subject: [PATCH 33/41] Reduce Grouped MLP Fuser CPU Overhead (#3410) * Reduce grouped MLP fuser CPU overhead Reuse fused operation plans when full activation recompute changes grad mode, and avoid redundant CUDA current-device discovery for grouped MLP stream lookups. Co-authored-by: Ting-Yang Kao Signed-off-by: Zhongbo Zhu * resolve comments Signed-off-by: Zhongbo Zhu * fix cutedsl wgrad crash Signed-off-by: tingyangk * resolve comments Signed-off-by: Zhongbo Zhu --------- Signed-off-by: Zhongbo Zhu Signed-off-by: tingyangk Co-authored-by: Ting-Yang Kao --- tests/pytorch/test_fusible_ops.py | 167 ++++++++++++++++++ .../pytorch/ops/fused/grouped_mlp.py | 4 +- transformer_engine/pytorch/ops/fuser.py | 84 +++++---- 3 files changed, 219 insertions(+), 36 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 6cd1fc3065..2adce717c9 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -146,6 +146,173 @@ def maybe_skip_quantization( pytest.skip("NVFP4 quantization is only supported with BF16 data") +def test_operation_fuser_caches_plans_by_grad_requirement(monkeypatch) -> None: + """Cache and restore fusion plans for checkpoint forward and recompute.""" + + # Count fusion-plan construction without depending on any particular real + # fusion implementation. Each distinct fusion configuration invokes this + # hook once, while a cache hit must bypass it entirely. + fusion_calls = 0 + + def track_fusion(ops, *, recipe): # pylint: disable=unused-argument + nonlocal fusion_calls + fusion_calls += 1 + # Preserve the operation list so this hook observes plan construction + # without changing the topology under test. + return ops + + # The fusion registries are class attributes shared by every OperationFuser. + # pytest's monkeypatch fixture restores all three after the test, preventing + # this synthetic fusion function from leaking into other tests. Keep only a + # joint forward-backward fusion hook so each plan build has one countable + # callback and no registered TE fusion can affect the result. + monkeypatch.setattr(OperationFuser, "forward_backward_fusion_functions", [track_fusion]) + monkeypatch.setattr(OperationFuser, "forward_fusion_functions", []) + monkeypatch.setattr(OperationFuser, "backward_fusion_functions", []) + + # One Identity op is enough to exercise the cache. With one basic op, + # first_op_requiring_backward has an intentionally simple interpretation: + # 0: backward starts at the Identity op; + # 1: the boundary is past the only op, so no backward work is required. + fuser = OperationFuser([te_ops.Identity()]) + x = torch.ones(1, requires_grad=True) + # maybe_fuse_ops expects one extra-input collection per basic op. Identity + # has no extra inputs, so its collection is an empty tuple. + extra_inputs = [()] + + # Phase 1: the original checkpointed forward runs with grad disabled. This + # is the first invocation, so the fuser must construct and cache the no-grad + # configuration. The runtime backward boundary is past the only op. + fuser.maybe_fuse_ops(False, None, x, extra_inputs) + assert fusion_calls == 1 + assert fuser.first_op_requiring_backward == 1 + no_grad_forward_ops = fuser._forward_ops + no_grad_backward_ops = fuser._backward_ops + + # Phase 2: backward replays the checkpointed region with grad enabled. The + # backward boundary is part of the fusion key, allowing future fusion rules + # to choose a training-specific topology. The first grad-enabled invocation + # therefore constructs and caches a second configuration. + fuser.maybe_fuse_ops(True, None, x, extra_inputs) + assert fusion_calls == 2 + assert fuser.first_op_requiring_backward == 0 + grad_forward_ops = fuser._forward_ops + grad_backward_ops = fuser._backward_ops + assert grad_forward_ops is not no_grad_forward_ops + assert grad_backward_ops is not no_grad_backward_ops + + # Phase 3: the next checkpointed forward must select the exact no-grad lists + # cached in phase 1. Before the cache was added, every boundary transition + # rebuilt the fused operations and called track_fusion again. + fuser.maybe_fuse_ops(False, None, x, extra_inputs) + assert fusion_calls == 2 + assert fuser.first_op_requiring_backward == 1 + assert fuser._forward_ops is no_grad_forward_ops + assert fuser._backward_ops is no_grad_backward_ops + + # Phase 4: another recomputation must likewise restore the grad-enabled + # lists from phase 2. The full alternating sequence has built only the two + # configurations represented by its two fusion keys. + fuser.maybe_fuse_ops(True, None, x, extra_inputs) + assert fusion_calls == 2 + assert fuser.first_op_requiring_backward == 0 + assert fuser._forward_ops is grad_forward_ops + assert fuser._backward_ops is grad_backward_ops + + +def test_operation_fuser_resets_recipe_state_independently_from_plan_cache(monkeypatch) -> None: + """Track recipe-state resets independently from fusion-plan construction.""" + + fusion_calls = 0 + + def track_fusion(ops, *, recipe): # pylint: disable=unused-argument + nonlocal fusion_calls + fusion_calls += 1 + return ops + + # Replace the process-wide fusion registries so one callback corresponds to + # one plan construction. monkeypatch restores the registries after the test. + monkeypatch.setattr(OperationFuser, "forward_backward_fusion_functions", [track_fusion]) + monkeypatch.setattr(OperationFuser, "forward_fusion_functions", []) + monkeypatch.setattr(OperationFuser, "backward_fusion_functions", []) + + op = te_ops.Identity() + reset_recipes = [] + first_forward_calls = 0 + + def track_recipe_reset(*, recipe): + reset_recipes.append(recipe) + + def track_first_forward(): + nonlocal first_forward_calls + first_forward_calls += 1 + + # Identity has no quantizers, so replace its state hooks with counters. This + # keeps the test CPU-only and isolates OperationFuser's reset decisions. + monkeypatch.setattr(op, "reset_recipe_state", track_recipe_reset) + monkeypatch.setattr(op, "pre_first_fuser_forward", track_first_forward) + + fuser = OperationFuser([op]) + x = torch.ones(1) + extra_inputs = [()] + + current_scaling = transformer_engine.common.recipe.Float8CurrentScaling(backward_override=None) + fuser.maybe_fuse_ops(False, current_scaling, x, extra_inputs) + assert reset_recipes == [current_scaling] + assert first_forward_calls == 1 + assert fusion_calls == 1 + + # A fresh but equivalent recipe does not invalidate state or the plan. + equivalent_current_scaling = transformer_engine.common.recipe.Float8CurrentScaling( + backward_override=None + ) + fuser.maybe_fuse_ops(False, equivalent_current_scaling, x, extra_inputs) + assert reset_recipes == [current_scaling] + assert first_forward_calls == 1 + assert fusion_calls == 1 + + # Backward override affects both recipe state and fusion topology, so it + # triggers one reset and constructs a distinct cached plan. + overridden_current_scaling = transformer_engine.common.recipe.Float8CurrentScaling( + backward_override="high_precision" + ) + fuser.maybe_fuse_ops(False, overridden_current_scaling, x, extra_inputs) + assert reset_recipes == [current_scaling, overridden_current_scaling] + assert first_forward_calls == 1 + assert fusion_calls == 2 + + delayed_scaling = transformer_engine.common.recipe.DelayedScaling( + amax_history_len=8, + backward_override=None, + ) + fuser.maybe_fuse_ops(False, delayed_scaling, x, extra_inputs) + assert reset_recipes == [current_scaling, overridden_current_scaling, delayed_scaling] + assert first_forward_calls == 1 + assert fusion_calls == 3 + + # Amax history length only affects delayed-scaling recipe state. Reset that + # state, but restore the existing DelayedScaling fusion plan from the cache. + resized_delayed_scaling = transformer_engine.common.recipe.DelayedScaling( + amax_history_len=16, + backward_override=None, + ) + fuser.maybe_fuse_ops(False, resized_delayed_scaling, x, extra_inputs) + assert reset_recipes == [ + current_scaling, + overridden_current_scaling, + delayed_scaling, + resized_delayed_scaling, + ] + assert first_forward_calls == 1 + assert fusion_calls == 3 + + # Repeating the exact recipe parameters performs neither operation again. + fuser.maybe_fuse_ops(False, resized_delayed_scaling, x, extra_inputs) + assert len(reset_recipes) == 4 + assert first_forward_calls == 1 + assert fusion_calls == 3 + + @torch.no_grad() def make_reference_and_test_tensors( shape: int | Iterable[int], diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 2128bae1ec..66c5bbb196 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1404,7 +1404,7 @@ def fuser_forward( alpha_tensor = get_cached_ones_tensor(num_groups, dtype, device) norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device) - current_stream = torch.cuda.current_stream().cuda_stream + current_stream = torch.cuda.current_stream(device.index).cuda_stream fc1_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc1_op) fc2_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc2_op) @@ -2095,7 +2095,7 @@ def fuser_backward( # Kernel scaling factors alpha_tensor = get_cached_ones_tensor(num_groups, dtype, device) norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device) - current_stream = torch.cuda.current_stream().cuda_stream + current_stream = torch.cuda.current_stream(device.index).cuda_stream unit_activation_scale = bool(getattr(fc1_ctx, "unit_activation_scale", False)) scales_f32 = None diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index fd66529ba8..2500002700 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -11,7 +11,7 @@ import torch -from ..quantization import FP8GlobalStateManager, Recipe, DelayedScaling +from ..quantization import FP8GlobalStateManager, Recipe from ..quantized_tensor import prepare_for_saving, restore_from_func_ctx from .op import ( BasicOperation, @@ -48,6 +48,8 @@ def _is_graph_capturing() -> bool: OperationFusionFunction: TypeAlias = ( "Callable[tuple[list[FusibleOperation], ...], list[FusibleOperation]]" ) +_FusedOpList: TypeAlias = list[tuple[FusibleOperation, list[int]]] +_FusionParams: TypeAlias = tuple[type, int, Optional[str]] class _OperationFuserAutogradFunction(torch.autograd.Function): @@ -535,15 +537,22 @@ def __init__( op._lock_extra_tensor_channels() # Ops for forward and backward pass, will be populated in maybe_fuse_ops - self._forward_ops: list[tuple[FusibleOperation, list[int]]] - self._backward_ops: list[tuple[FusibleOperation, list[int]]] + self._forward_ops: _FusedOpList + self._backward_ops: _FusedOpList + + # Fused operation configurations are reusable wrappers around the basic + # ops, so cache each configuration by the state that selected it. + self._fused_ops_cache: dict[_FusionParams, tuple[_FusedOpList, _FusedOpList]] = {} # Cache and detect change of state relevant for fusing operations self.recipe_type = None - self.first_op_requiring_backward = 0 self.backward_override = None self._last_amax_history_len = 0 + # Runtime backward boundary. Full activation recompute alternates this + # between the checkpointed forward and the grad-enabled recomputation. + self.first_op_requiring_backward = 0 + # Flatten list of parameters self._basic_op_params = [list(op.parameters()) for op in self._basic_ops] self._basic_op_num_params = list(map(len, self._basic_op_params)) @@ -626,36 +635,48 @@ def maybe_fuse_ops( first_op_requiring_backward = op_idx break - # Early exit if fusion parameters haven't changed - need_reset = False + # Update the runtime backward boundary on every invocation, including + # paths that reuse a cached fused operation configuration. + self.first_op_requiring_backward = first_op_requiring_backward + + # Check if recipe parameters don't match cached values. In this case, + # the recipe state in the basic ops might be invalid, so reset it. recipe_type = type(recipe) + need_to_reset_recipe_state = self.recipe_type != recipe_type + backward_override = recipe.backward_override if recipe is not None else None - fusion_params = (recipe_type, first_op_requiring_backward, backward_override) - if fusion_params != ( - self.recipe_type, - self.first_op_requiring_backward, - self.backward_override, - ): - # Recipe type, backward override, or grad requirements have changed - need_reset = True - elif ( + if backward_override != self.backward_override: + self.backward_override = backward_override + need_to_reset_recipe_state = True + + if ( recipe is not None and recipe.delayed() and self._last_amax_history_len != recipe.amax_history_len ): - # FP8 delayed scaling has changed amax history length - need_reset = True - if not need_reset: - return - - # Reset recipe state - for op in self._basic_ops: - op.reset_recipe_state(recipe=recipe) + self._last_amax_history_len = recipe.amax_history_len + need_to_reset_recipe_state = True - # Check if this is the first iteration - if self.recipe_type is None: + if need_to_reset_recipe_state: for op in self._basic_ops: - op.pre_first_fuser_forward() + op.reset_recipe_state(recipe=recipe) + + # Check if this is the first iteration + if self.recipe_type is None: + for op in self._basic_ops: + op.pre_first_fuser_forward() + + self.recipe_type = recipe_type + + # Training and inference may support different fusions. Keep the + # backward boundary in the key, but pay construction cost only once for + # each configuration. Full recompute therefore builds at most one + # no-grad plan and one grad-enabled plan for a stable recipe. + fusion_params = (recipe_type, first_op_requiring_backward, backward_override) + cached_ops = self._fused_ops_cache.get(fusion_params) + if cached_ops is not None: + self._forward_ops, self._backward_ops = cached_ops + return # Apply joint forward-backward fusions first joint_ops = OperationFuser._apply_fusions( @@ -682,14 +703,9 @@ def maybe_fuse_ops( self._basic_ops, ) - # Save current fusion params - self.recipe_type, self.first_op_requiring_backward, self.backward_override = fusion_params - - # Save amax history length - if isinstance(recipe, DelayedScaling): - self._last_amax_history_len = recipe.amax_history_len - else: - self._last_amax_history_len = 0 + # The FusedOperation contract excludes parameters and per-invocation + # state, so the mapped lists can be selected directly on cache hits. + self._fused_ops_cache[fusion_params] = (self._forward_ops, self._backward_ops) def __call__( self, From d1e9c33449ef583937532a003aa26abac2fe8981 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 4 Sep 2026 14:22:00 -0700 Subject: [PATCH 34/41] Use cuDNN's deterministic dprob in the fused grouped MLP (#3407) * [PyTorch] Ask cuDNN for a deterministic dprob under NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 The cuDNN grouped-GEMM dactivation backward that the CuTe DSL fused grouped MLP calls accumulates the scale gradient (dprob) with cross-CTA atomic adds, so its floating-point summation order follows the tile scheduler and varies run to run. Until now there was no way to switch that off, and NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 did not reach it: the run trained fine and was silently not reproducible. cuDNN frontend 1.28.0 (NVIDIA/cudnn-frontend#521) added a `deterministic` argument to grouped_gemm_dsrelu_wrapper_sm100 that parks each N-subtile's partial result in its own slot and sums the slots in a canonical order, for dprob and for dbias. Pass it from the TE flag. Passed as True or not at all, never as False. The wrapper's own default is None, which follows torch.use_deterministic_algorithms; sending an explicit False would override that and take determinism away from a caller who asked torch for it without setting the TE variable. The capability is reported per subclass rather than per environment variable, because grouped_gemm_dglu_wrapper_sm100 has no equivalent argument -- a GLU activation stays non-deterministic however new the installed front-end is. That case, and an SReLU op on a front-end older than 1.28.0, warn instead, once per distinct reason since the remedies differ. The warning is raised from where dprob is actually produced: with a unit activation scale the epilogue never runs its atomic accumulation, so there is nothing to make deterministic and nothing to warn about. Tests: TestGroupedMLPDeterminism covers the env-var parse, that only the SReLU op reports the capability and that it tracks the front-end version (no GPU or cuDNN needed for either), that the warning fires once per reason, and an MXFP8 end-to-end run under determinism for both SwiGLU and SReLU that checks numerics and pins which of the two arms warns. Signed-off-by: Zhiyu Li * Honor torch.use_deterministic_algorithms too, not just the env variable _deterministic_algorithms_required() copied the narrow check from transformer_engine.pytorch.triton.grouped_dbias_dscales, which reads NVTE_ALLOW_NONDETERMINISTIC_ALGO and nothing else. DotProductAttention takes the union instead -- the variable OR torch.use_deterministic_algorithms -- and that is the right precedent here. The two knobs answer different questions. The variable is set once in a job launcher, applies uniformly across ranks, and is the only one TE's C++ layer can read. The torch flag is the framework standard, is togglable at runtime, and is what a user who wants reproducibility usually reaches for; most have never heard of the variable. Keying on the variable alone left the torch flag half-honored. The SReLU path happened to come out right, but by delegation rather than by decision: TE passed nothing and the wrapper's own default read torch.are_deterministic_algorithms_enabled(). The GLU path did not -- TE stayed silent about an atomic dprob it cannot fix, for a user who had asked torch for reproducibility. That silence is the exact failure mode the warning exists to prevent, so it was the one case that most needed to warn. Passing the argument only as True, never as False, now needs a different justification than the one the first commit gave: with the union in place the two are equivalent, since the wrapper's default reads the same torch flag TE just read. The reason that survives is narrower and firmer -- the argument does not exist on the dGLU wrapper or on a front-end older than 1.28.0, where passing it at all, even as False, is a TypeError. Tests: the env-var parametrization becomes the two-knob truth table, including the row that motivates the change (torch flag set, NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -- the variable's default is the absence of a request, not a request for non-determinism, so the torch flag still wins). A fixture restores the process-global torch flag. Signed-off-by: Zhiyu Li * Test that dprob is actually bit-exact, not just within tolerance Review caught that nothing in the suite tested the property this change exists for. The end-to-end test runs the op once and checks numerics against a reference with rtol=0.125 / atol=0.25; reordering the same atomic adds moves dprob by about an ulp, so a run that is silently not reproducible passes it comfortably. The tolerance check proves the deterministic path is correct, which is worth keeping, but it cannot prove the path is deterministic. Add a second run. Same module, same inputs, grads cleared between passes, probs.grad compared with torch.equal. Three things the test has to get right to be worth having: * hidden_size 1024, not the 128 used elsewhere. dprob's reduction is over that extent and the tile is 256 wide, so 128 gives a single N-tile, one writer per token, and nothing to reorder -- the assertion would hold by construction and test nothing. * No bias. With an FC2 scale_bias the scale gradient is finished by the Triton grouped dbias/dscales kernel, which refuses to run under determinism, and probs.grad would stop being the dprob under test. * An assertion that the fusion happened, since dprob only comes from the cuDNN epilogue on the fused path. Skipped rather than xfailed on a front-end older than 1.28.0: there the kernel has no deterministic mode and is expected to vary, which is not a failure of this change. Weight gradients are deliberately left out of the comparison -- the CuTe DSL wgrad kernel has its own K-split atomics that this PR does not address. Signed-off-by: Zhiyu Li * Probe the dsrelu wrapper's signature instead of the frontend version _cudnn_frontend_supports_deterministic_dprob() gated on _cudnn_frontend_version_at_least("1.28.0"). That check is too coarse to answer the question it is asked, and would have raised at runtime on a build TE is actually run against. #521 merged after v1.27.0 was tagged, so `deterministic` ships in 1.28.0. But cudnn-frontend's develop branch has called itself 1.28.0 since shortly after that tag -- eleven days before the merge. Any front-end built from develop in that window reports 1.28.0 and does not accept the argument, so the version check passes, TE adds `deterministic=True` to the call, and the backward dies with TypeError: grouped_gemm_dsrelu_wrapper_sm100() got an unexpected keyword argument 'deterministic' This is not hypothetical, and not new. The same coarseness already bit use_single_group_runtime_offsets: a cuDNN reporting 1.27.0 that did not implement 1.27.0's arguments failed the identical way, in fuser_forward, before any backward code ran. Version numbers describe a release; they do not describe whatever happens to be installed. Ask the function instead. `"deterministic" in inspect.signature(...).parameters` is exact, cannot drift, and needs no maintenance when the release lands. The import is wrapped the way _grouped_gemm_dsrelu_backward_supported() already wraps it, so a missing cuDNN answers False rather than raising. Cached, since the call site runs every backward. This also removes the version constant from the code path entirely -- 1.28.0 now appears only in user-facing text, where a release number is the useful thing to say. Tests: a smoke test that the probe returns a bool without raising, with or without cuDNN installed, since reading a signature has more ways to fail than comparing two version strings. It deliberately does not assert which answer -- that depends on the installed front-end, and pinning it would only restate the implementation. Signed-off-by: Zhiyu Li * Revert the unrelated nccl-extensions submodule bump `git add -u` in the previous commit swept in a local 3rdparty/nccl-extensions pointer change that has nothing to do with this PR. Restore it to main's commit so the branch touches only the three files it means to. Signed-off-by: Zhiyu Li * Raise instead of warning, and cut the change down to what it needs Review asked for two things on the unsupported path: make it an error rather than a warning, and stop branching on self._cudnn_dact_func to pick a message. Both are right, and taking them removes most of the machinery this PR had accumulated. Raising matches what TE already does elsewhere: the Triton grouped dbias/dscales kernel refuses to run under determinism rather than running non-deterministically. It also matches what the variable documents -- "only deterministic algorithms are allowed" is not "prefer deterministic algorithms". A silently non-reproducible run is the failure this PR exists to prevent, so continuing past a request TE cannot honor was the wrong default. Checked that no existing determinism test hits this path: test_hybrid_quantization sets the variable for an attention recipe, and test_fusible_ops_with_userbuffers for linear ops. One message, no branch. The two cases did have different remedies, which is why the branch was there, but a single sentence states both facts -- "needs the scaled-SReLU activation and nvidia-cudnn-frontend 1.28.0 or later" -- without telling a SwiGLU user to go upgrade. What that let me delete: * _warn_nondeterministic_cudnn_dprob and its per-reason lru_cache, the two reason strings and the branch selecting them: 30 lines at the call site and above it, down to a single raise. * _cudnn_frontend_supports_deterministic_dprob as a standalone function. The probe now lives in GroupedMLP_CuTeGEMMUnary.grouped_gemm_dactivation_is_deterministic(), which reaches the wrapper through grouped_gemm_dactivation_kernel() -- the import and its ImportError handling already existed there, so folding it in dropped a duplicate import and an indirection. * The warn-once cache-clearing fixture in the tests, and the two tests that existed only to cover the warning. Tests: test_deterministic_dactivation_is_numerically_correct becomes test_determinism_either_runs_or_refuses -- it expects RuntimeError where the request cannot be honored and runs the full numerical check where it can, so both arms assert something either way. The bit-exactness and two-knob tests are unchanged in substance. Net: transformer_engine/pytorch/ops/fused/grouped_mlp.py goes from +106 to +68, all of it addition, no line of pre-existing code touched. Signed-off-by: Zhiyu Li * Apply suggestion from @vthumbe1503 Signed-off-by: vthumbe1503 * Match the feature-detection idiom main just landed The SiTU-GLU merge (#3402) brought _cudnn_frontend_supports_grouped_gemm_situglu() into this file, which asks inspect.signature(wrapper).parameters for the arguments it needs rather than comparing frontend versions -- the same conclusion this branch reached independently, now the house style. Two things to match. Guard the signature call with `except (TypeError, ValueError)`: a callable that is not introspectable answers "no" instead of raising out of a backward pass. I had left this out on the grounds that the wrapper is a plain undecorated function, which is true today but is not a property this code controls. And say "feature-detect" in the docstring summary, as the neighbor does. Also dropped the sentence about use_single_group_runtime_offsets from the docstring. The neighbor now demonstrates the pattern in the same file, so the cautionary tale is no longer what makes the choice legible. `import inspect` came in with the merge, so this branch no longer adds it. Signed-off-by: Zhiyu Li * Cut the comments down to the file's own register The new code carried multi-paragraph docstrings into a file whose 44 functions have a median docstring of one line. Measured before and after: grouped_mlp.py _deterministic_algorithms_required 10 -> 3 lines grouped_gemm_dactivation_is_deterministic (base) 5 -> 1 grouped_gemm_dactivation_is_deterministic (unary) 7 -> 1 test_grouped_mlp.py four new tests 4-6 -> 1-4 four inline comment blocks 2-3 -> 1 each Before this, the three new functions were the 2nd, 3rd and 5th longest docstrings in grouped_mlp.py; only fuse_grouped_mlp_ops, which has a full Parameters block, was longer. In the test file, 63 pre-existing tests have a median docstring of zero lines. Most of what came out was rationale, not explanation: why the union matches DotProductAttention, why feature detection beats a version compare, which cuDNN release window motivated it. That belongs in the commits that made those choices, where it already is, and it reads as noise next to _cudnn_frontend_supports_grouped_gemm_situglu -- the neighbor doing the very same feature detection in a one-line docstring with no rationale at all. What stayed is what the code cannot say itself: that the check sits inside the non-unit-scale branch because a unit scale produces no dprob; that hidden_size must exceed one N-tile or the bit-exactness test is vacuous; that bias would reroute probs.grad through Triton; that weight grads are excluded because wgrad has its own atomics. Each is now one line. No behavior change -- comments, docstrings and one local variable's reading order only. Signed-off-by: Zhiyu Li * Flatten the determinism check and the runs-or-refuses test Structural cleanups from the review pass. grouped_mlp.py: the check was nested two deep inside `if not unit_activation_scale`, and assigned deterministic_dactivation only to immediately test its own assignment. Hoisted to two flat statements right after unit_activation_scale is computed. `not unit_activation_scale and _deterministic_algorithms_required()` now says in the expression what the comment had to say in prose, and the separate `= False` initializer is gone. The local itself stays -- the kwargs dict is built about sixty lines further down. Also shortened the error: the tile-scheduler detail was not actionable, and "this activation's cuDNN dactivation kernel" is more accurate than naming the grouped-GEMM backward, since which kernel it is depends on the activation. test_grouped_mlp.py: fused_cls was derived from `activation` by a five-line conditional inside the test; it is now the second half of the parametrize pair. That also fixes the skip guard, which asked GroupedMLP_CuTeGEMMGLU.is_supported() on both parametrizations including the SReLU one -- the sibling test three functions down already gets this right. The _run closure existed only so an if/else could call it twice; a contextlib.nullcontext / pytest.raises choice removes the closure and the branch. nullcontext is used in ten test files here, so it is the local idiom rather than a new one. Not taken: dropping the `isinstance(..., bool)` assertion. It looks vacuous but it is the only coverage of the ImportError branch in the capability probe, which is the branch that runs on every machine without cuDNN -- including CI. Signed-off-by: Zhiyu Li * Close the second dprob producer, and stop discarding fp64 test tensors Two findings from the review pass. dprob has two producers in this backward, and the check only covered one. The cuDNN epilogue produces grad_scales at fuser_backward, and when scale_bias is set compute_grouped_dbias_dscales accumulates into it further down -- the Triton kernel that grouped_dbias_dscales.py documents as nondeterministic atomic adds. That kernel's own guard reads NVTE_ALLOW_NONDETERMINISTIC_ALGO and nothing else. So the hole opened exactly where this branch widened the trigger. With torch.use_deterministic_algorithms(True) and the variable unset -- the case the union exists to start honoring -- SReLU on a 1.28.0 front-end with scale_bias passed the new check, set deterministic=True, raised nothing, and then routed dprob through the nondeterministic path anyway. Env-var users were never exposed: the Triton guard fires for them. It was reachable only via the torch flag, which is to say only through what this branch added. The test picked bias=False and so never crossed it. scale_bias is computed ~130 lines earlier in the same scope, so the fix is to require both producers rather than one. Still one condition and one message, per review -- the message now lists all three requirements instead of two. Separately, the bit-exactness test built its tensors with make_reference_and_test_tensors and discarded the reference every time. That helper allocates an fp64 CPU companion, quantizes and dequantizes for MXFP8 representability, then copies back D2H with an implicit sync -- about 16 MB of host allocation across the two (1024, 1024) calls, for a test that compares run 1 against run 2 and never against a reference. Twelve of the file's other fifteen uses keep the reference; this one had no use for it. Plain uniform_ tensors instead. Also dropped a .item() sync for a token count already known in Python. Not taken, with reasons: * Hoisting _deterministic_algorithms_required into pytorch/utils.py so the Triton guard reads the same union. That is the deeper fix and it is correct, but broadening that guard changes behavior for callers this PR does not touch (ops/basic/grouped_linear.py, module/grouped_linear.py) -- users who set only the torch flag would start seeing RuntimeError where they now get silent nondeterminism. Worth doing deliberately, not as a side effect of this branch. * Extracting the signature-probe shared with _cudnn_frontend_supports_grouped_gemm_situglu. The overlap is about four lines and the two are not interchangeable; refactoring working code outside the diff to save them is not this PR's job. Signed-off-by: Zhiyu Li * Add the regression test for the scale_bias hole The previous commit fixed a real bug and shipped it with no test. Every test in the class used bias=False and every end-to-end one set the env var, so neither half of the bug was reachable: not scale_bias, and not the torch-flag-only trigger. Both halves are load-bearing. With the env var the Triton kernel raises on its own, so an env-var test would have passed before the fix as well as after and pinned nothing. Only torch.use_deterministic_algorithms with the variable unset reaches the state where this op's check said yes and the Triton reduction then ran nondeterministically. warn_only=True so torch's own enforcement cannot raise first and be mistaken for TE's refusal. are_deterministic_algorithms_enabled() still reports True in that mode -- the separate is_deterministic_algorithms_warn_only_enabled() getter exists precisely because the two are independent -- so the predicate under test sees what it should. Not executed: no GPU or torch on the machine this was written on. Formatting and syntax only, like the rest of the branch. Signed-off-by: Zhiyu Li * Make the bit-exactness test capable of failing Followed cudnn-frontend#521's own test work and found this test had the flaw its commit 88c7fab was written to fix, at the same config. That commit measured 16 launches per shape and found that at l=4 / [256]*4 / n=512 the NONDETERMINISTIC dprob is already bit-stable: the assertion cannot fail there, so a pass certifies nothing. It varies 15/15 at l=8 / [1024]*8 / n=2048. This test used l=4 / [256]*4 / n=1024 -- the vacuous shape, one power of two along n. Moved to the shape that actually varies. n > 256 was necessary but not sufficient, which is what the old comment got wrong. Spanning several N-tiles exercises the within-CTA subtile ordering; making the cross-CTA reduction unstable needs the larger token count and expert count too. Also took the rest of #521's discipline for these comparisons: * Repeat rather than compare a pair. The order determinism removes is set by the tile scheduler, so two runs can match by luck. Four by default, NVTE_TEST_DETERMINISM_REPEATS to raise it, matching that file's DETERMINISM_REPEATS. * Compare bytes, not values. torch.equal treats +0.0 and -0.0 as equal, and a change in reduction order produces exactly that; upstream's bitwise_bits views as uint8 for the same reason. * Assert the output is finite first, so a NaN run cannot be read as a determinism result. Not copied: asserting that the nondeterministic path *does* vary. It is the thing that makes the config meaningful, but as an assertion it is timing-dependent and would flake. Upstream settled this by measuring once and pinning the config; the comment now cites that measurement so the next person does not shrink the shape back. Not run yet: job 535935 is building the previous revision of this test. Signed-off-by: Zhiyu Li * Pick the bit-exactness config by measuring it, not by borrowing one Measured on GB300 across five shapes, determinism off, 8 launches each (job 538058), counting how many runs differ from run 0: l=8 tok/grp=1024 n=2048 2/7 max|d| 5.96e-08 l=8 tok/grp=1024 n=4096 2/7 max|d| 9.54e-07 l=16 tok/grp=1024 n=2048 7/7 max|d| 1.19e-07 l=8 tok/grp=2048 n=2048 5/7 max|d| 7.63e-06 l=4 tok/grp=512 n=8192 6/7 max|d| 1.91e-06 Moved to l=16, the only shape where every run differs, so the assertion cannot pass by luck. The previous choice, l=8, varies 2/7 -- an eight-run sample calls it stable often enough to be a poor detector, and an earlier control run (537313) did exactly that and reported 0/7 at this shape. I took that single sample as proof the config was vacuous and said so; it was a sampling artifact, and the shape does vary, just weakly. That earlier shape came from cudnn-frontend#521's own measurement, which was taken on its direct wrapper test. It does not transfer to TE's path -- different scheduler settings, different quantization -- so borrowing the number was the mistake underneath both errors. This config is measured through the fused grouped MLP itself. Two things the same job settled that are worth recording: * Without #521 the values genuinely move: 6e-08 to 8e-06 absolute across these shapes. Small, but nonzero every time, and the reason the refusal exists rather than a warning. * The refusal cannot be exercised against the stock 1.27.0 frontend on this image at all. TE's forward passes prob_tensor=None because _cudnn_frontend_version_at_least("1.27.0") reports optional-prob support that a stock 1.27.0 does not implement, so the op dies in fuser_forward with "prob_tensor is required" before any determinism code runs. Same version-gate-too-coarse failure this PR avoids for its own argument, on a gate it does not own. Signed-off-by: Zhiyu Li * Make the scale_bias half of the dprob check readable The condition was written as `not (A and not B)` with `scale_bias` as B, which gives a reader no way to tell why an FC2 bias flag decides whether a scale gradient is reproducible -- and the call that makes it relevant is ~250 lines further down. Same logic, named and nested: `dprob_is_deterministic` says what the conjunction means, and the comment names the mechanism instead of gesturing at it. dprob is finished by two kernels, not one -- the cuDNN dactivation epilogue writes it, and then, when scale_bias is set, fuser_backward hands it to compute_grouped_dbias_dscales as the `dscales` accumulator, which atomically adds into it. Its docstring is explicit: "Both outputs use fp32 atomic adds, so pre-populated tensors are accumulated into." No logic change; the truth table is identical. Signed-off-by: Zhiyu Li --------- Signed-off-by: Zhiyu Li Signed-off-by: vthumbe1503 Co-authored-by: vthumbe1503 --- tests/pytorch/test_grouped_mlp.py | 189 ++++++++++++++++++ .../pytorch/ops/fused/grouped_mlp.py | 54 +++++ 2 files changed, 243 insertions(+) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 1d173c76d8..15c1ff6c51 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -5,6 +5,7 @@ from __future__ import annotations from collections.abc import Iterable +import contextlib import functools import os import math @@ -2845,6 +2846,194 @@ def train_step( assert_close(graph_grad, param.grad, **tols) +class TestGroupedMLPDeterminism: + """Determinism coverage for the CuTe DSL fused grouped MLP. + + Only the dSReLU wrapper can make ``dprob`` bit-exact, and only from cuDNN FE 1.28.0 on. + Anything else must refuse a determinism request rather than run non-deterministically. + """ + + @pytest.fixture + def _restore_torch_determinism(self): + """``use_deterministic_algorithms`` is process-global, so put it back.""" + previous = torch.are_deterministic_algorithms_enabled() + yield + torch.use_deterministic_algorithms(previous) + + @pytest.mark.parametrize( + "allow_nondeterministic,torch_flag,expected", + ( + (None, False, False), # default: non-deterministic algorithms are allowed + ("1", False, False), + ("0", False, True), # the TE variable alone + (None, True, True), # the torch flag alone, which TE must not ignore + ("1", True, True), # ... including when the TE variable says otherwise + ("0", True, True), + ), + ) + def test_either_knob_requests_determinism( + self, + monkeypatch, + _restore_torch_determinism, + *, + allow_nondeterministic: Optional[str], + torch_flag: bool, + expected: bool, + ) -> None: + """``=1`` is the absence of a request, not a request for non-determinism.""" + if allow_nondeterministic is None: + monkeypatch.delenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", raising=False) + else: + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", allow_nondeterministic) + torch.use_deterministic_algorithms(torch_flag) + assert grouped_mlp_module._deterministic_algorithms_required() is expected + + def test_only_the_srelu_path_can_be_deterministic(self) -> None: + """The capability belongs to the wrapper, not the environment. Needs no GPU.""" + glu = grouped_mlp_module.GroupedMLP_CuTeGEMMGLU + unary = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary + assert glu.grouped_gemm_dactivation_is_deterministic() is False + assert isinstance(unary.grouped_gemm_dactivation_is_deterministic(), bool) + + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + @pytest.mark.parametrize( + "activation,fused_cls", + ( + ("scaled_srelu", grouped_mlp_module.GroupedMLP_CuTeGEMMUnary), + ("scaled_swiglu", grouped_mlp_module.GroupedMLP_CuTeGEMMGLU), + ), + ) + def test_determinism_either_runs_or_refuses( + self, monkeypatch, *, activation, fused_cls + ) -> None: + """A request TE cannot honor must fail loudly; one it can must still be correct.""" + if not fused_cls.is_supported(): + pytest.skip("MXFP8 fused grouped MLP is not supported on this system") + + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + expectation = ( + contextlib.nullcontext() + if fused_cls.grouped_gemm_dactivation_is_deterministic() + else pytest.raises(RuntimeError, match="dprob") + ) + with expectation: + TestGroupedMLPFusedOp().test_grouped_mlp( + bias=False, + hidden_size=128, + quantization="mxfp8", + single_grouped_weight=False, + activation=activation, + ) + + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_scale_bias_refuses_under_the_torch_flag( + self, monkeypatch, _restore_torch_determinism + ) -> None: + """``scale_bias`` finishes ``dprob`` in a Triton kernel that reads only the env var. + + So the torch flag alone is the combination that used to pass this op's own check and + then reduce nondeterministically anyway, on a front-end new enough to say yes. + """ + fused_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary + if not fused_cls.is_supported(): + pytest.skip("MXFP8 fused grouped MLP is not supported on this system") + + monkeypatch.delenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", raising=False) + # warn_only so torch's own enforcement cannot raise first and mask what TE does. + torch.use_deterministic_algorithms(True, warn_only=True) + with pytest.raises(RuntimeError, match="dprob"): + TestGroupedMLPFusedOp().test_grouped_mlp( + bias=True, + hidden_size=128, + quantization="mxfp8", + single_grouped_weight=False, + activation="scaled_srelu", + ) + + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None: + """Repeated identical runs must give a bit-identical ``dprob``. + + An ulp of reordering passes every tolerance in this file, so only an exact + comparison across runs can see it. + """ + fused_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary + if not fused_cls.is_supported(): + pytest.skip("MXFP8 fused grouped MLP is not supported on this system") + if not fused_cls.grouped_gemm_dactivation_is_deterministic(): + pytest.skip("dSReLU determinism needs cuDNN frontend 1.28.0 or later") + + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + + device = torch.device("cuda") + dtype = torch.bfloat16 + # Measured on GB300, determinism off, 8 launches per shape (job 538058): this shape + # gives 7/7 runs differing from run 0, so the assertion below can actually fail. + # Shapes matter more than they look -- l=8 with the same n and tokens/group varies + # only 2/7, which an 8-run sample reports as stable often enough to be useless, and + # cudnn-frontend#521 measured its own l=4 / [256]*4 / n=512 as never varying. + group_size = 16 + hidden_size = 2048 + tokens_per_group = 1024 + split_sizes = torch.tensor([tokens_per_group] * group_size, dtype=torch.int, device=device) + num_tokens = tokens_per_group * group_size + + recipe = make_recipe("mxfp8") + + # Plain random tensors, not make_reference_and_test_tensors: this test compares two + # runs against each other, never against a reference, so the fp64 companion and the + # MXFP8 representability round-trip would both be allocated and thrown away. + def _rand(*shape, requires_grad=True) -> torch.Tensor: + out = torch.empty(shape, dtype=dtype, device=device).uniform_(-0.25, 0.25) + return out.requires_grad_() if requires_grad else out + + x = _rand(num_tokens, hidden_size) + dy = _rand(num_tokens, hidden_size, requires_grad=False) + probs = _rand(num_tokens) + + # No bias, or probs.grad comes from the Triton dbias kernel instead of cuDNN. + with te.quantized_model_init(enabled=True, recipe=recipe): + module = te.ops.Sequential( + te.ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ), + te.ops.ScaledSReLU(), + te.ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ), + ) + + def _run() -> torch.Tensor: + x.grad = None + probs.grad = None + with te.autocast(enabled=True, recipe=recipe): + y = module(x, split_sizes, probs, split_sizes) + y.backward(dy) + return probs.grad.detach().clone() + + runs = [_run()] + # Without the fusion there is no cuDNN dprob and the comparison proves nothing. + forward_ops = module._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], fused_cls) + # More than two, as cudnn-frontend#521 does: the cross-CTA order that determinism + # removes is set by the scheduler, so two runs can agree by luck. + runs += [_run() for _ in range(int(os.getenv("NVTE_TEST_DETERMINISM_REPEATS", "4")) - 1)] + torch.cuda.synchronize() + + assert torch.isfinite(runs[0]).all(), "dprob is not finite; the comparison would be moot" + # Bytes, not values: torch.equal calls +0.0 and -0.0 equal, and a change in reduction + # order can produce exactly that. Weight grads are excluded from the comparison -- + # the CuTe DSL wgrad kernel has its own K-split atomics, which this change leaves. + for index, later in enumerate(runs[1:], start=1): + assert torch.equal( + runs[0].contiguous().view(torch.uint8), later.contiguous().view(torch.uint8) + ), ( + f"dprob differs between run 0 and run {index} under determinism; max |delta| =" + f" {(runs[0].float() - later.float()).abs().max().item()}" + ) + + def test_grouped_gemm_quant_cute_matches_mxfp8_quantized() -> None: if not mxfp8_available: pytest.skip(reason_for_no_mxfp8) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 66c5bbb196..58f22c874c 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -170,6 +170,17 @@ def _cudnn_frontend_supports_single_group_runtime_offsets( ) and _cudnn_frontend_version_at_least("1.27.0") +def _deterministic_algorithms_required() -> bool: + """Whether bit-exact reproducibility was asked for. Same union as ``DotProductAttention``. + + Uncached: both knobs can change during the process. + """ + return ( + not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) + or torch.are_deterministic_algorithms_enabled() + ) + + def _wrap_single_quantized_as_grouped( tensor: torch.Tensor, quantized: MXFP8Tensor | NVFP4Tensor | NVFP4TensorStorage, @@ -956,6 +967,11 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable: """Fused kernel for grouped GEMM, activation backward, and scale grad.""" raise NotImplementedError + @classmethod + def grouped_gemm_dactivation_is_deterministic(cls) -> bool: + """Whether this op's dactivation kernel can produce a bit-exact ``dprob``.""" + return False + @classmethod @functools.lru_cache(maxsize=None) def grouped_gemm_quant_kernel(cls) -> Callable: @@ -2098,6 +2114,28 @@ def fuser_backward( current_stream = torch.cuda.current_stream(device.index).cuda_stream unit_activation_scale = bool(getattr(fc1_ctx, "unit_activation_scale", False)) + # A unit activation scale produces no dprob, so there is nothing to make deterministic. + deterministic_dactivation = ( + not unit_activation_scale and _deterministic_algorithms_required() + ) + if deterministic_dactivation: + # Two kernels write dprob and both have to be exact. The cuDNN dactivation + # epilogue produces it below; then, when scale_bias is set, it is passed to + # compute_grouped_dbias_dscales as the ``dscales`` accumulator and atomically + # added into (see triton/grouped_dbias_dscales.py). That Triton kernel is never + # deterministic, so scale_bias rules out a bit-exact dprob on its own. + dprob_is_deterministic = ( + self.grouped_gemm_dactivation_is_deterministic() and not scale_bias + ) + if not dprob_is_deterministic: + raise RuntimeError( + "Deterministic execution was requested" + " (NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 or" + " torch.use_deterministic_algorithms), but the scale gradient (dprob) is" + " accumulated with nondeterministic atomics on this configuration." + " A bit-exact dprob requires the scaled-SReLU activation," + " nvidia-cudnn-frontend 1.28.0 or later, and an FC2 without scale_bias." + ) scales_f32 = None scales_tensor = None dscales_tensor = None @@ -2152,6 +2190,9 @@ def fuser_backward( "use_dynamic_sched": True, } dactivation_kernel = self.grouped_gemm_dactivation_kernel() + if deterministic_dactivation: + # Never passed to a wrapper that would reject it -- the check above raises first. + fc2_dactivation_kwargs["deterministic"] = True if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)): fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 if self._cudnn_dact_func is not None: @@ -2682,6 +2723,19 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable: return grouped_gemm_dsrelu_wrapper_sm100 + @classmethod + @functools.lru_cache(maxsize=None) + def grouped_gemm_dactivation_is_deterministic(cls) -> bool: + """Feature-detect the dSReLU wrapper's ``deterministic`` argument (cuDNN FE 1.28.0+).""" + try: + kernel = cls.grouped_gemm_dactivation_kernel() + except ImportError: + return False + try: + return "deterministic" in inspect.signature(kernel).parameters + except (TypeError, ValueError): + return False + def fuse_ops( ops: list[FusibleOperation], From 29229b5a030d1672bb79af61a735caf42ce1effc Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Sat, 5 Sep 2026 06:51:41 +0800 Subject: [PATCH 35/41] [PyTorch] Reduce CUDA graph memory retention (#3427) * [PyTorch] Release warmup outputs after their last use Signed-off-by: Robin Zhang * [PyTorch] Release buffer-reuse capture temporaries Signed-off-by: Robin Zhang * [PyTorch] Release per-callable state on reset Signed-off-by: Robin Zhang * [PyTorch] Bundle per-callable lifecycle helpers Signed-off-by: Robin Zhang --------- Signed-off-by: Robin Zhang --- tests/pytorch/test_cuda_graphs.py | 202 ++++++++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 129 +++++++++++++++--- 2 files changed, 310 insertions(+), 21 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 5a848dc0e8..8f85f57f32 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -5,6 +5,8 @@ from typing import Callable, Dict, Iterable, List, Tuple, Union import pytest import copy +import gc +import weakref import torch from transformer_engine.pytorch import ( @@ -994,6 +996,206 @@ def hook(module: torch.nn.Module) -> None: ] +def test_ordered_warmup_releases_consumed_outputs() -> None: + """Ordered warmup should only retain outputs until their corresponding backward.""" + + class OutputLifetimeModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.previous_output = None + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + is_warmup = not torch.cuda.is_current_stream_capturing() + if is_warmup and self.previous_output is not None: + assert self.previous_output() is None + output = input_ * 2 + if is_warmup: + self.previous_output = weakref.ref(output) + return output + + module = OutputLifetimeModule() + sample_args = tuple((torch.ones(4, 8, device="cuda", requires_grad=True),) for _ in range(2)) + graphed_callables = make_graphed_callables( + (module,), + sample_args, + num_warmup_iters=2, + _order=[1, -1, 1, -1], + _num_layers_per_chunk=[1], + ) + assert module.previous_output is not None + assert module.previous_output() is None + reset_graphs(graphed_callables) + + +def test_unordered_warmup_releases_consumed_outputs() -> None: + """Unordered warmup should release each output after its corresponding backward.""" + + class OutputLifetimeModule(torch.nn.Module): + def __init__(self, output_refs: list, module_idx: int) -> None: + super().__init__() + self.output_refs = output_refs + self.module_idx = module_idx + self.capture_started = False + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + output = input_ * 2 + if torch.cuda.is_current_stream_capturing(): + self.capture_started = True + else: + self.output_refs[self.module_idx] = weakref.ref(output) + return output + + output_refs = [None, None] + modules = tuple(OutputLifetimeModule(output_refs, module_idx) for module_idx in range(2)) + + def first_module_backward_pre_hook(_module: torch.nn.Module) -> None: + if not modules[0].capture_started: + assert output_refs[1] is not None + assert output_refs[1]() is None + + graphed_callables = make_graphed_callables( + modules, + tuple((torch.ones(4, 8, device="cuda", requires_grad=True),) for _ in modules), + num_warmup_iters=2, + capture_time_hooks=[ + {"backward_pre_hooks": {0: first_module_backward_pre_hook}}, + None, + ], + ) + assert all(output_ref is not None and output_ref() is None for output_ref in output_refs) + reset_graphs(graphed_callables) + + +def test_inference_warmup_does_not_retain_outputs() -> None: + """Inference warmup should release outputs as soon as each forward returns.""" + + class OutputLifetimeModule(torch.nn.Module): + def __init__(self, previous_output: list) -> None: + super().__init__() + self.previous_output = previous_output + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + is_warmup = not torch.cuda.is_current_stream_capturing() + if is_warmup and self.previous_output[0] is not None: + assert self.previous_output[0]() is None + output = input_ * 2 + if is_warmup: + self.previous_output[0] = weakref.ref(output) + return output + + previous_output = [None] + modules = tuple(OutputLifetimeModule(previous_output).eval() for _ in range(2)) + graphed_callables = make_graphed_callables( + modules, + tuple((torch.ones(4, 8, device="cuda"),) for _ in modules), + num_warmup_iters=2, + ) + assert previous_output[0] is not None + assert previous_output[0]() is None + reset_graphs(graphed_callables) + + +def test_reused_capture_buffers_release_outputs_after_backward() -> None: + """Capture locals must not keep weak-refed output buffers alive.""" + + class OutputLifetimeModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.previous_capture_output = None + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + if ( + torch.cuda.is_current_stream_capturing() + and self.previous_capture_output is not None + ): + assert self.previous_capture_output() is None + output = input_ * 2 + if torch.cuda.is_current_stream_capturing(): + self.previous_capture_output = weakref.ref(output) + return output + + module = OutputLifetimeModule() + sample_args = tuple((torch.ones(4, 8, device="cuda", requires_grad=True),) for _ in range(2)) + graphed_callables = make_graphed_callables( + (module,), + sample_args, + _order=[1, -1, 1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + ) + assert module.previous_capture_output is not None + assert module.previous_capture_output() is None + reset_graphs(graphed_callables) + + +def test_reset_releases_only_the_selected_callable() -> None: + """Reset releases one callable's graph state without retaining its peers.""" + + class CaptureOutputModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.capture_output = None + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + output = input_ * 2 + if torch.cuda.is_current_stream_capturing(): + self.capture_output = weakref.ref(output) + return output + + modules = tuple(CaptureOutputModule().cuda() for _ in range(2)) + graphed_callables = make_graphed_callables( + modules, + tuple((torch.ones(4, device="cuda", requires_grad=True),) for _ in modules), + ) + capture_outputs = tuple(module.capture_output for module in modules) + assert all(output is not None and output() is not None for output in capture_outputs) + + graphed_callables[0].reset() + graphed_callables[0].reset() + gc.collect() + assert capture_outputs[0]() is None + assert capture_outputs[1]() is not None + + output = graphed_callables[1](torch.randn(4, device="cuda", requires_grad=True)) + output.sum().backward() + del output + graphed_callables[1].reset() + gc.collect() + assert capture_outputs[1]() is None + + +@pytest.mark.parametrize("with_order", (False, True)) +def test_reset_rejects_all_replay_entry_points(with_order: bool) -> None: + """Reset is idempotent and terminal for forward and backward replay.""" + + class TestModule(torch.nn.Module): + def forward(self, input_: torch.Tensor) -> torch.Tensor: + return input_ * 2 + + module = TestModule().cuda() + sample_input = torch.ones(4, device="cuda", requires_grad=True) + graph_options = {} + if with_order: + graph_options = {"_order": [1, -1], "_num_layers_per_chunk": [1]} + graphed_callable = make_graphed_callables(module, (sample_input,), **graph_options) + output = graphed_callable(torch.randn_like(sample_input, requires_grad=True)) + torch.cuda.synchronize() + + graphed_callable.reset() + graphed_callable.reset() + if not with_order: + # The eager fallback for a different training state is invalid after reset too. + graphed_callable.eval() + + error = "has been reset and can no longer be used" + with pytest.raises(RuntimeError, match=error): + graphed_callable(torch.randn_like(sample_input, requires_grad=True)) + with pytest.raises(RuntimeError, match=error): + graphed_callable.backward_dw() + with pytest.raises(RuntimeError, match=error): + output.sum().backward() + + @pytest.mark.parametrize("with_order", (False, True)) def test_make_graphed_callables_with_capture_time_hooks(with_order: bool) -> None: """Test capture-time hooks around warmup and graph capture.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index b298b3d8ff..04fa56721d 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -8,7 +8,7 @@ import gc import warnings from math import ceil -from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union +from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple, TypeVar, Union import torch from torch.utils._pytree import tree_flatten as _tree_flatten @@ -44,6 +44,13 @@ ) +class _GraphedCallableHelpers(NamedTuple): + """Lifecycle helpers owned by one graphed callable invocation.""" + + ensure_not_reset: Callable[[], None] + release_static_state: Callable[[], None] + + def set_capture_start() -> None: """Record beginning of `make_graphed_callables`.""" global _IS_GRAPH_CAPTURING @@ -633,13 +640,17 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): warmup_outputs = [] for func_idx, func in zip(warmup_func_idx, warmup_func): outputs = _run_warmup_forward(func_idx, func, func_idx) - warmup_outputs.append((func_idx, func, outputs)) - if is_training: - for func_idx, func, outputs in reversed(warmup_outputs): - _run_warmup_backward(func_idx, func, outputs, warmup_iter, func_idx) + if is_training: + warmup_outputs.append((func_idx, func, outputs)) + else: + del outputs + while warmup_outputs: + func_idx, func, outputs = warmup_outputs.pop() + _run_warmup_backward(func_idx, func, outputs, warmup_iter, func_idx) + del outputs else: # Follow _order exactly, mirroring the capture phase. - per_fwd_outputs = {} # per_callable_fwd_idx -> flattened outputs + per_fwd_outputs = {} # per_callable_fwd_idx -> outstanding flattened outputs fwd_idx = [0] * num_model_chunks bwd_idx = [0] * num_model_chunks for c_id in _order: @@ -653,7 +664,10 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): ) + (fwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no) func = callables[callable_idx] outputs = _run_warmup_forward(per_callable_fwd_idx, func, callable_idx) - per_fwd_outputs[per_callable_fwd_idx] = outputs + if is_training: + per_fwd_outputs[per_callable_fwd_idx] = outputs + else: + del outputs fwd_idx[m_chunk] += 1 elif ceil(c_id) == c_id: # Backward pass for chunk -c_id. @@ -665,10 +679,11 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): _prefix_num_layers[m_chunk] * num_microbatches ) + (bwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no) func = callables[callable_idx] - outputs = per_fwd_outputs[per_callable_bwd_idx] + outputs = per_fwd_outputs.pop(per_callable_bwd_idx) _run_warmup_backward( per_callable_bwd_idx, func, outputs, warmup_iter, callable_idx ) + del outputs bwd_idx[m_chunk] += 1 if post_warmup_hook is not None: @@ -729,6 +744,7 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): per_callable_static_outputs[per_callable_fwd_idx] = tuple(flatten_outputs) per_callable_output_unflatten_spec[per_callable_fwd_idx] = spec graph_callables[per_callable_fwd_idx] = func + del outputs, flatten_outputs fwd_idx[m_chunk] += 1 else: # Capture backward graph for model chunk c_id, microbatch bwd_idx[-c_id-1] @@ -917,6 +933,11 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): per_callable_static_grad_inputs[idx] ) previous_chunk_last_callable_bwd_idx = per_callable_bwd_idx + + # The per-callable containers now own all tensors that must survive + # capture. Drop local strong references so weak-refed graph buffers can + # be returned to the shared CUDA graph pool before the next capture. + del static_outputs, static_grad_inputs, grad_inputs if ceil(c_id) == c_id: bwd_idx[m_chunk] += 1 else: @@ -1028,12 +1049,22 @@ def make_graphed_autograd_function( static_grad_inputs, returned_param_grad_clone_slots, ): + is_reset = False + + def ensure_not_reset(): + """Reject replay after this callable's graph state has been released.""" + if is_reset: + raise RuntimeError( + "This graphed callable has been reset and can no longer be used." + ) + class Graphed(torch.autograd.Function): """Autograd function for graph replay.""" @staticmethod def forward(ctx, skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *inputs): # pylint: disable=missing-function-docstring + ensure_not_reset() # Set flag for whether to update FP8 weight updates ctx.is_first_module = FP8GlobalStateManager.is_first_fp8_module() @@ -1071,6 +1102,7 @@ def forward(ctx, skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *i @torch.autograd.function.once_differentiable def backward(ctx, *grads): # pylint: disable=missing-function-docstring + ensure_not_reset() # Replay backward graph if len(grads) != len(static_grad_outputs): @@ -1119,6 +1151,7 @@ def backward(ctx, *grads): return (None, None, None) + tuple(grad_inputs) def functionalized(*user_args, **user_kwargs): + ensure_not_reset() # Decide whether to update FP8 weights skip_fp8_weight_update = None @@ -1170,16 +1203,43 @@ def functionalized(*user_args, **user_kwargs): ) return _tree_unflatten(out, output_unflatten_spec) - return functionalized - - def make_graphed_attribute_functions(graph_idx): - # Get te modules for current graph + def release_static_state(): + """Release per-callable state captured by replay closures.""" + nonlocal fwd_graph, bwd_graph, is_reset + nonlocal module_params + nonlocal static_input_surface, static_outputs + nonlocal static_grad_outputs, static_grad_inputs + + is_reset = True + + # Drop the per-callable references that can own graph-pool storage. + fwd_graph = None + bwd_graph = None + module_params = () + static_input_surface = () + static_outputs = () + static_grad_outputs = () + static_grad_inputs = () + + helpers = _GraphedCallableHelpers( + ensure_not_reset=ensure_not_reset, + release_static_state=release_static_state, + ) + return functionalized, helpers + + def make_graphed_attribute_functions(graph_idx, helpers): + # Snapshot per-callable state so returned closures do not retain the outer lists. + fwd_graph = fwd_graphs[graph_idx] + bwd_graph = bwd_graphs[graph_idx] + bwd_dw_graph = bwd_dw_graphs[graph_idx] + need_bwd_dw = need_bwd_dw_graph.get(graph_idx, False) te_modules = visited_te_modules.get(graph_idx, set()) # Attach backward_dw as an attribute to the graphed callable. def backward_dw(): - if need_bwd_dw_graph.get(graph_idx, False): - bwd_dw_graphs[graph_idx].replay() + helpers.ensure_not_reset() + if need_bwd_dw: + bwd_dw_graph.replay() # Trigger the grad accumulation hook for wgrad graphs. for module in te_modules: @@ -1191,16 +1251,24 @@ def backward_dw(): # Attach reset as an attribute to the graphed callable. def reset(): - fwd_graphs[graph_idx].reset() - bwd_graphs[graph_idx].reset() - bwd_dw_graphs[graph_idx].reset() + nonlocal fwd_graph, bwd_graph, bwd_dw_graph, te_modules + + for graph in (fwd_graph, bwd_graph, bwd_dw_graph): + if graph is not None: + graph.reset() + + fwd_graph = None + bwd_graph = None + bwd_dw_graph = None + te_modules = () + helpers.release_static_state() return backward_dw, reset # Put together the final graphed callables ret = [] for i in range(len(sample_args)): - graphed = make_graphed_autograd_function( + graphed, helpers = make_graphed_autograd_function( fwd_graphs[i], bwd_graphs[i], per_callable_module_params[i], @@ -1218,8 +1286,17 @@ def reset(): te_modules = visited_te_modules.get(i, set()) if isinstance(func, torch.nn.Module): - def make_graphed_forward(func, graph_training_state, graphed, orig_fwd, te_modules): + def make_graphed_forward( + func, + graph_training_state, + graphed, + orig_fwd, + te_modules, + helpers, + ): def new_fwd(*user_args, **user_kwargs): + helpers.ensure_not_reset() + # If the module's training-or-eval state matches what we graphed, # run the graph, otherwise run the original forward method if func.training == graph_training_state: @@ -1264,7 +1341,14 @@ def new_fwd(*user_args, **user_kwargs): return new_fwd - forward = make_graphed_forward(func, func.training, graphed, func.forward, te_modules) + forward = make_graphed_forward( + func, + func.training, + graphed, + func.forward, + te_modules, + helpers, + ) if _order is None: func.forward = forward ret.append(func) @@ -1273,7 +1357,10 @@ def new_fwd(*user_args, **user_kwargs): else: ret.append(graphed) - backward_dw_func, reset_func = make_graphed_attribute_functions(i) + backward_dw_func, reset_func = make_graphed_attribute_functions( + i, + helpers, + ) setattr(ret[-1], "backward_dw", backward_dw_func) setattr(ret[-1], "reset", reset_func) From 846379d3c4e8dee0d12fe8066c8c31dce3ef3ee2 Mon Sep 17 00:00:00 2001 From: Ravi Ghadia <40660742+ghadiaravi13@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:52:54 -0500 Subject: [PATCH 36/41] Relax runtime checks for activation recompute into Warnings (#3436) * Remove redundant runtime checks for activation recompute in MLP from _ScaledUnary class in activation.py Signed-off-by: Ravi Ghadia * Add warning for activation recompute in MLP outside fused path in _ScaledUnary class Signed-off-by: Ravi Ghadia * Add test for Scaled SReLU activation recompute warning outside fused MLP path Signed-off-by: Ravi Ghadia * Enhance activation recompute warning in ScaledSReLU: Update test to verify multiple warnings during backward passes and refactor warning mechanism to ensure it survives Dynamo tracing. Signed-off-by: Ravi Ghadia * Remove redundant logic to only warn once Signed-off-by: Tim Moon * Remove unhelpful test Signed-off-by: Tim Moon --------- Signed-off-by: Ravi Ghadia Signed-off-by: Tim Moon Co-authored-by: Tim Moon --- .../pytorch/ops/basic/activation.py | 28 ++++++++++--------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index 5c33c08b44..26d5261b13 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -16,7 +16,7 @@ from ...constants import DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...tensor.float8_tensor import Float8CurrentScalingQuantizer, Quantizer -from ...utils import clear_tensor_data +from ...utils import _compile_safe_warn, clear_tensor_data from ..op import BasicOperation, OperationContext from .._common import maybe_dequantize @@ -404,14 +404,14 @@ def fuser_forward( next_op_input_quantizer: Optional[Quantizer], # pylint: disable=unused-argument basic_op_kwargs: list[dict[str, Any]], # pylint: disable=unused-argument ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: + extra_input = basic_op_extra_inputs[0][0] + if self.activation_recompute_in_mlp: - raise RuntimeError( - f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " - "fused grouped MLP path." + _compile_safe_warn( + f"{self.__class__.__name__}(activation_recompute_in_mlp=True) is only supported " + "in the fused grouped MLP path." ) - extra_input = basic_op_extra_inputs[0][0] - if torch.is_autocast_enabled(): dtype = torch.get_autocast_dtype("cuda") elif isinstance(input_, torch.Tensor): @@ -447,18 +447,18 @@ def fuser_backward( ]: del basic_op_grad_extra_outputs - if self.activation_recompute_in_mlp: - raise RuntimeError( - f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " - "fused grouped MLP path." - ) - ctx = basic_op_ctxs[0] x, scales = ctx.saved_tensors x = maybe_dequantize(x.contiguous(), ctx.dtype) scales = maybe_dequantize(scales, ctx.dtype) grad_output = maybe_dequantize(grad_output.contiguous(), ctx.dtype) + if self.activation_recompute_in_mlp: + _compile_safe_warn( + f"{self.__class__.__name__}(activation_recompute_in_mlp=True) is only supported " + "in the fused grouped MLP path." + ) + grad_input, grad_extra_input = self._scaled_unary_backward( grad_output, x, @@ -483,7 +483,9 @@ class ScaledSReLU(_ScaledUnary): ---------- activation_recompute_in_mlp : bool, default = ``False`` Enable fused grouped MLP kernels to recompute activation outputs - during backward when supported instead of saving them. + during backward when supported instead of saving them. Outside the + fused grouped MLP path this option has no effect and a warning is + emitted. """ def _scaled_unary_forward( From a30aee5b5536d9e0edf8645eda6eabb66238e108 Mon Sep 17 00:00:00 2001 From: Andrew White Date: Fri, 4 Sep 2026 19:06:45 -0500 Subject: [PATCH 37/41] fix: Unreachable backend check after earlier backend skip (#3368) * fix: remove unreachable fused-attn backend skip and add regression test Root cause: FusedAttnRunner._check_configs skipped with 'Unsupported inputs combination or device compute capability.' unless the backend was NVTE_F16_arbitrary_seqlen. The later elif re-testing self.backend != NVTE_F16_arbitrary_seqlen could therefore never be reached, so its skip message ('B1SS, BHSS and 11SS bias shapes are only supported for the F16_arbitrary_seqlen backend') was dead code. Fix: drop the dead elif arm. The padding-mask skip in the sibling if arm is retained. A regression test locks the remaining behavior: a non-1HSS post-scale-bias config (BiasShape._B1SS) that passes _check_configs selects NVTE_F16_arbitrary_seqlen, proving the removal is behaviorally invisible and the earlier guard is the sole gate. Testing: not run locally - the JAX test stack (jax, transformer_engine_jax) is not installed on this machine. The suite runs in NVIDIA TransformerEngine CI. Contribution: tests/jax tests target real GPU/cuDNN fused-attention kernels and skip otherwise; the new test will follow that path via CI. Signed-off-by: andrewwhitecdw * Remove unnecessary test Signed-off-by: Przemyslaw Tredak --------- Signed-off-by: andrewwhitecdw Signed-off-by: Przemyslaw Tredak Co-authored-by: andrewwhitecdw Co-authored-by: Przemyslaw Tredak --- tests/jax/test_fused_attn.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index b6fc8f7794..1ca4121a4f 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -612,11 +612,6 @@ def _check_configs(self): pytest.skip( "B1SS, BHSS and 11SS bias shapes are only supported for non-padding mask" ) - elif self.backend != NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: - pytest.skip( - "B1SS, BHSS and 11SS bias shapes are only supported for " - "the F16_arbitrary_seqlen backend." - ) def _setup_inputs(self): self._check_configs() From 5f6105b900778068144f1b24e25a2e84066983cb Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Sat, 5 Sep 2026 15:13:46 +0200 Subject: [PATCH 38/41] [Common] Preserve shared-memory pointer provenance in TMA kernels (#3482) * Use the shmem alignment operator consistently Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Oleg Goncharov Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/common/cast/core/grouped_tma.cuh | 6 ------ .../fp8_blockwise/group_quantize_fp8_blockwise.cuh | 10 +++++----- .../common/cast/mxfp8/group_quantize_mxfp8.cuh | 4 ++-- .../common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh | 4 ++-- .../specialized/quantize_transpose_nvfp4_tuned_1D.cuh | 4 ++-- 5 files changed, 11 insertions(+), 17 deletions(-) diff --git a/transformer_engine/common/cast/core/grouped_tma.cuh b/transformer_engine/common/cast/core/grouped_tma.cuh index 61218d654a..8603fd1fd2 100644 --- a/transformer_engine/common/cast/core/grouped_tma.cuh +++ b/transformer_engine/common/cast/core/grouped_tma.cuh @@ -53,12 +53,6 @@ inline bool dimensions_supported_by_TMA(const Tensor *const t) { return cols % alignment_requirement == 0; } -__device__ __forceinline__ unsigned char *align_smem_ptr_per_TMA_requirements(unsigned char *p) { - size_t addr = reinterpret_cast(p); - addr = (addr + TMA_SHMEM_ALIGNMENT - 1) & ~(TMA_SHMEM_ALIGNMENT - 1); - return reinterpret_cast(addr); -} - // Copies the base tensor map to shmem, modifies the copy, stores the modified tensor map at index __device__ __forceinline__ void modify_base_tensor_map(const CUtensorMap base_tensor_map, CUtensorMap *global_tensor_map, diff --git a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh index 203f569471..31feaf833d 100644 --- a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh +++ b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh @@ -314,9 +314,9 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma // Dynamic smem holds the IType input tile (TMA dest, must be 128 B aligned). // warp_amaxes and tma_mbar are static smem. - extern __shared__ unsigned char smem_raw_2d_tma[]; - IType(*smem_in)[kTileDim] = reinterpret_cast( - common::align_smem_ptr_per_TMA_requirements(smem_raw_2d_tma)); + extern __shared__ char smem_raw_2d_tma[]; + IType(*smem_in)[kTileDim] = + reinterpret_cast(align_up(smem_raw_2d_tma, TMA_SHMEM_ALIGNMENT)); __shared__ CType warp_amaxes[kNumWarps]; __shared__ size_t warp_offset_partials[kNumWarps]; @@ -603,8 +603,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke // Dynamic smem: IType[kTileDim][kTileDim], 128 B aligned for TMA. Static smem // (smem_T when CW, tma_mbar) lives outside the dynamic region. - extern __shared__ unsigned char smem_raw_1d_tma[]; - unsigned char* smem_base = common::align_smem_ptr_per_TMA_requirements(smem_raw_1d_tma); + extern __shared__ char smem_raw_1d_tma[]; + char* smem_base = align_up(smem_raw_1d_tma, TMA_SHMEM_ALIGNMENT); IType(*smem)[kTileDim] = reinterpret_cast(smem_base); __shared__ uint64_t tma_mbar; diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 980a77db0a..b0383d95f3 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -704,8 +704,8 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_ constexpr size_t out_mem_rowwise = (ROWWISE_SCALING ? buff_size_aligned_out : 0); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned - extern __shared__ unsigned char dynamic_shmem[]; - unsigned char *dshmem = align_smem_ptr_per_TMA_requirements(dynamic_shmem); + extern __shared__ char dynamic_shmem[]; + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *sIn_ptr = reinterpret_cast(dshmem); diff --git a/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh index 24f84fa359..878fc93107 100644 --- a/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh @@ -209,8 +209,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_scaled_swiglu_mxfp8_k DIVUP_TO_MULTIPLE(CHUNK_DIM_Y * sizeof(float), TMA_SHMEM_ALIGNMENT); // shmem layout: [act input][gate input][colwise output][prob] - extern __shared__ unsigned char dynamic_shmem[]; - unsigned char *dshmem = align_smem_ptr_per_TMA_requirements(dynamic_shmem); + extern __shared__ char dynamic_shmem[]; + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); IType *sInAct_ptr = reinterpret_cast(dshmem); IType *sInGate_ptr = reinterpret_cast(dshmem + buff_size_aligned_in); diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh index cdd0d4916a..58af1f7938 100644 --- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh @@ -410,8 +410,8 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(nvfp4_scale_t), TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned - extern __shared__ unsigned char dynamic_shmem[]; - unsigned char *dshmem = common::align_smem_ptr_per_TMA_requirements(dynamic_shmem); + extern __shared__ char dynamic_shmem[]; + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); IType *sIn_ptr = reinterpret_cast(dshmem); fp4e2m1x2 *sOut_ptr = reinterpret_cast(dshmem + in_mem); From dc8909da2e2c9cebba3f2583ea2a47b531ff1bcb Mon Sep 17 00:00:00 2001 From: Evgeny Tsykunov Date: Tue, 8 Sep 2026 12:24:29 +0200 Subject: [PATCH 39/41] Update Linear docstring referring torch.Linear (#3491) * Warn Linear argument documentation Signed-off-by: Evgeny * Update docstring, remove warning Signed-off-by: Evgeny --------- Signed-off-by: Evgeny --- docs/api/pytorch.rst | 2 +- transformer_engine/pytorch/module/linear.py | 8 +++++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 54981c9086..739a9864b9 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -6,7 +6,7 @@ PyTorch ======= -.. autoapiclass:: transformer_engine.pytorch.Linear(in_features, out_features, bias=True, **kwargs) +.. autoapiclass:: transformer_engine.pytorch.Linear(in_features, out_features, **kwargs) :members: forward, set_tensor_parallel_group .. autoapiclass:: transformer_engine.pytorch.GroupedLinear(in_features, out_features, bias=True, **kwargs) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 55fc69ef7f..94de69e975 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1901,7 +1901,13 @@ def _linear_eager( class Linear(TransformerEngineBaseModule): """Applies a linear transformation to the incoming data :math:`y = xA^T + b` - On NVIDIA GPUs it is a drop-in replacement for ``torch.nn.Linear``. + On NVIDIA GPUs, this module implements the same linear transformation as + ``torch.nn.Linear``. + + .. note:: + + Its constructor signature differs from ``torch.nn.Linear``. Pass optional + arguments, including ``bias``, by keyword. Parameters ---------- From f2bec2314c10754d4ea5a1fd0cb0a8fa1442bf6d Mon Sep 17 00:00:00 2001 From: Evgeny Tsykunov Date: Tue, 8 Sep 2026 12:25:51 +0200 Subject: [PATCH 40/41] [PyTorch] Document fine-grained quantization recipes (#3336) * Fine-grained recipe docs Signed-off-by: Evgeny * Update docs/features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Evgeny Tsykunov * resolve comments Signed-off-by: Evgeny * Rework heterogeneous quantization docs into mixed-format quantization; add per-recipe Quantizer sections Signed-off-by: Pawel Gadzinski * Align mixed-format quantization diagrams with shared diagram-colors.css and dark mode Signed-off-by: Pawel Gadzinski * Rename mixed-format quantization docs to fine-grained quantization recipes Signed-off-by: Pawel Gadzinski * Simplify quantizer factory paragraph Signed-off-by: Pawel Gadzinski * Generalize the fallback-path description Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Evgeny Signed-off-by: Evgeny Tsykunov Signed-off-by: Pawel Gadzinski Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Pawel Gadzinski --- docs/_static/css/diagram-colors.css | 15 + docs/api/pytorch.rst | 18 + ...torch_fine_grained_quantization_example.py | 130 +++++++ .../fine_grained_quantization.rst | 364 ++++++++++++++++++ .../img/fine_grained_assignments.svg | 95 +++++ .../img/fine_grained_linear_mapping.svg | 59 +++ .../img/hybrid_columnwise_source.svg | 74 ++++ .../img/hybrid_quantizer.svg | 58 +++ .../fp8_blockwise_scaling.rst | 31 ++ .../fp8_current_scaling.rst | 53 ++- .../fp8_delayed_scaling.rst | 67 +++- .../features/low_precision_training/index.rst | 1 + .../introduction/introduction.rst | 68 ++++ .../low_precision_training/mxfp8/mxfp8.rst | 53 ++- .../low_precision_training/nvfp4/nvfp4.rst | 54 +++ .../performance_considerations.rst | 56 ++- transformer_engine/common/recipe/__init__.py | 12 +- .../pytorch/tensor/hybrid_tensor.py | 4 + .../pytorch/tensor/identity_tensor.py | 19 + 19 files changed, 1224 insertions(+), 7 deletions(-) create mode 100644 docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py create mode 100644 docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst create mode 100644 docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg create mode 100644 docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg create mode 100644 docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg create mode 100644 docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg diff --git a/docs/_static/css/diagram-colors.css b/docs/_static/css/diagram-colors.css index f5dc7da4dd..9ee5827bd1 100644 --- a/docs/_static/css/diagram-colors.css +++ b/docs/_static/css/diagram-colors.css @@ -279,3 +279,18 @@ html[data-theme="dark"] .subtitle, html[data-theme="dark"] .memory-label { fill: #e0e0e0; } html[data-theme="dark"] .connector { stroke: #bdbdbd; } + +/* fine_grained_quantization diagrams */ +html[data-theme="dark"] .fmt-mxfp8 { fill: #10375c; stroke: #64b5f6; } +html[data-theme="dark"] .fmt-nvfp4 { fill: #5c3a10; stroke: #ffb74d; } +html[data-theme="dark"] .fmt-bf16 { fill: #1e4620; stroke: #81c784; } +html[data-theme="dark"] .fmt-mxfp8-text { fill: #90caf9; } +html[data-theme="dark"] .fmt-nvfp4-text { fill: #ffcc80; } +html[data-theme="dark"] .fmt-bf16-text { fill: #a5d6a7; } +html[data-theme="dark"] .source { fill: #3a2f5c; stroke: #b39ddb; } +html[data-theme="dark"] .quantizer { fill: #1e4620; stroke: #81c784; } +html[data-theme="dark"] .representation { fill: #10375c; stroke: #64b5f6; } +html[data-theme="dark"] .dequantize { fill: #5c3a10; stroke: #ffb74d; } +html[data-theme="dark"] .rowlabel, +html[data-theme="dark"] .legend, +html[data-theme="dark"] .op { fill: #e0e0e0; } diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 739a9864b9..a6afa2d0cc 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -112,6 +112,12 @@ Communication-computation overlap :members: FP8, NONE +Fine-grained quantization recipes +--------------------------------- + +.. autoapiclass:: transformer_engine.pytorch.QuantizerRole(module_type="", tensor_type="", name="") + + Quantized tensors ----------------- @@ -129,6 +135,10 @@ Quantized tensors .. autoapiclass:: transformer_engine.pytorch.NVFP4TensorStorage(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, amax_rowwise, amax_columnwise, fp4_dtype, quantizer) +.. autoapiclass:: transformer_engine.pytorch.HybridQuantizedTensorStorage(*, rowwise_storage, columnwise_storage, quantizer, fake_dtype=None) + +.. autoapiclass:: transformer_engine.pytorch.IdentityTensorStorage(*, hp_data, fake_dtype=None, quantizer=None) + .. autoapiclass:: transformer_engine.pytorch.Float8Tensor(shape, dtype, data, fp8_scale_inv, fp8_dtype, requires_grad=False, data_transpose=None, quantizer=None) .. autoapiclass:: transformer_engine.pytorch.MXFP8Tensor(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, fp8_dtype, quantizer) @@ -137,6 +147,10 @@ Quantized tensors .. autoapiclass:: transformer_engine.pytorch.NVFP4Tensor(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, amax_rowwise, amax_columnwise, fp4_dtype, quantizer) +.. autoapiclass:: transformer_engine.pytorch.HybridQuantizedTensor(shape, dtype, *, rowwise_storage, columnwise_storage, quantizer, requires_grad=False, device=None) + +.. autoapiclass:: transformer_engine.pytorch.IdentityTensor(shape, dtype, *, hp_data, quantizer=None, requires_grad=False, device=None) + Quantizers ---------- @@ -153,6 +167,10 @@ Quantizers .. autoapiclass:: transformer_engine.pytorch.NVFP4Quantizer(fp4_dtype, *, rowwise=True, columnwise=True, **kwargs) +.. autoapiclass:: transformer_engine.pytorch.HybridQuantizer(*, rowwise_quantizer, columnwise_quantizer, columnwise_source="original") + +.. autoapiclass:: transformer_engine.pytorch.IdentityQuantizer(*, dtype=None, rowwise=True, columnwise=True) + Tensor saving and restoring functions ------------------------------------- diff --git a/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py b/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py new file mode 100644 index 0000000000..b511fc41dc --- /dev/null +++ b/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py @@ -0,0 +1,130 @@ +#!/usr/bin/env python3 +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Runnable fine-grained quantization recipe example. + +The factory assigns one precision to each ``demo.fc1`` Linear GEMM: + +* fprop: ``weight.row(MXFP8) x input.row(MXFP8)`` +* dgrad: ``weight.col(NVFP4) x grad_output.row(NVFP4)`` +* wgrad: ``input.col(original BF16) x grad_output.col(original BF16)`` + +``demo.fc2`` runs every GEMM in high precision. ``demo.output`` is not +special-cased and therefore exercises the MXFP8 base-factory fallback. + +Run from the Transformer Engine repository root:: + + python docs/examples/fine_grained_quantization/\ + pytorch_fine_grained_quantization_example.py +""" + +from __future__ import annotations + +import torch +import transformer_engine.pytorch as te + + +def require_supported_hardware() -> None: + """Fail early with TE's reason when either required format is unavailable.""" + + if not torch.cuda.is_available(): + raise SystemExit("This example requires a CUDA-capable NVIDIA GPU.") + + failures = [] + for name, check in ( + ("MXFP8", te.is_mxfp8_available), + ("NVFP4", te.is_nvfp4_available), + ): + available, reason = check(return_reason=True) + if not available: + failures.append(f"{name}: {reason}") + if failures: + raise SystemExit("Required formats are unavailable: " + "; ".join(failures)) + + +require_supported_hardware() + +# START_FINE_GRAINED_QUANTIZATION_EXAMPLE + +from typing import Optional + +import torch + +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import CustomRecipe +from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + mxfp8_factory, + nvfp4_factory, +) + + +THREE_FORMAT_MODULE = "demo.fc1" +HIGH_PRECISION_MODULE = "demo.fc2" +BASE_FACTORY = mxfp8_factory + + +def quantizer_factory(role: Optional[te.QuantizerRole]): + """Return a fresh quantizer for every role, including ``None``. + + ``BASE_FACTORY`` makes the factory total: unknown roles, future role values, + and untargeted modules all retain valid MXFP8 behavior. + """ + + if role is not None and role.name == THREE_FORMAT_MODULE: + # Constructing fresh child quantizers for every call is recommended. + if role.tensor_type == "input": + # Wgrad retains the original BF16 input. + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="original", + ) + if role.tensor_type == "weight": + # Dgrad uses NVFP4 quantized from the dequantized MXFP8 fprop weight. + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=nvfp4_factory(role), + columnwise_source="rowwise_dequantized", + ) + if role.tensor_type == "grad_output": + # Dgrad uses NVFP4 while wgrad retains the original BF16 gradient. + return te.HybridQuantizer( + rowwise_quantizer=nvfp4_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="original", + ) + + if role is not None and role.name == HIGH_PRECISION_MODULE: + return te.IdentityQuantizer() + + return BASE_FACTORY(role) + + +linear_options = {"bias": False, "params_dtype": torch.bfloat16, "device": "cuda"} +model = torch.nn.Sequential( + te.Linear(128, 256, name=THREE_FORMAT_MODULE, **linear_options), + torch.nn.GELU(), + te.Linear(256, 256, name=HIGH_PRECISION_MODULE, **linear_options), + torch.nn.GELU(), + te.Linear(256, 128, name="demo.output", **linear_options), +) +inputs = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) +recipe = CustomRecipe(qfactory=quantizer_factory) + +with te.autocast(enabled=True, recipe=recipe): + outputs = model(inputs) + +loss = outputs.float().square().mean() +loss.backward() + +# END_FINE_GRAINED_QUANTIZATION_EXAMPLE + +gradients = [inputs.grad, *(parameter.grad for parameter in model.parameters())] +assert all(gradient is not None for gradient in gradients) +assert all(torch.isfinite(gradient).all() for gradient in gradients) + +print(f"GPU: {torch.cuda.get_device_name()}") +print(f"TE Linear names: {[model[index].name for index in (0, 2, 4)]}") +print(f"loss: {loss.item():.6f}; forward and backward completed") diff --git a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst new file mode 100644 index 0000000000..301b655b3f --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst @@ -0,0 +1,364 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _fine-grained-quantization-recipes: +.. _heterogeneous-quantization-recipes: + +Fine-grained quantization recipes +================================= + +Standard TE recipes quantize the whole model the same way. That is often too +coarse: one sensitive layer may need BF16 while the rest runs in MXFP8, or a +gradient GEMM may tolerate a cheaper format than the forward pass. Fine-grained +recipes lift this restriction: you write a small factory function that picks a +quantizer for each slot TE asks about, and pass it via +:class:`~transformer_engine.common.recipe.CustomRecipe` to the usual +:class:`~transformer_engine.pytorch.autocast`. +"Fine-grained" refers to the granularity of that choice (per module, tensor +role, and GEMM direction), not to the block size of the scaling factors. + +.. warning:: + + Fine-grained recipes are currently available only in the PyTorch API of + TE. + +.. warning:: + + Fine-grained recipes and their construction APIs are experimental: API, + validation, and kernel coverage may change without notice. This guide does + not define a supported recipe or an expected accuracy/performance ordering. + + +Example: mixing MXFP8, NVFP4, and BF16 +-------------------------------------- + +The `runnable example `__ +makes the following assignments: + +.. raw:: html + :file: img/fine_grained_assignments.svg + +*Figure 1. Precision assignments per module and GEMM used throughout this +guide.* + +A minimal factory implementing these assignments, plugged into the standard +TE autocast path: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + import transformer_engine.pytorch as te + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + mxfp8_factory, + nvfp4_factory, + ) + + + def quantizer_factory(role): + if role is not None and role.name == "demo.fc1": + if role.tensor_type == "input": + # wgrad keeps the original BF16 input + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="original", + ) + if role.tensor_type == "weight": + # fprop in MXFP8, dgrad in NVFP4 + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=nvfp4_factory(role), + columnwise_source="rowwise_dequantized", + ) + if role.tensor_type == "grad_output": + # dgrad in NVFP4, wgrad keeps the original BF16 gradient + return te.HybridQuantizer( + rowwise_quantizer=nvfp4_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="original", + ) + if role is not None and role.name == "demo.fc2": + return te.IdentityQuantizer() # whole module stays in BF16 + return mxfp8_factory(role) # every other TE module in MXFP8 + + + recipe = CustomRecipe(qfactory=quantizer_factory) + + with te.autocast(enabled=True, recipe=recipe): + output = model(inputs) + +The complete, runnable version is available +`on GitHub `__ +(requires Blackwell or later); run it from the repository root after +installing TE: + +.. code-block:: bash + + python docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py + +CustomRecipe and quantizer factory +---------------------------------- + +:class:`~transformer_engine.common.recipe.CustomRecipe` is used like any +other TE recipe (``DelayedScaling``, ``MXFP8BlockScaling``, ...), but carries +no quantization logic of its own: TE asks your ``qfactory`` for a quantizer +whenever a module needs one. + +Each TE module defines an ordered role list for the forward and backward +quantizer slots it needs. When module recipe state is initialized or rebuilt, +a ``CustomRecipe`` calls ``qfactory(role)`` once for every slot in that list. +It does not call the factory on every unchanged forward. + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + # QuantizerRole describes the slot being configured (fields below): + # + # @dataclasses.dataclass(frozen=True) + # class QuantizerRole: + # module_type: str = "" + # tensor_type: str = "" + # name: str = "" + + + def quantizer_factory(role: Optional[te.QuantizerRole]): + # construct a fresh quantizer on every call + ... + # Boundary slots may pass role=None or a role with empty fields, so + # always end with a default that covers every remaining role. + return mxfp8_factory(role) + + + # The factory plugs into the standard TE autocast path: + recipe = CustomRecipe(qfactory=quantizer_factory) + + with te.autocast(enabled=True, recipe=recipe): + output = model(inputs) + + **Module type** + + The kind of TE module that owns the slot, filled in by TE itself: + + * ``"linear"`` — ``Linear``, ``LayerNormLinear``, ``fc1``/``fc2`` in + ``LayerNormMLP``, ``qkv``/``proj`` in ``MultiheadAttention``; + * ``"grouped_linear"`` — ``GroupedLinear``; + * ``"dpa"`` — ``DotProductAttention``. + + **Tensor type** + + Which tensor of that module the quantizer will process, also filled in + by TE. For ``"linear"`` and ``"grouped_linear"``: + + * ``"input"`` — the activation (fprop, wgrad); + * ``"weight"`` — (fprop, dgrad); + * ``"grad_output"`` — the incoming gradient (dgrad, wgrad). + + For ``"dpa"``: + + * ``"qkv"`` — the query/key/value tensor; + * ``"s"`` — the softmax output; + * ``"do"`` — the output gradient; + * ``"dp"`` — the gradient of ``"s"``. + + **Name** + + The identity of one concrete module instance, supplied by the caller: + ``te.Linear(..., name="decoder.39.fc2")``. Composite TE modules may + append suffixes such as ``.fc1``, ``.fc2``, and ``.proj``. + + The role vocabulary is experimental and may grow between releases — + one more reason to end the factory with a total default. Treat the role + strings as selectors, not a fixed enumeration. Prefer a module-level + function for the factory itself, so that launchers and checkpointing + setups can import or pickle it. + + TE provides factories for its native quantizers in + ``transformer_engine.pytorch.custom_recipes.quantizer_factories`` + (``mxfp8_factory``, ``nvfp4_factory``, ...). They can be used as + defaults or to construct ``HybridQuantizer`` children. Additional + specialized recipes are available in + ``transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo``. + + The factory is not limited to TE-native quantizers: it may return your + own :class:`~transformer_engine.pytorch.Quantizer` subclass, and custom + quantizers can also serve as ``HybridQuantizer`` children. The GEMMs + still need to receive representations in formats they support. + +HybridQuantizer +--------------- + +During training, each tensor of a ``Linear`` or ``GroupedLinear`` layer feeds +two different GEMMs: its rowwise representation feeds one, its columnwise +representation the other (the exact operand layout is described in the +:doc:`Introduction <../introduction/introduction>`). Since those two GEMMs may +want different formats, the tensor needs a quantizer per direction: +:class:`~transformer_engine.pytorch.HybridQuantizer` composes a rowwise and a +columnwise quantizer, and its output, +:class:`~transformer_engine.pytorch.HybridQuantizedTensor`, composes the +corresponding representations. + +.. tabs:: + + .. tab:: PyTorch + + The following is pseudocode illustrating the composition: + + .. code-block:: text + + quantizer = te.HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=DType.kFloat8E4M3), + columnwise_quantizer=NVFP4Quantizer(), + columnwise_source="original", # or "rowwise_dequantized" + ) + + # Quantization yields a HybridQuantizedTensor whose rowwise + # representation is MXFP8 and columnwise representation is NVFP4; + # each GEMM consumes the representation it needs. + qtensor = quantizer(tensor) + +.. raw:: html + :file: img/hybrid_quantizer.svg + +*Figure 2. HybridQuantizer composes a rowwise and a columnwise quantizer; each +representation of the result feeds a different GEMM.* + +Choosing the columnwise source +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +``columnwise_source`` is a separate numerical recipe choice that controls the +source for the columnwise representation: + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Value + - Columnwise source + * - ``"original"`` + - The original high-precision tensor. + * - ``"rowwise_dequantized"`` + - Dequantized rowwise representation. + +.. raw:: html + :file: img/hybrid_columnwise_source.svg + +*Figure 3. The columnwise representation can be derived from the original +high-precision tensor or from the dequantized rowwise representation.* + +For forward inputs and weights, ``"rowwise_dequantized"`` derives the backward +representation from the value consumed in the forward direction. This +can improve forward/backward numerical consistency and may affect convergence. +It does not recover information discarded by rowwise quantization. +``"original"`` instead derives both representations from the original tensor. +Choose the provenance as part of the numerical recipe. + +IdentityQuantizer +----------------- + +:class:`~transformer_engine.pytorch.IdentityQuantizer` stores its input in the +held compute dtype, typically BF16, FP16, or FP32. It can keep a complete slot +in high precision or act as one child of a ``HybridQuantizer``: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + # whole slot in high precision (e.g. a module kept in BF16) + quantizer = te.IdentityQuantizer() + + # one direction in high precision, the other quantized + quantizer = te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + +Note that in the second example the columnwise direction is high precision but +holds the value reconstructed from MXFP8, not the original input — see +`Choosing the columnwise source`_ above. + +Example: one format per GEMM +---------------------------- + +A natural way to design a recipe is to pick one format for each GEMM. To +translate that into quantizers, look at what each GEMM consumes — both of its +operands must be in that GEMM's format: + +* **fprop** consumes ``input.rowwise`` and ``weight.rowwise``; +* **dgrad** consumes ``grad_output.rowwise`` and ``weight.columnwise``; +* **wgrad** consumes ``input.columnwise`` and ``grad_output.columnwise``. + +Reading the same table per tensor gives the ``HybridQuantizer`` for each role. +For the example assignments (fprop in MXFP8, dgrad in NVFP4, wgrad in BF16): + +.. code-block:: text + + input = HybridQuantizer(rowwise=MXFP8, columnwise=BF16) # fprop | wgrad + weight = HybridQuantizer(rowwise=MXFP8, columnwise=NVFP4) # fprop | dgrad + grad_output = HybridQuantizer(rowwise=NVFP4, columnwise=BF16) # dgrad | wgrad + +.. raw:: html + :file: img/fine_grained_linear_mapping.svg + +*Figure 4. Each GEMM consumes one representation of each of its two operand +tensors; giving both operands the same format sets that GEMM's precision.* + +If two directions use the same quantizer configuration, a plain quantizer may +replace the corresponding hybrid; one factory may return both plain and hybrid +quantizers. +The two operands of each GEMM still need a combination supported by that GEMM +backend. TE may reject incompatible quantizer pairs or unsupported layouts. + +.. note:: + + On supported hardware these recipes run TE's regular quantized kernels: the + tensors are quantized on the GPU and the GEMMs execute in the selected + low-precision formats. TE does not fall back to fake quantization + (quantize-dequantize followed by a high-precision GEMM). + + +Validating and optimizing a recipe +---------------------------------- + +The factory API can express more recipes than TE has kernels for, so any +assignment lands in one of three buckets: + +* **Fast** — quantization hits TE's fused kernels and every GEMM runs a + native low-precision implementation. +* **Correct but potentially unoptimized** — the recipe executes, but some + selected paths may not have fused or optimized implementations in the + current TE release. For example, ``HybridQuantizer`` may produce its rowwise + and columnwise representations in separate kernel launches; future releases + may fuse this work. +* **Rejected** — the two operands of some GEMM end up in a combination of + formats or layouts that no GEMM backend supports, and TE raises an error. + This can happen with plain and hybrid quantizers alike. + +Before adopting a recipe for a real workload, check that: + +* it executes at all on the target GPU, software version, and modules; +* it runs on optimized kernels rather than fallback paths; +* accuracy and convergence hold on the target model and distributed setup; +* throughput and memory actually improve on the target workload. + +The unoptimized paths are still useful: accuracy and convergence experiments can run +on them before dedicated kernels exist, so the precision of each GEMM can be +treated as an accuracy/performance trade-off to explore. + +API reference +------------- + +See the :doc:`PyTorch API <../../../api/pytorch>` for ``QuantizerRole``, +``HybridQuantizer``, ``IdentityQuantizer``, and their returned tensor types. +See the :doc:`Common API <../../../api/common>` for ``CustomRecipe``. diff --git a/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg new file mode 100644 index 0000000000..b9aa48333d --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg @@ -0,0 +1,95 @@ + + + Precision assignment by tensor role and module + Each tensor role provides a rowwise and a columnwise representation, consumed by fprop, dgrad, and wgrad GEMMs. demo.fc1: input is MXFP8 rowwise and original BF16 columnwise; weight is MXFP8 rowwise and NVFP4 columnwise; grad_output is NVFP4 rowwise and original BF16 columnwise. demo.fc2 keeps every tensor in BF16. Other TE modules use MXFP8 everywhere. + + + + + Precision assignment by tensor role and module + + demo.fc1 + demo.fc2 + Other TE modules + + + + input + rowwise (fprop) + + MXFP8 + + BF16 + + MXFP8 + + columnwise (wgrad) + + BF16 (original) + + BF16 + + MXFP8 + + + + + weight + rowwise (fprop) + + MXFP8 + + BF16 + + MXFP8 + + columnwise (dgrad) + + NVFP4 + + BF16 + + MXFP8 + + + + + grad_output + rowwise (dgrad) + + NVFP4 + + BF16 + + MXFP8 + + columnwise (wgrad) + + BF16 (original) + + BF16 + + MXFP8 + + + + + MXFP8 + + NVFP4 + + BF16 (high precision) + + diff --git a/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg new file mode 100644 index 0000000000..0e17ffdfa4 --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg @@ -0,0 +1,59 @@ + + + Per-GEMM formats and the operands each GEMM consumes + Three GEMM cards. Fprop consumes input.rowwise and weight.rowwise, both MXFP8. Dgrad consumes weight.columnwise and grad_output.rowwise, both NVFP4. Wgrad consumes input.columnwise and grad_output.columnwise, both BF16. + + + + + + + fprop + format: MXFP8 + + input.rowwise + MXFP8 + × + + weight.rowwise + MXFP8 + + + + + dgrad + format: NVFP4 + + grad_output.rowwise + NVFP4 + × + + weight.columnwise + NVFP4 + + + + + wgrad + format: BF16 + + input.columnwise + BF16 + × + + grad_output.columnwise + BF16 + + diff --git a/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg new file mode 100644 index 0000000000..ebc077eb04 --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg @@ -0,0 +1,74 @@ + + + Hybrid quantizer columnwise source choices + With original provenance, both quantizers consume the original high-precision tensor. With rowwise-dequantized provenance, the columnwise quantizer consumes the dequantized rowwise representation. + + + + + + + + Choosing the columnwise source + + + + columnwise_source="original" + + + High-precision tensor + + + + same original source + + + Rowwise quantizer + + Columnwise quantizer + + + + + Rowwise + representation + + Columnwise + representation + + + + + columnwise_source="rowwise_dequantized" + + + High-precision tensor + + + + Rowwise quantizer + + + + Rowwise + representation + + + + Dequantize + + + + Columnwise quantizer + + + Columnwise + representation + + diff --git a/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg new file mode 100644 index 0000000000..6e542306c7 --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg @@ -0,0 +1,58 @@ + + + HybridQuantizer data flow + A high-precision tensor enters a HybridQuantizer whose rowwise child is an MXFP8 quantizer and columnwise child is an NVFP4 quantizer. The result is a HybridQuantizedTensor with an MXFP8 rowwise representation and an NVFP4 columnwise representation, each consumed by a different GEMM. + + + + + + + + + + tensor + high precision (BF16) + + + + + + + HybridQuantizer + + rowwise_quantizer + MXFP8Quantizer + + columnwise_quantizer + NVFP4Quantizer + + + + + + + HybridQuantizedTensor + + rowwise + MXFP8 + + columnwise + NVFP4 + + + + + GEMM 1 + GEMM 2 + diff --git a/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst b/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst index 48d17db8d5..557a4b09e8 100644 --- a/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst +++ b/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst @@ -180,6 +180,37 @@ Blackwell and later (SM >= 10.0) – the recipe is emulated with MXFP8. Note tha ---- + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + Blockwise scaling uses + :class:`~transformer_engine.pytorch.Float8BlockQuantizer`. Each block of + the tensor gets its own power-of-two scale: ``block_scaling_dim=1`` + scales 1x128 blocks, ``block_scaling_dim=2`` (the default) scales + 128x128 blocks. This recipe is not available in TE/JAX. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.Float8BlockQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=1, + ) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + Developer Notes --------------- diff --git a/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst b/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst index cac3792194..2436a07566 100644 --- a/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst +++ b/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst @@ -164,6 +164,57 @@ Here's how to use FP8 Current Scaling recipe in PyTorch and JAX: ---- + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + Current scaling uses + :class:`~transformer_engine.pytorch.Float8CurrentScalingQuantizer`. It + needs no external state: at each call it computes the amax of the input + tensor, derives the scale from it, and then quantizes. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.Float8CurrentScalingQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + device="cuda", + ) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + + .. tab:: JAX + + Current scaling uses ``CurrentScaleQuantizer``. At each call it computes + the amax of the input tensor, derives the scale from it, and then + quantizes. + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() + Developer Notes --------------- @@ -177,4 +228,4 @@ On Blackwell and later, rowwise and columnwise tensors share the same memory lay so all-gather of columnwise tensors is directly supported. For Hopper and Ada, all-gather of transposed FP8 tensors is not supported. -The rowwise tensor is gathered first, then transposed to columnwise format. \ No newline at end of file +The rowwise tensor is gathered first, then transposed to columnwise format. diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst index d39787f6f5..99a379eed1 100644 --- a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst +++ b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst @@ -160,4 +160,69 @@ However, amax reduction works slightly differently in different frameworks. Supported devices ----------------- -Ada and later (SM 8.9+) \ No newline at end of file +Ada and later (SM 8.9+) + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + Delayed scaling uses + :class:`~transformer_engine.pytorch.Float8Quantizer`. It does not + compute the scaling factor from the current tensor: one-element + ``scale`` and ``amax`` buffers are supplied at construction. + Quantization applies the given scale and records the tensor's amax into + the ``amax`` buffer. + + During training both buffers are views into the recipe state: ``scale`` + into its per-quantizer scale vector, ``amax`` into the current row of + its ``(amax_history_len, num_quantizers)`` amax history. At the end of + each step the recipe state computes a new scale from the history (its + max or most recent entry, per ``amax_compute_algo``), rolls the history + by one slot, and zeroes the current row — all in place, so the views + held by the quantizer stay valid for the whole training run. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.Float8Quantizer( + scale=torch.ones(1, device="cuda"), + amax=torch.zeros(1, device="cuda"), + fp8_dtype=te.DType.kFloat8E4M3, + ) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + + .. tab:: JAX + + Delayed scaling uses ``DelayedScaleQuantizer``. The ``scale`` and the + ``amax_history`` (1024 entries by default) are fields of the quantizer + itself, carried through JAX transformations as its pytree state. Each + ``quantize()`` call applies the current ``scale``, then updates the + state: the tensor's amax is written into the history, a new scale is + computed from the history (max or most-recent entry, per + ``amax_compute_algo``), and the history is rolled by one slot. + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.DELAYED_TENSOR_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() diff --git a/docs/features/low_precision_training/index.rst b/docs/features/low_precision_training/index.rst index 0a798f1364..b9649c00a4 100644 --- a/docs/features/low_precision_training/index.rst +++ b/docs/features/low_precision_training/index.rst @@ -15,4 +15,5 @@ Low precision training fp8_blockwise_scaling/fp8_blockwise_scaling.rst mxfp8/mxfp8.rst nvfp4/nvfp4.rst + fine_grained_quantization/fine_grained_quantization.rst speedups.rst diff --git a/docs/features/low_precision_training/introduction/introduction.rst b/docs/features/low_precision_training/introduction/introduction.rst index fba7796ece..2255308b04 100644 --- a/docs/features/low_precision_training/introduction/introduction.rst +++ b/docs/features/low_precision_training/introduction/introduction.rst @@ -283,3 +283,71 @@ so GEMM with tensors ``A`` and ``B`` returns ``B * A^T``. :file: img/fp8_linear_flow.svg *Figure 4: Forward pass of a Linear layer with low precision data flow.* + +Quantizers +---------- + +Every recipe implements its quantization logic in a **quantizer** — an object +that converts a high-precision tensor into a quantized one. TE modules create +and use quantizers internally according to the active recipe, but a quantizer +can also be used directly: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + qtensor = quantizer(tensor) # quantize + roundtrip = qtensor.dequantize() # back to high precision + + The main parts of the interface are: + + * ``quantize(tensor)`` — quantizes a high-precision tensor and returns a + ``QuantizedTensor``; calling the quantizer (``quantizer(tensor)``) is + a shorthand; + * ``update_quantized(src, dst)`` — quantizes ``src`` in place into an + already-allocated quantized tensor ``dst``; + * ``make_empty(shape)`` — allocates an uninitialized quantized tensor to + be filled later; + * ``rowwise_usage`` / ``columnwise_usage`` — flags selecting which of + the two GEMM-oriented representations the produced tensor holds; + * the returned ``QuantizedTensor`` supports ``dequantize()`` back to + high precision. + + .. tab:: JAX + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() + + The main parts of the interface are: + + * ``quantize(x, is_rowwise=..., is_colwise=...)`` — quantizes a tensor + and returns a ``ScaledTensor`` holding the requested representations + (the default comes from the quantizer's ``q_layout``); + * the returned ``ScaledTensor`` supports ``dequantize()`` back to high + precision; + * quantizers are registered pytrees, so they can be passed through JAX + transformations. + +Each recipe section ends with a short description of that recipe's quantizer. diff --git a/docs/features/low_precision_training/mxfp8/mxfp8.rst b/docs/features/low_precision_training/mxfp8/mxfp8.rst index 1fbcc43af9..1827d42cb6 100644 --- a/docs/features/low_precision_training/mxfp8/mxfp8.rst +++ b/docs/features/low_precision_training/mxfp8/mxfp8.rst @@ -152,6 +152,57 @@ SM 10.0, SM 10.3 ---- + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + MXFP8 uses :class:`~transformer_engine.pytorch.MXFP8Quantizer`. Every + 32-element block shares one power-of-two (E8M0) scale, computed from the + block's amax at quantization time; no external state is needed. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + + .. tab:: JAX + + MXFP8 uses ``BlockScaleQuantizer`` — the JAX quantizer for block-based + scaling, selected by ``ScalingMode.MXFP8_1D_SCALING``. Instead of one + scale per tensor, the tensor is split along the quantization axis into + 32-element blocks and each block gets its own power-of-two (E8M0) scale, + computed from that block's amax at quantization time. Because the scale + is derived from the current data, no external state (scale buffers or + amax history) is needed. + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() + Developer Notes --------------- @@ -210,4 +261,4 @@ All-gather of columnwise tensors All-gather of columnwise tensors is supported and necessary because: - columnwise quantized tensors cannot be computed from rowwise quantized ones, -- gathering high-precision tensors is avoided in most cases for performance reasons. \ No newline at end of file +- gathering high-precision tensors is avoided in most cases for performance reasons. diff --git a/docs/features/low_precision_training/nvfp4/nvfp4.rst b/docs/features/low_precision_training/nvfp4/nvfp4.rst index 900789b0d3..26d798651d 100644 --- a/docs/features/low_precision_training/nvfp4/nvfp4.rst +++ b/docs/features/low_precision_training/nvfp4/nvfp4.rst @@ -250,6 +250,60 @@ Supported devices ---- + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + NVFP4 uses :class:`~transformer_engine.pytorch.NVFP4Quantizer`. It + implements the two-level scaling described above: an FP8 (E4M3) scale + per 16-element block plus one FP32 scale per tensor. Further keyword + options select the recipe variations from this page (random Hadamard + transforms, stochastic rounding, 2D weight scaling); they are internal + knobs and may change without notice. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + ) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + + .. tab:: JAX + + NVFP4 uses its own ``NVFP4Quantizer``, with the same two-level scaling. + ``ScalingMode.NVFP4_1D_SCALING`` selects per-block scaling only, + ``ScalingMode.NVFP4_2D_SCALING`` adds 2D weight scaling. + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.NVFP4_1D_SCALING, + q_dtype=jnp.float4_e2m1fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() + Developer Notes --------------- diff --git a/docs/features/low_precision_training/performance_considerations/performance_considerations.rst b/docs/features/low_precision_training/performance_considerations/performance_considerations.rst index 2c21799dd6..afa5d16c86 100644 --- a/docs/features/low_precision_training/performance_considerations/performance_considerations.rst +++ b/docs/features/low_precision_training/performance_considerations/performance_considerations.rst @@ -143,6 +143,61 @@ Transformer Engine chooses the best possible fusion internally taking the recipe *Figure 3: Three scenarios of producing quantized tensors in rowwise and columnwise usages.* +**Usages in the quantizer API** + +The usages are visible directly in the quantizer API: + +.. tabs:: + + .. tab:: PyTorch + + At quantization time, the quantizer's ``rowwise_usage`` and + ``columnwise_usage`` flags select which representations ``quantize()`` + produces; when both are set, the representations are computed together + in one fused kernel (scenario 1 above). + + After quantization, ``update_usage()`` on the quantized tensor removes a + representation or, when supported by the format, generates a missing one. + Passing ``rowwise_usage=False`` after the forward pass frees the rowwise + data while keeping the columnwise data for backward. Some formats also + support ``columnwise_usage=True`` to create the columnwise representation + from the data already present (e.g. by a transpose on Hopper — scenario 3 + above); unsupported requests raise an error. Arguments left as ``None`` + preserve the current state. + + .. code-block:: python + + quantizer = te.MXFP8Quantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + ) + + qtensor = quantizer(tensor) # both representations, one fused kernel + + qtensor.update_usage(rowwise_usage=False) # drop rowwise, keep columnwise + + .. tab:: JAX + + The usages are selected when the tensor is quantized: the quantizer's + ``q_layout`` (``QuantizeLayout.ROWWISE``, ``COLWISE``, or + ``ROWWISE_COLWISE``) sets the default, and ``quantize()`` accepts + ``is_rowwise``/``is_colwise`` overrides. Requesting both usages returns + a ``ScaledTensor2x`` holding the two representations. There is no + in-place ``update_usage()``: JAX arrays are immutable, so a + representation is not added or dropped later — unneeded ones are simply + not requested and get dropped by XLA's dead-code elimination. + + .. code-block:: python + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE_COLWISE, + ) + + qtensor = quantizer.quantize(x) # ScaledTensor2x, both representations + rowwise_only = quantizer.quantize(x, is_rowwise=True, is_colwise=False) Memory usage @@ -470,4 +525,3 @@ Actual behavior depends on the recipe and module configuration. *Figure 5: All-gather of quantized tensors for input and gradient tensors. This is one possible scenario — actual behavior varies depending on the recipe and module configuration.* - diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 5d5ce1f6cf..128e8280bb 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -639,13 +639,18 @@ class CustomRecipe(Recipe): ---------- qfactory : Callable Factory callable that returns a quantizer instance *or* a - ``QuantizerRequest`` subclass for a given ``QuantizerRole``. + ``QuantizerRequest`` subclass for a given optional ``QuantizerRole``. The callable is invoked as:: qfactory( - role: QuantizerRole, + role: Optional[QuantizerRole], ) -> Union[Quantizer, QuantizerRequest] + Boundary slots may provide ``None`` or a role with empty fields. The + factory must return a valid object for every call. Return an + ``IdentityQuantizer`` for an intentional high-precision slot instead + of returning ``None``. + ``QuantizerRole`` is a frozen dataclass with the following fields: - ``module_type`` (str): module type (empty string when not set), e.g. @@ -663,7 +668,8 @@ class CustomRecipe(Recipe): See ``transformer_engine.pytorch.quantization.QuantizerRole`` and ``transformer_engine.pytorch.quantization.DelayedScalingRequest`` - for full documentation. + for API details. See :ref:`heterogeneous-quantization-recipes` for + construction rules and direction mapping. backward_override : {None, 'high_precision', 'dequantized'}, default = None Backward precision mode. None does not modify backward behavior, diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 8df2ec8b4b..26d0798b92 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -19,6 +19,10 @@ class HybridQuantizer(Quantizer): """Quantizer that composes rowwise and columnwise representations. + .. warning:: + **EXPERIMENTAL**: ``HybridQuantizer`` is under active development and + its API is subject to change without notice. + When both representations are requested, applies ``rowwise_quantizer`` to produce the rowwise representation and ``columnwise_quantizer`` to produce the columnwise representation. The results are wrapped in a diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index 8310afc653..ec171564fe 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -26,6 +26,10 @@ class IdentityQuantizer(Quantizer): """Quantizer that produces a high-precision passthrough representation. + .. warning:: + **EXPERIMENTAL**: ``IdentityQuantizer`` is under active development and + its API is subject to change without notice. + Returns an :class:`IdentityTensorStorage` (or :class:`IdentityTensor`) holding the tensor directly, without a low-precision encoding. ``general_gemm`` materializes it as a plain tensor, so a GEMM consumes it @@ -174,6 +178,21 @@ class IdentityTensor(IdentityTensorStorage, QuantizedTensor): Presents as a standard tensor of its nominal dtype; internally it just holds data directly in that dtype, without a low-precision encoding. + + Parameters + ---------- + shape : iterable of int + Tensor dimensions. + dtype : torch.dtype + Logical tensor datatype. + hp_data : torch.Tensor + Held high-precision data. + quantizer : IdentityQuantizer, optional + Quantizer that produced the tensor. + requires_grad : bool, default = False + Whether to compute gradients for this tensor. + device : torch.device, optional + Device containing the tensor. """ def __repr__(self, *, tensor_contents=None): From cd245040282ba06d98afc3aa09ef7e71d06ac30a Mon Sep 17 00:00:00 2001 From: Evgeny Tsykunov Date: Tue, 8 Sep 2026 12:29:02 +0200 Subject: [PATCH 41/41] Fix hybrid fp8 test guard (#3492) Signed-off-by: Evgeny --- tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py | 4 ++++ tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 08e762045a..fc482ce5b2 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -1733,6 +1733,10 @@ def test_fused_adam_hybrid_scale_uniform_across_shards(hybrid_recipe_name): ), f"missing hybrid current-scaling directions: {checked}" +@pytest.mark.skipif( + not te.is_fp8_available(), + reason=te.is_fp8_available(return_reason=True)[1], +) def test_fused_adam_hybrid_identity_fp8_master_weights(): """FSDP2 + FusedAdam with Hybrid(FP8 current rowwise, Identity columnwise). diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 39d825e701..c2a9df765e 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -510,6 +510,10 @@ def _hybrid_param_count(): _check_fp8_fsdp2_allgather(model, tols=dict(atol=5e-4, rtol=5e-3)) +@pytest.mark.skipif( + not te.is_fp8_available(), + reason=te.is_fp8_available(return_reason=True)[1], +) def test_distributed_hybrid_identity_all(): """FSDP2 training/all-gather with an all-Identity CustomRecipe.