diff --git a/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml index a342485f4..6bb6b5b75 100644 --- a/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml @@ -3,6 +3,16 @@ user_name: ${PRIMUS_USER:root} exp_name: ${PRIMUS_EXP_NAME:qwen3_30B_A3B-pretrain} workspace: ${PRIMUS_WORKSPACE:./output} +# Fused packed-QKV split + Q/K RMSNorm + RoPE. head_dim 128 needs a Primus-Turbo +# with QK_RMSNORM_ROPE_HEAD_DIMS; older Turbo keeps the unfused path. +# +# Full DeepEP fence at turbo_deepep_num_cu>=160: the cheap fence (#539) was only +# row-validated at <=80 CUs; the full opt stack hit intermittent NaN forward loss +# at 160/192 (~1/13). DISABLE=1 selects the release/acquire fence (~2-3% step). +env: + PRIMUS_FUSED_QK_RMSNORM_ROPE: "1" + PRIMUS_TURBO_DEEPEP_DISABLE_CHEAP_FENCE: "1" + modules: pre_trainer: framework: megatron @@ -83,16 +93,21 @@ modules: # Turbo enable_primus_turbo: true use_turbo_attention: false - use_turbo_grouped_gemm: false + use_turbo_gemm: true + use_turbo_grouped_gemm: true use_turbo_rms_norm: false # bug + turbo_fused_grouped_gemm: true + use_turbo_fused_act_with_probs: true + # needs a Primus-Turbo whose moe_permute takes quantize_dtype + turbo_fp8_permute: true # deepep use_turbo_deepep: true moe_shared_expert_overlap: false moe_router_dtype: fp32 - # 64 or 80 for ep8, 32 for ep16-64 is best practice - turbo_deepep_num_cu: 80 + # 160 for ep8 on MI355X (192 gives a NaN loss, 256 is slower), 32 for ep16-64 + turbo_deepep_num_cu: 160 turbo_deepep_use_comm_stream: false # sync-free moe support stage 1-2, 0 means not use sync-free moe @@ -105,5 +120,6 @@ modules: # enable fp8 training fp8: hybrid - moe_use_legacy_grouped_gemm: true + fp8_recipe: tensorwise + moe_use_legacy_grouped_gemm: false gradient_accumulation_fusion: true diff --git a/primus/backends/megatron/core/extensions/primus_turbo.py b/primus/backends/megatron/core/extensions/primus_turbo.py index 6b8addc88..b5d5716ab 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo.py +++ b/primus/backends/megatron/core/extensions/primus_turbo.py @@ -975,6 +975,7 @@ def __init__( self._num_heads_for_sinks = self.config.num_attention_heads self.offload = args.offload and "attn" in args.offload_ops + self.use_fp8_attention = args.enable_turbo_attention_float8 if args.enable_turbo_attention_float8: self.attn = ( primus_turbo_torch.ops.flash_attn_fp8_usp_func @@ -1129,6 +1130,15 @@ def forward( key = key.permute(0, 2, 1, 3) value = value.permute(0, 2, 1, 3) + if self.use_fp8_attention: + # flash_attn_fp8_func permutes non-bshd *storage* as if it were the + # logical layout, swapping b and s, so hand it bshd-contiguous q/k/v. + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + # flash_attn_fp8_func takes no ``sink`` argument. + sink_kwargs = {"sink": sink_tensor} if sink_tensor is not None else {} o = self.attn( query, key, @@ -1142,7 +1152,7 @@ def forward( deterministic=self.deterministic_mode, return_lse=False, return_attn_probs=False, - sink=sink_tensor, # PR 208: pass sink tensor to Primus-Turbo + **sink_kwargs, **self.attn_kwargs, ) @@ -2630,6 +2640,23 @@ def __init__( self.moe_router_force_load_balancing_type = getattr( args, "moe_router_force_load_balancing_type", "uniform" ) + self.turbo_fp8_permute = getattr(args, "turbo_fp8_permute", False) + + def _fp8_permute_dtypes(self): + """(input, grad) FP8 dtypes the fused grouped MLP expects, or (None, None). + + Only tensorwise current scaling: that is the one recipe whose scale is a single scalar, + so quantizing the dispatched tokens before permute is exact. + """ + if not self.turbo_fp8_permute or not PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled(): + return None, None + quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() + if not quant_config.current_scaling(): + return None, None + from primus_turbo.pytorch.ops.utils import _get_fp8_dtype + + fmt = quant_config.data().format + return _get_fp8_dtype(fmt, True), _get_fp8_dtype(fmt, False) def dispatch_preprocess( self, hidden_states: torch.Tensor, routing_map: torch.Tensor, probs: torch.Tensor @@ -2711,9 +2738,15 @@ def dispatch_postprocess(self, hidden_states: torch.Tensor, probs: torch.Tensor) Returns: A tuple of permuted tokens, token counts per expert, and permuted probabilities. """ - permuted_input, tokens_per_expert, permuted_probs = self.deepep_dispatcher._post_dispatch( - hidden_states, probs - ) + quantize_dtype, _ = self._fp8_permute_dtypes() + if quantize_dtype is None: + permuted_input, tokens_per_expert, permuted_probs = self.deepep_dispatcher._post_dispatch( + hidden_states, probs + ) + else: + permuted_input, tokens_per_expert, permuted_probs = self.deepep_dispatcher._post_dispatch( + hidden_states, probs, quantize_dtype=quantize_dtype + ) if self.config.moe_router_dtype == "fp64": permuted_probs = permuted_probs.to(torch.float64) return permuted_input, tokens_per_expert, permuted_probs @@ -2724,7 +2757,13 @@ def combine_preprocess(self, hidden_states: torch.Tensor): This method restores the hidden states to their original ordering before expert processing by using the communication manager's restoration function. """ - hidden_states = self.deepep_dispatcher._pre_combine(hidden_states) + _, grad_quantize_dtype = self._fp8_permute_dtypes() + if grad_quantize_dtype is None: + hidden_states = self.deepep_dispatcher._pre_combine(hidden_states) + else: + hidden_states = self.deepep_dispatcher._pre_combine( + hidden_states, grad_quantize_dtype=grad_quantize_dtype + ) return hidden_states def token_combine( diff --git a/primus/backends/megatron/patches/args/rocm_arg_validation.py b/primus/backends/megatron/patches/args/rocm_arg_validation.py index c00621a4e..c8c189e4c 100644 --- a/primus/backends/megatron/patches/args/rocm_arg_validation.py +++ b/primus/backends/megatron/patches/args/rocm_arg_validation.py @@ -108,6 +108,19 @@ def validate_turbo_grouped_gemm_without_padding(args) -> None: raise ValueError(f"{option}=True requires moe_router_padding_for_quantization=False.") +def validate_turbo_fp8_permute(args) -> None: + """The FP8 permute handoff only has one consumer: the fused grouped MLP behind Turbo DeepEP.""" + option = "turbo_fp8_permute" + if not getattr(args, option, False): + return + for required in ("enable_primus_turbo", "use_turbo_deepep", "turbo_fused_grouped_gemm"): + if not getattr(args, required, False): + raise ValueError(f"{option}=True requires {required}=True.") + recompute_modules = getattr(args, "recompute_modules", None) or [] + if getattr(args, "recompute_granularity", None) == "selective" and "moe_act" in recompute_modules: + raise ValueError(f"{option}=True does not support recompute of moe_act.") + + def validate_args_on_rocm(args): # Primus-Turbo auto-tuning use_turbo_autotune = getattr(args, "use_turbo_autotune", False) @@ -216,6 +229,7 @@ def validate_args_on_rocm(args): ) validate_turbo_grouped_gemm_without_padding(args) + validate_turbo_fp8_permute(args) # turbo deepep if args.use_turbo_deepep: diff --git a/primus/backends/megatron/patches/turbo/qk_rmsnorm_rope_patches.py b/primus/backends/megatron/patches/turbo/qk_rmsnorm_rope_patches.py index a8193444c..a7aa6f151 100644 --- a/primus/backends/megatron/patches/turbo/qk_rmsnorm_rope_patches.py +++ b/primus/backends/megatron/patches/turbo/qk_rmsnorm_rope_patches.py @@ -4,7 +4,7 @@ # See LICENSE for license information. ############################################################################### -"""Opt-in GPT-OSS packed-QKV RMSNorm + RoPE integration. +"""Opt-in packed-QKV RMSNorm + RoPE integration (GPT-OSS head_dim 64, Qwen3 head_dim 128). ``PRIMUS_FUSED_QK_RMSNORM_ROPE=1`` replaces the training-only sequence @@ -12,8 +12,8 @@ with Primus-Turbo's packed FlyDSL operator. The patch is intentionally narrow: TP=1, self attention, SBHD, non-interleaved full RoPE, no packed sequences, -and PrimusTurboRMSNorm for both Q and K. Unsupported calls continue through -the original Megatron path. +a head_dim the installed Turbo supports, and TE RMSNorm (or PrimusTurboRMSNorm) +for both Q and K. Unsupported calls continue through the original Megatron path. """ from __future__ import annotations @@ -79,7 +79,7 @@ def _enabled(_ctx: PatchContext) -> bool: "megatron.turbo.qk_rmsnorm_rope", backend="megatron", phase="before_train", - description="Fuse GPT-OSS packed QKV split, Q/K RMSNorm, and RoPE with FlyDSL", + description="Fuse packed QKV split, Q/K RMSNorm, and RoPE with FlyDSL", condition=_enabled, priority=70, ) @@ -92,8 +92,14 @@ def patch_qk_rmsnorm_rope(_ctx: PatchContext): qk_rmsnorm_rope_shape_error, ) from primus_turbo.pytorch.ops.rope import fused_qkv_rmsnorm_rope + from transformer_engine.pytorch.module.rmsnorm import RMSNorm as TERMSNorm - from primus.backends.megatron.core.extensions.primus_turbo import PrimusTurboRMSNorm + try: + from primus_turbo.pytorch.kernels.rope.qk_rmsnorm_rope_impl import ( + QK_RMSNORM_ROPE_HEAD_DIMS, + ) + except ImportError: + QK_RMSNORM_ROPE_HEAD_DIMS = (64,) if getattr(attention_module, "_primus_qk_rmsnorm_rope_installed", False): return @@ -121,15 +127,16 @@ def _eligible(self, rotary_pos_emb, inference_context, packed_seq_params) -> boo and rotary_pos_emb is not None and not no_rope and self.attention_type == "self" - and self.hidden_size_per_attention_head == 64 + and self.hidden_size_per_attention_head in QK_RMSNORM_ROPE_HEAD_DIMS and not getattr(self.config, "rotary_interleaved", False) and getattr(self.config, "rotary_percent", 1.0) == 1.0 and attention_module._yarn_get_concentration_factor_from_config(self.config) == 1.0 and not getattr(self.config, "attention_output_gate", False) and getattr(self, "world_size", 1) == 1 and not getattr(self, "offload_qkv_linear", False) - and isinstance(q_norm, PrimusTurboRMSNorm) - and isinstance(k_norm, PrimusTurboRMSNorm) + # Only weight / eps are consumed, so TE RMSNorm and its PrimusTurboRMSNorm subclass both qualify. + and isinstance(q_norm, TERMSNorm) + and isinstance(k_norm, TERMSNorm) and _supported_forward_pre_hooks(q_norm) is not None and _supported_forward_pre_hooks(k_norm) is not None and not getattr(q_norm, "zero_centered_gamma", False) @@ -278,6 +285,5 @@ def _apply_rotary_pos_emb(t, *args, **kwargs): attention_module.apply_rotary_pos_emb = _apply_rotary_pos_emb attention_module._primus_qk_rmsnorm_rope_installed = True log_rank_0( - "[Patch:megatron.turbo.qk_rmsnorm_rope] Installed GPT-OSS FlyDSL " - "packed QKV + Q/K RMSNorm + RoPE fusion" + "[Patch:megatron.turbo.qk_rmsnorm_rope] Installed FlyDSL packed QKV + Q/K RMSNorm + RoPE fusion" ) diff --git a/primus/configs/modules/megatron/primus_turbo.yaml b/primus/configs/modules/megatron/primus_turbo.yaml index a09fa893c..3d1b56460 100644 --- a/primus/configs/modules/megatron/primus_turbo.yaml +++ b/primus/configs/modules/megatron/primus_turbo.yaml @@ -61,6 +61,11 @@ use_turbo_grouped_gemm: false # Fuse the grouped MLP chain (fc1 + GLU + fc2, including WGRAD) into Primus-Turbo's # fused grouped MLP op. turbo_fused_grouped_gemm: false +# Hand the fused grouped MLP its input and output gradient already in FP8 (tensorwise current +# scaling only): the DeepEP dispatcher quantizes the dispatched tokens before permute, and the +# combine gradient before un-permute's backward, instead of their topk-times larger permuted +# copies. Bit-identical. Requires turbo_fused_grouped_gemm and use_turbo_deepep. +turbo_fp8_permute: false # let Turbo grouped GEMMs consume non-aligned expert token counts directly, # without TE Fp8Padding/Fp8Unpadding turbo_grouped_gemm_without_padding: false diff --git a/tests/unit_tests/backends/megatron/test_rocm_arg_validation.py b/tests/unit_tests/backends/megatron/test_rocm_arg_validation.py index dd52c5bd3..bcb868b12 100644 --- a/tests/unit_tests/backends/megatron/test_rocm_arg_validation.py +++ b/tests/unit_tests/backends/megatron/test_rocm_arg_validation.py @@ -9,6 +9,7 @@ import pytest from primus.backends.megatron.patches.args.rocm_arg_validation import ( + validate_turbo_fp8_permute, validate_turbo_grouped_gemm_without_padding, ) @@ -62,3 +63,49 @@ def test_turbo_grouped_gemm_without_padding_disabled_is_noop(): def test_turbo_grouped_gemm_without_padding_rejects_unsupported_config(overrides, message): with pytest.raises(ValueError, match=message): validate_turbo_grouped_gemm_without_padding(_no_padding_args(**overrides)) + + +def _fp8_permute_args(**overrides): + values = { + "turbo_fp8_permute": True, + "enable_primus_turbo": True, + "use_turbo_deepep": True, + "turbo_fused_grouped_gemm": True, + } + values.update(overrides) + return SimpleNamespace(**values) + + +@pytest.mark.parametrize( + "overrides", + [ + {}, + {"recompute_granularity": "full", "recompute_modules": ["moe_act"]}, + {"recompute_granularity": "selective", "recompute_modules": ["core_attn"]}, + {"recompute_granularity": "selective", "recompute_modules": None}, + ], +) +def test_turbo_fp8_permute_accepts_supported_config(overrides): + validate_turbo_fp8_permute(_fp8_permute_args(**overrides)) + + +def test_turbo_fp8_permute_disabled_is_noop(): + validate_turbo_fp8_permute(SimpleNamespace()) + validate_turbo_fp8_permute(SimpleNamespace(turbo_fp8_permute=False, use_turbo_deepep=False)) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"enable_primus_turbo": False}, "requires enable_primus_turbo=True"), + ({"use_turbo_deepep": False}, "requires use_turbo_deepep=True"), + ({"turbo_fused_grouped_gemm": False}, "requires turbo_fused_grouped_gemm=True"), + ( + {"recompute_granularity": "selective", "recompute_modules": ["core_attn", "moe_act"]}, + "does not support recompute of moe_act", + ), + ], +) +def test_turbo_fp8_permute_rejects_unsupported_config(overrides, message): + with pytest.raises(ValueError, match=message): + validate_turbo_fp8_permute(_fp8_permute_args(**overrides)) diff --git a/tests/unit_tests/backends/megatron/test_turbo_fp8_permute_dtypes.py b/tests/unit_tests/backends/megatron/test_turbo_fp8_permute_dtypes.py new file mode 100644 index 000000000..4fc7e3b12 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_turbo_fp8_permute_dtypes.py @@ -0,0 +1,69 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Unit tests for the FP8 dtypes the Turbo DeepEP dispatcher asks the permute for under +``turbo_fp8_permute``. Quantizing before permute is only exact for tensorwise current +scaling, so every other state must keep the bf16 permute. +""" + +from types import SimpleNamespace + +import pytest + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus_turbo.pytorch.core.low_precision import ( # noqa: E402 isort:skip + Format, + ScalingGranularity, + float8_e4m3, + float8_e5m2, +) + +from primus.backends.megatron.core.extensions.primus_turbo import ( # noqa: E402 isort:skip + PrimusTurboDeepEPTokenDispatcher, + PrimusTurboLowPrecisionGlobalStateManager, + PrimusTurboQuantConfig, +) + + +def _dtypes(monkeypatch, *, flag=True, turbo_fp8=True, quant_config=None): + monkeypatch.setattr(PrimusTurboLowPrecisionGlobalStateManager, "PRIMUS_TURBO_FP8_ENABLED", turbo_fp8) + monkeypatch.setattr(PrimusTurboLowPrecisionGlobalStateManager, "PRIMUS_TURBO_QUANT_CONFIG", quant_config) + dispatcher = SimpleNamespace(turbo_fp8_permute=flag) + return PrimusTurboDeepEPTokenDispatcher._fp8_permute_dtypes(dispatcher) + + +@pytest.mark.parametrize( + ("fmt", "expected"), + [ + (Format.HYBRID, (float8_e4m3, float8_e5m2)), + (Format.E4M3, (float8_e4m3, float8_e4m3)), + ], +) +def test_tensorwise_returns_input_and_grad_dtypes(monkeypatch, fmt, expected): + quant_config = PrimusTurboQuantConfig(format=fmt, granularity=ScalingGranularity.TENSORWISE) + assert _dtypes(monkeypatch, quant_config=quant_config) == expected + + +def test_flag_off_keeps_bf16_permute(monkeypatch): + quant_config = PrimusTurboQuantConfig(format=Format.HYBRID, granularity=ScalingGranularity.TENSORWISE) + assert _dtypes(monkeypatch, flag=False, quant_config=quant_config) == (None, None) + + +def test_turbo_fp8_off_keeps_bf16_permute(monkeypatch): + """A layer outside Turbo FP8 (e.g. pinned to BF16) must not get FP8 tokens.""" + quant_config = PrimusTurboQuantConfig(format=Format.HYBRID, granularity=ScalingGranularity.TENSORWISE) + assert _dtypes(monkeypatch, turbo_fp8=False, quant_config=quant_config) == (None, None) + + +def test_blockwise_keeps_bf16_permute(monkeypatch): + quant_config = PrimusTurboQuantConfig( + format=Format.HYBRID, granularity=ScalingGranularity.BLOCKWISE, block_size=128 + ) + assert _dtypes(monkeypatch, quant_config=quant_config) == (None, None)