diff --git a/docs/03-configuration-reference/environment-variables.md b/docs/03-configuration-reference/environment-variables.md index 61aca7f1b..4b0586f29 100644 --- a/docs/03-configuration-reference/environment-variables.md +++ b/docs/03-configuration-reference/environment-variables.md @@ -121,6 +121,7 @@ Primus seeds many of these in `runner/helpers/envs/base_env.sh`. RCCL honors NCC | `HSA_NO_SCRATCH_RECLAIM` | `1` | `base_env.sh`, container passthrough | ROCm runtime; documented for MoE stability | `1` keeps scratch allocated (often used for MoE stability). See [ROCR environment](https://rocm.docs.amd.com/projects/ROCR-Runtime/en/docs-7.1.1/environment_variables.html). | | `HIP_VISIBLE_DEVICES` | `0..GPUS_PER_NODE-1` | `base_env.sh` | ROCm device visibility | Restricts which GPU indices ROCm exposes. | | `ROCBLAS_DEFAULT_ATOMICS_MODE` | (unset) | User | `primus/backends/megatron/patches/args/rocm_arg_validation.py` | Read for deterministic / accuracy-sensitive GEMM behavior. | +| `TORCH_COMPILE_CACHE_KEY_TAG` | (unset) | `primus/core/patches/triton_bufops_war_patches.py`, on ROCm only | torch.compile caching | Gets `primus-bufops-war-v` appended so graphs compiled before (or under a different version of) the workaround are not reused. An existing value is appended to, not replaced. | --- diff --git a/docs/04-technical-guides/performance-tuning.md b/docs/04-technical-guides/performance-tuning.md index 321ccf47d..b02fc5aa3 100644 --- a/docs/04-technical-guides/performance-tuning.md +++ b/docs/04-technical-guides/performance-tuning.md @@ -233,6 +233,35 @@ From `trainer_base.yaml` and model settings: --- +## 8. Known ROCm correctness workarounds + +The AMD Triton backend can emit a `buffer_store_dwordx4` whose data registers are redefined by a +later instruction with no `s_waitcnt vmcnt` in between, so the store writes whatever the +clobbering instruction left behind. Nothing faults and nothing warns; a few percent of the +kernel's output elements simply hold garbage. Inductor's SwiGLU backward fusion trips it and +poisons the MLP weight gradients, which diverges TorchTitan training. Megatron is exposed to the +same kernels, since its fused activations are inductor-generated as well (`megatron/core/jit.py` +sets `jit_fuser = torch.compile` on torch 2.2+). + +`primus/core/patches/triton_bufops_war_patches.py` handles this automatically on ROCm: every +kernel compiles normally, and only those whose emitted AMDGCN contains the hazard are recompiled +with buffer ops disabled. It has no switch: the recompile is a correctness fix, and a kernel that +does not need one is left alone anyway. + +Two things about the design are deliberate. It selects per kernel rather than setting +`AMDGCN_USE_BUFFER_OPS=0` globally, because the hazardous kernels lose nothing without buffer +addressing while the Primus-Turbo grouped-GEMM kernels spill without it and carry no hazard — the +global switch costs ~18% on MoE recipes for kernels that never needed fixing. And it decides from +the emitted machine code rather than from a list of affected architectures, so a newly shipped +architecture cannot silently fall out of coverage. + +It also appends to `TORCH_COMPILE_CACHE_KEY_TAG`, because on an FX graph cache hit inductor never +calls `triton.compile`; without that, a cache filled before the patch existed would keep serving +hazardous binaries. Note that every Triton build tried so far still emits the pattern, so this is +not something an image upgrade removes. + +--- + ## Quick checklist 1. **GEMMs:** run HipBLASLt stages 1–3 or use `offline_tune_gemm.py` for custom workflows. diff --git a/primus/backends/torchtitan/patches/__init__.py b/primus/backends/torchtitan/patches/__init__.py index 5f4b8c157..0bf9fd46a 100644 --- a/primus/backends/torchtitan/patches/__init__.py +++ b/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, diff --git a/primus/backends/torchtitan/patches/dsv3_v022_perf_patches.py b/primus/backends/torchtitan/patches/dsv3_v022_perf_patches.py index edd3e7d41..71314f88e 100644 --- a/primus/backends/torchtitan/patches/dsv3_v022_perf_patches.py +++ b/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)", ) diff --git a/primus/backends/torchtitan/patches/inductor_precision_casts_patches.py b/primus/backends/torchtitan/patches/inductor_precision_casts_patches.py new file mode 100644 index 000000000..f0364e24c --- /dev/null +++ b/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") diff --git a/primus/backends/torchtitan/patches/turbo/__init__.py b/primus/backends/torchtitan/patches/turbo/__init__.py index dfc2b79bc..5c61bf1ad 100644 --- a/primus/backends/torchtitan/patches/turbo/__init__.py +++ b/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, diff --git a/primus/backends/torchtitan/patches/turbo/aiter_mla_capability_patches.py b/primus/backends/torchtitan/patches/turbo/aiter_mla_capability_patches.py new file mode 100644 index 000000000..930da89de --- /dev/null +++ b/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") diff --git a/primus/backends/torchtitan/patches/turbo/attention_patches.py b/primus/backends/torchtitan/patches/turbo/attention_patches.py index 8dbf1fe05..abc693550 100644 --- a/primus/backends/torchtitan/patches/turbo/attention_patches.py +++ b/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: """ diff --git a/primus/backends/torchtitan/patches/turbo/attention_safety.py b/primus/backends/torchtitan/patches/turbo/attention_safety.py new file mode 100644 index 000000000..2b874ed82 --- /dev/null +++ b/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) + ) diff --git a/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py b/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py index 17cf70b8d..2e9365ff7 100644 --- a/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py +++ b/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] " diff --git a/primus/backends/torchtitan/primus_turbo_extensions/primus_turbo_converter.py b/primus/backends/torchtitan/primus_turbo_extensions/primus_turbo_converter.py index 14d864c3a..c3ebeb81b 100644 --- a/primus/backends/torchtitan/primus_turbo_extensions/primus_turbo_converter.py +++ b/primus/backends/torchtitan/primus_turbo_extensions/primus_turbo_converter.py @@ -16,6 +16,10 @@ register_model_converter, ) +from primus.backends.torchtitan.patches.turbo.attention_safety import ( + should_use_turbo_attention, +) + def replace_turbo_attention_modules(model: torch.nn.Module, fp8_config): from primus_turbo.pytorch.modules import TurboAttention # TODO: import Check @@ -33,13 +37,20 @@ def replace_turbo_attention_modules(model: torch.nn.Module, fp8_config): class PrimusTubroConverter(ModelConverter): def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims): + self.primus_turbo_config = job_config.primus_turbo + self.enabled = should_use_turbo_attention( + job_config.model, + self.primus_turbo_config, + ) + self.fp8_config = None + if not self.enabled: + return + from primus_turbo.pytorch.core.low_precision import ( Float8QuantConfig, ScalingGranularity, ) - self.enabled = True - self.primus_turbo_config = job_config.primus_turbo self.fp8_config = ( Float8QuantConfig( granularity=ScalingGranularity.BLOCKWISE, diff --git a/primus/core/patches/__init__.py b/primus/core/patches/__init__.py index c40568ad8..8e9f1999c 100644 --- a/primus/core/patches/__init__.py +++ b/primus/core/patches/__init__.py @@ -64,3 +64,8 @@ def fix_deepseek_moe(ctx: PatchContext): "run_patches", "version_matches", ] + +# Backend-independent patches. Unlike the per-backend collections, nothing walks +# this package for "*_patches" modules, so a patch that has to register for every +# backend is imported here explicitly (after the names above exist). +import primus.core.patches.triton_bufops_war_patches # noqa: F401, E402 diff --git a/primus/core/patches/triton_bufops_war_patches.py b/primus/core/patches/triton_bufops_war_patches.py new file mode 100644 index 000000000..611889385 --- /dev/null +++ b/primus/core/patches/triton_bufops_war_patches.py @@ -0,0 +1,216 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Triton buffer-store WAR miscompile workaround +============================================= + +The AMD Triton backend can emit a ``buffer_store_dwordx4`` whose data VGPRs are +redefined by a later instruction with no intervening ``s_waitcnt vmcnt``, so the +store writes the clobbering value instead of the computed one. The corruption +is silent: it drops Inf/NaN on TorchTitan and stays inside the normal numeric +range on Megatron, so a finite loss does not prove a clean run. + +``AMDGCN_USE_BUFFER_OPS=0`` avoids it globally but costs ~18% on MoE recipes, +because the kernels that spill without buffer addressing (the Primus-Turbo +grouped-GEMM ones) are not the kernels that miscompile. So compile every +kernel normally, look for the hazard in the emitted AMDGCN, and recompile only +the affected kernels with buffer ops off. + +Keying on the machine code means neither a kernel nor an architecture allowlist +has to be maintained, and it follows the toolchain: every Triton build tried so +far (upstream 3.7.0, ROCm 3.7.0/3.7.1/3.8.0) still emits the pattern on at least +one production kernel. +""" + +import functools +import os +import re +import sys + +from primus.core.patches.context import PatchContext +from primus.core.patches.patch_registry import register_patch +from primus.core.utils.module_utils import log_rank_0 + +# Bump whenever the detection below changes, so that torch.compile caches +# populated by an older guard are not reused. See _break_inductor_caches. +_GUARD_VERSION = "1" + +_STORE = re.compile(r"^\s*buffer_store_dwordx4\s+v\[?(\d+)(?::(\d+))?\]?") +_WRITE = re.compile(r"^\s*[a-z_0-9]+\s+v\[?(\d+)(?::(\d+))?\]?\s*,") +_WAIT = re.compile(r"^\s*s_waitcnt\b") +_VMCNT = re.compile(r"vmcnt\((\d+)\)") +_STOP = re.compile(r"^\s*s_(endpgm|branch|cbranch|setpc|swappc)\b") + +_stats = {"scanned": 0, "patched": 0, "residual": 0, "names": {}} + +_PREFIX = "[Patch:triton.buffer_store_war]" + + +def _log(msg: str) -> None: + """Log without ever raising: this runs inside the compiler, which can be + reached before the Primus logger exists.""" + try: + log_rank_0(f"{_PREFIX} {msg}") + except Exception: + print(f"{_PREFIX} {msg}", file=sys.stderr) + + +def has_hazard(asm: str) -> bool: + """True if a buffer_store_dwordx4 can have its data VGPRs clobbered. + + Walks forward from each store to the first instruction that either waits on + vmcnt (the store's registers are safe from there on) or redefines a + register the store is reading (the miscompile). + """ + if "buffer_store_dwordx4" not in asm: + return False + lines = asm.splitlines() + for i, line in enumerate(lines): + m = _STORE.match(line) + if not m: + continue + lo = int(m.group(1)) + hi = int(m.group(2)) if m.group(2) else lo + for j in range(i + 1, len(lines)): + nxt = lines[j] + if not nxt.strip() or nxt.lstrip().startswith((";", ".", "//")): + continue + if _STOP.match(nxt): + break + if _WAIT.match(nxt): + if _VMCNT.search(nxt): + break + continue + w = _WRITE.match(nxt) + if w: + wlo = int(w.group(1)) + whi = int(w.group(2)) if w.group(2) else wlo + # Both operands are contiguous VGPR ranges, so they overlap iff + # neither ends before the other begins. + if wlo <= hi and lo <= whi: + return True + break + return False + + +def _on_rocm() -> bool: + """True on any ROCm GPU. + + Deliberately not an architecture allowlist. Wrong values were measured on + gfx942; gfx950 emits the same pattern but was measured to tolerate it. Since + the compiler is omitting a required ``s_waitcnt`` on both, which chips happen + to tolerate that is not a safe thing to encode -- and an allowlist would have + to be widened for every new architecture, leaving it unprotected until + somebody remembers. Scanning decides instead, at a substring test per + compiled kernel. + """ + try: + import torch + + return torch.cuda.is_available() and torch.version.hip is not None + except Exception: + return False + + +def _break_inductor_caches() -> None: + """Make torch.compile caches from before this guard unreachable. + + On an FX graph cache hit inductor never calls ``triton.compile``, so a cache + filled before the guard existed keeps serving hazardous binaries -- and does + so invisibly, since ``_stats`` only counts kernels that reach the compiler. + """ + tag = f"primus-bufops-war-v{_GUARD_VERSION}" + existing = os.environ.get("TORCH_COMPILE_CACHE_KEY_TAG", "") + if tag in existing.split(","): + return + combined = f"{existing},{tag}" if existing else tag + os.environ["TORCH_COMPILE_CACHE_KEY_TAG"] = combined + try: + from torch.compiler import config as compiler_config + + # The env var is only read when torch.compiler.config is first imported, + # which may already have happened. + compiler_config.cache_key_tag = combined + except Exception as exc: + _log( + f"could not set cache_key_tag ({exc}); wipe the torch.compile cache " + f"directory before trusting this run" + ) + + +@register_patch( + "triton.buffer_store_war", + backend=None, # both Megatron and TorchTitan compile Triton kernels + phase="build_args", # earliest phase, well before the first kernel compile + condition=lambda ctx: _on_rocm(), +) +def patch_triton_buffer_store_war(ctx: PatchContext) -> None: + """Disable buffer ops for the individual kernels that hit the store WAR bug.""" + import triton.knobs as knobs + from triton._C.libtriton import get_cache_invalidating_env_vars + from triton.compiler import compiler as _cc + + if getattr(_cc.compile, "_primus_bufops_war", False): + return + + _break_inductor_caches() + orig = _cc.compile + + @functools.wraps(orig) + def guarded(src, target=None, options=None, _env_vars=None): + kernel = orig(src, target=target, options=options, _env_vars=_env_vars) + try: + if kernel.metadata.target.backend != "hip": + return kernel + asm = kernel.asm["amdgcn"] + except Exception: + return kernel + + _stats["scanned"] += 1 + if not has_hazard(asm): + return kernel + + name = getattr(kernel, "name", "?") + # Record the override in the cache key so flipping the policy off cannot + # resurrect an artifact compiled under the other setting. + env = dict(get_cache_invalidating_env_vars() if _env_vars is None else _env_vars) + env["AMDGCN_USE_BUFFER_OPS"] = "0" + with knobs.amd.scope(): + knobs.amd.use_buffer_ops = False + try: + fixed = orig(src, target=target, options=options, _env_vars=env) + except Exception as exc: # a broken workaround must not break the build + _stats["residual"] += 1 + _log(f"{name}: recompile failed ({exc}); KEEPING HAZARDOUS KERNEL") + return kernel + + _stats["patched"] += 1 + _stats["names"][name] = _stats["names"].get(name, 0) + 1 + if has_hazard(fixed.asm["amdgcn"]): + _stats["residual"] += 1 + _log(f"{name}: hazard survives with buffer ops off; the workaround does not cover it") + return fixed + + guarded._primus_bufops_war = True + _cc.compile = guarded + # JITFunction.create_binder resolves `from ..compiler import compile` lazily, + # so the package attribute is the binding that actually reaches the JIT. + for modname in ("triton.compiler", "triton", "triton.runtime.jit"): + mod = sys.modules.get(modname) + if mod is not None and getattr(mod, "compile", None) is orig: + mod.compile = guarded + + _log("installed") + + +def get_stats() -> dict: + """Kernels seen, recompiled, and still hazardous afterwards. + + ``residual`` must be zero. It only counts kernels that reached the + compiler, not ones served from a torch.compile cache. + """ + return {k: (dict(v) if isinstance(v, dict) else v) for k, v in _stats.items()} diff --git a/tests/unit_tests/backends/torchtitan/test_aiter_mla_capability_patch.py b/tests/unit_tests/backends/torchtitan/test_aiter_mla_capability_patch.py new file mode 100644 index 000000000..6ec52c775 --- /dev/null +++ b/tests/unit_tests/backends/torchtitan/test_aiter_mla_capability_patch.py @@ -0,0 +1,94 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tests for the gfx942 DeepSeek MLA fmha_v3 capability correction.""" + +from unittest.mock import MagicMock + +import pytest +import torch + +import primus.backends.torchtitan.patches.turbo.aiter_mla_capability_patches as capability_patch +from primus.core.patches.patch_registry import PatchRegistry + +PATCH_ID = "torchtitan.primus_turbo.aiter_mla_capability" + + +def test_patch_registered_for_torchtitan_setup(): + patch = PatchRegistry.get(PATCH_ID) + assert patch is not None + assert patch.backend == "torchtitan" + assert patch.phase == "setup" + + +def _kwargs(qk_head_dim=128, v_head_dim=128): + q = torch.empty(1, 8, 2, qk_head_dim) + k = torch.empty_like(q) + v = torch.empty(1, 8, 2, v_head_dim) + return { + "dout": torch.empty_like(v), + "q": q, + "k": k, + "v": v, + "out": torch.empty_like(v), + "softmax_lse": torch.empty(1, 2, 8), + "dq": torch.empty_like(q), + "dk": torch.empty_like(k), + "dv": torch.empty_like(v), + "dbias": None, + "dropout_p": 0.0, + "softmax_scale": 0.1, + "causal": True, + "window_size_left": -1, + "window_size_right": -1, + "bias": None, + "alibi_slopes": None, + "deterministic": False, + "rng_state": torch.tensor([3, 7], dtype=torch.int64), + "sink": None, + "dsink": None, + "qkv_format": "bshd", + } + + +@pytest.mark.parametrize("qk_head_dim", [128, 192]) +def test_deepseek_mla_shapes_route_to_ck(monkeypatch, qk_head_dim): + monkeypatch.setattr( + capability_patch, + "_is_gfx942_tensor", + lambda tensor: True, + ) + original = MagicMock(return_value="original") + mha = MagicMock() + mha.mha_bwd.return_value = "softmax_d" + execute = capability_patch._make_execute_wrapper(original, mha) + kwargs = _kwargs(qk_head_dim=qk_head_dim) + + result = execute(**kwargs) + assert result == ( + "softmax_d", + kwargs["dq"], + kwargs["dk"], + kwargs["dv"], + None, + None, + ) + original.assert_not_called() + assert mha.mha_bwd.call_args.args[0] is kwargs["dout"] + + +def test_other_shapes_keep_original_backend(monkeypatch): + monkeypatch.setattr( + capability_patch, + "_is_gfx942_tensor", + lambda tensor: True, + ) + original = MagicMock(return_value="original") + execute = capability_patch._make_execute_wrapper(original, MagicMock()) + kwargs = _kwargs(qk_head_dim=64, v_head_dim=64) + + assert execute(**kwargs) == "original" + original.assert_called_once_with(**kwargs) diff --git a/tests/unit_tests/backends/torchtitan/test_dsv3_v022_perf_patch.py b/tests/unit_tests/backends/torchtitan/test_dsv3_v022_perf_patch.py new file mode 100644 index 000000000..ad38be52f --- /dev/null +++ b/tests/unit_tests/backends/torchtitan/test_dsv3_v022_perf_patch.py @@ -0,0 +1,127 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Regression tests for the TorchTitan v0.2.2 DeepSeek memory patch. + +DeepSeek keeps low-fragment whole-block compilation for dense layers, but MoE +TransformerBlocks and ``MoE.forward`` must remain eager. Wrapping a MoE block +silently produced a NaN forward loss even with a graph break around its MoE. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import torch.nn as nn + +import primus.backends.torchtitan.patches.dsv3_v022_perf_patches as dsv3_patch +from primus.core.patches import PatchContext +from primus.core.patches.patch_registry import PatchRegistry + +COMBINE_PATCH_ID = "torchtitan.dsv3.moe_bf16_combine" +COMPILE_PATCH_ID = "torchtitan.dsv3.whole_block_compile" + + +def _ctx(model_name, compile_enable=True): + params = SimpleNamespace(compile=SimpleNamespace(enable=compile_enable)) + module_config = SimpleNamespace(params=params) + return PatchContext( + backend="torchtitan", + phase="setup", + model_name=model_name, + extra={"module_config": module_config}, + ) + + +class TestDsv3V022PatchRegistration: + def test_safe_whole_block_compile_is_registered(self): + patch = PatchRegistry.get(COMPILE_PATCH_ID) + assert patch is not None + assert patch.backend == "torchtitan" + assert patch.phase == "setup" + + def test_bf16_combine_remains_registered(self): + patch = PatchRegistry.get(COMBINE_PATCH_ID) + assert patch is not None + assert patch.backend == "torchtitan" + assert patch.phase == "setup" + + +class TestDsv3V022PatchCondition: + def test_compile_patch_requires_deepseek_and_compile(self): + patch = PatchRegistry.get(COMPILE_PATCH_ID) + assert patch is not None + assert patch.condition(_ctx("deepseek_v3", True)) is True + assert patch.condition(_ctx("deepseek_v3", False)) is False + assert patch.condition(_ctx("llama3", True)) is False + + def test_bf16_combine_applies_to_deepseek(self): + patch = PatchRegistry.get(COMBINE_PATCH_ID) + assert patch is not None + assert patch.condition(_ctx("deepseek_v3")) is True + assert patch.condition(_ctx(SimpleNamespace(name="DeepSeek-V3"))) is True + + def test_bf16_combine_does_not_apply_to_other_models(self): + patch = PatchRegistry.get(COMBINE_PATCH_ID) + assert patch is not None + assert patch.condition(_ctx("llama3")) is False + assert patch.condition(_ctx(None)) is False + + +class _Block(nn.Module): + def __init__(self, moe_enabled): + super().__init__() + self.moe_enabled = moe_enabled + + +class _Model(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleDict( + { + "dense": _Block(moe_enabled=False), + "moe": _Block(moe_enabled=True), + } + ) + + +def test_safe_compile_only_compiles_dense_blocks(monkeypatch): + # Load TorchTitan before replacing torch.compile; its attention module + # compiles a helper at import time with a different signature. + import torchtitan.tools.logging # noqa: F401 + + model = _Model() + calls = [] + + def fake_compile(module, *, backend, fullgraph): + calls.append((module.moe_enabled, backend, fullgraph)) + return module + + monkeypatch.setattr(dsv3_patch.torch, "compile", fake_compile) + dsv3_patch._apply_dense_only_compile( + model, + SimpleNamespace(backend="inductor"), + ep_enabled=True, + ) + + assert calls == [(False, "inductor", True)] + + +def test_moe_forward_is_installed_behind_compiler_disable(monkeypatch): + import torchtitan.models.moe.moe as moe_module + + def eager_boundary(*args, **kwargs): + return args, kwargs + + monkeypatch.setattr( + dsv3_patch.torch.compiler, + "disable", + lambda fn: eager_boundary, + ) + monkeypatch.setattr(dsv3_patch, "log_rank_0", lambda *args, **kwargs: None) + + with patch.object(moe_module.MoE, "forward", moe_module.MoE.forward): + dsv3_patch.patch_moe_bf16_combine(_ctx("deepseek_v3")) + assert moe_module.MoE.forward is eager_boundary diff --git a/tests/unit_tests/backends/torchtitan/test_inductor_precision_casts_patch.py b/tests/unit_tests/backends/torchtitan/test_inductor_precision_casts_patch.py new file mode 100644 index 000000000..e00990635 --- /dev/null +++ b/tests/unit_tests/backends/torchtitan/test_inductor_precision_casts_patch.py @@ -0,0 +1,107 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Unit tests for the inductor precision-cast emulation patch. + +Regression focus: the patch has to reach inductor before the first compile and +only for compiled runs, and it must not clobber a value somebody else already +set -- TorchTitan turns the same flag on for float8 rowwise recipes. + +The last test pins the reason no cache busting is needed: the flag is part of +the inductor FX graph cache key. If a torch upgrade ever moves it into one of +the ignore lists, caches populated before the patch would keep serving the +kernels that produce the NaN, silently, and this patch would need the same +``TORCH_COMPILE_CACHE_KEY_TAG`` treatment as the Triton buffer-store one. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import primus.backends.torchtitan.patches.inductor_precision_casts_patches as pc_patch +from primus.core.patches import PatchContext +from primus.core.patches.patch_registry import PatchRegistry + +PATCH_ID = "torchtitan.torch.inductor_precision_casts" + + +def _ctx(compile_enable=True): + params = SimpleNamespace(compile=SimpleNamespace(enable=compile_enable)) + module_config = SimpleNamespace(params=params) + return PatchContext(backend="torchtitan", phase="setup", extra={"module_config": module_config}) + + +class TestRegistration: + def test_patch_registered(self): + assert PATCH_ID in PatchRegistry.list_ids() + p = PatchRegistry.get(PATCH_ID) + assert p is not None + assert p.backend == "torchtitan" + # Must land before parallelize_fn calls apply_compile. + assert p.phase == "setup" + + +class TestCondition: + def test_enabled_when_compiling(self): + import torch + + with patch.object(torch.version, "hip", "7.15"): + assert pc_patch._compile_enabled(_ctx(compile_enable=True)) is True + + def test_disabled_off_rocm(self): + import torch + + with patch.object(torch.version, "hip", None): + assert pc_patch._compile_enabled(_ctx(compile_enable=True)) is False + + def test_disabled_without_compile(self): + assert pc_patch._compile_enabled(_ctx(compile_enable=False)) is False + + def test_disabled_when_config_absent(self): + ctx = PatchContext(backend="torchtitan", phase="setup", extra={}) + assert pc_patch._compile_enabled(ctx) is False + + +class _AlreadyOn: + def __bool__(self): + return True + + +class TestApply: + # Patch the attribute on the real config module: swapping the module out of + # sys.modules makes torch re-import torch._inductor.test_operators, which + # fails on a duplicate TORCH_LIBRARY registration. + def test_turns_the_flag_on(self): + import torch._inductor.config as inductor_config + + with patch.object(inductor_config, "emulate_precision_casts", False), patch.object( + pc_patch, "log_rank_0" + ): + pc_patch.patch_inductor_precision_casts(_ctx()) + assert inductor_config.emulate_precision_casts is True + + def test_leaves_an_existing_true_alone(self): + # TorchTitan sets this for float8 rowwise; re-setting it must stay a no-op + # rather than racing with whoever owns the value. A truthy marker instead + # of True makes the early return observable. + import torch._inductor.config as inductor_config + + marker = _AlreadyOn() + with patch.object(inductor_config, "emulate_precision_casts", marker), patch.object( + pc_patch, "log_rank_0" + ): + pc_patch.patch_inductor_precision_casts(_ctx()) + assert inductor_config.emulate_precision_casts is marker + + +class TestCacheKeyParticipation: + def test_flag_is_part_of_the_inductor_cache_key(self): + import torch._inductor.config as inductor_config + + ignore = set(getattr(inductor_config, "_save_config_ignore", ()) or ()) + prefixes = tuple(getattr(inductor_config, "_cache_config_ignore_prefix", ()) or ()) + assert "emulate_precision_casts" not in ignore + assert not any("emulate_precision_casts".startswith(p) for p in prefixes) diff --git a/tests/unit_tests/backends/torchtitan/test_primus_turbo_converter.py b/tests/unit_tests/backends/torchtitan/test_primus_turbo_converter.py new file mode 100644 index 000000000..da0e3c2c4 --- /dev/null +++ b/tests/unit_tests/backends/torchtitan/test_primus_turbo_converter.py @@ -0,0 +1,151 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tests for TorchTitan's Primus-Turbo model converter gating.""" + +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +import primus.backends.torchtitan.patches.turbo.attention_patches as attention_patch +import primus.backends.torchtitan.patches.turbo.attention_safety as attention_safety +import primus.backends.torchtitan.primus_turbo_extensions.primus_turbo_converter as converter_module +from primus.core.patches import PatchContext +from primus.core.patches.patch_registry import PatchRegistry + + +def _job_config( + *, + model_name="llama3", + enable_turbo=True, + use_turbo_attention=True, + attention_float8=False, + use_turbo_float8_linear=False, + use_moe_fp8=False, + use_classic_attention=False, +): + turbo = SimpleNamespace( + enable_primus_turbo=enable_turbo, + use_turbo_attention=use_turbo_attention, + enable_attention_float8=attention_float8, + use_turbo_float8_linear=use_turbo_float8_linear, + use_moe_fp8=use_moe_fp8, + use_classic_attention=use_classic_attention, + ) + return SimpleNamespace( + model=SimpleNamespace(name=model_name), + primus_turbo=turbo, + ) + + +def _fake_low_precision_module(): + module = ModuleType("primus_turbo.pytorch.core.low_precision") + module.ScalingGranularity = SimpleNamespace(BLOCKWISE="blockwise") + module.Float8QuantConfig = lambda **kwargs: ("fp8", kwargs) + return module + + +@pytest.mark.parametrize( + "enable_turbo,use_turbo_attention", + [(False, True), (True, False), (False, False)], +) +def test_converter_skips_attention_replacement_when_disabled( + monkeypatch, + enable_turbo, + use_turbo_attention, +): + replace = MagicMock() + monkeypatch.setattr(converter_module, "replace_turbo_attention_modules", replace) + + converter = converter_module.PrimusTubroConverter( + _job_config( + enable_turbo=enable_turbo, + use_turbo_attention=use_turbo_attention, + ), + parallel_dims=None, + ) + model = object() + + assert converter.enabled is False + assert converter.fp8_config is None + assert converter.convert(model) is None + replace.assert_not_called() + + +def test_converter_replaces_attention_when_both_flags_enabled(monkeypatch): + monkeypatch.setitem( + sys.modules, + "primus_turbo.pytorch.core.low_precision", + _fake_low_precision_module(), + ) + replace = MagicMock() + monkeypatch.setattr(converter_module, "replace_turbo_attention_modules", replace) + + converter = converter_module.PrimusTubroConverter( + _job_config(attention_float8=True), + parallel_dims=None, + ) + model = object() + + assert converter.enabled is True + assert converter.fp8_config == ( + "fp8", + {"granularity": "blockwise", "block_size": 64}, + ) + assert converter.convert(model) is model + replace.assert_called_once_with(model, converter.fp8_config) + + +def _deepseek_fp8_config(): + return _job_config( + model_name="deepseek_v3", + use_turbo_float8_linear=True, + use_moe_fp8=True, + use_classic_attention=False, + ) + + +def test_converter_keeps_turbo_attention_enabled_for_gfx942_deepseek_fp8( + monkeypatch, +): + monkeypatch.setattr(attention_safety, "_is_gfx942", lambda: True) + monkeypatch.setitem( + sys.modules, + "primus_turbo.pytorch.core.low_precision", + _fake_low_precision_module(), + ) + replace = MagicMock() + monkeypatch.setattr(converter_module, "replace_turbo_attention_modules", replace) + + converter = converter_module.PrimusTubroConverter( + _deepseek_fp8_config(), + parallel_dims=None, + ) + model = object() + + assert converter.enabled is True + assert converter.convert(model) is model + replace.assert_called_once_with(model, converter.fp8_config) + + +def test_setup_patch_uses_the_same_runtime_safety_gate(monkeypatch): + monkeypatch.setattr(attention_safety, "_is_gfx942", lambda: True) + monkeypatch.setattr(attention_patch, "log_rank_0", lambda *args, **kwargs: None) + config = _deepseek_fp8_config() + ctx = PatchContext( + backend="torchtitan", + phase="setup", + model_name=config.model, + extra={"module_config": SimpleNamespace(params=config)}, + ) + patch = PatchRegistry.get("torchtitan.primus_turbo.turbo_attention") + assert patch is not None + assert patch.condition(ctx) is True + + monkeypatch.setattr(attention_safety, "_is_gfx942", lambda: False) + assert patch.condition(ctx) is True diff --git a/tests/unit_tests/core/patches/test_triton_bufops_war.py b/tests/unit_tests/core/patches/test_triton_bufops_war.py new file mode 100644 index 000000000..452680988 --- /dev/null +++ b/tests/unit_tests/core/patches/test_triton_bufops_war.py @@ -0,0 +1,175 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Detection of the buffer-store WAR hazard, and the gate that enables it. + +Pure text analysis of AMDGCN, so these run anywhere; no GPU and no Triton. +The patch that uses this only recompiles kernels the detector flags, which +makes a false negative a silently wrong training run and a false positive +merely a slower one. +""" + +import inspect +import sys +import types + +import primus.core.patches.triton_bufops_war_patches as war + + +def _asm(*lines: str) -> str: + return "\n".join(f"\t{line}" for line in lines) + + +class TestHasHazard: + def test_store_then_immediate_clobber(self): + """The bug: v8 is redefined while the store that reads it is in flight.""" + assert war.has_hazard( + _asm( + "buffer_store_dwordx4 v[8:11], v4, s[0:3], 0 offen", + "v_add_f32_e32 v8, v12, v13", + ) + ) + + def test_vmcnt_wait_clears_it(self): + assert not war.has_hazard( + _asm( + "buffer_store_dwordx4 v[8:11], v4, s[0:3], 0 offen", + "s_waitcnt vmcnt(0)", + "v_add_f32_e32 v8, v12, v13", + ) + ) + + def test_wait_without_vmcnt_does_not_clear_it(self): + """Only a vmcnt wait orders the store; lgkmcnt says nothing about it.""" + assert war.has_hazard( + _asm( + "buffer_store_dwordx4 v[8:11], v4, s[0:3], 0 offen", + "s_waitcnt lgkmcnt(0)", + "v_add_f32_e32 v8, v12, v13", + ) + ) + + def test_clobber_inside_register_range(self): + """The store reads v[8:11]; redefining any one of them is enough.""" + assert war.has_hazard( + _asm( + "buffer_store_dwordx4 v[8:11], v4, s[0:3], 0 offen", + "v_mul_f32_e32 v10, v12, v13", + ) + ) + + def test_partially_overlapping_write_range(self): + """A wide write only has to straddle one end of the store's range.""" + assert war.has_hazard( + _asm( + "buffer_store_dwordx4 v[8:11], v4, s[0:3], 0 offen", + "v_lshlrev_b64 v[10:13], 2, v[2:3]", + ) + ) + + def test_adjacent_write_range_is_clean(self): + assert not war.has_hazard( + _asm( + "buffer_store_dwordx4 v[8:11], v4, s[0:3], 0 offen", + "v_lshlrev_b64 v[12:15], 2, v[2:3]", + ) + ) + + def test_unrelated_register_is_clean(self): + assert not war.has_hazard( + _asm( + "buffer_store_dwordx4 v[8:11], v4, s[0:3], 0 offen", + "v_add_f32_e32 v20, v12, v13", + ) + ) + + def test_narrow_stores_are_ignored(self): + """Only dwordx4 has been observed to miscompile, and buffer ops are kept + wherever possible because dropping them is what costs throughput.""" + assert not war.has_hazard( + _asm( + "buffer_store_dwordx2 v[8:9], v4, s[0:3], 0 offen", + "v_add_f32_e32 v8, v12, v13", + ) + ) + + def test_global_stores_are_ignored(self): + assert not war.has_hazard( + _asm( + "global_store_dwordx4 v4, v[8:11], s[0:1]", + "v_add_f32_e32 v8, v12, v13", + ) + ) + + def test_end_of_program_terminates_the_scan(self): + assert not war.has_hazard( + _asm( + "buffer_store_dwordx4 v[8:11], v4, s[0:3], 0 offen", + "s_endpgm", + "v_add_f32_e32 v8, v12, v13", + ) + ) + + def test_comments_and_directives_are_skipped(self): + assert war.has_hazard( + _asm( + "buffer_store_dwordx4 v[8:11], v4, s[0:3], 0 offen", + "; %bb.1:", + ".p2align 6", + "v_add_f32_e32 v8, v12, v13", + ) + ) + + def test_second_store_is_still_checked(self): + """A clean first store must not stop the scan.""" + assert war.has_hazard( + _asm( + "buffer_store_dwordx4 v[4:7], v0, s[0:3], 0 offen", + "s_waitcnt vmcnt(0)", + "buffer_store_dwordx4 v[8:11], v0, s[0:3], 0 offen", + "v_add_f32_e32 v9, v12, v13", + ) + ) + + def test_kernel_without_buffer_stores(self): + assert not war.has_hazard(_asm("v_add_f32_e32 v8, v12, v13", "s_endpgm")) + + def test_empty_input(self): + assert not war.has_hazard("") + + +class TestEnableGate: + """The gate must not narrow to an architecture allowlist. + + Wrong values were measured on gfx942; gfx950 emits the same pattern and was + measured to tolerate it. Since the omitted `s_waitcnt` is a compiler defect + on both, which chips tolerate it is not something to encode in a gate. + Whether a kernel is affected is decided by its emitted machine code, so the + gate only asks whether this is a ROCm GPU at all. + """ + + def _fake_torch(self, hip): + return types.SimpleNamespace( + cuda=types.SimpleNamespace(is_available=lambda: True), + version=types.SimpleNamespace(hip=hip), + ) + + def test_enabled_on_any_rocm_gpu(self, monkeypatch): + monkeypatch.setitem(sys.modules, "torch", self._fake_torch("6.4.0")) + assert war._on_rocm() + + def test_disabled_on_a_cuda_build(self, monkeypatch): + monkeypatch.setitem(sys.modules, "torch", self._fake_torch(None)) + assert not war._on_rocm() + + def test_disabled_when_torch_is_missing(self, monkeypatch): + """The module also imports on machines with no torch at all.""" + monkeypatch.setitem(sys.modules, "torch", None) + assert not war._on_rocm() + + def test_gate_does_not_inspect_the_architecture_name(self): + """A gate reading gcnArchName is how the gfx950 blind spot happened.""" + assert "gcnArchName" not in inspect.getsource(war._on_rocm) diff --git a/tests/unit_tests/test_finite_training_metrics.py b/tests/unit_tests/test_finite_training_metrics.py new file mode 100644 index 000000000..57b747371 --- /dev/null +++ b/tests/unit_tests/test_finite_training_metrics.py @@ -0,0 +1,59 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Detection of non-finite per-step metrics in trainer logs. + +A diverged run still exits 0 and still prints the completion marker, so this +check is what keeps such a run from passing as green. It fails open by design: +a log it cannot parse reports zero checked steps rather than erroring, which +makes a silent parsing regression the failure mode worth testing for. +""" + +import pytest + +from tests.utils import assert_finite_training_metrics + +TORCHTITAN = "step: {step} loss: {loss} grad_norm: {grad} memory: 30.5GiB(38.10%)" +MEGATRON = ( + " iteration {step}/ 3 | consumed samples: 24 | elapsed time per iteration (ms): 1234.5 |" + " lm loss: {loss} | grad norm: {grad} | number of skipped iterations: 0 |" +) + + +@pytest.mark.parametrize("fmt", [TORCHTITAN, MEGATRON], ids=["torchtitan", "megatron"]) +class TestFiniteTrainingMetrics: + def test_clean_run_is_parsed(self, fmt): + checked = assert_finite_training_metrics("t", fmt.format(step=1, loss="1.17E+01", grad="5.885")) + assert checked == 1, "log format no longer recognized; divergence would go undetected" + + @pytest.mark.parametrize("bad", ["nan", "inf", "-inf"]) + def test_non_finite_loss_fails(self, fmt, bad): + with pytest.raises(AssertionError, match="non-finite"): + assert_finite_training_metrics("t", fmt.format(step=2, loss=bad, grad="5.885")) + + def test_non_finite_grad_norm_fails(self, fmt): + with pytest.raises(AssertionError, match="non-finite"): + assert_finite_training_metrics("t", fmt.format(step=3, loss="1.17E+01", grad="nan")) + + def test_only_the_diverged_step_is_reported(self, fmt): + log = "\n".join( + ( + fmt.format(step=1, loss="1.17E+01", grad="5.885"), + fmt.format(step=2, loss="nan", grad="5.885"), + ) + ) + with pytest.raises(AssertionError, match=r"step 2: loss=nan"): + assert_finite_training_metrics("t", log) + + +def test_ansi_colored_log_is_still_parsed(): + colored = "\x1b[32m" + TORCHTITAN.format(step=1, loss="nan", grad="5.885") + "\x1b[0m" + with pytest.raises(AssertionError, match="non-finite"): + assert_finite_training_metrics("t", colored) + + +def test_unrecognized_format_is_not_a_failure(): + assert assert_finite_training_metrics("t", "some launcher output\nwithout metrics\n") == 0 diff --git a/tests/utils.py b/tests/utils.py index 83490432b..c285de0ed 100755 --- a/tests/utils.py +++ b/tests/utils.py @@ -5,7 +5,9 @@ ############################################################################### +import math import os +import re import subprocess import sys import time @@ -16,6 +18,58 @@ TRAINING_COMPLETED_MARKER = "Training completed." +_ANSI_ESCAPE_RE = re.compile(r"\x1b\[[0-9;]*m") + +# Per-step metric lines, one pattern per backend log format: +# torchtitan: "step: 1 loss: 12.78468 grad_norm: nan memory: ..." +# megatron: "iteration 3/ 3 | ... | lm loss: 1.17E+01 | ... | grad norm: 5.885 | ..." +_STEP_METRIC_RES = ( + re.compile(r"\bstep:\s*(?P\d+)\b.*?\bloss:\s*(?P\S+).*?\bgrad_norm:\s*(?P\S+)"), + re.compile( + r"\biteration\s*(?P\d+)\s*/.*?\blm loss:\s*(?P\S+).*?\bgrad norm:\s*(?P\S+)" + ), +) + + +def assert_finite_training_metrics(tag: str, log_text: str) -> int: + """Fail when any logged per-step loss or grad norm is NaN/Inf. + + A training run that diverges numerically still exits 0 and still prints the + "Training completed." marker, so without this check such a run passes as a + green test. Returns the number of steps that were checked; 0 means the log + had no recognizable metric lines, which is not treated as a failure so that + backends with other log formats keep working. + """ + plain = _ANSI_ESCAPE_RE.sub("", log_text) + + checked = 0 + bad: list[str] = [] + for line in plain.splitlines(): + if "loss:" not in line: # the vast majority of log lines; skip the regexes + continue + for pattern in _STEP_METRIC_RES: + match = pattern.search(line) + if match is None: + continue + checked += 1 + for field in ("loss", "grad"): + raw = match.group(field).rstrip("|,") + try: + value = float(raw) + except ValueError: + continue + if not math.isfinite(value): + bad.append(f"step {match.group('step')}: {field}={raw}") + break + + if bad: + raise AssertionError( + f"[{tag}] Training reported non-finite metrics, so the run diverged even " + f"though the process exited 0: {', '.join(bad[:8])}" + ) + + return checked + def skip_if_no_cuda(reason: str = "requires GPU (primus_turbo initializes CUDA at import)") -> None: """Skip the calling test module at collection time when CUDA is unavailable. @@ -97,6 +151,7 @@ def run_training_script( cmd: list[str], train_log_path: str, env: Optional[dict] = None, + check_metrics: bool = True, ) -> tuple[str, str]: """Execute a training command and validate that training completed successfully. @@ -110,6 +165,7 @@ def run_training_script( cmd: Command to execute (passed to subprocess.run). train_log_path: Path to the training log file written by the launcher. env: Environment variables for the subprocess. + check_metrics: Also require every logged loss and grad norm to be finite. Returns: (stdout_output, stderr_output) tuple where stdout_output is the @@ -151,6 +207,9 @@ def run_training_script( f"Log file: {train_log_path}" ) + if check_metrics: + assert_finite_training_metrics(tag, stdout_output) + return stdout_output, "" except subprocess.CalledProcessError as e: diff --git a/tools/auto_benchmark/metrics.py b/tools/auto_benchmark/metrics.py index a38a98a71..dbfcc67ba 100644 --- a/tools/auto_benchmark/metrics.py +++ b/tools/auto_benchmark/metrics.py @@ -2,6 +2,7 @@ import argparse import csv +import math import os import re import shutil @@ -16,8 +17,20 @@ WARMUP_SKIP = 5 ERROR_STATUS = "Error found - check log" +NONFINITE_STATUS = "Diverged (non-finite loss) - invalid" OK_STATUS = "OK" +# Per-step metric lines, one pattern per backend log format: +# torchtitan: "step: 1 loss: 12.78468 grad_norm: nan memory: ..." +# megatron: "iteration 3/ 3 | ... | lm loss: 1.17E+01 | ... | grad norm: 5.885 | ..." +# Same parsing as tests/utils.py:assert_finite_training_metrics. +STEP_METRIC_REGEXES = ( + re.compile(r"\bstep:\s*(?P\d+)\b.*?\bloss:\s*(?P\S+).*?\bgrad_norm:\s*(?P\S+)"), + re.compile( + r"\biteration\s*(?P\d+)\s*/.*?\blm loss:\s*(?P\S+).*?\bgrad norm:\s*(?P\S+)" + ), +) + ANSI_ESCAPE_REGEX = re.compile(r"\x1b\[[0-9;]*m") LOG_EXIT_CODE_REGEX = re.compile(r"primus launcher exited with code (\d+)", re.IGNORECASE) @@ -65,6 +78,8 @@ - Timing/throughput fields use the current value before "/" (e.g. 5896.1/5913.1 -> 5896.1). - Warm-up: the first five iterations are excluded before averaging. - "Status" is "Error found - check log" when the log contains errors or metrics could not be computed. +- "Status" is "Diverged (non-finite loss) - invalid" when any step logged a NaN/Inf loss or grad + norm. Such a run exits 0 and still prints throughput, but the numbers are meaningless. - Numeric fields may contain commas in logs; commas are removed before averaging. """ @@ -76,6 +91,8 @@ - "Steps" is the number of training steps used after dropping the first five warm-up steps. - TPS and TFLOPS values may contain commas in logs; commas are removed before averaging. - "Status" is "Error found - check log" when the log contains errors or metrics could not be computed. +- "Status" is "Diverged (non-finite loss) - invalid" when any step logged a NaN/Inf loss or grad + norm. Such a run exits 0 and still prints throughput, but the numbers are meaningless. """ MEGATRON_NUM = r"[\d,]+(?:\.\d+)?" @@ -311,9 +328,35 @@ def strip_ansi(text): return ANSI_ESCAPE_REGEX.sub("", text) -def log_has_error(path): +def find_nonfinite_metric(plain): + if "loss:" not in plain: # the vast majority of log lines; skip the regexes + return None + for pattern in STEP_METRIC_REGEXES: + m = pattern.search(plain) + if m is None: + continue + for field in ("loss", "grad"): + raw = m.group(field).rstrip("|,") + try: + value = float(raw) + except ValueError: + continue + if not math.isfinite(value): + return f"step {m.group('step')} {field}={raw}" + return None + return None + + +def scan_log(path): + """Single pass over a run log for launcher errors and non-finite metrics. + + A diverged run still exits 0 and still prints throughput, so without the + second check a NaN run reports a (inflated) tps as if it were a valid result. + Returns (has_error, nonfinite) where nonfinite lists the offending steps. + """ exit_code = None saw_error_line = False + nonfinite = [] with open(path, "r", errors="ignore") as f: for line in f: @@ -323,15 +366,18 @@ def log_has_error(path): if m: exit_code = int(m.group(1)) + bad = find_nonfinite_metric(plain) + if bad: + nonfinite.append(bad) + if any(pattern.search(plain) for pattern in LOG_ERROR_EXCLUSIONS): continue if any(pattern.search(plain) for pattern in LOG_ERROR_PATTERNS): saw_error_line = True - if exit_code not in (None, 0): - return True - return saw_error_line + has_error = True if exit_code not in (None, 0) else saw_error_line + return has_error, nonfinite def render_results(backend, headers, rows, note_text): @@ -349,6 +395,14 @@ def render_results(backend, headers, rows, note_text): " Open the corresponding log file for details." ) + diverged_rows = [row for row in rows if len(row) > 2 and row[2] == NONFINITE_STATUS] + if diverged_rows: + labels = ", ".join(f"{row[0]} {row[1]}" for row in diverged_rows) + print( + f"\n{len(diverged_rows)} run(s) logged a non-finite loss or grad norm and are" + f" reported without throughput: {labels}" + ) + table_width = sum(len(str(header)) for header in headers) + 3 * len(headers) if table_width > terminal_width(): print("\n(Table may be wider than the terminal. Open the CSV for the full view.)") @@ -444,17 +498,17 @@ def megatron_collect_rows(): continue path = os.path.join(log_dir, fname) - has_error = log_has_error(path) + has_error, nonfinite = scan_log(path) records = megatron_parse_log_file(path) stats = megatron_compute_averages(records) bs, seq, gbs = load_training_params("megatron", meta["device"], meta["model"], fname, log_dir) - if has_error or not stats: + if has_error or nonfinite or not stats: if gbs == "-" and records: gbs = records[0]["gbs"] values = [ - ERROR_STATUS, + NONFINITE_STATUS if nonfinite and not has_error else ERROR_STATUS, "megatron", meta["device"], bs, @@ -615,7 +669,7 @@ def torchtitan_collect_rows(): continue path = os.path.join(log_dir, fname) - has_error = log_has_error(path) + has_error, nonfinite = scan_log(path) steps = torchtitan_parse_log_file(path) stats = torchtitan_compute_averages(steps) @@ -629,9 +683,9 @@ def torchtitan_collect_rows(): if log_seq is not None: seq = log_seq - if has_error or not stats: + if has_error or nonfinite or not stats: values = [ - ERROR_STATUS, + NONFINITE_STATUS if nonfinite and not has_error else ERROR_STATUS, "torchtitan", meta["device"], bs,