fix(megatron): preserve diffusion resume state

docs/04-technical-guides/diffusion-models/mxfp4_training.md

@@ -8,7 +8,10 @@ MXFP4 stores activations and weights in 4-bit microscale floating-point with one
 
 - Uses a **local spec** (`PrimusTurboMXFP4LocalSpecProvider`) with **no Transformer Engine dependency**—MXFP4 linear layers are self-contained autograd `Function`s that call Primus-Turbo's `gemm_fp4_impl` directly, so the path is `torch.compile`-friendly with minimal graph breaks.
 - Keeps **attention, optimizer state / main params, and inter-rank communication in BF16**. Only the MMA inputs of the column- and row-parallel linears are quantized.
-- Supports two backward modes via `mxfp4_backward_precision`: pure **MXFP4** (default) or **FP8** hybrid (E5M2 backward with tensorwise scaling on HipBLASLt).
+- Selects forward and backward precision independently. Forward supports
+  **MXFP4** (default), tensorwise **FP8 E4M3**, or **BF16** through
+  `mxfp4_forward_precision`; backward supports MXFP4 or tensorwise FP8 E5M2
+  through `mxfp4_backward_precision`.
 - Dispatches the FP4 GEMM through Primus-Turbo's pluggable backend layer, which can route to either AITER (recommended for MI355X) or HipBLASLt.
 
 ## Table of contents
@@ -67,6 +70,7 @@ The relevant overrides in [`examples/megatron/configs/MI355X/diffusion/flux_12b_
 # MXFP4 precision
 fp4: "mxfp4"
 fp4_recipe: "mxfp4"              # default is "nvfp4" in trainer_base.yaml; must override
+mxfp4_forward_precision: "mxfp4"  # "mxfp4", "fp8", or "bf16"
 mxfp4_backward_precision: "mxfp4"  # "mxfp4" (pure) or "fp8" (hybrid)
 
 # Local spec + Primus-Turbo
@@ -86,6 +90,7 @@ gradient_accumulation_fusion: false
 |------|--------|-------|
 | `fp4` | `"mxfp4"` | Top-level switch to enable FP4. |
 | `fp4_recipe` | `"mxfp4"` for this guide | Default in [`primus/configs/modules/megatron/trainer_base.yaml`](../../../primus/configs/modules/megatron/trainer_base.yaml) is `nvfp4`; the MXFP4 config overrides it. |
+| `mxfp4_forward_precision` | `"mxfp4"`, `"fp8"`, or `"bf16"` | Selects only the forward linear GEMMs. FP8 uses dynamic tensorwise E4M3 on HipBLASLt; BF16 uses the native linear operation. MXFP4 tensors are still prepared for backward when `mxfp4_backward_precision: "mxfp4"`. |
 | `mxfp4_backward_precision` | `"mxfp4"` or `"fp8"` | Exhaustive set (checked by branch in [`primus_turbo_mxfp4_local.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py)). `"fp8"` uses E5M2 with tensorwise HipBLASLt for backward. |
 | `mxfp4_gradient_stochastic_rounding` | `true` / `false` | Optional. Enables SR on FP4 gradient quantization. |
 
@@ -198,7 +203,7 @@ Formal A/B benchmarks vs BF16 and FP8 (delayed and tensorwise) are pending and w
 - MXFP4 spec provider: [`primus/backends/megatron/core/extensions/primus_turbo_local_spec.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_local_spec.py) (`PrimusTurboMXFP4LocalSpecProvider`).
 - MXFP4 linear-layer autograd / fwd-bwd: [`primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py).
 - Config schema defaults: [`primus/configs/modules/megatron/trainer_base.yaml`](../../../primus/configs/modules/megatron/trainer_base.yaml).
-- Dataclass field `mxfp4_backward_precision`: [`primus/backends/megatron/core/models/diffusion/common/config.py`](../../../primus/backends/megatron/core/models/diffusion/common/config.py).
+- Dataclass fields `mxfp4_forward_precision` and `mxfp4_backward_precision`: [`primus/backends/megatron/core/models/diffusion/common/config.py`](../../../primus/backends/megatron/core/models/diffusion/common/config.py).
 - FP4 backend selection (Primus-Turbo): `primus_turbo/common/constants.py`, `primus_turbo/pytorch/core/backend.py`, `primus_turbo/pytorch/kernels/gemm/gemm_fp4_impl.py`.
 - AITER tuned-config loader: `aiter/jit/core.py` (`AITER_CONFIG_GEMM_A4W4`).
 - AITER A4W4 dispatch + hit/miss logging: `aiter/ops/gemm_op_a4w4.py`.

examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml

@@ -141,6 +141,7 @@ modules:
 
       fp4: "mxfp4"
       fp4_recipe: "mxfp4"
+      mxfp4_forward_precision: "mxfp4"  # "mxfp4", "fp8", or "bf16"
       mxfp4_backward_precision: "mxfp4"  # "mxfp4" (pure) or "fp8" (hybrid)
 
       empty_unused_memory_level: 0

primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py

@@ -12,7 +12,8 @@
 Key properties:
 - Uses setup_context pattern with primitive-only args so torch.compile
   can trace through without graph breaks.
-- Two backward modes: pure MXFP4 or hybrid (FP4 fwd / FP8 bwd).
+- Independently selects MXFP4, FP8, or BF16 forward precision and
+  MXFP4 or FP8 backward precision.
 - gemm_fp4_impl and gemm_fp8_impl are torch.library.custom_op with register_fake.
 - Zero TransformerEngine dependencies.
 - Requires tensor_model_parallel_size=1, no GAF, no sequence_parallel.
@@ -229,6 +230,14 @@ def _quantize_mxfp4_dual_backward(ctx, *grad_outputs):
 _FP4_DTYPE = torch.float4_e2m1fn_x2
 _GRAN_VALUE = ScalingGranularity.MX_BLOCKWISE.value
 _DEFAULT_BACKEND = BackendType.HIPBLASLT.value
+_FORWARD_PRECISION_MXFP4 = 0
+_FORWARD_PRECISION_FP8 = 1
+_FORWARD_PRECISION_BF16 = 2
+_FORWARD_PRECISION_VALUES = {
+    "mxfp4": _FORWARD_PRECISION_MXFP4,
+    "fp8": _FORWARD_PRECISION_FP8,
+    "bf16": _FORWARD_PRECISION_BF16,
+}
 
 
 def _quantize_input_dual(input_2d, preshuffle):
@@ -291,9 +300,13 @@ def _quantize_grad_dual(grad_2d, preshuffle, use_sr=True):
 class MXFP4LinearFunction(torch.autograd.Function):
     """MXFP4 linear (Y = X @ W^T) with MX block-of-32 scaling.
 
-    Two modes via backward_is_fp8 bool primitive:
-    - Pure MXFP4: forward + backward both use FP4 quantization + gemm_fp4_impl
-    - Hybrid: forward uses FP4, backward re-quantizes saved BF16 to FP8 tensorwise
+    Forward and backward precision are independent:
+    - Forward: MXFP4, dynamic tensorwise FP8 E4M3, or BF16.
+    - Backward: MXFP4 or dynamic tensorwise FP8 E5M2.
+
+    MXFP4 backward always uses the transposed MXFP4 tensors prepared during
+    forward. This keeps its memory and compute path identical when only the
+    forward GEMM moves to higher precision.
 
     Uses setup_context pattern with primitive-only args for torch.compile.
     """
@@ -308,6 +321,10 @@ def forward(
         fp8_gran_value,
         fp8_backend_value,
         use_gradient_sr,
+        forward_precision_value=_FORWARD_PRECISION_MXFP4,
+        fp8_fwd_dtype=None,
+        fp8_fwd_gran_value=0,
+        fp8_fwd_backend_value=0,
     ):
         out_dtype = input.dtype
         orig_shape = input.shape
@@ -316,19 +333,39 @@ def forward(
         a_fp4, a_scale, a_t_fp4, a_t_scale = _quantize_input_dual(input_2d, preshuffle)
         b_fp4, b_scale, b_t_fp4, b_t_scale = _quantize_weight_dual(weight, preshuffle)
 
-        output = gemm_fp4_impl(
-            a_fp4,
-            a_scale,
-            False,
-            b_fp4,
-            b_scale,
-            True,
-            out_dtype,
-            False,
-            granularity=_GRAN_VALUE,
-            default_backend=_DEFAULT_BACKEND,
-            preshuffled=preshuffle,
-        )
+        if forward_precision_value == _FORWARD_PRECISION_MXFP4:
+            output = gemm_fp4_impl(
+                a_fp4,
+                a_scale,
+                False,
+                b_fp4,
+                b_scale,
+                True,
+                out_dtype,
+                False,
+                granularity=_GRAN_VALUE,
+                default_backend=_DEFAULT_BACKEND,
+                preshuffled=preshuffle,
+            )
+        elif forward_precision_value == _FORWARD_PRECISION_FP8:
+            a_fp8, a_scale_inv = _quantize_fp8_tw(input_2d, fp8_fwd_dtype)
+            b_fp8, b_scale_inv = _quantize_fp8_tw(weight, fp8_fwd_dtype)
+            output = gemm_fp8_impl(
+                a_fp8,
+                a_scale_inv,
+                False,
+                b_fp8,
+                b_scale_inv,
+                True,
+                out_dtype,
+                False,
+                granularity=fp8_fwd_gran_value,
+                default_backend=fp8_fwd_backend_value,
+            )
+        elif forward_precision_value == _FORWARD_PRECISION_BF16:
+            output = torch.nn.functional.linear(input_2d, weight)
+        else:
+            raise ValueError(f"Unsupported MXFP4 forward precision value: {forward_precision_value}")
         output = output.reshape(*orig_shape[:-1], output.shape[-1])
 
         if backward_is_fp8:
@@ -348,6 +385,7 @@ def forward(
 
     @staticmethod
     def setup_context(ctx, inputs, output):
+        ctx.num_inputs = len(inputs)
         (
             _,
             _,
@@ -357,7 +395,7 @@ def setup_context(ctx, inputs, output):
             fp8_gran_value,
             fp8_backend_value,
             use_gradient_sr,
-        ) = inputs
+        ) = inputs[:8]
 
         ctx.preshuffle = preshuffle
         ctx.backward_is_fp8 = backward_is_fp8
@@ -458,7 +496,7 @@ def backward(ctx, grad_output, *_):
                 preshuffled=preshuffle,
             )
 
-        return grad_input, grad_weight, None, None, None, None, None, None
+        return (grad_input, grad_weight) + (None,) * (ctx.num_inputs - 2)
 
 
 # ---------------------------------------------------------------------------
@@ -492,9 +530,27 @@ def __init__(self, *args, **kwargs):
 
         self._preshuffle = _enable_preshuffle()
         _assert_preshuffle_contract(self.config, self._preshuffle)
+        self._forward_precision = getattr(self.config, "mxfp4_forward_precision", "mxfp4")
+        if self._forward_precision not in _FORWARD_PRECISION_VALUES:
+            raise ValueError(
+                "mxfp4_forward_precision must be one of "
+                f"{tuple(_FORWARD_PRECISION_VALUES)}, got {self._forward_precision!r}"
+            )
+        self._forward_precision_value = _FORWARD_PRECISION_VALUES[self._forward_precision]
         self._backward_is_fp8 = getattr(self.config, "mxfp4_backward_precision", "mxfp4") == "fp8"
         self._use_gradient_sr = getattr(self.config, "mxfp4_gradient_stochastic_rounding", False)
 
+        if self._forward_precision == "fp8":
+            from primus_turbo.pytorch.core.low_precision import float8_e4m3
+
+            self._fp8_fwd_dtype = float8_e4m3
+            self._fp8_fwd_gran_value = ScalingGranularity.TENSORWISE.value
+            self._fp8_fwd_backend_value = BackendType.HIPBLASLT.value
+        else:
+            self._fp8_fwd_dtype = None
+            self._fp8_fwd_gran_value = 0
+            self._fp8_fwd_backend_value = 0
+
         if self._backward_is_fp8:
             from primus_turbo.pytorch.core.low_precision import float8_e5m2
 
@@ -518,6 +574,10 @@ def _forward_impl(self, input, weight, *args, **kwargs):
             self._fp8_gran_value,
             self._fp8_backend_value,
             self._use_gradient_sr,
+            self._forward_precision_value,
+            self._fp8_fwd_dtype,
+            self._fp8_fwd_gran_value,
+            self._fp8_fwd_backend_value,
         )
         output = result[0]
 
@@ -552,9 +612,27 @@ def __init__(self, *args, **kwargs):
 
         self._preshuffle = _enable_preshuffle()
         _assert_preshuffle_contract(self.config, self._preshuffle)
+        self._forward_precision = getattr(self.config, "mxfp4_forward_precision", "mxfp4")
+        if self._forward_precision not in _FORWARD_PRECISION_VALUES:
+            raise ValueError(
+                "mxfp4_forward_precision must be one of "
+                f"{tuple(_FORWARD_PRECISION_VALUES)}, got {self._forward_precision!r}"
+            )
+        self._forward_precision_value = _FORWARD_PRECISION_VALUES[self._forward_precision]
         self._backward_is_fp8 = getattr(self.config, "mxfp4_backward_precision", "mxfp4") == "fp8"
         self._use_gradient_sr = getattr(self.config, "mxfp4_gradient_stochastic_rounding", False)
 
+        if self._forward_precision == "fp8":
+            from primus_turbo.pytorch.core.low_precision import float8_e4m3
+
+            self._fp8_fwd_dtype = float8_e4m3
+            self._fp8_fwd_gran_value = ScalingGranularity.TENSORWISE.value
+            self._fp8_fwd_backend_value = BackendType.HIPBLASLT.value
+        else:
+            self._fp8_fwd_dtype = None
+            self._fp8_fwd_gran_value = 0
+            self._fp8_fwd_backend_value = 0
+
         if self._backward_is_fp8:
             from primus_turbo.pytorch.core.low_precision import float8_e5m2
 
@@ -578,6 +656,10 @@ def _forward_impl(self, input, weight, *args, **kwargs):
             self._fp8_gran_value,
             self._fp8_backend_value,
             self._use_gradient_sr,
+            self._forward_precision_value,
+            self._fp8_fwd_dtype,
+            self._fp8_fwd_gran_value,
+            self._fp8_fwd_backend_value,
         )
         output = result[0]
 

primus/backends/megatron/core/models/diffusion/common/config.py

@@ -32,6 +32,8 @@ class BaseDiffusionConfig(TransformerConfig):
         fp8_scaling_strategy: FP8 scaling strategy for local spec provider (default: 'dynamic')
         fp8_force_nt_layout: FP8 backward GEMM layout (default: False)
         fp8_reduce_amax: Whether to allreduce amax across ranks (default: False)
+        mxfp4_forward_precision: MXFP4 forward precision, 'mxfp4', 'fp8', or 'bf16'
+            (default: 'mxfp4')
         mxfp4_backward_precision: MXFP4 backward precision, 'mxfp4' or 'fp8' (default: 'mxfp4')
         mxfp4_gradient_stochastic_rounding: Stochastic rounding on gradients (default: False)
         sensitive_layers_enabled: Enable sensitive layer configuration (default: False)
@@ -71,6 +73,11 @@ class BaseDiffusionConfig(TransformerConfig):
     # Whether to allreduce amax across DP/TP ranks for delayed FP8 scaling
     fp8_reduce_amax: bool = False
 
+    # MXFP4 forward precision: "mxfp4" (pure), "fp8", or "bf16".
+    # The latter two retain MXFP4 backward unless the backward knob below is
+    # independently changed.
+    mxfp4_forward_precision: str = "mxfp4"
+
     # MXFP4 backward precision: "mxfp4" (pure) or "fp8" (hybrid)
     mxfp4_backward_precision: str = "mxfp4"
 
@@ -168,6 +175,17 @@ def validate(self):
         Raises:
             ValueError: If configuration is invalid
         """
+        if self.mxfp4_forward_precision not in {"mxfp4", "fp8", "bf16"}:
+            raise ValueError(
+                "mxfp4_forward_precision must be 'mxfp4', 'fp8', or 'bf16', "
+                f"got {self.mxfp4_forward_precision!r}"
+            )
+
+        if self.mxfp4_backward_precision not in {"mxfp4", "fp8"}:
+            raise ValueError(
+                "mxfp4_backward_precision must be 'mxfp4' or 'fp8', " f"got {self.mxfp4_backward_precision!r}"
+            )
+
         if self.in_channels <= 0:
             raise ValueError(f"in_channels must be positive, got {self.in_channels}")
 

primus/backends/megatron/core/models/diffusion/flux/config.py

@@ -78,6 +78,8 @@ class FluxConfig(BaseDiffusionConfig):
         apply_rope_fusion: Whether to apply RoPE fusion optimization (default: False)
         add_qkv_bias: Whether to add bias to QKV projections (default: True)
         single_block_bias: Whether to add bias to single block linear layers (default: True)
+        hetereogenous_dist_checkpoint: Keep joint and single blocks in distinct
+            per-layer distributed-checkpoint namespaces (required, always True)
         activation_func: Activation function (default: openai_gelu_no_jit)
         use_te_rng_tracker: Whether to use Transformer Engine RNG tracker (default: False)
     """
@@ -91,6 +93,10 @@ class FluxConfig(BaseDiffusionConfig):
     # Architecture: Number of layers
     num_joint_layers: int = 19  # Default: Flux 12B
     num_single_layers: int = 38  # Default: Flux 12B
+    # Megatron preserves this misspelling in its public TransformerConfig API.
+    # Flux's joint and single blocks have different parameter sets and shapes,
+    # so they cannot share the homogeneous layer-stacked checkpoint namespace.
+    hetereogenous_dist_checkpoint: bool = True
 
     # Architecture: Dimensions (Flux standard: 3072)
     hidden_size: int = 3072
@@ -214,6 +220,12 @@ def __post_init__(self):
 
         super().__post_init__()  # BaseDiffusionConfig (mapping) -> TransformerConfig (validation)
 
+        if not self.hetereogenous_dist_checkpoint:
+            raise ValueError(
+                "Flux requires hetereogenous_dist_checkpoint=True because its joint "
+                "and single transformer blocks have different checkpoint schemas"
+            )
+
         # Xavier uniform for Megatron parallel linear layers (common Flux reference default).
         self.init_method = nn.init.xavier_uniform_
         self.output_layer_init_method = nn.init.xavier_uniform_

primus/backends/megatron/core/models/diffusion/flux/model.py

@@ -604,8 +604,8 @@ def sharded_state_dict(
         indexed sharded keys (``transformer.layers.<i>.*``) are required here:
         the homogeneous layer-stacked path would leave unclaimed slots for the
         params absent in single blocks and raise a CheckpointingException at
-        save time. TransformerBlock.sharded_state_dict provides this
-        automatically, so no config toggle is needed.
+        save time. ``FluxConfig.hetereogenous_dist_checkpoint`` enables
+        TransformerBlock's supported per-layer checkpoint namespace.
 
         Args:
             prefix: Prefix for state dict keys (e.g., 'module.')
@@ -617,8 +617,8 @@ def sharded_state_dict(
         """
         sharded_state_dict = {}
 
-        # Delegate transformer layers to TransformerBlock
-        # Handles heterogeneous layers automatically
+        # Delegate transformer layers to TransformerBlock. FluxConfig requires
+        # its heterogeneous per-layer checkpoint namespace.
         sharded_state_dict.update(
             self.transformer.sharded_state_dict(f"{prefix}transformer.", sharded_offsets, metadata)
         )