fix(torchtitan): harden ROCm compiled training against silent NaNs

Preserve eager precision-cast semantics in Inductor graphs, keep DeepSeek MoE
blocks on TorchTitan's mixed compile strategy, and route gfx942 DeepSeek FP8
MLA backward through aiter CK instead of the invalid fmha_v3 ASM path. Also
bind grouped-MM patches to the imported module so CI no longer depends on
import order.

Author: WangLingxun

primus/backends/torchtitan/patches/__init__.py

@@ -31,6 +31,7 @@
     embedding_amp_patches,
     flex_attention_patches,
     fsdp_weight_tying_patches,
+    inductor_precision_casts_patches,
     logger_patches,
     metrics_output_format,
     mock_dataset_patches,

primus/backends/torchtitan/patches/dsv3_v022_perf_patches.py

@@ -10,14 +10,22 @@
 Reduce DSv3-16B HBM on MI355X with torchtitan v0.2.2 without modifying the
 upstream submodule. The balanced-routing field migration
 (``training.debug_moe_force_load_balance`` -> ``debug.moe_force_load_balance``)
-is already handled in the DeepSeek configs on main, so this module only carries
-the two memory fixes:
-
-1. Replace v0.2.2 ``apply_compile`` (capture_scalar_outputs + per-submodule
-   compile) with whole-block compile matching v0.1.0 behavior (drops the
-   fragmented reserved HBM re-introduced by per-submodule compilation).
-2. Replace MoE combine fp32 ``bmm`` with a bf16 weighted sum (drops the fp32
-   activation copy retained across the MoE layers).
+is already handled in the DeepSeek configs on main.
+
+Whole-block compilation is still needed to avoid the fragmented reserved HBM
+caused by v0.2.2's per-submodule compilation. Compiling the entire MoE forward
+is unsafe, however: it pulls expert-parallel token dispatch/combine and
+GroupedExperts FSDP hooks into one inductor graph. On MI300X that silently
+produces a NaN forward loss from the first step.
+
+Use a safe dense-only boundary instead: dense blocks compile as a full graph,
+while MoE TransformerBlocks remain eager. Compiling even the outer MoE block
+with a graph break around ``MoE.forward`` still produces the NaN, whereas
+v0.2.2's per-submodule policy reintroduces the fragmented reserved HBM this
+patch was created to avoid. Keeping MoE blocks eager avoids both failure modes;
+the expensive attention and grouped-GEMM work still uses Primus-Turbo kernels.
+The MoE forward replacement also changes its fp32 ``bmm`` combine into a bf16
+weighted sum, dropping the fp32 activation copy retained across MoE layers.
 """
 
 from __future__ import annotations
@@ -52,36 +60,38 @@ def _compile_enabled(ctx: PatchContext) -> bool:
     return bool(get_param(ctx, "compile.enable", False))
 
 
+def _apply_dense_only_compile(model: nn.Module, compile_config: Any, ep_enabled: bool) -> None:
+    """Whole-block compile dense layers and leave MoE TransformerBlocks eager."""
+    from torchtitan.tools.logging import logger
+
+    del ep_enabled  # MoE blocks stay eager, including EP dispatch.
+    for layer_id, transformer_block in model.layers.named_children():
+        if not transformer_block.moe_enabled:
+            transformer_block = torch.compile(
+                transformer_block,
+                backend=compile_config.backend,
+                fullgraph=True,
+            )
+        model.layers.register_module(layer_id, transformer_block)
+
+    logger.info("Compiling dense TransformerBlocks; leaving MoE TransformerBlocks eager (Primus patch)")
+
+
 @register_patch(
     "torchtitan.dsv3.whole_block_compile",
     backend="torchtitan",
     phase="setup",
-    description="Whole-block torch.compile; leave capture_scalar_outputs disabled",
+    description="Whole-block compile dense layers; leave MoE TransformerBlocks eager",
     condition=lambda ctx: _is_deepseek_model(ctx) and _compile_enabled(ctx),
 )
 def patch_whole_block_compile(ctx: PatchContext) -> None:
-    """Replace v0.2.2 apply_compile with v0.1.0-style whole TransformerBlock compile."""
-    from torchtitan.config.job_config import Compile as CompileConfig
+    """Install low-fragment dense-only compilation."""
     from torchtitan.models.llama4.infra import parallelize as parallelize_module
-    from torchtitan.tools.logging import logger
-
-    def apply_compile_patched(model: nn.Module, compile_config: CompileConfig, ep_enabled: bool) -> None:
-        # Match v0.1.0: do NOT set capture_scalar_outputs=True (avoids fragmentation).
-        for layer_id, transformer_block in model.layers.named_children():
-            fullgraph = not transformer_block.moe_enabled
-            transformer_block = torch.compile(
-                transformer_block,
-                backend=compile_config.backend,
-                fullgraph=fullgraph,
-            )
-            model.layers.register_module(layer_id, transformer_block)
-
-        logger.info("Compiling each TransformerBlock with torch.compile (Primus whole-block patch)")
 
-    parallelize_module.apply_compile = apply_compile_patched
+    parallelize_module.apply_compile = _apply_dense_only_compile
     log_rank_0(
         "[Patch:torchtitan.dsv3.whole_block_compile] "
-        "Patched torchtitan.models.llama4.infra.parallelize.apply_compile",
+        "Patched apply_compile with eager MoE TransformerBlocks",
     )
 
 
@@ -143,8 +153,10 @@ def patch_moe_bf16_combine(ctx: PatchContext) -> None:
     """Replace MoE.forward to drop the fp32 bmm copy in the combine step."""
     import torchtitan.models.moe.moe as moe_module
 
-    moe_module.MoE.forward = _moe_forward_bf16_combine
+    # Keep this explicit eager boundary in case a caller wraps a larger parent
+    # module in torch.compile outside the DeepSeek parallelization path.
+    moe_module.MoE.forward = torch.compiler.disable(_moe_forward_bf16_combine)
     log_rank_0(
         "[Patch:torchtitan.dsv3.moe_bf16_combine] "
-        "Patched torchtitan.models.moe.moe.MoE.forward (bf16 weighted combine)",
+        "Patched MoE.forward (eager boundary, bf16 weighted combine)",
     )

primus/backends/torchtitan/patches/inductor_precision_casts_patches.py

@@ -0,0 +1,53 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Preserve eager low-precision casts in ROCm Inductor graphs.
+
+On the pinned ROCm torch 2.12 build, compiled FSDP BF16 llama/qwen recipes
+produce a finite first loss but NaN gradients unless
+``emulate_precision_casts`` is enabled. Eager and ``aot_eager`` remain finite.
+
+This is not attributed to pytorch#150859: that issue covered float8 rowwise
+training with selective activation checkpointing and is no longer reproducible
+on upstream 2.12 nightly. The plain-BF16 ROCm reproducer here is different even
+though the same compatibility flag fixes it.
+
+The flag is part of Inductor's cache key, so no manual cache invalidation is
+required.
+"""
+
+from primus.core.patches import PatchContext, get_param, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_PREFIX = "[Patch:torchtitan.torch.inductor_precision_casts]"
+
+
+def _compile_enabled(ctx: PatchContext) -> bool:
+    if not bool(get_param(ctx, "compile.enable", False)):
+        return False
+
+    import torch
+
+    return torch.version.hip is not None
+
+
+@register_patch(
+    "torchtitan.torch.inductor_precision_casts",
+    backend="torchtitan",
+    phase="setup",  # before parallelize_fn reaches apply_compile
+    description="Emulate eager precision casts in ROCm Inductor BF16/FSDP graphs",
+    condition=_compile_enabled,
+)
+def patch_inductor_precision_casts(ctx: PatchContext) -> None:
+    """Make inductor round intermediates the way eager does."""
+    import torch._inductor.config as inductor_config
+
+    if inductor_config.emulate_precision_casts:
+        log_rank_0(f"{_PREFIX} already enabled; leaving it alone")
+        return
+
+    inductor_config.emulate_precision_casts = True
+    log_rank_0(f"{_PREFIX} torch._inductor.config.emulate_precision_casts = True")

primus/backends/torchtitan/patches/turbo/__init__.py

@@ -13,6 +13,7 @@
 """
 
 from primus.backends.torchtitan.patches.turbo import (  # noqa: F401
+    aiter_mla_capability_patches,
     async_tp_patches,
     attention_patches,
     deepseek_v3_classic_attention_patches,

primus/backends/torchtitan/patches/turbo/aiter_mla_capability_patches.py

@@ -0,0 +1,117 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Route the failing gfx942 DeepSeek MLA backward directly to aiter CK."""
+
+from typing import Any
+
+from primus.backends.torchtitan.patches.turbo.attention_safety import (
+    requires_ck_mla_backward,
+)
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_LOG_PREFIX = "[Patch:torchtitan.primus_turbo.aiter_mla_capability]"
+_WRAPPED_ATTR = "_primus_gfx942_mla_ck_backward"
+
+
+def _can_patch(ctx: PatchContext) -> bool:
+    args = get_args(ctx)
+    return requires_ck_mla_backward(ctx.model_name, args.primus_turbo)
+
+
+def _is_gfx942_tensor(tensor) -> bool:
+    try:
+        import torch
+
+        arch = torch.cuda.get_device_properties(tensor.device).gcnArchName or ""
+        return "gfx942" in arch
+    except Exception:  # pragma: no cover - defensive
+        return False
+
+
+def _should_use_ck(kwargs: dict[str, Any]) -> bool:
+    q = kwargs["q"]
+    v = kwargs["v"]
+    return bool(
+        kwargs.get("sink") is None
+        and kwargs.get("qkv_format", "bshd") == "bshd"
+        and q.ndim == 4
+        and v.ndim == 4
+        and q.shape[-1] in (128, 192)
+        and v.shape[-1] == 128
+        and kwargs["dropout_p"] == 0.0
+        and kwargs.get("bias") is None
+        and kwargs.get("alibi_slopes") is None
+        and kwargs.get("dbias") is None
+        and _is_gfx942_tensor(q)
+    )
+
+
+def _run_ck_backward(mha, kwargs: dict[str, Any]):
+    softmax_d = mha.mha_bwd(
+        kwargs["dout"],
+        kwargs["q"],
+        kwargs["k"],
+        kwargs["v"],
+        kwargs["out"],
+        kwargs["softmax_lse"],
+        kwargs["dropout_p"],
+        kwargs["softmax_scale"],
+        kwargs["causal"],
+        kwargs["window_size_left"],
+        kwargs["window_size_right"],
+        kwargs["deterministic"],
+        kwargs["dq"],
+        kwargs["dk"],
+        kwargs["dv"],
+        kwargs["dbias"],
+        kwargs["bias"],
+        kwargs["alibi_slopes"],
+        kwargs["rng_state"],
+        None,
+        kwargs["sink"],
+        kwargs["dsink"],
+    )
+    return (
+        softmax_d,
+        kwargs["dq"],
+        kwargs["dk"],
+        kwargs["dv"],
+        kwargs["dbias"],
+        kwargs["dsink"],
+    )
+
+
+def _make_execute_wrapper(original, mha):
+    def execute(*args, **kwargs):
+        if not args and _should_use_ck(kwargs):
+            return _run_ck_backward(mha, kwargs)
+        return original(*args, **kwargs)
+
+    setattr(execute, _WRAPPED_ATTR, True)
+    return execute
+
+
+@register_patch(
+    "torchtitan.primus_turbo.aiter_mla_capability",
+    backend="torchtitan",
+    phase="setup",
+    description="Route failing gfx942 DeepSeek MLA backward from v3 to aiter CK",
+    condition=_can_patch,
+    priority=51,
+)
+def patch_aiter_mla_capability(ctx: PatchContext) -> None:
+    import aiter.ops.mha as mha
+    from primus_turbo.pytorch.kernels.attention import attention_aiter_impl
+
+    backend = attention_aiter_impl.AttnBwdAiterBackend
+    original = backend.execute
+    if getattr(original, _WRAPPED_ATTR, False):
+        return
+
+    backend.execute = staticmethod(_make_execute_wrapper(original, mha))
+    log_rank_0(f"{_LOG_PREFIX} routing gfx942 DeepSeek MLA backward directly to aiter CK")

primus/backends/torchtitan/patches/turbo/attention_patches.py

@@ -15,19 +15,30 @@
 expressed as a backend patch so it can be managed via the Primus patch system.
 """
 
-from primus.core.patches import PatchContext, get_param, register_patch
+from primus.backends.torchtitan.patches.turbo.attention_safety import (
+    requires_ck_mla_backward,
+    should_use_turbo_attention,
+)
+from primus.core.patches import PatchContext, get_args, register_patch
 from primus.core.utils.module_utils import log_rank_0
 
 
+def _can_patch_turbo_attention(ctx: PatchContext) -> bool:
+    args = get_args(ctx)
+    if requires_ck_mla_backward(ctx.model_name, args.primus_turbo):
+        log_rank_0(
+            "[Patch:torchtitan.primus_turbo.turbo_attention] "
+            "Using Turbo attention with CK MLA backward on gfx942."
+        )
+    return should_use_turbo_attention(ctx.model_name, args.primus_turbo)
+
+
 @register_patch(
     "torchtitan.primus_turbo.turbo_attention",
     backend="torchtitan",
     phase="setup",
     description="Use Primus-Turbo Attention kernels for supported models",
-    condition=lambda ctx: (
-        get_param(ctx, "primus_turbo.enable_primus_turbo", False)
-        and get_param(ctx, "primus_turbo.use_turbo_attention", False)
-    ),
+    condition=_can_patch_turbo_attention,
 )
 def patch_turbo_attention(ctx: PatchContext) -> None:
     """

primus/backends/torchtitan/patches/turbo/attention_safety.py

@@ -0,0 +1,43 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Runtime safety gate for TorchTitan Primus-Turbo attention."""
+
+from typing import Any
+
+
+def _model_name(model: Any) -> str:
+    if isinstance(model, str):
+        return model
+    return str(getattr(model, "name", model) or "")
+
+
+def _is_gfx942() -> bool:
+    try:
+        import torch
+
+        if not torch.cuda.is_available():
+            return False
+        return "gfx942" in (torch.cuda.get_device_properties(0).gcnArchName or "")
+    except Exception:  # pragma: no cover - fail open outside supported GPU runs
+        return False
+
+
+def requires_ck_mla_backward(model: Any, turbo_config: Any) -> bool:
+    """Whether gfx942 MLA must reject the unsupported fmha_v3 backward path."""
+    is_deepseek = "deepseek" in _model_name(model).lower()
+    is_fp8_recipe = bool(
+        getattr(turbo_config, "use_turbo_float8_linear", False) or getattr(turbo_config, "use_moe_fp8", False)
+    )
+    uses_nonclassic_attention = not bool(getattr(turbo_config, "use_classic_attention", False))
+    return is_deepseek and is_fp8_recipe and uses_nonclassic_attention and _is_gfx942()
+
+
+def should_use_turbo_attention(model: Any, turbo_config: Any) -> bool:
+    return bool(
+        getattr(turbo_config, "enable_primus_turbo", False)
+        and getattr(turbo_config, "use_turbo_attention", False)
+    )

primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py

@@ -41,7 +41,7 @@ def patch_torchtitan_moe(ctx: PatchContext) -> None:
     """
     Patch TorchTitan MoE to use Primus-Turbo grouped_mm implementation.
     """
-    import torchtitan.models.moe.moe
+    import torchtitan.models.moe.moe as moe_module
 
     from primus.backends.torchtitan.models.moe.moe import _run_experts_grouped_mm
 
@@ -64,7 +64,7 @@ def _run_experts_grouped_mm_dynamic(*args, **kwargs):
     _run_experts_grouped_mm_dynamic.__qualname__ = _ALREADY_PATCHED_SENTINEL
     _run_experts_grouped_mm_dynamic.__name__ = _ALREADY_PATCHED_SENTINEL
 
-    torchtitan.models.moe.moe._run_experts_grouped_mm = _run_experts_grouped_mm_dynamic
+    moe_module._run_experts_grouped_mm = _run_experts_grouped_mm_dynamic
 
     log_rank_0(
         "[Patch:torchtitan.primus_turbo.moe_grouped_mm] "