Enable native-TE MXFP4 training for Flux on MI355X

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

@@ -0,0 +1,204 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — DDP + TE Spec + MXFP4 (native-TE autocast) (MI355X)
+#
+# Combines Megatron DDP + distributed optimizer with the TransformerEngine spec,
+# routed through TE's native FP4 autocast (MXFP4BlockScaling) for MXFP4 training.
+#
+# Key configuration:
+#   - Megatron DDP with overlap_grad_reduce + overlap_param_gather
+#   - TEColumnParallelLinear / TERowParallelLinear / TEDotProductAttention / TENorm
+#   - Fused CK attention (attention_backend: fused), ~2.5x faster bwd than FlashAttn
+#   - MXFP4 (E2M1 + E8M0 block-of-32 scales) via TE-native fp4_use_native_te_autocast
+#   - Primus Turbo disabled (native TE linears consume the model-level FP4 autocast)
+#   - torch.compile disabled (TE FP4 autocast global state breaks graph capture)
+#   - Energon pre-encoded dataset with stored VAE mean/logvar (resample mode)
+#
+# Set these environment variables before launching:
+#   export NVTE_USE_CAST_TRANSPOSE_TRITON=0
+#   export NVTE_FUSED_ATTN=1
+#   export NVTE_FUSED_ATTN_CK=1
+#   export NVTE_USE_HIPBLASLT=1
+#   export USE_HIPBLASLT=1
+#   export TORCH_BLAS_PREFER_HIPBLASLT=1
+#   export NVTE_FP4_USE_TUNED_GEMM=1
+#   export NVTE_FP4_LOG_GEMM_SHAPES=0
+#   unset NVTE_FLASH_ATTN 
+#   unset NVTE_UNFUSED_ATTN
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_te_mxfp4}
+
+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
+      # ==========================================
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+      vae_latent_mode: resample
+
+      # ==========================================
+      # RoPE Fusion
+      # ==========================================
+      rotary_interleaved: true
+      apply_rope_fusion: true
+      position_embedding_type: rope
+
+      # ==========================================
+      # TransformerEngine Spec
+      # ==========================================
+      transformer_impl: "transformer_engine"
+
+      # ==========================================
+      # Attention backend — FUSED (TE CK fused attention)
+      # ==========================================
+      attention_backend: fused
+
+      # ==========================================
+      # MXFP4 — via TE-native MXFP4BlockScaling autocast (no Primus-Turbo)
+      # ==========================================
+      fp4: "mxfp4"
+      fp4_recipe: "mxfp4"
+      fp4_use_native_te_autocast: true
+      mxfp4_backward_precision: "mxfp4"
+
+      # ==========================================
+      # Energon Dataset
+      # ==========================================
+      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
+      # ==========================================
+      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
+      num_workers: 8
+      prefetch_factor: 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: false
+
+      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
+
+      # ==========================================
+      # Memory / Misc
+      # ==========================================
+      empty_unused_memory_level: 0
+
+      # Manual GC — align GC timing across ranks to avoid stragglers
+      manual_gc: true
+      manual_gc_interval: 1000
+
+      distributed_timeout_minutes: 60
+      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
+      log_throughput: true
+      wall_clock_step_timer: true
+
+      # Profiler — disabled
+      profile: false
+
+      # ==========================================
+      # NATIVE TE — Primus-Turbo disabled
+      # ==========================================
+      # The native-TE MXFP4 path uses TE's own linears + fp8_autocast(MXFP4BlockScaling).
+      # Primus-Turbo is disabled, so the model builds native TE linears that consume the
+      # model-level FP4 autocast returned by Flux.get_fp8_context().
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      seed: 2025
+      te_rng_tracker: true
+
+      # torch.compile — DISABLED for native-TE MXFP4.
+      # TE's FP4 autocast toggles process-global quantization state, which breaks
+      # torch.compile graph capture (graph breaks / stale guards). Run eager, matching
+      # the validated native-TE MXFP4 runs.
+      torch_compile:
+        enable: false

primus/backends/megatron/core/fp4_utils.py

@@ -144,7 +144,13 @@ def get_fp4_context(config: TransformerConfig, layer_no: int = -1, is_init: bool
                 fp4_context = nullcontext()
             else:
                 fp4_recipe, fp4_recipe_none_reason = get_fp4_recipe(config)
-                turbo_enabled = _primus_turbo_enabled()
+                # fp4_use_native_te_autocast forces the TE-native autocast branch
+                # (TE fp8_autocast + MXFP4BlockScaling -> AITER a4w4), bypassing
+                # Primus-Turbo entirely -- even if the Turbo autocast is otherwise
+                # enabled. This is the pure-TE MXFP4 path (enable_primus_turbo=false).
+                turbo_enabled = _primus_turbo_enabled() and not getattr(
+                    config, "fp4_use_native_te_autocast", False
+                )
 
                 global WARN_ONCE
                 if WARN_ONCE:

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

@@ -174,6 +174,15 @@ class FluxConfig(BaseDiffusionConfig):
     # quantization instead of the default inline implementation.
     use_cpp_fp8_quantize: bool = False
 
+    # FP4 autocast backend (TE-spec only). When True, the model-level FP4 autocast
+    # uses TE's native fp8_autocast with the MXFP4BlockScaling recipe (TE -> AITER
+    # a4w4) instead of Primus-Turbo's primus_turbo_fp4_autocast, forcing
+    # get_fp4_context onto the TE-native branch. Enables pure-TE MXFP4
+    # (enable_primus_turbo=false); requires use_turbo_gemm=false (native TE linears)
+    # because PrimusTurboLinear quantizes only under the Turbo autocast and would
+    # otherwise fall back to BF16. Enforced in validate_args_on_rocm.
+    fp4_use_native_te_autocast: bool = False
+
     # CUDA graph support
     enable_cuda_graph: bool = False
     cuda_graph_scope: Optional[str] = None  # Options: "full", "full_iteration"

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

@@ -634,25 +634,46 @@ def sharded_state_dict(
 
     def get_fp8_context(self):
         """
-        Return FP8 context manager for FP8 training.
+        Return the active low-precision (FP8 or FP4) quantization context manager.
 
-        When transformer_impl="local", FP8 is handled per-module at init time
-        (inside Float8ColumnParallelLinear / Float8RowParallelLinear), so no
-        global FP8 context manager is needed. This avoids mutable global state
-        that would cause torch.compile graph breaks.
+        For the TransformerEngine spec, a TE linear only issues an FP8/FP4 GEMM
+        while the corresponding autocast context is active. The Flux forward wraps
+        every compile strategy in this single method, so it must return whichever
+        quantization context the config requests:
+
+        - FP8 (config.fp8): TE FP8 autocast via megatron.core.fp8_utils.
+        - FP4 (config.fp4, e.g. MXFP4): TE FP4 autocast via megatron.core.fp4_utils.
+          Without this branch a pure-FP4 run would get nullcontext and silently
+          fall back to BF16 GEMMs (no AITER FP4 routing).
+
+        FP8 takes precedence over FP4 when both are set, mirroring stock Megatron's
+        transformer_block (`if config.fp8 ... elif config.fp4 ...`).
+
+        When transformer_impl="local", FP8/FP4 is handled per-module at init time
+        (inside Float8/MXFP4 ColumnParallelLinear / RowParallelLinear), so no global
+        context manager is needed. This avoids mutable global state that would cause
+        torch.compile graph breaks.
 
         Returns:
-            Context manager for FP8 operations (or nullcontext if FP8 disabled
-            or using local spec)
+            Context manager for the active quantization mode (or nullcontext if
+            quantization is disabled or using local spec).
         """
         from contextlib import nullcontext
 
         if self.config.transformer_impl == "local":
             return nullcontext()
 
-        from megatron.core.fp8_utils import get_fp8_context as get_fp8_context_util
+        if getattr(self.config, "fp8", None):
+            from megatron.core.fp8_utils import get_fp8_context as get_fp8_context_util
+
+            return get_fp8_context_util(self.config)
+
+        if getattr(self.config, "fp4", None):
+            from megatron.core.fp4_utils import get_fp4_context as get_fp4_context_util
+
+            return get_fp4_context_util(self.config)
 
-        return get_fp8_context_util(self.config)
+        return nullcontext()
 
     def forward(
         self,

primus/backends/megatron/flux_pretrain_trainer.py

@@ -466,6 +466,30 @@ def _build_flux_config_from_yaml(self):
         transformer_impl = getattr(params, "transformer_impl", "transformer_engine")
         config_params["transformer_impl"] = transformer_impl
 
+        # Convert attention_backend (string) to the AttnBackend enum. A YAML value
+        # like "fused" arrives as a plain string, but _set_attention_backend()
+        # compares it against AttnBackend enum members -- a raw string matches no
+        # branch, so the backend is silently left unset and falls back to
+        # FlashAttention. AttnBackend["fused"] returns AttnBackend.fused.
+        from megatron.core.transformer.enums import AttnBackend
+
+        # The getattr default only applies when the key is absent; an explicit
+        # `attention_backend: null` in YAML yields None, so coerce that (and an
+        # empty string) to auto rather than crashing on attn_backend.name below.
+        attn_backend = getattr(params, "attention_backend", AttnBackend.auto)
+        if attn_backend is None or (isinstance(attn_backend, str) and not attn_backend.strip()):
+            attn_backend = AttnBackend.auto
+        if isinstance(attn_backend, str):
+            try:
+                attn_backend = AttnBackend[attn_backend.strip()]
+            except KeyError as exc:
+                raise ValueError(
+                    f"Unknown attention_backend '{attn_backend}'. Choose from: "
+                    f"{[b.name for b in AttnBackend]}."
+                ) from exc
+        config_params["attention_backend"] = attn_backend
+        log_rank_0(f"Attention backend: {attn_backend.name}")
+
         # Precision settings from params (not in model YAML)
         config_params.update(
             {
@@ -505,6 +529,7 @@ def _build_flux_config_from_yaml(self):
                 "fp4": fp4_enabled,
                 "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),
             }
         )
 

primus/backends/megatron/patches/args/rocm_arg_validation.py

@@ -143,6 +143,23 @@ def validate_args_on_rocm(args):
             args.fp4_recipe in support_fp4_recipe
         ), f"{args.fp4_recipe} recipe is not support when enable `use_turbo_gemm`."
 
+    # Turbo FP4 autocast check: PrimusTurboLinear quantizes only under the Turbo
+    # autocast, which the TE-native path never enters. Local specs quantize
+    # in-module rather than off that state, so they are exempt.
+    if (
+        args.fp4
+        and use_turbo_gemm
+        and getattr(args, "fp4_use_native_te_autocast", False)
+        and getattr(args, "transformer_impl", "transformer_engine") != "local"
+    ):
+        raise ValueError(
+            "fp4_use_native_te_autocast=True is incompatible with use_turbo_gemm=True. "
+            "Native TE FP4 autocast does not set the Primus-Turbo FP4 state that "
+            "PrimusTurboLinear reads, so Turbo GEMM linears would silently run BF16 "
+            "instead of FP4. Set use_turbo_gemm=false to use native TE MXFP4 linears, "
+            "or set fp4_use_native_te_autocast=false to use the Primus-Turbo FP4 path."
+        )
+
     # NOTE: mxfp8 environment variable must be set to 1 to enable mxfp8 recipe on ROCm.
     if args.fp8_recipe == "mxfp8":
         assert (