diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4_to_fp8.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4_to_fp8.yaml new file mode 100644 index 000000000..53b56e593 --- /dev/null +++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4_to_fp8.yaml @@ -0,0 +1,267 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell Training — DDP + Local Spec + MXFP4 -> FP8 switch (MI355X) +# +# Starts in MXFP4 and flips every MXFP4 linear to dynamic tensorwise FP8 at +# mxfp4_to_fp8_switch_iter, with no checkpoint, no weight conversion and no +# optimizer state remap (neither precision stores a quantized weight: both keep a +# BF16 nn.Parameter and quantize inside forward). Checkpoint-mediated transitions +# are forbidden under MLPerf, which is why the switch is iteration-triggered. +# +# REQUIRED ENVIRONMENT — per-precision Primus-Turbo GEMM backends: +# +# export PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER,FP8:FLYDSL +# # leave PRIMUS_TURBO_AUTO_TUNE unset +# +# One process has to serve both precisions, and their usual recipes conflict: +# MXFP4 needs FP4 pinned to AITER with autotune off (otherwise the linear refuses +# to construct -- the AITER preshuffle fast path is the only one that understands +# the shuffled layout), while the FP8 guide asks for autotune on and the backend +# unset. PRIMUS_TURBO_GEMM_BACKEND is per-precision, so pinning each slot +# satisfies both: FP8 gets an explicit fast backend and therefore never needs the +# autotuner MXFP4 forbids. +# +# Getting this wrong fails in two very different ways. Dropping the FP4 pin raises +# at model construction, which is loud and fine. Dropping the FP8 pin silently +# runs post-switch GEMMs on whatever the dispatcher picks, which quietly +# invalidates any step-time comparison against the FP8 baseline. +# +# CONSTRAINT: FlyDSL FP8 is gfx950-only (the kernel uses mfma_f32_16x16x128_f8f6f4), +# so this config is MI355X-only. On gfx942/MI300X the FP8 pin is unusable. +# +# CONSTRAINT: a pinned backend is strict. When a pin is set the dispatcher raises +# if can_handle is False rather than falling back, so every FP8 GEMM in the +# process must satisfy FlyDSL's tensorwise constraints: K > 128 (its pipeline +# needs two K tiles), scalar per-tensor scales, and any layout but TT. Flux clears +# these everywhere -- forward K is hidden or ffn, and the weight-gradient K is the +# token count -- but the margin is per-GEMM, not global, so a genuinely small FP8 +# matmul added anywhere in this process would hard-fail instead of degrading. +# Measured on MI355X: the whole model's FP8 arm is one extra compiled graph, and +# saved operands go from 0.53 B/element (packed FP4 plus E8M0 block scales) to +# 1.0 B/element, so budget the recurring activation peak at ~1.9x its MXFP4 value. +# +# Note for local runs: exporting this pin also applies to the unit tests, where +# some FP8 cases use small shapes that FlyDSL declines. Run the test suite without +# it (tests pin backends in-code where they need to). + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_local_mxfp4_to_fp8} + +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_12b.yaml + + trainer_class: FluxPretrainTrainer + + overrides: + model_type: flux_schnell + + # ========================================== + # Flux Schnell training recipe + # ========================================== + # Short LR warmup ramp over the first optimizer steps for stability. + nemo_aligned_lr_warmup: true + # Also hosts the FP8 graph pre-warm: the only place a real grad-enabled step + # runs at the production micro_batch_size, which is what the pre-warmed cache + # entry has to be guarded on. + warmup_train_steps: 2 + + timestep_sampling_strategy: "direct_uniform" + activation_func: "openai_gelu" + cfg_dropout_prob: 0.1 + + # VAE latent normalization (required for resample mode) + vae_scale: 0.3611 + vae_shift: 0.1159 + + # Resample mode: re-draw latents from mean+logvar each step + vae_latent_mode: resample + + # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout + rotary_interleaved: true + + # ========================================== + # PrimusTurboMXFP4LocalSpecProvider + # ========================================== + transformer_impl: "local" + + # ========================================== + # DATASET CONFIGURATION + # ========================================== + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Training iterations (example value; increase for a full training run) + train_iters: 1000 + eval_interval: 10000 + eval_iters: 0 + log_interval: 1 + save_interval: 10000 + + # Batch configuration + micro_batch_size: 64 + global_batch_size: 512 + seq_length: 512 # 256 img tokens + 256 text tokens (schnell) + + # ========================================== + # BF16 + MXFP4 Precision + # ========================================== + bf16: true + fp16: false + params_dtype: bfloat16 + grad_reduce_in_bf16: true + + # Optimizer settings + optimizer: adam + lr: 2.0e-4 + min_lr: 2.0e-4 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-8 + clip_grad: 1.0 + + # Learning rate scheduler (warmup-hold, no decay) + lr_warmup_iters: 1600 + lr_decay_iters: 4000 + lr_decay_style: constant + + # DataLoader settings + num_workers: 8 + dataloader_type: external + max_samples_per_sequence: null + + # ========================================== + # Megatron DDP + Distributed Optimizer + # ========================================== + + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + + use_torch_fsdp2: false + use_megatron_fsdp: false + + use_distributed_optimizer: true + overlap_grad_reduce: true + overlap_param_gather: true + overlap_param_gather_with_optimizer_step: false + use_precision_aware_optimizer: true + main_params_dtype: fp32 + main_grads_dtype: bf16 + exp_avg_dtype: bf16 + exp_avg_sq_dtype: bf16 + + ddp_bucket_size: 256000000 + use_fsdp2_fp32_param_optimizer: false + + ckpt_format: torch_dist + + gradient_accumulation_fusion: false + + check_for_nan_in_loss_and_grad: true + + # ========================================== + # MXFP4 Configuration — Block-scaled via Primus Turbo + AITER + # ========================================== + + use_flash_attn: true + + fp4: "mxfp4" + fp4_recipe: "mxfp4" + mxfp4_backward_precision: "mxfp4" # "mxfp4" (pure) or "fp8" (hybrid) + + # ========================================== + # MXFP4 -> FP8 runtime switch + # ========================================== + # Fires between two iterations, after the optimizer step and before the next + # forward. The saved-activation pool is empty at that point (no pipelining, + # so backward has already consumed everything), so the two formats never + # coexist and there is nothing to discard or copy. What changes is the + # recurring steady-state peak: MXFP4 saves ~0.53 bytes/element, FP8 + # tensorwise ~1.0, so the high-water mark roughly doubles from here on. That + # is inherent to any design whose end state is 100% FP8; the pre-warm below + # measures it before the measured run starts. + mxfp4_to_fp8_switch_iter: 600 + + # Trace the FP8 graph during warmup so the switch is a guard-driven cache hit + # rather than a recompile at the boundary. Also the startup check that the + # switch is not a silent no-op: if Dynamo does not guard _fp8_mode, no new + # graph is traced and the run fails here instead of quietly training on in + # MXFP4 while logging a successful switch. + mxfp4_to_fp8_prewarm: true + + # 0 converts every layer in one loop at the boundary. The 19 MMDiT and 38 + # single blocks share compiled graphs (Dynamo keys its cache on the code + # object), so the whole-model flip costs ~2 recompiles rather than 57 -- and + # pre-warming removes even those. Set nonzero only if the pre-warm shows full + # FP8 does not fit, in which case the end state is a partial mix. + mxfp4_to_fp8_layers_per_iter: 0 + mxfp4_to_fp8_order: "deep_to_shallow" + + empty_unused_memory_level: 0 + + distributed_timeout_minutes: 180 + distributed_backend: nccl + + # Checkpointing (disabled) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging + tensorboard_dir: output/tensorboard/flux_12b_ddp_local_mxfp4_to_fp8 + wandb_project: flux_12b_ddp_local_mxfp4_to_fp8 + log_throughput: true + wall_clock_step_timer: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + log_memory_to_tensorboard: true + + # PyTorch Profiler + profile: false + use_pytorch_profiler: true + profile_ranks: [0] + profile_step_start: 9 + profile_step_end: 11 + torch_profiler_record_shapes: true + torch_profiler_with_stack: true + torch_profiler_use_gzip: false + disable_profiler_activity_cpu: false + + # Primus Turbo Configuration + enable_primus_turbo: true + use_turbo_attention: true + + seed: 2025 + per_step_rng_reseed: false + nemo_chimera_init: false + + # Torch Compile — compatible with MXFP4 local spec (per-module FP4) + torch_compile: + enable: true + strategy: "per_block" + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: false + # true here, overriding the MXFP4 config's false, because this is a GLOBAL + # Inductor flag applied during tracing rather than at decoration time. The + # FP8 arm is traced during pre-warm and the MXFP4 arm at model build, so + # flipping it in between would give the two arms different BF16 cast + # semantics. Setting it true from the start keeps them consistent and makes + # post-switch numerics comparable to the FP8 MLPerf baseline, which also + # sets it true. + emulate_precision_casts: true + fused_ln_modulate: true diff --git a/primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py b/primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py index 73f5ed68b..fb527296e 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py +++ b/primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py @@ -34,7 +34,10 @@ from primus_turbo.pytorch.kernels.gemm.gemm_fp4_impl import gemm_fp4_impl from primus_turbo.pytorch.kernels.gemm.gemm_fp8_impl import gemm_fp8_impl -from .primus_turbo_float8_local import _quantize_fp8_tw +from .primus_turbo_float8_local import ( + OpaqueFP8LinearTensorwiseFunction, + _quantize_fp8_tw, +) # The AITER MXFP4 preshuffle fast path used to be implemented here as a # monkey patch on GEMMFP4AITERBackend.execute. It now lives in Primus-Turbo @@ -68,10 +71,19 @@ def _enable_preshuffle() -> bool: (pre Primus-Turbo PR #383) so MXFP4LinearFunction can keep passing a bool into its custom ops. """ - return ( - GlobalBackendManager.get_gemm_backend(PrecisionType.FP4) == BackendType.AITER - and not GlobalBackendManager.auto_tune_enabled() - ) + choice = GlobalBackendManager.get_gemm_backend(PrecisionType.FP4) + + # Primus-Turbo PR #447 ("support auto tune on different ops") changed this + # getter from returning a bare BackendType to returning a BackendChoice that + # carries its own auto_tune flag. Both shapes have to work here: the Flux local + # spec needs a Turbo new enough to expose the merged FlashAttnFunc, which + # postdates #447, while older pinned Turbos are still in use elsewhere. + # Comparing the BackendChoice against BackendType silently yields False, which + # would turn every MXFP4 run into a preshuffle-contract failure. + backend = getattr(choice, "backend", choice) + auto_tune = getattr(choice, "auto_tune", False) or GlobalBackendManager.auto_tune_enabled() + + return backend == BackendType.AITER and not auto_tune def _assert_preshuffle_contract(config, preshuffle: bool) -> None: @@ -461,6 +473,72 @@ def backward(ctx, grad_output, *_): return grad_input, grad_weight, None, None, None, None, None, None +# --------------------------------------------------------------------------- +# MXFP4 -> FP8 runtime precision switch +# +# Neither precision stores a quantized weight -- both inherit a plain BF16 +# nn.Parameter from Megatron and quantize inside forward -- so switching is +# purely a change of which autograd Function _forward_impl dispatches to. No +# weight conversion, no parameter re-identification, and therefore nothing for +# DDP's bucket groups or the optimizer's identity-keyed state to notice. +# +# FlyDSL is the FP8 backend. It handles NT/NN/TN natively for tensorwise, so +# force_nt is False: normalizing to NT exists only to steer hipBLASLt onto its +# fast TN kernel, and it costs a pre-transposed copy of both operands. +# +# Resolved leniently. A Primus-Turbo without FLYDSL must still build pure-MXFP4 +# models, so a missing backend degrades to 0 ("let the dispatcher choose") here +# and is rejected by the switch patch, which is the only caller that knows the +# run actually intends to switch. +# --------------------------------------------------------------------------- + +_SWITCH_FP8_BACKEND = getattr(BackendType, "FLYDSL", None) + + +def _init_fp8_switch_state(module) -> None: + """Initialise the FP8 arm of an MXFP4 linear, left dormant at ``_fp8_mode=False``. + + The attribute names are deliberately distinct from the ``_fp8_bwd_dtype`` / + ``_fp8_gran_value`` / ``_fp8_backend_value`` triple that backs + ``mxfp4_backward_precision='fp8'``. That triple is passed into + MXFP4LinearFunction on every call including the pure-MXFP4 path, where it is + ``None``/``0``/``0``; reusing it would change traced constants for a path this + switch must leave bit-identical. + """ + from primus_turbo.pytorch.core.low_precision import float8_e4m3, float8_e5m2 + + module._fp8_mode = False + module._switch_fp8_fwd_dtype = float8_e4m3 + module._switch_fp8_bwd_dtype = float8_e5m2 + module._switch_fp8_gran_value = ScalingGranularity.TENSORWISE.value + module._switch_fp8_backend_value = 0 if _SWITCH_FP8_BACKEND is None else _SWITCH_FP8_BACKEND.value + module._switch_force_nt = False + + +def _fp8_forward(module, input, weight): + """Dynamic tensorwise FP8 linear, landing on the same Function the FP8 spec uses. + + ``_quantize_fp8_tw`` is called directly rather than via ``_extract_fp8_weight``: + the latter only exists to unwrap an ``FP8UnshardedWeightTensor`` produced by an + FSDP2 all-gather, and MXFP4 requires DDP, so that subclass cannot appear. This + keeps an isinstance check and a function-local import out of the traced region. + """ + weight_fp8, weight_scale_inv = _quantize_fp8_tw(weight, module._switch_fp8_fwd_dtype) + + result = OpaqueFP8LinearTensorwiseFunction.apply( + input, + weight, + weight_fp8, + weight_scale_inv, + module._switch_fp8_fwd_dtype, + module._switch_fp8_bwd_dtype, + module._switch_fp8_gran_value, + module._switch_fp8_backend_value, + module._switch_force_nt, + ) + return result[0] + + # --------------------------------------------------------------------------- # MXFP4-aware parallel linear layers # --------------------------------------------------------------------------- @@ -506,20 +584,25 @@ def __init__(self, *args, **kwargs): self._fp8_gran_value = 0 self._fp8_backend_value = 0 + _init_fp8_switch_state(self) + def _forward_impl(self, input, weight, *args, **kwargs): bias = kwargs.get("bias", None) - result = MXFP4LinearFunction.apply( - input, - weight, - self._preshuffle, - self._backward_is_fp8, - self._fp8_bwd_dtype, - self._fp8_gran_value, - self._fp8_backend_value, - self._use_gradient_sr, - ) - output = result[0] + if self._fp8_mode: + output = _fp8_forward(self, input, weight) + else: + result = MXFP4LinearFunction.apply( + input, + weight, + self._preshuffle, + self._backward_is_fp8, + self._fp8_bwd_dtype, + self._fp8_gran_value, + self._fp8_backend_value, + self._use_gradient_sr, + ) + output = result[0] if bias is not None: output = output + bias @@ -566,20 +649,25 @@ def __init__(self, *args, **kwargs): self._fp8_gran_value = 0 self._fp8_backend_value = 0 + _init_fp8_switch_state(self) + def _forward_impl(self, input, weight, *args, **kwargs): bias = kwargs.get("bias", None) - result = MXFP4LinearFunction.apply( - input, - weight, - self._preshuffle, - self._backward_is_fp8, - self._fp8_bwd_dtype, - self._fp8_gran_value, - self._fp8_backend_value, - self._use_gradient_sr, - ) - output = result[0] + if self._fp8_mode: + output = _fp8_forward(self, input, weight) + else: + result = MXFP4LinearFunction.apply( + input, + weight, + self._preshuffle, + self._backward_is_fp8, + self._fp8_bwd_dtype, + self._fp8_gran_value, + self._fp8_backend_value, + self._use_gradient_sr, + ) + output = result[0] if bias is not None: output = output + bias diff --git a/primus/backends/megatron/core/models/diffusion/common/config.py b/primus/backends/megatron/core/models/diffusion/common/config.py index de4074744..a33f326b6 100644 --- a/primus/backends/megatron/core/models/diffusion/common/config.py +++ b/primus/backends/megatron/core/models/diffusion/common/config.py @@ -14,6 +14,10 @@ from megatron.core.enums import Fp8Recipe from megatron.core.transformer.transformer_config import TransformerConfig +# Order the MXFP4 -> FP8 ramp walks layers in. Only consulted by the ramp +# fallback; the default single-boundary switch converts everything at once. +MXFP4_TO_FP8_ORDERS = ("deep_to_shallow", "shallow_to_deep") + @dataclass class BaseDiffusionConfig(TransformerConfig): @@ -77,6 +81,23 @@ class BaseDiffusionConfig(TransformerConfig): # Stochastic rounding on MXFP4 gradients (paper Section 4.4) mxfp4_gradient_stochastic_rounding: bool = False + # Iteration at which MXFP4 linears flip to dynamic tensorwise FP8; 0 disables. + # These four are declared here for validation and for the startup config dump. + # The switch patch itself reads the Primus YAML params namespace, so treat this + # dataclass as a cross-check rather than the source of truth. + mxfp4_to_fp8_switch_iter: int = 0 + + # Trace the FP8 graph during warmup so the switch is a cache hit, not a recompile. + mxfp4_to_fp8_prewarm: bool = True + + # Ramp fallback: layers converted per iteration. 0 converts every layer in one + # loop at the boundary, which is the default now that block instances share a + # compiled graph. Nonzero only helps if full FP8 does not fit in memory. + mxfp4_to_fp8_layers_per_iter: int = 0 + + # Ramp order; see MXFP4_TO_FP8_ORDERS. Irrelevant to the single-boundary switch. + mxfp4_to_fp8_order: str = "deep_to_shallow" + # Sensitive layer configuration (clean naming, maps to Megatron internals) sensitive_layers_enabled: bool = False sensitive_layers_start: int = 0 @@ -116,6 +137,29 @@ def __post_init__(self): self.num_layers_at_start_in_bf16 = self.sensitive_layers_start self.num_layers_at_end_in_bf16 = self.sensitive_layers_end + # Validated before super() so a misconfigured switch fails on its own terms. + # The switch never assigns self.fp8 -- it sets the FP8 dtypes straight onto the + # module -- so Megatron's "fp4 and fp8 cannot coexist" check never sees it, and + # this is the only place the fp4 cross-check can happen. + if self.mxfp4_to_fp8_switch_iter < 0: + raise ValueError(f"mxfp4_to_fp8_switch_iter must be >= 0, got {self.mxfp4_to_fp8_switch_iter}.") + if self.mxfp4_to_fp8_switch_iter > 0 and not getattr(self, "fp4", None): + raise ValueError( + f"mxfp4_to_fp8_switch_iter={self.mxfp4_to_fp8_switch_iter} requires fp4 to be " + "set (e.g. fp4: mxfp4). The switch flips MXFP4 linears to FP8, so with no " + "MXFP4 linears there is nothing to switch." + ) + if self.mxfp4_to_fp8_layers_per_iter < 0: + raise ValueError( + "mxfp4_to_fp8_layers_per_iter must be >= 0 (0 = convert every layer at the " + f"boundary), got {self.mxfp4_to_fp8_layers_per_iter}." + ) + if self.mxfp4_to_fp8_order not in MXFP4_TO_FP8_ORDERS: + raise ValueError( + f"Unknown mxfp4_to_fp8_order '{self.mxfp4_to_fp8_order}'. " + f"Choose from: {list(MXFP4_TO_FP8_ORDERS)}." + ) + if self.sensitive_layers_enabled and self.sensitive_layer_precision == "tw_fp8": _deferred_fp8 = "e4m3" if self.fp8 is None else None _deferred_fp8_recipe = ( diff --git a/primus/backends/megatron/core/models/diffusion/flux/attention.py b/primus/backends/megatron/core/models/diffusion/flux/attention.py index b444b527c..8e279bf0d 100644 --- a/primus/backends/megatron/core/models/diffusion/flux/attention.py +++ b/primus/backends/megatron/core/models/diffusion/flux/attention.py @@ -292,12 +292,15 @@ def get_query_key_value_tensors( # Split into Q, K, V query, key, value = self._split_qkv(mixed_qkv) - # Apply optional Q/K normalization + # Apply optional Q/K normalization. The cast back to value's dtype matters: + # the norm can return a wider dtype than it was given, and every dense + # flash-attention backend refuses a mixed-precision (q, k, v) triple rather + # than casting for us. Reference Flux does the same cast inside QKNorm. if self.q_layernorm is not None: - query = self.q_layernorm(query) + query = self.q_layernorm(query).to(value.dtype) if self.k_layernorm is not None: - key = self.k_layernorm(key) + key = self.k_layernorm(key).to(value.dtype) return query, key, value @@ -322,10 +325,10 @@ def get_added_query_key_value_tensors( # Apply optional Q/K normalization if self.added_q_layernorm is not None: - query = self.added_q_layernorm(query) + query = self.added_q_layernorm(query).to(value.dtype) if self.added_k_layernorm is not None: - key = self.added_k_layernorm(key) + key = self.added_k_layernorm(key).to(value.dtype) return query, key, value @@ -544,6 +547,10 @@ def forward( # Get Q, K, V query, key, value = self.get_query_key_value_tensors(hidden_states, key_value_states) + # The inherited Q/K norm can widen q/k past value's dtype, and every dense + # flash-attention backend refuses a mixed-precision (q, k, v) triple rather + # than casting for us. Reference Flux does the same cast inside QKNorm. + query, key = query.to(value.dtype), key.to(value.dtype) # Adjust for inference query, key, value, rotary_pos_emb, attn_mask_type, *_ = self._adjust_key_value_for_inference( diff --git a/primus/backends/megatron/flux_pretrain_trainer.py b/primus/backends/megatron/flux_pretrain_trainer.py index 602a21d85..daedcfc57 100644 --- a/primus/backends/megatron/flux_pretrain_trainer.py +++ b/primus/backends/megatron/flux_pretrain_trainer.py @@ -530,6 +530,10 @@ def _build_flux_config_from_yaml(self): "fp4_recipe": fp4_recipe, "mxfp4_backward_precision": getattr(params, "mxfp4_backward_precision", "mxfp4"), "fp4_use_native_te_autocast": getattr(params, "fp4_use_native_te_autocast", False), + "mxfp4_to_fp8_switch_iter": getattr(params, "mxfp4_to_fp8_switch_iter", 0), + "mxfp4_to_fp8_prewarm": getattr(params, "mxfp4_to_fp8_prewarm", True), + "mxfp4_to_fp8_layers_per_iter": getattr(params, "mxfp4_to_fp8_layers_per_iter", 0), + "mxfp4_to_fp8_order": getattr(params, "mxfp4_to_fp8_order", "deep_to_shallow"), } ) @@ -728,6 +732,10 @@ def _log_flux_config(self, config, args): "fp4_recipe", "mxfp4_backward_precision", "mxfp4_gradient_stochastic_rounding", + "mxfp4_to_fp8_switch_iter", + "mxfp4_to_fp8_prewarm", + "mxfp4_to_fp8_layers_per_iter", + "mxfp4_to_fp8_order", "sensitive_layers_enabled", "sensitive_layers_start", "sensitive_layers_end", diff --git a/primus/backends/megatron/patches/mlperf_warmup_patches.py b/primus/backends/megatron/patches/mlperf_warmup_patches.py index 5ac532978..4e6e79f3e 100644 --- a/primus/backends/megatron/patches/mlperf_warmup_patches.py +++ b/primus/backends/megatron/patches/mlperf_warmup_patches.py @@ -327,6 +327,36 @@ def _hooked_train_step( ) _log(f"Completed {warmup_steps} warmup steps") + # ---- 4b. Pre-warm the FP8 graph for the MXFP4 -> FP8 switch ---- + # Hosted here because this is the only place that runs a real grad-enabled + # step at the production micro_batch_size. automatic_dynamic_shapes is on, + # so pre-warming anywhere else (a small dummy batch) would guard on + # different shapes, mark dims dynamic, and change the graph for everyone. + # The optimizer is already neutered and parameters are restored below, so + # the extra step cannot perturb the measured run. + if int(getattr(primus_args, "mxfp4_to_fp8_switch_iter", 0) or 0) > 0 and getattr( + primus_args, "mxfp4_to_fp8_prewarm", True + ): + from primus.backends.megatron.patches.mxfp4_to_fp8_switch_patches import ( + prewarm_fp8_graphs, + ) + + _log("Pre-warming the FP8 graph for the MXFP4 -> FP8 switch") + prewarm_fp8_graphs( + models, + lambda: _wrapped_chain( + forward_step_func, + synthetic_iter, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=iteration, + ), + order=getattr(primus_args, "mxfp4_to_fp8_order", "deep_to_shallow"), + ) + # ---- 5. Restore optimizer ---- _restore_optimizer(optimizer, saved_opt) _reset_optimizer_state(optimizer) diff --git a/primus/backends/megatron/patches/mxfp4_to_fp8_switch_patches.py b/primus/backends/megatron/patches/mxfp4_to_fp8_switch_patches.py new file mode 100644 index 000000000..541cc9958 --- /dev/null +++ b/primus/backends/megatron/patches/mxfp4_to_fp8_switch_patches.py @@ -0,0 +1,389 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Runtime MXFP4 -> FP8 precision switch for Flux (Megatron local spec). + +WHAT IT DOES +------------ +At a preset iteration every MXFP4 linear flips its ``_fp8_mode`` flag and starts +dispatching to the dynamic tensorwise FP8 autograd Function instead of the MXFP4 +one. Neither precision stores a quantized weight -- both keep a plain BF16 +``nn.Parameter`` and quantize inside forward -- so the flip needs no weight +conversion, no optimizer state remap, and no checkpoint. That last point is the +reason this exists: MLPerf forbids a checkpoint-mediated transition. + +WHY ONE BOUNDARY AND NOT A RAMP +------------------------------- +The two Flux block classes are instantiated 19 and 38 times, but Dynamo keys its +code cache on the *code object*, so with ``inline_inbuilt_nn_modules`` the +instances share compiled graphs. Flipping the whole model therefore costs on the +order of two recompiles, not 57, and pre-warming those two during warmup removes +even that. ``mxfp4_to_fp8_layers_per_iter`` keeps the layer-at-a-time ramp +available for the case where full FP8 does not fit in memory. + +THE CONVERTED SET IS A PURE FUNCTION OF THE ITERATION +----------------------------------------------------- +Never a call counter. Three things depend on this: + +1. ``mlperf_warmup`` (priority 95) wraps this patch (priority 46) and re-enters + the inner chain ``warmup_steps + 1`` times with the *same* iteration. A + counter-driven switch would fire during warmup. +2. Rank agreement. Every rank must flip the identical set at the identical + iteration. If the decision took input from loss, grad norms or memory, ranks + would desynchronize and the next collective would *hang* rather than fail. + Peak memory is logged here but must never feed the decision. +3. Resumability, with no switch progress to persist anywhere. +""" + +from __future__ import annotations + +import re +from typing import Callable, Dict, List, Optional, Sequence, Tuple + +import torch + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + +_PATCH_KEY = "megatron.mxfp4.to_fp8_switch" + +# Matches the owning transformer layer index in a module's qualified name. Left +# unanchored on purpose so it survives the DDP / Float16Module name prefixes +# without any unwrapping. +_LAYER_RE = re.compile(r"transformer\.layers\.(\d+)\.") + +LayerPlan = List[Tuple[int, List[torch.nn.Module]]] + + +def _log(msg: str) -> None: + log_rank_0(f"[Patch:mxfp4_to_fp8_switch] {msg}") + + +def _needs_mxfp4_to_fp8_switch(ctx: PatchContext) -> bool: + args = get_args(ctx) + if args is None: + return False + return int(getattr(args, "mxfp4_to_fp8_switch_iter", 0) or 0) > 0 + + +def _assert_fp8_backend_available() -> None: + """Reject a run that intends to switch on a Primus-Turbo without FlyDSL. + + ``_init_fp8_switch_state`` resolves the backend leniently so a stack without + FLYDSL can still build pure-MXFP4 models. This is the callsite that knows the + run actually means to switch, so this is where a missing backend is fatal + rather than a silent downgrade to whatever the dispatcher picks. + """ + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + _SWITCH_FP8_BACKEND, + ) + + if _SWITCH_FP8_BACKEND is None: + raise RuntimeError( + "mxfp4_to_fp8_switch_iter > 0 requires BackendType.FLYDSL in Primus-Turbo, " + "which this build does not expose. Upgrade Primus-Turbo, or unset " + "mxfp4_to_fp8_switch_iter to keep the run in MXFP4." + ) + + # A GEMM backend pinned *in code* is stored as a plain dict that + # get_gemm_backend indexes directly, so a dict covering only FP4 -- exactly what + # set_gemm_backend(AITER, PrecisionType.FP4) builds, and what the MXFP4 + # preshuffle error recommends -- raises KeyError on the first FP8 GEMM, deep + # inside a custom op long after the switch fired. Probing it here turns that + # into a startup failure naming the fix. The env-var path is read with .get and + # degrades to the module's default backend, so it needs no such probe. + from primus_turbo.pytorch.core.backend import GlobalBackendManager, PrecisionType + + try: + GlobalBackendManager.get_gemm_backend(PrecisionType.FP8) + except KeyError as exc: + raise RuntimeError( + "mxfp4_to_fp8_switch_iter > 0 but the in-code Primus-Turbo GEMM backend " + "map has no FP8 entry, so the post-switch FP8 GEMMs would fail with " + f"KeyError({exc}). Pin both precisions -- " + "GlobalBackendManager.set_gemm_backend(BackendType.FLYDSL, PrecisionType.FP8) " + "alongside the FP4 pin, or drop the in-code pin and export " + "PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER,FP8:FLYDSL instead." + ) from exc + + +def build_layer_plan( + models: Sequence[torch.nn.Module], + order: str = "deep_to_shallow", +) -> Tuple[LayerPlan, List[torch.nn.Module]]: + """Group every MXFP4 linear by its owning transformer layer index. + + Derived from the live module tree rather than from ``num_layers``: the 19 + 38 + figure is only a ``FluxConfig`` default, and it is wrong whenever + ``sensitive_layers_enabled`` is set, since those layers are built as + ``Float8*ParallelLinear`` or plain linears and have no ``_fp8_mode`` to flip. + + Returns the ordered ``(layer_index, linears)`` plan plus any MXFP4 linears that + sit outside the layer stack. For Flux the second list is empty -- every MXFP4 + linear comes from the block specs -- but it is collected and converted anyway + so an unexpected one cannot leave the model in a silently mixed end state. + """ + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4ColumnParallelLinear, + MXFP4RowParallelLinear, + ) + + mxfp4_types = (MXFP4ColumnParallelLinear, MXFP4RowParallelLinear) + by_layer: Dict[int, List[torch.nn.Module]] = {} + extras: List[torch.nn.Module] = [] + + for chunk in models: + for name, module in chunk.named_modules(): + if not isinstance(module, mxfp4_types): + continue + match = _LAYER_RE.search(name) + if match is None: + extras.append(module) + else: + by_layer.setdefault(int(match.group(1)), []).append(module) + + indices = sorted(by_layer) + if order == "deep_to_shallow": + indices.reverse() + + return [(idx, by_layer[idx]) for idx in indices], extras + + +def _linear_class_names(models: Sequence[torch.nn.Module]) -> set: + """Class names of every linear-like module, for the empty-plan error message.""" + names = set() + for chunk in models: + for module in chunk.modules(): + cls = type(module).__name__ + if "Linear" in cls: + names.add(cls) + return names + + +def target_layer_count( + iteration: Optional[int], + switch_iter: int, + layers_per_iter: int, + plan_len: int, +) -> int: + """How many layers must be in FP8 at ``iteration``. Pure; no side effects. + + ``layers_per_iter <= 0`` converts everything at the boundary. A rate at or + above the plan length degenerates to the same thing, so one code path covers + both the single-boundary switch and the ramp fallback. + """ + if iteration is None or iteration < switch_iter: + return 0 + if layers_per_iter <= 0: + return plan_len + steps_taken = iteration - switch_iter + 1 + return min(plan_len, steps_taken * layers_per_iter) + + +def set_fp8_mode(plan: LayerPlan, extras: Sequence[torch.nn.Module], count: int) -> None: + """Put the first ``count`` planned layers into FP8 and the rest in MXFP4. + + Written as an absolute assignment rather than an incremental flip so that + calling it repeatedly with the same ``count`` is a genuine no-op -- which is + what makes warmup's re-entry harmless. + """ + for position, (_, linears) in enumerate(plan): + fp8 = position < count + for linear in linears: + linear._fp8_mode = fp8 + + # Non-layer linears ride along with the first flip: once any conversion has + # started they are FP8, so the end state is never partially converted. + for linear in extras: + linear._fp8_mode = count > 0 + + +def _unique_graph_count() -> int: + """Dynamo's compiled-graph counter, or -1 when it cannot be read.""" + try: + from torch._dynamo.utils import counters + + return int(counters["stats"].get("unique_graphs", 0)) + except Exception: # pragma: no cover - torch internal moved + return -1 + + +def prewarm_fp8_graphs( + models: Sequence[torch.nn.Module], + run_step: Callable[[], None], + order: str = "deep_to_shallow", +) -> None: + """Trace the FP8 arm once at production shapes, then restore MXFP4. + + This is what makes the switch seamless, and it is also the only runtime check + that the switch is not a silent no-op. ``run_step`` must be a real + grad-enabled training step at the production ``micro_batch_size``: + + - Grad enabled, so AOTAutograd compiles the backward partition too. A + ``no_grad`` pre-warm would leave the backward to compile at the boundary. + - Production shapes, because ``automatic_dynamic_shapes`` is on -- pre-warming + at a different batch size would mark dims dynamic and change the graph for + everyone. + + The check is structural: ``unique_graphs`` must increase, proving Dynamo + guarded ``_fp8_mode`` and traced a distinct graph rather than reusing the + MXFP4 entry. It deliberately does *not* check numerics; there is no eager FP8 + reference mid-warmup, and a second forward at production shapes would be + expensive. Numerical equivalence to the production Float8 path is covered by + the unit tests instead, which is why those are the load-bearing ones. + """ + plan, extras = build_layer_plan(models, order) + if not plan and not extras: + _log("Pre-warm skipped: no MXFP4 linears found.") + return + + before = _unique_graph_count() + set_fp8_mode(plan, extras, len(plan)) + try: + run_step() + finally: + set_fp8_mode(plan, extras, 0) + + after = _unique_graph_count() + if before < 0 or after < 0: + _log("WARNING: could not read Dynamo's unique_graphs counter; pre-warm unverified.") + return + if after <= before: + raise RuntimeError( + "MXFP4 -> FP8 pre-warm traced no new graph " + f"(unique_graphs stayed at {after}). Dynamo did not guard _fp8_mode, so the " + "switch would be a silent no-op: the run would log a successful switch while " + "still training in MXFP4. Make the dispatch Dynamo-observable before using " + "this config." + ) + _log(f"Pre-warm traced the FP8 graph (unique_graphs {before} -> {after}); restored MXFP4.") + + +def _report_switch(args, iteration: Optional[int], converted: int, total: int) -> None: + """Log the conversion. Peak memory is reported, never fed back into the decision.""" + peak_gib = torch.cuda.max_memory_allocated() / (1024**3) if torch.cuda.is_available() else 0.0 + _log(f"iteration={iteration}: {converted}/{total} layers now FP8; peak allocated {peak_gib:.2f} GiB") + + # The MLPerf logger only exists when mlperf_mode is set, and the MXFP4 configs + # this derives from do not set it. An unguarded emission would crash the run at + # the one iteration that must not fail. + if not getattr(args, "mlperf_mode", False): + return + try: + from mlperf_logging import mllog + + mllog.get_mllogger().event( + key="mxfp4_to_fp8_switch", + value={"iteration": iteration, "layers_converted": converted, "layers_total": total}, + ) + except Exception as exc: # pragma: no cover - logger not initialised + _log(f"mllog emission skipped ({exc}).") + + +@register_patch( + "megatron.mxfp4.to_fp8_switch", + backend="megatron", + phase="before_train", + description=( + "Flip MXFP4 linears to dynamic tensorwise FP8 at a preset iteration, " + "on pre-warmed torch.compile graphs." + ), + priority=46, + condition=_needs_mxfp4_to_fp8_switch, +) +def patch_mxfp4_to_fp8_switch(ctx: PatchContext) -> None: + """Wrap ``train_step`` so the switch lands on an iteration boundary. + + Priority 46 sits after the delayed-scaling preamble (40) and the FSDP2 FP8 + cache refresh (45), and before ``empty_cache_interval`` (50). It is *inside* + ``mlperf_warmup`` (95), so warmup does re-enter it; the pure-function form + above is what makes that harmless, which makes the placement a convenience + rather than a correctness argument. + """ + import megatron.training.training as megatron_training + + from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched + + if is_patched(megatron_training, _PATCH_KEY): + _log("Already applied; skipping re-wrap.") + return + + args = get_args(ctx) + switch_iter = int(getattr(args, "mxfp4_to_fp8_switch_iter", 0) or 0) + layers_per_iter = int(getattr(args, "mxfp4_to_fp8_layers_per_iter", 0) or 0) + order = getattr(args, "mxfp4_to_fp8_order", "deep_to_shallow") + + _assert_fp8_backend_available() + + original_train_step = megatron_training.train_step + state: Dict[str, object] = {"plan": None, "extras": None, "applied": 0} + + # *args/**kwargs passthrough rather than the explicit eight-arg signature the + # delayed-FP8 patch uses. Restating the signature means *synthesizing* an + # `iteration=` on the way down, which breaks against train_step_seq_split + # (parallelism/train_step_patches.py, priority 40) whose wrapper takes seven + # positionals and no iteration at all. Forwarding exactly what arrived adds no + # such coupling. + def _patched_train_step(*fn_args, **fn_kwargs): + model = fn_args[2] if len(fn_args) > 2 else fn_kwargs.get("model") + + if state["plan"] is None and model is not None: + models = model if isinstance(model, (list, tuple)) else [model] + plan, extras = build_layer_plan(models, order) + if not plan and not extras: + # An empty plan is the one failure this patch cannot survive: every + # later step is a no-op, so the run would train to completion in + # MXFP4 while reporting a successful switch. Name what was actually + # found, since the usual cause is a model built from some other + # spec provider than PrimusTurboMXFP4LocalSpecProvider. + raise RuntimeError( + "mxfp4_to_fp8_switch_iter > 0 but no MXFP4 linear was found in " + f"the model, so the switch would be a no-op. Linear-like modules " + f"present: {sorted(_linear_class_names(models))}. Check that fp4 " + "is set and that transformer_impl selects the MXFP4 local spec." + ) + state["plan"] = plan + state["extras"] = extras + mode = "all at once" if layers_per_iter <= 0 else f"{layers_per_iter}/iter" + _log( + f"Planned {len(plan)} MXFP4 layers ({len(extras)} unindexed linears); " + f"switch at iteration {switch_iter}, {mode}, order={order}" + ) + + plan = state["plan"] or [] + extras = state["extras"] or [] + + # The iteration kwarg is authoritative. args.iteration only moves on + # checkpoint load/save, so it must not be used as the gate; curr_iteration + # is the live value Megatron sets just before calling train_step. + iteration = fn_kwargs.get("iteration") + if iteration is None: + from megatron.training import get_args as megatron_get_args + + iteration = getattr(megatron_get_args(), "curr_iteration", None) + + target = target_layer_count(iteration, switch_iter, layers_per_iter, len(plan)) + if target != state["applied"]: + set_fp8_mode(plan, extras, target) + state["applied"] = target + + # The caching allocator is holding blocks sized for FP4 tensors while + # the next forward asks for differently sized FP8 ones, so it may + # cudaMalloc fresh segments with unusable cached ones sitting idle. + # This is the only real transient at the switch -- no saved activations + # are live here, the pool is empty between the optimizer step and the + # next forward. Deliberately bypasses the empty_cache_interval throttle, + # which exists because empty_cache is expensive; at one call per switch + # the cost is irrelevant. + if torch.cuda.is_available(): + torch.cuda.empty_cache() + _report_switch(args, iteration, target, len(plan)) + + return original_train_step(*fn_args, **fn_kwargs) + + megatron_training.train_step = _patched_train_step + mark_patched(megatron_training, _PATCH_KEY) + _log(f"Wrapped train_step (switch_iter={switch_iter}, priority=46)") diff --git a/runner/helpers/envs/base_env.sh b/runner/helpers/envs/base_env.sh index 2060dac02..c15ca4759 100755 --- a/runner/helpers/envs/base_env.sh +++ b/runner/helpers/envs/base_env.sh @@ -126,8 +126,12 @@ log_exported_vars "Python Path and Data Paths" \ # NCCL and Network Configuration # ============================================================================= -# Set visible GPUs for the current node (0 to GPUS_PER_NODE-1) -HIP_VISIBLE_DEVICES=$(seq -s, 0 $((GPUS_PER_NODE - 1))) +# Set visible GPUs for the current node (0 to GPUS_PER_NODE-1). A value supplied +# by the caller wins, so a specific device set or ordering can be pinned from +# outside -- needed to tell a rank-specific software fault apart from a faulty GPU. +if [ -z "${HIP_VISIBLE_DEVICES:-}" ]; then + HIP_VISIBLE_DEVICES=$(seq -s, 0 $((GPUS_PER_NODE - 1))) +fi export HIP_VISIBLE_DEVICES # Keep ROCm libraries ahead of any system-provided HSA runtime. diff --git a/runner/primus-cli-direct.sh b/runner/primus-cli-direct.sh index c90a0d86d..99ebf6e42 100755 --- a/runner/primus-cli-direct.sh +++ b/runner/primus-cli-direct.sh @@ -698,8 +698,15 @@ elif [[ "$RUN_MODE" == "torchrun" ]]; then FILTERS+=($((${GPUS_PER_NODE:-8} - 1))) fi - # Build filter argument (only if FILTERS is non-empty) - if [ "${#FILTERS[@]}" -gt 0 ]; then + # Build filter argument (only if FILTERS is non-empty). + # PRIMUS_LOCAL_RANKS_FILTER overrides the default rank-0-only console filter, + # which otherwise discards the stderr of a crashing non-zero rank. Set it to + # "all" to keep every rank's output. + if [ "${PRIMUS_LOCAL_RANKS_FILTER:-}" = "all" ]; then + FILTER_ARG=() + elif [ -n "${PRIMUS_LOCAL_RANKS_FILTER:-}" ]; then + FILTER_ARG=(--local-ranks-filter "$PRIMUS_LOCAL_RANKS_FILTER") + elif [ "${#FILTERS[@]}" -gt 0 ]; then LOCAL_FILTER=$(IFS=,; echo "${FILTERS[*]}") FILTER_ARG=(--local-ranks-filter "$LOCAL_FILTER") else diff --git a/tests/unit_tests/backends/megatron/diffusion/test_mxfp4_to_fp8_switch.py b/tests/unit_tests/backends/megatron/diffusion/test_mxfp4_to_fp8_switch.py new file mode 100644 index 000000000..bd6503d17 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_mxfp4_to_fp8_switch.py @@ -0,0 +1,771 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for the MXFP4 -> FP8 runtime precision switch. + +Two tiers: + +- Schedule and plan mechanics, which need neither a GPU nor Primus-Turbo. These + cover the invariant the whole design rests on: the converted set is a pure + function of the iteration, so ``mlperf_warmup``'s re-entry is harmless and every + rank flips identically (a memory- or loss-driven decision would desynchronize + ranks and *hang* the next collective rather than fail). +- Numerics and tracing, gated on MXFP4 hardware. The equivalence test is the + load-bearing one: the runtime pre-warm check is structural only (it proves a + distinct graph was traced, not that it computes the right thing), so proving the + FP8 arm lands exactly on the production Float8 path is this file's job. +""" + +import functools + +import pytest +import torch + +from tests.unit_tests.backends.megatron.conftest import requires_mxfp4 + + +def _init_method(): + return functools.partial(torch.nn.init.xavier_uniform_) + + +class _FakeLinear(torch.nn.Module): + """Stand-in for an MXFP4 linear: the schedule only touches ``_fp8_mode``.""" + + def __init__(self): + super().__init__() + self._fp8_mode = False + + +def _fake_plan(num_layers, per_layer=2): + return [(idx, [_FakeLinear() for _ in range(per_layer)]) for idx in range(num_layers)] + + +def _modes(plan): + return [all(lin._fp8_mode for lin in linears) for _, linears in plan] + + +# --------------------------------------------------------------------------- +# Schedule: a pure function of the iteration +# --------------------------------------------------------------------------- + + +class TestTargetLayerCount: + """The switch schedule, which must depend on nothing but the iteration.""" + + def test_no_conversion_before_switch_iter(self): + from primus.backends.megatron.patches.mxfp4_to_fp8_switch_patches import ( + target_layer_count, + ) + + for iteration in (0, 1, 599): + assert target_layer_count(iteration, 600, 0, 57) == 0 + + def test_single_boundary_converts_everything_at_once(self): + from primus.backends.megatron.patches.mxfp4_to_fp8_switch_patches import ( + target_layer_count, + ) + + assert target_layer_count(600, 600, 0, 57) == 57 + assert target_layer_count(999, 600, 0, 57) == 57 + + def test_missing_iteration_converts_nothing(self): + """A None iteration must not be read as 0 and trip the switch.""" + from primus.backends.megatron.patches.mxfp4_to_fp8_switch_patches import ( + target_layer_count, + ) + + assert target_layer_count(None, 0, 0, 57) == 0 + + def test_ramp_grows_by_rate_and_clamps(self): + from primus.backends.megatron.patches.mxfp4_to_fp8_switch_patches import ( + target_layer_count, + ) + + assert target_layer_count(600, 600, 4, 57) == 4 + assert target_layer_count(601, 600, 4, 57) == 8 + assert target_layer_count(613, 600, 4, 57) == 56 + # Clamped, not run past the end. + assert target_layer_count(614, 600, 4, 57) == 57 + assert target_layer_count(9999, 600, 4, 57) == 57 + + def test_rate_at_or_above_plan_length_degenerates_to_single_boundary(self): + """One code path has to cover both schedules.""" + from primus.backends.megatron.patches.mxfp4_to_fp8_switch_patches import ( + target_layer_count, + ) + + assert target_layer_count(600, 600, 57, 57) == 57 + assert target_layer_count(600, 600, 999, 57) == 57 + + def test_idempotent_across_repeated_calls(self): + """mlperf_warmup re-enters the inner chain with the *same* iteration. + + A call-counter-driven schedule would advance the ramp during warmup; this + is the regression guard for that. + """ + from primus.backends.megatron.patches.mxfp4_to_fp8_switch_patches import ( + target_layer_count, + ) + + counts = {target_layer_count(600, 600, 4, 57) for _ in range(10)} + assert counts == {4} + + def test_order_independent(self): + """Evaluating iterations out of order must not change the answer.""" + from primus.backends.megatron.patches.mxfp4_to_fp8_switch_patches import ( + target_layer_count, + ) + + forward = [target_layer_count(i, 600, 4, 57) for i in range(598, 615)] + backward = [target_layer_count(i, 600, 4, 57) for i in reversed(range(598, 615))] + assert forward == list(reversed(backward)) + + +class TestSetFp8Mode: + """Absolute assignment, so repeated application is a genuine no-op.""" + + def test_converts_prefix_of_plan(self): + from primus.backends.megatron.patches.mxfp4_to_fp8_switch_patches import ( + set_fp8_mode, + ) + + plan = _fake_plan(5) + set_fp8_mode(plan, [], 2) + assert _modes(plan) == [True, True, False, False, False] + + def test_repeated_application_is_a_no_op(self): + from primus.backends.megatron.patches.mxfp4_to_fp8_switch_patches import ( + set_fp8_mode, + ) + + plan = _fake_plan(5) + for _ in range(3): + set_fp8_mode(plan, [], 3) + assert _modes(plan) == [True, True, True, False, False] + + def test_reset_to_zero_restores_mxfp4(self): + """What the pre-warm relies on to leave no trace.""" + from primus.backends.megatron.patches.mxfp4_to_fp8_switch_patches import ( + set_fp8_mode, + ) + + plan = _fake_plan(4) + extras = [_FakeLinear()] + set_fp8_mode(plan, extras, 4) + set_fp8_mode(plan, extras, 0) + assert _modes(plan) == [False] * 4 + assert extras[0]._fp8_mode is False + + def test_extras_convert_with_the_first_flip(self): + """Linears outside the layer stack must never be left behind.""" + from primus.backends.megatron.patches.mxfp4_to_fp8_switch_patches import ( + set_fp8_mode, + ) + + plan = _fake_plan(4) + extras = [_FakeLinear(), _FakeLinear()] + set_fp8_mode(plan, extras, 0) + assert all(not lin._fp8_mode for lin in extras) + set_fp8_mode(plan, extras, 1) + assert all(lin._fp8_mode for lin in extras) + + +class TestEmptyPlanDetection: + """A model with no MXFP4 linears has to be caught, not planned around. + + This is not hypothetical: the Flux spec selection swallows ImportError when + resolving the MXFP4 provider and silently falls back to BF16 linears, which is + what happens whenever the installed Primus-Turbo is out of step with the + attention API this repo expects. Without a check the run trains to completion + in the wrong precision while reporting a successful switch. + """ + + def test_plan_is_empty_for_a_model_without_mxfp4_linears(self): + from primus.backends.megatron.patches.mxfp4_to_fp8_switch_patches import ( + build_layer_plan, + ) + + model = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8)) + plan, extras = build_layer_plan([model]) + assert plan == [] and extras == [] + + def test_linear_class_names_reports_what_was_found(self): + from primus.backends.megatron.patches.mxfp4_to_fp8_switch_patches import ( + _linear_class_names, + ) + + model = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.ReLU()) + assert _linear_class_names([model]) == {"Linear"} + + +# --------------------------------------------------------------------------- +# Config validation +# --------------------------------------------------------------------------- + + +def _diffusion_config(**overrides): + from primus.backends.megatron.core.models.diffusion.common.config import ( + BaseDiffusionConfig, + ) + + defaults = dict(hidden_size=256, num_attention_heads=8, num_layers=2) + defaults.update(overrides) + return BaseDiffusionConfig(**defaults) + + +class TestConfigValidation: + """Validated before super().__post_init__() so it fails on its own terms.""" + + def test_default_is_disabled(self): + config = _diffusion_config() + assert config.mxfp4_to_fp8_switch_iter == 0 + assert config.mxfp4_to_fp8_prewarm is True + assert config.mxfp4_to_fp8_layers_per_iter == 0 + assert config.mxfp4_to_fp8_order == "deep_to_shallow" + + def test_switch_without_fp4_is_rejected(self): + """Nothing to switch without MXFP4 linears, so say so rather than no-op.""" + with pytest.raises(ValueError, match="requires fp4"): + _diffusion_config(mxfp4_to_fp8_switch_iter=600) + + def test_switch_with_fp4_is_accepted(self): + config = _diffusion_config(fp4="mxfp4", mxfp4_to_fp8_switch_iter=600) + assert config.mxfp4_to_fp8_switch_iter == 600 + # The switch must never set config.fp8: Megatron rejects fp4 and fp8 + # together, and the FP8 dtypes are set straight onto the module instead. + assert not config.fp8 + + def test_negative_switch_iter_is_rejected(self): + with pytest.raises(ValueError, match="must be >= 0"): + _diffusion_config(mxfp4_to_fp8_switch_iter=-1) + + def test_negative_rate_is_rejected(self): + with pytest.raises(ValueError, match="layers_per_iter must be >= 0"): + _diffusion_config(fp4="mxfp4", mxfp4_to_fp8_layers_per_iter=-1) + + def test_unknown_order_is_rejected(self): + with pytest.raises(ValueError, match="Unknown mxfp4_to_fp8_order"): + _diffusion_config(fp4="mxfp4", mxfp4_to_fp8_order="sideways") + + +# --------------------------------------------------------------------------- +# Numerics and tracing (MXFP4 hardware) +# --------------------------------------------------------------------------- + +# Token count of the test inputs. Not arbitrary: the weight-gradient GEMM reduces +# over tokens, so this is its K. FlyDSL tensorwise FP8 requires K > 128 (its +# software pipeline needs at least two K tiles) and MX-blockwise requires +# K % 128 == 0 and K >= 256, so anything smaller makes the backend refuse the +# inputs rather than exercise the path. Production token counts are far larger, +# so this only keeps the test off a cliff production never approaches. +_TOKENS = 256 + + +def _mxfp4_config(**overrides): + """Build a TransformerConfig for MXFP4 linears. + + The mxfp4_* knobs are Flux-side fields that the linears read off the config + with ``getattr``, so they are not TransformerConfig constructor arguments. + Anything not a declared field is attached afterwards, which is how the real + FluxConfig presents them. + """ + import dataclasses + + from megatron.core.transformer.transformer_config import TransformerConfig + + defaults = dict( + hidden_size=256, + num_attention_heads=8, + num_layers=1, + params_dtype=torch.bfloat16, + fp4="mxfp4", + fp4_recipe="mxfp4", + ) + declared = {f.name for f in dataclasses.fields(TransformerConfig)} + extras = {k: v for k, v in overrides.items() if k not in declared} + defaults.update({k: v for k, v in overrides.items() if k in declared}) + + config = TransformerConfig(**defaults) + for key, value in extras.items(): + setattr(config, key, value) + return config + + +def _fp8_config(**overrides): + from megatron.core.transformer.transformer_config import TransformerConfig + + defaults = dict( + hidden_size=256, + num_attention_heads=8, + num_layers=1, + params_dtype=torch.bfloat16, + fp8="hybrid", + fp8_recipe="tensorwise", + ) + defaults.update(overrides) + config = TransformerConfig(**defaults) + config.fp8_scaling_strategy = "dynamic" + config.fp8_force_nt_layout = False + return config + + +def _pin_gemm_backends(monkeypatch): + """Pin FP4 to AITER and FP8 to FlyDSL, autotune off -- the production recipe. + + FP4 must be AITER with autotune off or MXFP4 ``__init__`` fails its preshuffle + contract. FP8 is pinned too so the switch's arm runs on the backend production + pins rather than whatever the dispatcher would pick. + + Goes through the public setter instead of writing ``_gemm_backend`` directly, + because the entries are version-dependent: Primus-Turbo used to store a bare + ``BackendType`` per precision and now stores a ``BackendChoice``. Building the + table by hand pins one shape and breaks against the other. ``monkeypatch`` + still records the original attribute, so the setter's writes are undone on + teardown. + """ + import os + + from primus_turbo.pytorch.core.backend import ( + BackendType, + GlobalBackendManager, + PrecisionType, + ) + + if os.environ.get("PRIMUS_TURBO_GEMM_BACKEND", None) == "": + monkeypatch.delenv("PRIMUS_TURBO_GEMM_BACKEND", raising=False) + + monkeypatch.setattr(GlobalBackendManager, "_gemm_backend", None) + monkeypatch.setattr(GlobalBackendManager, "_auto_tune", False, raising=False) + GlobalBackendManager.set_gemm_backend(BackendType.AITER, PrecisionType.FP4) + GlobalBackendManager.set_gemm_backend(BackendType.FLYDSL, PrecisionType.FP8) + + +class TestSwitchState: + """The dormant FP8 state installed on every MXFP4 linear at construction. + + Plain pytest class rather than ``PrimusUT``: the latter is a + ``unittest.TestCase``, into whose test methods pytest cannot inject + function-scoped fixtures or ``parametrize`` arguments. The logger these tests + need comes from the session-scoped autouse fixture in the megatron conftest. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state, monkeypatch): + _pin_gemm_backends(monkeypatch) + + @requires_mxfp4 + def test_switch_state_defaults(self): + from primus_turbo.pytorch.core.low_precision import ( + ScalingGranularity, + float8_e4m3, + float8_e5m2, + ) + + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4ColumnParallelLinear, + ) + + linear = MXFP4ColumnParallelLinear( + input_size=256, + output_size=512, + config=_mxfp4_config(), + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + + assert linear._fp8_mode is False, "the switch must start dormant" + assert linear._switch_fp8_fwd_dtype == float8_e4m3 + assert linear._switch_fp8_bwd_dtype == float8_e5m2 + assert linear._switch_fp8_gran_value == ScalingGranularity.TENSORWISE.value + # FlyDSL handles NT/NN/TN natively, so normalizing to NT would only buy a + # pre-transposed copy of both operands. + assert linear._switch_force_nt is False + + @requires_mxfp4 + def test_flydsl_is_the_default_fp8_backend(self): + from primus_turbo.pytorch.core.backend import BackendType + + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4ColumnParallelLinear, + ) + + linear = MXFP4ColumnParallelLinear( + input_size=256, + output_size=512, + config=_mxfp4_config(), + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + assert linear._switch_fp8_backend_value == BackendType.FLYDSL.value + + @requires_mxfp4 + def test_switch_attributes_do_not_disturb_the_mxfp4_triple(self): + """The pure-MXFP4 path must stay bit-identical, which means its traced + constants must not move. Reusing the backward-precision triple would.""" + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4ColumnParallelLinear, + ) + + linear = MXFP4ColumnParallelLinear( + input_size=256, + output_size=512, + config=_mxfp4_config(mxfp4_backward_precision="mxfp4"), + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + assert linear._fp8_bwd_dtype is None + assert linear._fp8_gran_value == 0 + assert linear._fp8_backend_value == 0 + + +class TestFp8ArmEquivalence: + """The load-bearing test: the FP8 arm must be the production Float8 path. + + The runtime pre-warm check only proves a distinct graph was traced; it cannot + catch an FP8 arm that traces its own graph and computes the wrong thing. This + is what covers that. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state, monkeypatch): + _pin_gemm_backends(monkeypatch) + + @staticmethod + def _build_pair(mxfp4_cls, fp8_cls, input_size=256, output_size=512): + kwargs = dict( + input_size=input_size, + output_size=output_size, + init_method=_init_method(), + bias=False, + skip_bias_add=False, + is_expert=False, + ) + if mxfp4_cls.__name__.startswith("MXFP4Column"): + extra_mx, extra_fp8 = {"gather_output": False}, {"gather_output": False} + else: + extra_mx = extra_fp8 = {"input_is_parallel": False} + + mxfp4_linear = mxfp4_cls(config=_mxfp4_config(), **kwargs, **extra_mx) + fp8_linear = fp8_cls(config=_fp8_config(), **kwargs, **extra_fp8) + + fp8_linear.weight.data.copy_(mxfp4_linear.weight.data) + mxfp4_linear._fp8_mode = True + # Compare the dispatch path, not two GEMM backends: different backends can + # legitimately differ in accumulation order, which would make a bitwise + # claim meaningless. FlyDSL-as-default is asserted separately above. + mxfp4_linear._switch_fp8_backend_value = fp8_linear._fp8_backend_value + return mxfp4_linear, fp8_linear + + @requires_mxfp4 + @pytest.mark.parametrize("variant", ["column", "row"]) + def test_forward_and_backward_match_float8_linear(self, variant): + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + Float8ColumnParallelLinear, + Float8RowParallelLinear, + ) + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4ColumnParallelLinear, + MXFP4RowParallelLinear, + ) + + if variant == "column": + mxfp4_cls, fp8_cls = MXFP4ColumnParallelLinear, Float8ColumnParallelLinear + else: + mxfp4_cls, fp8_cls = MXFP4RowParallelLinear, Float8RowParallelLinear + + mxfp4_linear, fp8_linear = self._build_pair(mxfp4_cls, fp8_cls) + + x = torch.randn(_TOKENS, 256, dtype=torch.bfloat16, device="cuda") + x_mx = x.clone().requires_grad_(True) + x_fp8 = x.clone().requires_grad_(True) + + out_mx = mxfp4_linear(x_mx)[0] + out_fp8 = fp8_linear(x_fp8)[0] + torch.testing.assert_close(out_mx, out_fp8, rtol=0, atol=0) + + grad = torch.randn_like(out_mx) + out_mx.backward(grad) + out_fp8.backward(grad.clone()) + torch.testing.assert_close(x_mx.grad, x_fp8.grad, rtol=0, atol=0) + torch.testing.assert_close(mxfp4_linear.weight.grad, fp8_linear.weight.grad, rtol=0, atol=0) + + @requires_mxfp4 + def test_flipping_the_flag_changes_the_result(self): + """The unit-level version of the pre-warm no-op check.""" + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4ColumnParallelLinear, + ) + + linear = MXFP4ColumnParallelLinear( + input_size=256, + output_size=512, + config=_mxfp4_config(), + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + x = torch.randn(_TOKENS, 256, dtype=torch.bfloat16, device="cuda") + + mxfp4_out = linear(x)[0] + linear._fp8_mode = True + fp8_out = linear(x)[0] + + assert not torch.equal( + mxfp4_out, fp8_out + ), "MXFP4 and FP8 produced identical bits, which means the flag did not change the dispatch" + + +class TestNoGraphBreaks: + """Both arms must trace cleanly; the switch happens under per_block compile.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state, monkeypatch): + _pin_gemm_backends(monkeypatch) + + @requires_mxfp4 + @pytest.mark.parametrize("fp8_mode", [False, True]) + def test_zero_graph_breaks(self, fp8_mode): + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4ColumnParallelLinear, + ) + + linear = MXFP4ColumnParallelLinear( + input_size=256, + output_size=512, + config=_mxfp4_config(), + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + linear._fp8_mode = fp8_mode + + x = torch.randn(_TOKENS, 256, dtype=torch.bfloat16, device="cuda") + explanation = torch._dynamo.explain(lambda inp: linear(inp)[0])(x) + + assert explanation.graph_break_count == 0, ( + f"_fp8_mode={fp8_mode}: expected 0 graph breaks, got " + f"{explanation.graph_break_count}. Reasons: {explanation.break_reasons}" + ) + + @requires_mxfp4 + def test_flipping_the_flag_traces_a_new_graph(self): + """The mechanism the pre-warm assertion depends on. + + If a bool attribute read in an ``if`` is not guarded, the switch is a + silent no-op and pre-warm is the only thing standing between that and a + production run that logs success while training in MXFP4. + """ + from torch._dynamo.utils import counters + + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4ColumnParallelLinear, + ) + + linear = MXFP4ColumnParallelLinear( + input_size=256, + output_size=512, + config=_mxfp4_config(), + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + + torch._dynamo.reset() + counters.clear() + compiled = torch.compile(lambda inp: linear(inp)[0]) + x = torch.randn(_TOKENS, 256, dtype=torch.bfloat16, device="cuda") + + compiled(x) + after_mxfp4 = counters["stats"].get("unique_graphs", 0) + + linear._fp8_mode = True + compiled(x) + after_fp8 = counters["stats"].get("unique_graphs", 0) + + assert after_fp8 > after_mxfp4, ( + "flipping _fp8_mode did not fail a Dynamo guard " + f"(unique_graphs stayed at {after_fp8}); the switch would be a no-op " + "under torch.compile" + ) + + +class _TwoLinearBlock(torch.nn.Module): + """Stand-in for a transformer block: several instances, one code object.""" + + def __init__(self, config, hidden=256, ffn=512): + super().__init__() + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4ColumnParallelLinear, + MXFP4RowParallelLinear, + ) + + self.up = MXFP4ColumnParallelLinear( + input_size=hidden, + output_size=ffn, + config=config, + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + self.down = MXFP4RowParallelLinear( + input_size=ffn, + output_size=hidden, + config=config, + init_method=_init_method(), + bias=False, + skip_bias_add=False, + is_expert=False, + input_is_parallel=False, + ) + + def forward(self, x): + return self.down(self.up(x)[0])[0] + + +class TestGraphSharing: + """Whether the compiled-graph cost of the switch scales with layer count. + + This is what decides the architecture. If every block instance holds its own + cache entry, flipping all of them at one iteration pays N compiles inside a + single step, and the switch has to be spread over many iterations + (``mxfp4_to_fp8_layers_per_iter``). If instances share entries keyed on the + code object, the whole model can flip at one boundary for the price of one + compile, which is the design the default config assumes. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state, monkeypatch): + _pin_gemm_backends(monkeypatch) + + @requires_mxfp4 + def test_instances_share_graphs_across_the_switch(self): + from torch._dynamo.utils import counters + + n_blocks = 4 + config = _mxfp4_config() + blocks = [_TwoLinearBlock(config) for _ in range(n_blocks)] + # One torch.compile per block, as torch_compile_scope=per_block does. Each + # call returns its own wrapper, so any sharing comes from the Dynamo cache + # being keyed on the shared forward code object, not from reusing a wrapper. + compiled = [torch.compile(b) for b in blocks] + + torch._dynamo.reset() + counters.clear() + x = torch.randn(_TOKENS, 256, dtype=torch.bfloat16, device="cuda") + + for c in compiled: + c(x) + mxfp4_graphs = counters["stats"].get("unique_graphs", 0) + + for b in blocks: + b.up._fp8_mode = True + b.down._fp8_mode = True + for c in compiled: + c(x) + total_graphs = counters["stats"].get("unique_graphs", 0) + fp8_graphs = total_graphs - mxfp4_graphs + + # Two code objects per block (the block forward plus the inlined linears + # collapse into one graph each), so the bar is "constant in n_blocks", + # not an exact count. + assert mxfp4_graphs < n_blocks, ( + f"{n_blocks} blocks traced {mxfp4_graphs} MXFP4 graphs; instances are " + "not sharing cache entries, so a single-boundary switch would pay one " + "compile per block. Use mxfp4_to_fp8_layers_per_iter to spread it." + ) + assert fp8_graphs < n_blocks, ( + f"the flip traced {fp8_graphs} new graphs for {n_blocks} blocks; the " + "FP8 arm is being compiled per instance." + ) + assert fp8_graphs > 0, "the flip traced no new graph; the switch is a no-op" + + +class TestSavedActivationMemory: + """The memory cost of the switch, which is the reason it is a decision at all. + + MXFP4 stores its saved operands packed two-to-a-byte plus E8M0 block scales; + tensorwise FP8 stores a byte per element and a scalar scale. The switch + therefore raises the recurring steady-state activation peak, and a run that + fits in MXFP4 can fail after the flip. This pins the direction and bounds the + size of that increase. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state, monkeypatch): + _pin_gemm_backends(monkeypatch) + + @staticmethod + def _saved_bytes_per_element(linear, x): + """Bytes saved for backward, per element of the two GEMM operands. + + Measured as what stays allocated while the autograd graph is alive, minus + the output, so the result isolates the saved operand encoding instead of + being diluted by the bf16 output tensor. + """ + import gc + + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + out = linear(x)[0] + torch.cuda.synchronize() + held = torch.cuda.memory_allocated() - before + saved = held - out.numel() * out.element_size() + del out + return saved / (x.numel() + linear.weight.numel()) + + @requires_mxfp4 + def test_saved_operand_encoding_matches_the_memory_model(self): + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4ColumnParallelLinear, + ) + + linear = MXFP4ColumnParallelLinear( + input_size=256, + output_size=512, + config=_mxfp4_config(), + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + x = torch.randn(_TOKENS, 256, dtype=torch.bfloat16, device="cuda") + + linear._fp8_mode = False + mxfp4_per_elem = self._saved_bytes_per_element(linear, x) + linear._fp8_mode = True + fp8_per_elem = self._saved_bytes_per_element(linear, x) + + # MXFP4: half a byte of packed E2M1 plus one E8M0 byte per 32-element + # block, so 0.5 + 1/32 = 0.53125. FP8 tensorwise: one byte plus two scalar + # scales, so ~1.0. Bounds are tight on purpose -- a saved bf16 copy + # appearing on either arm would move these and silently invalidate the + # activation budget the switch is planned against. + assert 0.50 <= mxfp4_per_elem <= 0.60, ( + f"MXFP4 saved {mxfp4_per_elem:.3f} B/element, expected ~0.53 " + "(packed FP4 plus E8M0 block scales)" + ) + assert 0.95 <= fp8_per_elem <= 1.15, ( + f"FP8 saved {fp8_per_elem:.3f} B/element, expected ~1.0 " + "(one byte per element plus scalar scales)" + )