feat(mxfp6): MXFP6 (E2M3) Flux linears, a fused MLP, and the MLPerf recipe

.gitignore

@@ -38,3 +38,11 @@ pp_simulation_result
 .triton_cache_shared/
 .jax_cache/
 .cursor/
+
+# Throwaway run recipes. Configs worth keeping get a real name; an underscore
+# prefix marks a copy made to pin one experiment's knobs, which carries that
+# run's exp_name, dataset path and profiler window and is stale immediately.
+examples/megatron/configs/**/_*.yaml
+
+# Unit-test output directory.
+ut_out/

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

@@ -0,0 +1,226 @@
+# 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 + MXFP6 (MI355X)
+#
+# The MXFP4 sibling config of this file is the template; the only substantive
+# differences are the fp6 block near the bottom and the notes below.
+#
+# Key configuration:
+#   - Megatron DDP with overlap_grad_reduce + overlap_param_gather
+#   - PrimusTurboMXFP6LocalSpecProvider (NO TransformerEngine dependency)
+#   - MXFP6 (E2M3 + E8M0 block-of-32 scales) via Primus Turbo + AITER A6W6
+#   - Primus Turbo attention
+#   - torch.compile enabled (per_block strategy, compatible with local spec)
+#   - Energon pre-encoded dataset with stored VAE mean/logvar (resample mode)
+#
+# Requirements beyond the MXFP4 config:
+#   - gfx950 (MI350/MI355) only. The A6W6 kernels are gfx950 assembly, so unlike
+#     MXFP4 there is no fallback path on older parts.
+#   - An aiter carrying the A6W6 kernels (ROCm/aiter#4859). With an aiter that
+#     lacks them, startup fails loudly rather than quietly training in bf16.
+#
+# Shape constraint: every MXFP6 GEMM needs M, N and K to be multiples of 256. K is
+# included because the backward GEMMs use it as an output dimension. Flux 12B
+# satisfies this at every replaced linear (K is 3072 or 12288, N is 3072, 9216 or
+# 12288), and M = micro_batch_size * tokens stays aligned for the batch sizes here.
+# Note the spec only replaces the transformer block linears, so the VAE input
+# projection (K=64) is untouched.
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_local_mxfp6}
+
+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
+      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
+
+      # ==========================================
+      # PrimusTurboMXFP6LocalSpecProvider
+      # ==========================================
+      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 + MXFP6 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
+
+      # Swept, not inherited -- see the MLPerf MXFP6 recipe for the reasoning. The
+      # Megatron default of 256000000 elements leaves collective time exposed that a
+      # larger bucket hides: -17.4 ms per 512 images at this MBS=64/GBS=512, -27.9 ms
+      # at MBS=32/GBS=256, with peak memory unchanged. Both ends of the sweep lose, so
+      # 1024000000 is a measured optimum rather than a "bigger is better" guess.
+      ddp_bucket_size: 1024000000
+      use_fsdp2_fp32_param_optimizer: false
+
+      ckpt_format: torch_dist
+
+      # Required: the A6W6 entry point has no beta=1 accumulate epilogue, so wgrad
+      # cannot write main_grad in place.
+      gradient_accumulation_fusion: false
+
+      check_for_nan_in_loss_and_grad: true
+
+      # ==========================================
+      # MXFP6 Configuration — Block-scaled via Primus Turbo + AITER
+      # ==========================================
+
+      use_flash_attn: true
+
+      # Primus-owned field: Megatron has no notion of a 6-bit format, so fp6 must not
+      # be combined with fp4 or fp8 (rejected in BaseDiffusionConfig.__post_init__).
+      fp6: "mxfp6"
+      mxfp6_backward_precision: "mxfp6"  # "mxfp6" (pure) or "fp8" (hybrid)
+
+      # The MLP folds its bias-add + GELU into the MXFP6 packer, which is worth about 5%
+      # of step time here (74.1 -> 78.0 images/s/GPU on 8x MI355X at micro_batch_size 64).
+      # It engages on its own; there is no config key. Any configuration it cannot
+      # reproduce exactly falls back to Megatron's MLP per module, and
+      # PRIMUS_MXFP6_FUSED_MLP=off disables it outright while =on turns the fallback into
+      # an error, which is what an A/B measurement wants.
+
+      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_mxfp6
+      wandb_project: flux_12b_ddp_local_mxfp6
+      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 MXFP6 local spec (per-module FP6)
+      torch_compile:
+        enable: true
+        strategy: "per_block"
+        backend: "inductor"
+        mode: "default"
+        fullgraph: false
+        compile_optimizer: false
+        emulate_precision_casts: false
+        fused_ln_modulate: true