diff --git a/.gitignore b/.gitignore index 71d1928b9..6bd7a89f2 100644 --- a/.gitignore +++ b/.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/ diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp6.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp6.yaml new file mode 100644 index 000000000..7765fe303 --- /dev/null +++ b/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 diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp6_mlperf.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp6_mlperf.yaml new file mode 100644 index 000000000..e9e68df89 --- /dev/null +++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp6_mlperf.yaml @@ -0,0 +1,274 @@ +# 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 (MLPerf Mode, MI355X) +# +# The submission recipe for MXFP6. It is the FP8 MLPerf config with the +# precision block swapped, so that a comparison between the two is a comparison +# of numerics and nothing else; everything the ruleset pins -- AdamW +# 0.9/0.95/1e-8, weight decay 0.1, gradient clip 1.0, eval every 512 steps -- +# is identical on purpose and must not be tuned. For MXFP6 development work +# without the MLPerf machinery, use +# flux_12b_ddp_energon_schnell_resample_local_spec_mxfp6.yaml. +# +# Hardware: gfx950 (MI350/MI355) only. The A6W6 kernels are gfx950 assembly, so +# unlike MXFP4 there is no fallback path on older parts, and startup fails +# loudly on an aiter that lacks them rather than quietly training in bf16. +# +# Shape constraint: every MXFP6 GEMM needs M, N and K to be multiples of 256. +# Flux 12B satisfies this at every replaced linear, and M = micro_batch_size * +# tokens stays aligned at the batch sizes here. +# +# Before this produces a submittable log: the compliance checker accepts a +# fixed vocabulary for lowest_numerical_precision_in_linear and mxfp6 is not in +# it. A run under this recipe is a valid run whose log the checker rejects on +# that key alone, until the format is accepted upstream. + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_local_mxfp6_mlperf} + +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 + + # ========================================== + # MLPerf Mode + # ========================================== + mlperf_mode: true + warmup_train_steps: 2 + target_val_loss: 0.586 + + # Left explicit because the FP8 MLPerf recipe and the MXFP6 development + # recipe disagree on it, and the difference is a different LR trajectory + # through warmup rather than a stylistic one. This matches the FP8 MLPerf + # recipe, which is the configuration the convergence evidence so far was + # collected under; flipping it is a convergence experiment, not a cleanup. + nemo_aligned_lr_warmup: false + + # ========================================== + # MLPerf Training v5.1 Alignment + # ========================================== + 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 + vae_latent_mode: resample + + # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout + rotary_interleaved: true + + # ========================================== + # PrimusTurboFloat8LocalSpecProvider + # ========================================== + transformer_impl: "local" + + # ========================================== + # DATA — Real Energon data (not mock) + # ========================================== + mock_data: false + dataloader_type: external + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + num_workers: 16 + prefetch_factor: 4 + max_samples_per_sequence: null + + # Training iterations + # + # This is a safety cap, not a target: the run stops at the quality + # threshold, and reaching it is the result. The reference convergence + # points for flux1 at GBS 512 (mlperf_logging rcp_checker + # training_6.0.0/rcps_flux1.json, 20 NVIDIA BF16 runs) span 7077888 to + # 7602176 samples, i.e. 13824 to 14848 steps. The previous 5000 steps + # was 2560000 samples -- roughly a third of the way -- so every run + # exhausted its budget below target and no run could ever have counted. + # 16000 steps (8192000 samples) clears the slowest reference run with + # margin for a run that converges later than the reference spread. + train_iters: 16000 + eval_interval: 512 + # Cover the whole MLPerf validation set. eval_iters is derived from this + # (29696 / 512 = 58); setting both is rejected. The previous eval_iters of + # 10 read 5120 samples, and even at 58 the training num_workers of 16 + # would have left 1920 of them unread, so val_num_workers is pinned to 0. + eval_samples: 29696 + val_num_workers: 0 + # MLPerf assigns each val image one fixed timestep, carried in the Arrow + # 'timestep' column. Requires a val split ingested with that column; shards + # ingested earlier carry {"key": ...} only and will now fail loudly rather + # than silently fall back to injecting t = index % 8. + eval_timestep_source: dataset + log_interval: 10 + save_interval: 10000 + + # Batch configuration + micro_batch_size: 64 + global_batch_size: 512 + seq_length: 512 + + # ========================================== + # BF16 + FP8 Precision + # ========================================== + bf16: true + fp16: false + params_dtype: bfloat16 + grad_reduce_in_bf16: true + + # Optimizer settings (MLPerf v5.1) + 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 + # Inert while lr_decay_style is constant; kept equal to the budget so + # it cannot be misread as a decay that ends a quarter of the way in. + lr_decay_iters: 16000 + lr_decay_style: constant + + # ========================================== + # 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 + # FP32 optimizer states (m, v), matching the FP8 MLPerf recipe rather than + # the BF16 states the MXFP6 development recipe uses. Time to train is what + # is being measured, and FP32 states cost almost nothing in throughput here + # while removing a source of convergence difference from the FP8 comparison. + # + # Now measured rather than asserted, since the bucket sweep produced the + # matched arms for free. Against the development recipe's BF16 m/v with + # use_precision_aware_optimizer true, at MBS=64/GBS=512 and equal bucket + # size: +6.9 ms per 512 images, +0.85%, at both 256000000 and 1024000000. + # The throughput claim holds. What it does cost is memory, +16.2 GB of peak + # (267.0 against 250.8 GB), which is the number to watch if this recipe ever + # moves to a larger micro-batch, because 267 of 288 GB is most of the card. + # + # Note that main_grads_dtype below is inert: it belongs to the precision-aware + # optimizer and is ignored when that flag is false. The gradients DDP reduces + # are bf16 in this recipe exactly as in the development one, which is why the + # bucket sweep above transfers between them. + use_precision_aware_optimizer: false + main_params_dtype: fp32 + main_grads_dtype: fp32 + exp_avg_dtype: fp32 + exp_avg_sq_dtype: fp32 + + # Swept in this recipe's own optimizer regime rather than inherited from the + # development one. The Megatron default of 256000000 elements is the worst of + # the four sizes measured here: 1024000000 gives 801.9 ms per 512 images + # against 819.3, worth -17.4 ms, with peak memory unchanged (266.0 GB against + # 267.0). 512000000, the value the MLPerf v6.0 reference moved to, is an + # improvement but not the best one at 806.4. + # + # The collective path is identical to the development recipe's -- both reduce + # bf16 gradients (grad_reduce_in_fp32=False) and both build 43 buckets with the + # same element counts at this bucket size -- so the same 1024000000 optimum was + # expected here, and the win is the same -17.4 ms to the decimal. + # + # Do not push this further in one direction. Both ends lose: small buckets pay + # per-call overhead and get thin per-call bandwidth, and buckets spanning most + # of the model cannot start their reduce-scatter until far into backward, which + # re-exposes the tail. See mbs32_e2e/README.md in tiger-training-internal. + ddp_bucket_size: 1024000000 + use_fsdp2_fp32_param_optimizer: false + + ckpt_format: torch_dist + + gradient_accumulation_fusion: false + check_for_nan_in_loss_and_grad: false + + # ========================================== + # MXFP6 — Block-scaled (E2M3 + E8M0 over 32) 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, worth about 5% of + # step time (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. A submission run should be + # launched with PRIMUS_MXFP6_FUSED_MLP=on so that a configuration the fusion + # cannot reproduce exactly becomes an error rather than a silent per-module + # fallback to Megatron's MLP -- otherwise the measured throughput and the + # disclosed implementation can quietly diverge. + + empty_unused_memory_level: 0 + + distributed_timeout_minutes: 60 + distributed_backend: nccl + + # Checkpointing (disabled for MLPerf) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging — suppressed by mlperf_mode + tensorboard_dir: null + wandb_project: null + log_throughput: false + wall_clock_step_timer: true + log_timers_to_tensorboard: false + log_batch_size_to_tensorboard: false + log_learning_rate_to_tensorboard: false + log_memory_to_tensorboard: false + + # Profiler — disabled for MLPerf + profile: false + + # Primus Turbo + enable_primus_turbo: true + use_turbo_attention: true + use_dual_fp8_output_projection: false + + # Each run in a submission campaign needs its own seed, supplied from + # outside so the value in the log is the value the launcher chose. + seed: ${PRIMUS_SEED:2025} + # MLPerf-aligned per-step CUDA RNG reseed (defaults off elsewhere; MLPerf + # reproduction must opt in for run-to-run determinism). + per_step_rng_reseed: true + + # Torch Compile — per_block: compile each transformer layer individually + torch_compile: + enable: true + strategy: "per_block" + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: false + emulate_precision_casts: false + fused_ln_modulate: true diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp6_mlperf_2n.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp6_mlperf_2n.yaml new file mode 100644 index 000000000..fabf3d4df --- /dev/null +++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp6_mlperf_2n.yaml @@ -0,0 +1,279 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell — DDP + Local Spec + MXFP6, MLPerf mode, TWO NODES (16 ranks) +# +# The two-node counterpart of +# flux_12b_ddp_energon_schnell_resample_local_spec_mxfp6_mlperf.yaml. Nine seeds +# converged under this configuration (20131, 21139, 22147, 23155, 24163, 25171, +# 26179, 27187, 28195), every one of them reaching val_loss <= 0.586, at a mean +# of 7,718,684 samples and roughly 2.4 h per seed on 2x8 MI355X. +# +# THE BATCH SHAPE ASSUMES EXACTLY 16 RANKS. micro_batch_size 32 at +# global_batch_size 512 is one microbatch per rank on two nodes of eight. On one +# node it becomes two, and on four it does not divide. Changing node count means +# changing the batch configuration, and the learning rate and warmup below go +# with it -- see the reference-point note further down. +# +# WHY GBS 512 RATHER THAN 1024. Each batch size carries its own MLPerf reference +# convergence points, and flux_ref_512's olympic mean is 7,223,524 samples +# against 8,490,554 at GBS 1024 -- the smaller batch is 15% cheaper in samples +# for the same quality target. Wall clock per seed is close to a wash, because +# gradient accumulation drops to 1 and communication exposure per sample +# doubles: measured 565.0 ms/step with evaluation here against 999.9 at GBS 1024. +# +# LEARNING RATE AND WARMUP BELONG TO THE BATCH SIZE. lr 2.0e-4 and +# lr_warmup_iters 1600 are the flux_ref_512 values; at GBS 1024 they are 2.5e-4 +# and 800. The warmup field counts in samples as iters x GBS, so 1600 x 512 and +# 800 x 1024 are the same 819,200-sample ramp rather than a longer one. Neither +# field is constrained by the ruleset, so a GBS 512 run carrying the GBS 1024 +# hyperparameters passes every compliance check and quietly forfeits the +# reference points it is being scored against. There is no log line for this. +# +# Hardware: gfx950 (MI350/MI355) only. The A6W6 kernels are gfx950 assembly, so +# unlike MXFP4 there is no fallback path on older parts, and startup fails +# loudly on an aiter that lacks them rather than quietly training in bf16. +# +# Before this produces a submittable log: the compliance checker accepts a fixed +# vocabulary for lowest_numerical_precision_in_linear and mxfp6 is not in it. A +# run under this recipe is a valid run whose log the checker rejects on that key +# alone, until the format is accepted upstream. +# +# LAUNCHING THIS NEEDS MORE THAN THE RECIPE. Inter-node bandwidth, the eight +# rails, the MLLOG disclosure variables and the pinned container are all outside +# this file, and a run missing any of them either dies at startup or produces a +# result that is not comparable to the nine above. The full procedure is in +# tiger-training-internal docs/mlperf/flux_mxfp6_twonode_reproduce.md. + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_local_mxfp6_mlperf_2n} + +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 + + # ========================================== + # MLPerf Mode + # ========================================== + mlperf_mode: true + target_val_loss: 0.586 + warmup_train_steps: 2 + nemo_aligned_lr_warmup: true + + # ========================================== + # MLPerf Training v6.1 Alignment + # ========================================== + 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 + vae_latent_mode: resample + + # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout + rotary_interleaved: true + + # ========================================== + # PrimusTurboFloat8LocalSpecProvider + # ========================================== + transformer_impl: "local" + + # ========================================== + # DATA — real Energon data, train and val in one dataset + # ========================================== + # The default is the combined set, which is the data every converged seed + # here was trained on. The original mlperf_flux1 has no val split carrying + # the fixed-timestep column that eval_timestep_source below requires. + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/data/mlperf_flux1_combined_20260828} + train_data_path: [] + mock_data: false + dataloader_type: external + num_workers: 8 + max_samples_per_sequence: null + + # A safety cap, not a target: the run stops at the quality threshold and + # reaching it is the result. The slowest of the nine converged seeds took + # 15,360 steps, so 20000 clears the observed spread by 30%. + train_iters: 20000 + + eval_interval: 512 + # Megatron gates evaluation on `args.eval_interval and iteration % + # eval_interval == 0 and args.do_valid`. do_valid is normally set inside + # build_train_valid_test_data_loaders, which the flux path never calls + # because it builds Energon dataloaders itself, so it stays at its False + # default and evaluation never fires. A two-node run reached step 710 and + # logged zero validation lines. Set it explicitly. + do_valid: true + # eval_samples, NOT eval_iters. The megatron.args.eval_samples patch + # derives eval_iters from this and asserts coverage is exact: 29696 / 512 + # = 58 here. Setting eval_iters directly bypasses that patch, and a run + # doing so evaluated zero times in 400 steps. val_num_workers is pinned to + # 0 because a nonzero training num_workers leaves part of the set unread. + eval_samples: 29696 + val_num_workers: 0 + # MLPerf assigns each validation image one fixed timestep, carried in the + # Arrow 'timestep' column. Requires a val split ingested with that column; + # shards ingested earlier carry {"key": ...} only and fail loudly rather + # than silently falling back to injecting t = index % 8. Without this the + # run logs "Validation timestep source: equidistant" and is not on the + # evaluation path the nine converged seeds were measured with. + eval_timestep_source: dataset + + log_interval: 10 + + # ========================================== + # Batch configuration — 16 ranks, one microbatch per rank + # ========================================== + micro_batch_size: 32 + global_batch_size: 512 + seq_length: 512 + + # ========================================== + # BF16 + MXFP6 precision + # ========================================== + bf16: true + fp16: false + params_dtype: bfloat16 + grad_reduce_in_bf16: true + + # Optimizer settings (MLPerf-pinned; must not be tuned) + 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 + # Inert while lr_decay_style is constant. Left at the value every recorded + # run carried rather than tidied to train_iters, so a log from this recipe + # and the runs it reproduces agree field for field. + lr_decay_iters: 4000 + lr_decay_style: constant + + # ========================================== + # 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 + # Both overlaps are what keeps the inter-node reduce-scatter and param + # all-gather off the critical path. overlap_param_gather_with_optimizer_step + # is unavailable here: it asserts interleaved virtual pipelining, which + # does not exist at TP=PP=1. + 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 + # 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 + + use_flash_attn: true + + # ========================================== + # MXFP6 — block-scaled (E2M3 + E8M0 over 32) via Primus Turbo + AITER A6W6 + # ========================================== + # 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, worth about 5% + # of step time. It engages on its own; there is no config key. Launch with + # PRIMUS_MXFP6_FUSED_MLP=on so that a configuration the fusion cannot + # reproduce exactly becomes an error rather than a silent per-module + # fallback to Megatron's MLP -- otherwise measured throughput and the + # disclosed implementation can quietly diverge. + + empty_unused_memory_level: 0 + # Two nodes rendezvous over c10d and the first collective follows dataset + # startup, so the default 10 minutes is not enough headroom on a cold + # page cache. + distributed_timeout_minutes: 180 + distributed_backend: nccl + + # ========================================== + # Checkpointing — off + # ========================================== + # Every recorded run has save: null. A run that cannot finish inside its + # allocation therefore produces nothing, which is a deliberate trade + # against paying checkpoint I/O on a time-to-train measurement. + ckpt_format: torch_dist + save: null + save_interval: 100000 + disable_last_saving: true + load: null + finetune: false + use_checkpoint_args: false + auto_continue_train: false + + # ========================================== + # Logging + # ========================================== + tensorboard_dir: ${PRIMUS_TB_DIR:output/tensorboard/flux_12b_ddp_local_mxfp6_mlperf_2n} + wandb_project: null + 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 + + profile: false + use_pytorch_profiler: false + + enable_primus_turbo: true + use_turbo_attention: true + + # ========================================== + # Seed + # ========================================== + # Each run in a submission campaign needs its own seed. The stack is + # bit-deterministic, so repeating a seed reproduces its validation curve + # exactly and adds no convergence evidence. + seed: ${PRIMUS_SEED:20131} + # MLPerf-aligned per-step CUDA RNG reseed (defaults off elsewhere). + per_step_rng_reseed: true + nemo_chimera_init: false + + torch_compile: + enable: true + strategy: "per_block" + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: false + emulate_precision_casts: false + fused_ln_modulate: true diff --git a/examples/mlperf/flux1/megatron/README.md b/examples/mlperf/flux1/megatron/README.md index 48a668a90..b9301ac62 100644 --- a/examples/mlperf/flux1/megatron/README.md +++ b/examples/mlperf/flux1/megatron/README.md @@ -59,6 +59,20 @@ result. `evaluation_frequency` at exactly 262,144 samples (`eval_interval: 512` at GBS 512). Changing any of them in the recipe makes the log fail the checker. +## Running MXFP6 + +```bash +EXP=examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp6_mlperf.yaml \ +MLLOG_LOWEST_NUMERICAL_PRECISION_IN_LINEAR=mxfp6 \ +PRIMUS_MXFP6_FUSED_MLP=on \ +RESULTS_DIR=/results bash examples/mlperf/flux1/megatron/run_campaign.sh +``` + +`PRIMUS_MXFP6_FUSED_MLP=on` turns a configuration the fused MLP cannot +reproduce exactly into an error rather than a silent per-module fallback to +Megatron's MLP, so the throughput being measured and the implementation being +disclosed cannot drift apart. + ## MXFP6 and the disclosure vocabulary `training_6.0.0/common.yaml` accepts a fixed set of values for diff --git a/examples/run_pretrain.sh b/examples/run_pretrain.sh index 7878cc903..b117478d5 100755 --- a/examples/run_pretrain.sh +++ b/examples/run_pretrain.sh @@ -224,8 +224,85 @@ if [ "$USING_AINIC" == "1" ]; then export RCCL_HOME_DIR=${RCCL_HOME_DIR:-"/workspace/rccl"} export MPI_HOME_DIR=${MPI_HOME_DIR:-"/opt/ompi"} export NCCL_MAX_P2P_CHANNELS=56 - export NCCL_DMABUF_ENABLE=0 + export NCCL_DMABUF_ENABLE=${NCCL_DMABUF_ENABLE:-0} export NCCL_IB_QPS_PER_CONNECTION=1 + + # Without a topology file RCCL discovers the fabric itself and plans 2 channels inter-node, + # which is 46 GB/s; the Crusoe XML declares which rail is local to which GPU and the plan + # widens to 8, which is 118. See docs/mlperf/rccl_gdr_escalation.md, 2026-08-31. + # + # Resolved by GPU PCI device id, never by filename. The mi350x and mi355x files differ only + # there (0x75a0 against 0x75a3) and a mismatched file still parses, so RCCL applies it and + # silently loses the affinity the file exists to declare -- a wrong filename is worse than + # no file, because it looks configured. Read from sysfs rather than lspci, which is a + # package and not present in every image. /etc/crusoe is not mounted in every container, so + # a staged copy is accepted as a fallback. + # PRIMUS_RCCL_TOPO_DISABLE=1 restores the pre-2026-08-31 behaviour exactly: no topology + # file and no GDR level of our choosing. It exists so the A/B stays measurable from outside + # the launcher, and as an escape hatch if a future node ships a file that plans worse than + # RCCL's own discovery. + if [ "${PRIMUS_RCCL_TOPO_DISABLE:-0}" = "1" ]; then + LOG_INFO_RANK0 "RCCL topology: disabled by PRIMUS_RCCL_TOPO_DISABLE=1" + elif [ -z "${NCCL_TOPO_FILE:-}" ]; then + _gpu_devid="" + for _d in /sys/bus/pci/devices/*; do + [ "$(cat "$_d/vendor" 2>/dev/null)" = "0x1002" ] || continue + _gpu_devid="$(cat "$_d/device" 2>/dev/null)" + [ -n "$_gpu_devid" ] && break + done + for _f in /etc/crusoe/rccl_topo/*.xml /opt/rccl_topo_node.xml; do + [ -f "$_f" ] || continue + if [ -n "$_gpu_devid" ] && grep -q "device=\"${_gpu_devid}\"" "$_f"; then + export NCCL_TOPO_FILE="$_f" + break + fi + done + if [ -n "${NCCL_TOPO_FILE:-}" ]; then + LOG_INFO_RANK0 "RCCL topology: $NCCL_TOPO_FILE (matches GPU $_gpu_devid)" + else + LOG_INFO_RANK0 "RCCL topology: none matches GPU ${_gpu_devid:-unknown}; inter-node will plan 2 channels" + fi + unset _gpu_devid _d _f + fi + + # GPU-Direct RDMA is one switch because the four settings below only work together, and + # the failure mode of getting it partly right is that every rank dies during connection + # setup rather than falling back. + # + # dmabuf is the only registration path available here: the peerdirect client is absent + # from ib_core, and RCCL's own dmabuf gate reads a kernel config file this image does not + # ship, hence the force flag. Registration also has to come off the VMM allocator. With + # the default allocator the buffers are not dmabuf-exportable, RCCL falls back to + # ibv_reg_mr_iova2, and that returns EINVAL for every buffer on this fabric. That EINVAL + # was previously read as the fabric refusing GDR outright, and the level was pinned to LOC + # because of it -- but it is the allocator, not the fabric. NCCL_CUMEM_ENABLE=1 removes it. + # + # Two same-pod nodes at MBS=32/GBS=512, 2026-09-04: 537.0 -> 502.8 ms/step, a 34 ms + # saving against a 3.3 ms spread between repeats, final loss identical. Needs the + # topology file above, which is what makes RCCL attempt GDR at all; the separate 2.56x + # from the wider channel plan does not depend on GDR. + # + # Off by default because it is only safe once the dataloader can avoid fork: a process + # holding dmabuf-registered GPU memory segfaults inside os.fork(). That needs Energon + # carrying dev/patches/energon-7.3.2-no-fork.patch and ENERGON_MP_CONTEXT set away from + # fork, so refuse rather than hand back a segfault with no explanation. + if [ "${PRIMUS_RCCL_GDR:-0}" = "1" ]; then + if [ "${ENERGON_MP_CONTEXT:-fork}" = "fork" ]; then + LOG_ERROR "PRIMUS_RCCL_GDR=1 requires ENERGON_MP_CONTEXT=forkserver (or spawn)." + LOG_ERROR "Under fork the dataloader forks with GPU memory registered for GDR and" + LOG_ERROR "its workers segfault immediately. See dev/patches/README.md." + exit 1 + fi + # Set outright, not with :-, because NCCL_DMABUF_ENABLE has already been defaulted to + # 0 above and a :- default would silently keep that 0 and break registration. + export NCCL_NET_GDR_LEVEL="${PRIMUS_RCCL_GDR_LEVEL:-SYS}" + export NCCL_DMABUF_ENABLE=1 + export RCCL_FORCE_ENABLE_DMABUF=1 + export NCCL_CUMEM_ENABLE=1 + LOG_INFO_RANK0 "RCCL GDR: enabled, level $NCCL_NET_GDR_LEVEL (dmabuf via cumem)" + elif [ "${PRIMUS_RCCL_TOPO_DISABLE:-0}" != "1" ]; then + export NCCL_NET_GDR_LEVEL=${NCCL_NET_GDR_LEVEL:-LOC} + fi path_append_unique LD_LIBRARY_PATH \ /usr/lib/x86_64-linux-gnu \ /usr/lib/x86_64-linux-gnu/libibverbs \ diff --git a/primus/backends/megatron/core/extensions/primus_turbo_local_spec.py b/primus/backends/megatron/core/extensions/primus_turbo_local_spec.py index 0b903ef14..2d15873d3 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo_local_spec.py +++ b/primus/backends/megatron/core/extensions/primus_turbo_local_spec.py @@ -256,6 +256,38 @@ def core_attention(self) -> type: return PrimusTurboLocalAttention +class PrimusTurboMXFP6LocalSpecProvider(LocalSpecProvider): + """ + Compile-friendly MXFP6 (E2M3) spec: Primus Turbo attention + MXFP6 linear layers. + NO TransformerEngine. NO global FP6 state in forward path. + Requires tensor_model_parallel_size=1, and M/N/K all multiples of 256. + """ + + def column_parallel_linear(self) -> type: + from .primus_turbo_mxfp6_local import MXFP6ColumnParallelLinear + + return MXFP6ColumnParallelLinear + + def row_parallel_linear(self) -> type: + from .primus_turbo_mxfp6_local import MXFP6RowParallelLinear + + return MXFP6RowParallelLinear + + def mlp_module(self) -> type: + """MLP that folds its bias-add + GELU into the MXFP6 packer. + + Not part of ``BackendSpecProvider``; specs that can use it look for this method and + fall back to Megatron's ``MLP``. The subclass itself falls back per-module for any + configuration it cannot reproduce, so returning it unconditionally is safe. + """ + from .primus_turbo_mxfp6_local import MXFP6FusedMLP + + return MXFP6FusedMLP + + def core_attention(self) -> type: + return PrimusTurboLocalAttention + + class PrimusTurboFloat8LocalSpecProvider(LocalSpecProvider): """ Compile-friendly FP8 spec: Primus Turbo attention + FP8 linear layers. diff --git a/primus/backends/megatron/core/extensions/primus_turbo_mxfp6_local.py b/primus/backends/megatron/core/extensions/primus_turbo_mxfp6_local.py new file mode 100644 index 000000000..4d3d3537f --- /dev/null +++ b/primus/backends/megatron/core/extensions/primus_turbo_mxfp6_local.py @@ -0,0 +1,738 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Compile-friendly MXFP6 (E2M3) linear layers for Megatron local spec. + +The MXFP4 sibling of this module (``primus_turbo_mxfp4_local``) is the template, and +this one is deliberately smaller, because MXFP6 removes most of MXFP4's configuration +surface rather than because anything is left unfinished: + +- **No preshuffle contract.** The A6W6 kernels read AITER's packed C0/C1 tile blob + directly, so there is no unshuffled layout and no fast path to opt into. MXFP4's + ``_enable_preshuffle`` / ``_assert_preshuffle_contract`` dance, and the whole class of + misconfiguration it guards against, simply does not exist here. +- **No ScalingRecipe flags.** The 32-point Hadamard rotation is mandatory and fused into + the packer (the GEMM depends on it cancelling between the two operands), scaling is + strictly per-1x32 along the contraction axis so ``use_2d_block`` is meaningless, and + stochastic rounding is not implemented. MXFP4 threads twelve booleans through its + quantize op; MXFP6 has none to thread. +- **No local custom-op registration.** Primus-Turbo already exposes + ``primus_turbo::quantize_mxfp6_dual_impl`` as a ``torch.library.custom_op`` with a + correct fake, so unlike MXFP4 there is nothing to re-wrap in order to bypass + recipe construction. + +Retained from the MXFP4 design: the ``setup_context`` pattern with primitive-only +arguments so ``torch.compile`` traces without graph breaks, the two backward modes +(pure MXFP6, or hybrid MXFP6-forward / FP8-backward), and zero TransformerEngine +dependencies. + +Shape constraint worth knowing: MXFP6 needs the linear's M, N **and** K to be multiples +of 256. K is included because the backward GEMMs use it as an output dimension. This is +enforced inside Primus-Turbo's ``gemm_fp6``; here it means a hidden size or sequence +length that is only 128-aligned will be rejected at the first forward. +""" + +import functools +import os +import warnings + +import torch +import torch.nn.functional as F +from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear +from megatron.core.transformer.mlp import MLP +from primus_turbo.pytorch.core.backend import BackendType +from primus_turbo.pytorch.core.low_precision import ( + MXFP6_PROLOGUE_BIAS_GELU, + MXFP6_PROLOGUE_BIAS_GELU_BACKWARD, + MXFP6_PROLOGUE_IDENTITY, + ScalingGranularity, +) +from primus_turbo.pytorch.kernels.gemm.gemm_fp6_impl import gemm_fp6_impl, gemm_fp6_out_impl +from primus_turbo.pytorch.kernels.gemm.gemm_fp8_impl import gemm_fp8_impl +from primus_turbo.pytorch.kernels.quantization.mxfp6_pack import check_mxfp6_support + +from .primus_turbo_float8_local import _quantize_fp8_tw + +_GRAN_VALUE = ScalingGranularity.MX_BLOCKWISE.value + +# Registered by Primus-Turbo with a pure-arithmetic fake, so it is safe to trace. +_quantize_mxfp6_dual = torch.ops.primus_turbo.quantize_mxfp6_dual_impl +_quantize_mxfp6_fused_dual = torch.ops.primus_turbo.quantize_mxfp6_fused_dual_impl + + +def _wgrad_into_main_grad(weight, g_col, g_col_scale, a_col, a_col_scale, n, k, m): + """Store the weight gradient directly into ``weight.main_grad``. + + Saves the round trip the unfused path forces: a freshly allocated wgrad, handed to + autograd as ``weight.grad``, which Megatron's DDP backward hook then adds into + ``main_grad`` and frees. Here the A6W6 asm writes ``main_grad`` itself and the hook + has nothing left to do. + + Correct only because the store has beta=0 and there is exactly one microbatch per + optimizer step; ``_init_mxfp6_linear`` enforces that. + + ``grad_added_to_main_grad`` is *not* set here -- see ``_claim_main_grad``, which has to + do it from the forward instead. + + Returns the placeholder that has to go back to autograd as the weight gradient. It is + never read: with ``grad_added_to_main_grad`` set, the hook skips its ``add_`` and + immediately drops ``param.grad``. It exists only because the hook asserts + ``param.grad is not None`` whenever ``overlap_grad_reduce`` is on. This mirrors what + Megatron's own ``gradient_accumulation_fusion`` path returns. Note for plan item 0.2b: + this allocation sits inside what a full-iteration CUDA graph would capture. + """ + main_grad = getattr(weight, "main_grad", None) + if main_grad is None: + raise RuntimeError( + "MXFP6 wgrad fusion needs weight.main_grad, which Megatron DDP allocates. " + "Either wrap the model in DistributedDataParallel or set " + "mxfp6_fused_wgrad_accum=False." + ) + if main_grad.dtype != torch.bfloat16: + raise TypeError( + f"MXFP6 wgrad fusion writes bf16 only, but main_grad is {main_grad.dtype}. " + "Set main_grads_dtype=bf16 or mxfp6_fused_wgrad_accum=False." + ) + + gemm_fp6_out_impl(g_col, g_col_scale, a_col, a_col_scale, main_grad, n, k, m, _GRAN_VALUE) + return torch.empty_like(weight) + + +def _claim_main_grad(*weights) -> None: + """Tell Megatron's DDP hook that these weights' gradients are already in main_grad. + + Called from the forward, which reads oddly, because the natural place -- right after + the backward's store -- is not available: dynamo refuses to trace a mutation of state + owned outside an autograd.Function ("HOP: Unsafe side effect"), and rather than fail it + breaks the graph around every MXFP6 linear. On the MBS=32 Flux 12B arm that fragmented + one compiled block into hundreds and cost 42.7 ms of eager elementwise work per 512 + images, against the ~16 ms of ``add_`` the fusion removes -- a net regression. From the + forward the same assignment is ordinary traced code, which dynamo records as a side + effect and replays. + + Setting it before the store rather than after is safe in one direction only, and this + is that direction: the flag is read by the DDP backward hook, which cannot run until + the backward has stored, and ``zero_grad_buffer`` clears it at the top of every step. + A forward with no backward (eval, ``no_grad``) leaves it set with no hook to read it. + + Megatron does the same thing for its own reason -- ``zero_grad_buffer`` skips the reset + under TE CUDA graphs precisely because the capture "no longer has the opportunity to + set it back to True". + """ + for weight in weights: + weight.grad_added_to_main_grad = True + + +class MXFP6LinearFunction(torch.autograd.Function): + """MXFP6 linear (Y = X @ W^T) with MX block-of-32 scaling along the contraction axis. + + Two modes via the ``backward_is_fp8`` bool primitive: + + - Pure MXFP6: forward and backward both quantize to MXFP6 and call gemm_fp6_impl. + - Hybrid: forward is MXFP6, backward re-quantizes the saved BF16 to tensorwise FP8. + + In the pure path the forward returns the column-direction blobs as extra outputs so + ``setup_context`` can save them; they are already uint8, so unlike MXFP4 there is no + dtype-view juggling needed to keep the autograd engine from trying to allocate zero + gradients in an unsupported dtype. + + The bias is an input rather than something the caller adds afterwards, so that the + backward owns the bias gradient and can take it from the packer's column sums instead + of paying for a separate reduction over ``grad_output``. This costs nothing in the + forward: the biased tensor is a saved activation for the QK-norm and RoPE backward, so + it is materialized either way, and Inductor was already doing the add as a standalone + in-place pass over the GEMM output rather than fusing it into anything. + """ + + @staticmethod + def forward( + input, + weight, + bias, + backward_is_fp8, + fp8_bwd_dtype, + fp8_gran_value, + fp8_backend_value, + fuse_wgrad_accum, + ): + out_dtype = input.dtype + orig_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + + m, k = input_2d.shape + n = weight.shape[0] + + a_row, a_row_scale, a_col, a_col_scale = _quantize_mxfp6_dual(input_2d) + b_row, b_row_scale, b_col, b_col_scale = _quantize_mxfp6_dual(weight) + + # Bias goes into the GEMM's store epilogue, where it is free: the epilogue is bound by + # its scatter store rather than by VALU, so the add hides completely. Handing it to the + # GEMM rather than adding afterwards deletes a whole pass over the output, worth 7.3 ms + # per step at MBS=32, and rounds once instead of twice. + # + # Passed unconditionally. Whether the installed aiter can actually fold it is Turbo's to + # answer -- it probes, and adds the separate pass itself when it cannot -- so there is + # nothing to gate here and no aiter version for this layer to know about. + output = gemm_fp6_impl( + a_row, + a_row_scale, + b_row, + b_row_scale, + m, + n, + k, + out_dtype, + _GRAN_VALUE, + bias, + ) + output = output.reshape(*orig_shape[:-1], output.shape[-1]) + + if backward_is_fp8: + return output, input_2d.view_as(input_2d), weight.view_as(weight) + return output, a_col, a_col_scale, b_col, b_col_scale + + @staticmethod + def setup_context(ctx, inputs, output): + ( + input, + weight, + bias, + backward_is_fp8, + fp8_bwd_dtype, + fp8_gran_value, + fp8_backend_value, + fuse_wgrad_accum, + ) = inputs + + ctx.backward_is_fp8 = backward_is_fp8 + ctx.fuse_wgrad_accum = fuse_wgrad_accum + ctx.out_dtype = input.dtype + ctx.orig_shape = input.shape + # The packed blobs carry no shape, so the logical dims have to be saved too. + ctx.m = input.numel() // input.shape[-1] + ctx.k = input.shape[-1] + ctx.n = weight.shape[0] + + if backward_is_fp8: + _, input_2d_saved, weight_saved = output + ctx.save_for_backward(input_2d_saved, weight_saved) + ctx.fp8_bwd_dtype = fp8_bwd_dtype + ctx.fp8_gran_value = fp8_gran_value + ctx.fp8_backend_value = fp8_backend_value + else: + _, a_col, a_col_scale, b_col, b_col_scale = output + # The weight rides along only when the backward has to reach through it to + # weight.main_grad; save_for_backward hands back the same Parameter object. + if fuse_wgrad_accum: + ctx.save_for_backward(a_col, a_col_scale, b_col, b_col_scale, weight) + else: + ctx.save_for_backward(a_col, a_col_scale, b_col, b_col_scale) + ctx.mark_non_differentiable(a_col, a_col_scale, b_col, b_col_scale) + + @staticmethod + def backward(ctx, grad_output, *_): + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + + grad_2d = grad_output.reshape(-1, grad_output.shape[-1]) + m, n, k = ctx.m, ctx.n, ctx.k + want_bias_grad = ctx.needs_input_grad[2] + + if ctx.backward_is_fp8: + input_2d, weight = ctx.saved_tensors + + grad_fp8, grad_scale_inv = _quantize_fp8_tw(grad_2d, ctx.fp8_bwd_dtype) + a_fp8, a_scale_inv = _quantize_fp8_tw(input_2d, ctx.fp8_bwd_dtype) + b_fp8, b_scale_inv = _quantize_fp8_tw(weight, ctx.fp8_bwd_dtype) + + grad_input = gemm_fp8_impl( + grad_fp8, + grad_scale_inv, + False, + b_fp8, + b_scale_inv, + False, + ctx.out_dtype, + False, + granularity=ctx.fp8_gran_value, + default_backend=ctx.fp8_backend_value, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + grad_weight = gemm_fp8_impl( + a_fp8, + a_scale_inv, + True, + grad_fp8, + grad_scale_inv, + False, + ctx.out_dtype, + True, + granularity=ctx.fp8_gran_value, + default_backend=ctx.fp8_backend_value, + ) + # No packer runs on this path, so the bias gradient pays for its own reduction. + grad_bias = grad_2d.sum(0).to(ctx.out_dtype) if want_bias_grad else None + else: + if ctx.fuse_wgrad_accum: + a_col, a_col_scale, b_col, b_col_scale, weight = ctx.saved_tensors + else: + a_col, a_col_scale, b_col, b_col_scale = ctx.saved_tensors + + # The bias gradient is a reduction over exactly the tensor the packer is + # already streaming, so it rides along as a side output. Identity because + # there is no activation to undo here, unlike the MLP's fc1. + g_row, g_row_scale, g_col, g_col_scale, b_partial = _quantize_mxfp6_fused_dual( + grad_2d, None, None, MXFP6_PROLOGUE_IDENTITY, want_bias_grad + ) + grad_bias = b_partial.sum(0).to(ctx.out_dtype) if want_bias_grad else None + + # grad_input[M, K] = grad[M, N] @ weight[N, K], contracting N. b_col is the + # weight packed along N, i.e. logically [K, N] contracting N. + grad_input = gemm_fp6_impl( + g_row, + g_row_scale, + b_col, + b_col_scale, + m, + k, + n, + ctx.out_dtype, + _GRAN_VALUE, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + # grad_weight[N, K] = grad.T[N, M] @ input[M, K], contracting M. + if ctx.fuse_wgrad_accum: + grad_weight = _wgrad_into_main_grad( + weight, g_col, g_col_scale, a_col, a_col_scale, n, k, m + ) + else: + grad_weight = gemm_fp6_impl( + g_col, + g_col_scale, + a_col, + a_col_scale, + n, + k, + m, + ctx.out_dtype, + _GRAN_VALUE, + ) + + return grad_input, grad_weight, grad_bias, None, None, None, None, None + + +def _resolve_wgrad_fusion(module, name: str) -> bool: + """Whether this linear may write its wgrad straight into ``main_grad``. + + Gated on the Primus-owned ``mxfp6_fused_wgrad_accum`` rather than Megatron's + ``gradient_accumulation_fusion``, because the latter also moves every plain linear -- + see the field's own comment for what that costs Flux's AdaLN projections. + + Two hard requirements, both raised rather than silently downgraded so that a + misconfiguration does not read as a performance result: + + - One microbatch per optimizer step. The A6W6 asm stores with beta=0, so a second + microbatch would overwrite the first one's gradient instead of adding to it. + Megatron's own fused path uses a beta=1 accumulate kernel and has no such limit. + - The pure-MXFP6 backward. The FP8-backward mode forms its wgrad with + ``gemm_fp8_impl``, which has no caller-provided-output variant. + """ + if not getattr(module.config, "mxfp6_fused_wgrad_accum", False): + return False + + if module._backward_is_fp8: + raise ValueError( + f"{name} cannot combine mxfp6_fused_wgrad_accum=True with " + "mxfp6_backward_precision='fp8': the FP8 backward has no out-variant GEMM to " + "write main_grad with. Use mxfp6_backward_precision='mxfp6' or turn the " + "fusion off." + ) + + try: + from megatron.core.num_microbatches_calculator import get_num_microbatches + + num_microbatches = get_num_microbatches() + except (ImportError, AttributeError): + # Calculator not up yet (unit tests build these modules standalone). The config + # is the only authority available, and the store stays correct as long as the + # caller honours the one-microbatch rule. + num_microbatches = None + + if num_microbatches is not None and num_microbatches > 1: + raise ValueError( + f"{name} requires mxfp6_fused_wgrad_accum=False when there is more than one " + f"microbatch per step (got {num_microbatches}). The A6W6 store has no beta=1 " + "accumulate epilogue, so it would overwrite earlier microbatches." + ) + + return True + + +def _init_mxfp6_linear(module) -> None: + """Shared __init__ tail for both MXFP6 parallel linears. + + MXFP4 duplicates this block between its column and row classes; there is no reason + for the MXFP6 copy to inherit the duplication. + """ + name = type(module).__name__ + + if module.config.tensor_model_parallel_size != 1: + raise ValueError( + f"{name} requires tensor_model_parallel_size=1. " + f"Got {module.config.tensor_model_parallel_size}." + ) + if module.gradient_accumulation_fusion: + # Megatron's fused path needs a beta=1 accumulate epilogue, which the A6W6 entry + # point has not got. The MXFP6 equivalent is mxfp6_fused_wgrad_accum, which does + # not disturb the plain linears. + raise ValueError( + f"{name} requires gradient_accumulation_fusion=False. To fuse the MXFP6 " + "weight gradient into main_grad, set mxfp6_fused_wgrad_accum=True instead." + ) + if module.sequence_parallel: + raise ValueError(f"{name} requires sequence_parallel=False.") + + supported, reason = check_mxfp6_support() + if not supported: + raise RuntimeError(f"MXFP6 not supported on this device: {reason}") + + module._backward_is_fp8 = getattr(module.config, "mxfp6_backward_precision", "mxfp6") == "fp8" + module._fuse_wgrad_accum = _resolve_wgrad_fusion(module, name) + + if module._backward_is_fp8: + from primus_turbo.pytorch.core.low_precision import float8_e5m2 + + module._fp8_bwd_dtype = float8_e5m2 + module._fp8_gran_value = ScalingGranularity.TENSORWISE.value + module._fp8_backend_value = BackendType.HIPBLASLT.value + else: + module._fp8_bwd_dtype = None + module._fp8_gran_value = 0 + module._fp8_backend_value = 0 + + +def _mxfp6_forward_impl(module, input, weight, **kwargs): + bias = kwargs.get("bias", None) + + if module._fuse_wgrad_accum: + _claim_main_grad(weight) + + # The bias goes through the Function rather than being added here, so that the + # backward can lift the bias gradient out of the packer's column sums. See + # MXFP6LinearFunction's docstring for why this is free in the forward. + result = MXFP6LinearFunction.apply( + input, + weight, + bias, + module._backward_is_fp8, + module._fp8_bwd_dtype, + module._fp8_gran_value, + module._fp8_backend_value, + module._fuse_wgrad_accum, + ) + return result[0] + + +class MXFP6ColumnParallelLinear(ColumnParallelLinear): + """ColumnParallelLinear with per-module MXFP6. torch.compile friendly. + + Requires: tensor_model_parallel_size=1, gradient_accumulation_fusion=False, + sequence_parallel=False. ``mxfp6_fused_wgrad_accum=True`` is supported at one + microbatch per step; see ``_resolve_wgrad_fusion``. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + _init_mxfp6_linear(self) + + def _forward_impl(self, input, weight, *args, **kwargs): + return _mxfp6_forward_impl(self, input, weight, **kwargs) + + +class MXFP6RowParallelLinear(RowParallelLinear): + """RowParallelLinear with per-module MXFP6. torch.compile friendly. + + Requires: tensor_model_parallel_size=1, gradient_accumulation_fusion=False, + sequence_parallel=False. ``mxfp6_fused_wgrad_accum=True`` is supported at one + microbatch per step; see ``_resolve_wgrad_fusion``. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + _init_mxfp6_linear(self) + + def _forward_impl(self, input, weight, *args, **kwargs): + return _mxfp6_forward_impl(self, input, weight, **kwargs) + + +# --------------------------------------------------------------------------- +# Whole-MLP fusion. +# +# Splitting the MLP into two independent autograd Functions forces the activation to +# exist: fc2's forward has to receive a real tensor, and fc1's backward has to receive +# one. Owning fc1 -> epilogue -> fc2 in a single Function is what lets the packer take +# the epilogue as a prologue instead, in both directions -- the activation is packed +# straight out of LDS and the pre-activation gradient is never assembled at all. +# +# What that removes, per Flux 12B step at the profiled shapes: the bias-add + GELU kernel's +# read and write and the packer's read back of it in the forward, and the same round-trip +# plus the bias gradient's own reduction pass in the backward. Measured on 8x MI355X at +# micro_batch_size 64, against this same module with the fusion switched off: 75.3 ms/step of +# epilogue and reduction kernels go away, 35.5 ms/step of added prologue cost inside the +# packer replaces them, and the step's GPU busy time falls 856.5 -> 810.9 ms, 5.0% off wall +# clock (863.6 -> 820.3 ms/step). The step is GPU bound at 96.8% busy, so that lands as +# throughput: 74.1 -> 78.0 images/s/GPU. The pre-activation y1 is still saved, but it was +# already being saved for the activation's own backward, so peak allocated memory only grows +# by the column-sum buffer, 4 MB of 244 GB. Reserved memory grows more, 249.0 -> 250.2 GB, +# because the freed epilogue temporaries leave differently shaped holes in the caching +# allocator; that is the number the driver reports, so it is what a memory ceiling will see. +# +# The backward is where the win is, ~0.40 ms per call against ~0.11 for the forward, and the +# reason is worth knowing before trying to improve this. The packer is bandwidth bound at +# 3.8 TB/s without a prologue, so fusing work into it only pays while it stays that way. The +# forward prologue removes a 0.26 ms kernel and adds 0.14 ms of arithmetic to a 0.38 ms pack; +# the backward removes two kernels totalling 0.73 ms and its extra read of the incoming +# gradient is traffic it would have done anyway. An early version of the prologue used a libm +# tanh and a per-element bounds branch and cost 0.38 ms of arithmetic instead of 0.14, which +# made the forward a net regression and cost most of the win. +# --------------------------------------------------------------------------- + + +class MXFP6MLPFunction(torch.autograd.Function): + """fc1 GEMM, bias+GELU, fc2 GEMM as one op, with the activation never in HBM. + + Follows ``MXFP6LinearFunction``'s conventions: the column-direction blobs leave as + extra outputs so ``setup_context`` can save them and mark them non-differentiable, and + all non-tensor state lands on ``ctx`` as primitives so ``torch.compile`` traces cleanly. + + Only the pure-MXFP6 backward is supported. The FP8-backward mode re-quantizes saved + BF16 activations, which would put the activation back in HBM and defeat the point; + ``_fused_mlp_unusable_reason`` rejects that configuration before we get here. + """ + + @staticmethod + def forward(hidden_states, w1, b1, w2, fuse_wgrad_accum): + out_dtype = hidden_states.dtype + orig_shape = hidden_states.shape + x = hidden_states.reshape(-1, orig_shape[-1]) + + m, k = x.shape + f = w1.shape[0] + h = w2.shape[0] + + x_row, x_row_s, x_col, x_col_s = _quantize_mxfp6_dual(x) + w1_row, w1_row_s, w1_col, w1_col_s = _quantize_mxfp6_dual(w1) + + # Pre-activation. Saved for backward, where the epilogue is recomputed from it + # rather than its output being stashed -- the same bytes are held either way. + y1 = gemm_fp6_impl(x_row, x_row_s, w1_row, w1_row_s, m, f, k, out_dtype, _GRAN_VALUE) + + # gelu(y1 + b1), packed in both directions without ever being written out. + a_row, a_row_s, a_col, a_col_s, _ = _quantize_mxfp6_fused_dual( + y1, None, b1, MXFP6_PROLOGUE_BIAS_GELU, False + ) + w2_row, w2_row_s, w2_col, w2_col_s = _quantize_mxfp6_dual(w2) + + output = gemm_fp6_impl(a_row, a_row_s, w2_row, w2_row_s, m, h, f, out_dtype, _GRAN_VALUE) + output = output.reshape(*orig_shape[:-1], h) + + return output, y1, x_col, x_col_s, a_col, a_col_s, w1_col, w1_col_s, w2_col, w2_col_s + + @staticmethod + def setup_context(ctx, inputs, output): + hidden_states, w1, b1, w2, fuse_wgrad_accum = inputs + + ctx.fuse_wgrad_accum = fuse_wgrad_accum + ctx.out_dtype = hidden_states.dtype + ctx.orig_shape = hidden_states.shape + # The packed blobs carry no shape, so the logical dims have to be saved too. + ctx.m = hidden_states.numel() // hidden_states.shape[-1] + ctx.k = hidden_states.shape[-1] + ctx.f = w1.shape[0] + ctx.h = w2.shape[0] + + _, y1, x_col, x_col_s, a_col, a_col_s, w1_col, w1_col_s, w2_col, w2_col_s = output + blobs = (x_col, x_col_s, a_col, a_col_s, w1_col, w1_col_s, w2_col, w2_col_s) + # b1 is a leaf parameter, so saving it costs nothing, and the backward needs it to + # rebuild the pre-activation for the GELU derivative. w1 and w2 ride along only + # when the backward has to reach through them to their main_grad buffers. + extra = (w1, w2) if fuse_wgrad_accum else () + ctx.save_for_backward(y1, b1, *blobs, *extra) + ctx.mark_non_differentiable(*blobs) + + @staticmethod + def backward(ctx, grad_output, *_): + ( + y1, + b1, + x_col, + x_col_s, + a_col, + a_col_s, + w1_col, + w1_col_s, + w2_col, + w2_col_s, + *fused_weights, + ) = ctx.saved_tensors + m, k, f, h = ctx.m, ctx.k, ctx.f, ctx.h + out_dtype = ctx.out_dtype + + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + g2 = grad_output.reshape(-1, h) + + g2_row, g2_row_s, g2_col, g2_col_s = _quantize_mxfp6_dual(g2) + + # fc2 dgrad: [m, f] = g2[m, h] @ w2[h, f], contracting h. + grad_a = gemm_fp6_impl(g2_row, g2_row_s, w2_col, w2_col_s, m, f, h, out_dtype, _GRAN_VALUE) + # fc2 wgrad: [h, f] = g2.T[h, m] @ a[m, f], contracting m. + if ctx.fuse_wgrad_accum: + grad_w2 = _wgrad_into_main_grad( + fused_weights[1], g2_col, g2_col_s, a_col, a_col_s, h, f, m + ) + else: + grad_w2 = gemm_fp6_impl(g2_col, g2_col_s, a_col, a_col_s, h, f, m, out_dtype, _GRAN_VALUE) + + # The GELU derivative is applied while staging, so grad_y1 is never assembled. Its + # column sums come back as a side output because the bias gradient is a reduction + # over exactly the tensor that no longer exists. + want_bias_grad = ctx.needs_input_grad[2] + g1_row, g1_row_s, g1_col, g1_col_s, b1_partial = _quantize_mxfp6_fused_dual( + y1, grad_a, b1, MXFP6_PROLOGUE_BIAS_GELU_BACKWARD, want_bias_grad + ) + + # fc1 dgrad: [m, k] = grad_y1[m, f] @ w1[f, k], contracting f. + grad_x = gemm_fp6_impl(g1_row, g1_row_s, w1_col, w1_col_s, m, k, f, out_dtype, _GRAN_VALUE) + grad_x = grad_x.reshape(ctx.orig_shape) + # fc1 wgrad: [f, k] = grad_y1.T[f, m] @ x[m, k], contracting m. + if ctx.fuse_wgrad_accum: + grad_w1 = _wgrad_into_main_grad( + fused_weights[0], g1_col, g1_col_s, x_col, x_col_s, f, k, m + ) + else: + grad_w1 = gemm_fp6_impl(g1_col, g1_col_s, x_col, x_col_s, f, k, m, out_dtype, _GRAN_VALUE) + + grad_b1 = b1_partial.sum(0).to(out_dtype) if want_bias_grad else None + + return grad_x, grad_w1, grad_b1, grad_w2, None + + +def _is_tanh_gelu(fn) -> bool: + """Whether ``fn`` is exactly ``F.gelu(approximate="tanh")``. + + The packer's prologue implements that function and only that one, to within a rounding of + the tanh. This check is not paranoia: ``FluxConfig``'s default activation is + ``openai_gelu_no_jit``, the same mathematical function written as + ``beta * x * (1 + kappa * x^2)`` rather than ``beta * (x + kappa * x^3)``. Those disagree + well above the prologue's own rounding, and the YAML key that selects between them + (``activation_func: openai_gelu`` maps to the fused ATen one) + makes it easy to land on the other branch without noticing. + """ + if isinstance(fn, functools.partial): + return fn.func is F.gelu and fn.keywords.get("approximate") == "tanh" + return False + + +def _fused_mlp_unusable_reason(mlp) -> str: + """Why the fused MLP cannot be used for this module, or ``""`` if it can. + + Every branch corresponds to a path in ``MLP.forward`` that the fused Function does not + reproduce. Returning a reason rather than silently deferring keeps a misconfiguration + from looking like a performance result. + """ + config = mlp.config + + if not hasattr(torch.ops.primus_turbo, "quantize_mxfp6_fused_dual_impl"): + return "this Primus-Turbo build has no fused MXFP6 prologue packer" + if config.gated_linear_unit: + return "gated_linear_unit splits the fc1 output, which the prologue does not do" + if getattr(config, "bias_activation_fusion", False): + return "bias_activation_fusion routes the epilogue through Megatron's own fused kernel" + if getattr(config, "use_te_activation_func", False): + return "use_te_activation_func replaces the activation with a TE module" + if not _is_tanh_gelu(mlp.activation_func): + name = getattr(mlp.activation_func, "__name__", type(mlp.activation_func).__name__) + return ( + "the fused prologue implements F.gelu(approximate='tanh') only, but the " + f"activation is {name!r}" + ) + + for name in ("linear_fc1", "linear_fc2"): + linear = getattr(mlp, name) + if not isinstance(linear, (MXFP6ColumnParallelLinear, MXFP6RowParallelLinear)): + return f"{name} is {type(linear).__name__}, not an MXFP6 linear" + if getattr(linear, "_backward_is_fp8", False): + return ( + "mxfp6_backward_precision='fp8' saves the activation for backward " + "requantization, which the fusion removes" + ) + if not linear.skip_bias_add: + return f"{name} adds its own bias, so the epilogue is not the MLP's to fuse" + + return "" + + +def _fused_mlp_mode() -> str: + mode = os.environ.get("PRIMUS_MXFP6_FUSED_MLP", "").strip().lower() or "auto" + assert mode in ( + "auto", + "on", + "off", + ), f"PRIMUS_MXFP6_FUSED_MLP must be auto, on or off, got {mode!r}" + return mode + + +class MXFP6FusedMLP(MLP): + """MLP whose bias-add + GELU is folded into the MXFP6 packer, in both directions. + + Drop-in for ``MLP``: same submodules, same ``(output, output_bias)`` return, same + parameters and state dict. Only ``forward`` differs, and it defers to ``MLP.forward`` + for anything the fused path does not cover, so a configuration it cannot handle is + slow rather than wrong. + + ``PRIMUS_MXFP6_FUSED_MLP=off`` forces the stock path for A/B comparison; ``on`` makes + an unusable configuration an error instead of a silent fallback. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + mode = _fused_mlp_mode() + reason = "disabled by environment" if mode == "off" else _fused_mlp_unusable_reason(self) + self._fused_epilogue = reason == "" + + if not self._fused_epilogue and mode == "on": + raise RuntimeError(f"PRIMUS_MXFP6_FUSED_MLP=on but the fused MLP is unusable: {reason}") + if not self._fused_epilogue and mode == "auto": + warnings.warn( + f"MXFP6 fused MLP epilogue disabled, falling back to the stock MLP: {reason}", + stacklevel=2, + ) + + def forward(self, hidden_states, per_token_scale=None): + # per_token_scale scales the activation after the epilogue, which is exactly the + # tensor the fusion refuses to materialise. Only MoE experts pass it. + if not self._fused_epilogue or per_token_scale is not None: + return super().forward(hidden_states, per_token_scale=per_token_scale) + + # Both linears resolve the flag from the same config field, so fc1's answer covers + # fc2's weight too. + fuse_wgrad_accum = self.linear_fc1._fuse_wgrad_accum + if fuse_wgrad_accum: + _claim_main_grad(self.linear_fc1.weight, self.linear_fc2.weight) + + output = MXFP6MLPFunction.apply( + hidden_states, + self.linear_fc1.weight, + self.linear_fc1.bias, + self.linear_fc2.weight, + fuse_wgrad_accum, + )[0] + + # fc2 is built with skip_bias_add=True, so MLP's contract is to hand its bias back + # unadded for the caller to fuse into a residual. + return output, self.linear_fc2.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..952d779e2 100644 --- a/primus/backends/megatron/core/models/diffusion/common/config.py +++ b/primus/backends/megatron/core/models/diffusion/common/config.py @@ -14,6 +14,13 @@ from megatron.core.enums import Fp8Recipe from megatron.core.transformer.transformer_config import TransformerConfig +# Accepted values for `fp6`. Only one 6-bit format exists on gfx950 (E2M3, via AITER's +# A6W6 kernels), but keeping this a tuple mirrors how `fp4` names its format. +MXFP6_FORMATS = ("mxfp6",) + +# Accepted values for `mxfp6_backward_precision`. +MXFP6_BACKWARD_PRECISIONS = ("mxfp6", "fp8") + @dataclass class BaseDiffusionConfig(TransformerConfig): @@ -34,6 +41,9 @@ class BaseDiffusionConfig(TransformerConfig): fp8_reduce_amax: Whether to allreduce amax across ranks (default: False) mxfp4_backward_precision: MXFP4 backward precision, 'mxfp4' or 'fp8' (default: 'mxfp4') mxfp4_gradient_stochastic_rounding: Stochastic rounding on gradients (default: False) + fp6: Set to 'mxfp6' to run linears in MXFP6 (E2M3). None disables (default: None) + mxfp6_backward_precision: MXFP6 backward precision, 'mxfp6' or 'fp8' (default: 'mxfp6') + mxfp6_fused_wgrad_accum: MXFP6 wgrad writes weight.main_grad in place (default: False) sensitive_layers_enabled: Enable sensitive layer configuration (default: False) sensitive_layers_start: Number of sensitive layers at start (default: 0) sensitive_layers_end: Number of sensitive layers at end (default: 0) @@ -77,6 +87,30 @@ class BaseDiffusionConfig(TransformerConfig): # Stochastic rounding on MXFP4 gradients (paper Section 4.4) mxfp4_gradient_stochastic_rounding: bool = False + # MXFP6 (E2M3). Declared here rather than on TransformerConfig because Megatron has + # no notion of a 6-bit format, so unlike `fp4` this is Primus-owned -- which also + # means Megatron's "fp4 and fp8 cannot coexist" validation never sees it and the + # cross-checks below are the only place those combinations can be rejected. + fp6: Optional[str] = None + + # MXFP6 backward precision: "mxfp6" (pure) or "fp8" (hybrid), mirroring + # mxfp4_backward_precision. + mxfp6_backward_precision: str = "mxfp6" + + # Have the MXFP6 wgrad GEMM write weight.main_grad itself, replacing the elementwise + # add Megatron's DDP hook would otherwise run over every gradient. + # + # Deliberately not Megatron's `gradient_accumulation_fusion`: that flag is read by + # every plain linear too, and switching it on routes Flux's 76 AdaLN projections + # through `wgrad_gemm_accum_fp16`, which at their M=32 shapes is slower than the + # separate add it replaces -- measured at +11% step time on 8x MI355X, swamping the + # saving on the MXFP6 linears. This field moves only the MXFP6 ones. + # + # The A6W6 store has no beta=1 accumulate epilogue, so it overwrites main_grad and is + # only valid at one microbatch per optimizer step. Enforced per module, not here, + # because the microbatch count is not known at config time. + mxfp6_fused_wgrad_accum: bool = False + # Sensitive layer configuration (clean naming, maps to Megatron internals) sensitive_layers_enabled: bool = False sensitive_layers_start: int = 0 @@ -116,6 +150,55 @@ 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 + # MXFP6 cross-checks. Since `fp6` is Primus-owned, nothing downstream would + # notice a nonsense combination -- the layer spec would just pick one provider + # and silently ignore the other request. + if self.fp6 is not None: + if self.fp6 not in MXFP6_FORMATS: + raise ValueError(f"Unknown fp6 '{self.fp6}'. Choose from: {list(MXFP6_FORMATS)}.") + if getattr(self, "fp4", None) is not None: + raise ValueError( + f"fp4 ('{self.fp4}') and fp6 ('{self.fp6}') cannot both be set: the " + "layer spec selects one linear implementation per model." + ) + if self.fp8 is not None: + raise ValueError( + f"fp6 ('{self.fp6}') and fp8 ('{self.fp8}') cannot both be set, " + "mirroring Megatron's fp4/fp8 exclusion. For an MXFP6 forward with " + "an FP8 backward use mxfp6_backward_precision='fp8' instead." + ) + # Read through getattr because the MXFP4 -> FP8 switch is a separate change + # that may not be present; the check has to hold once both are, without + # making this branch depend on it. + switch_iter = int(getattr(self, "mxfp4_to_fp8_switch_iter", 0) or 0) + if switch_iter > 0: + # The switch patch walks the model for MXFP4ColumnParallelLinear / + # MXFP4RowParallelLinear and *skips* anything else, so with fp6 it would + # build a plan over zero layers and quietly never switch. Reject the + # combination rather than extend the patch: its prewarm and ramp logic + # are written around MXFP4 and there is no verified MXFP6 equivalent. + raise ValueError( + f"fp6 ('{self.fp6}') cannot be combined with mxfp4_to_fp8_switch_iter=" + f"{switch_iter}. The switch only converts MXFP4 " + "linears, of which an MXFP6 model has none. Use " + "mxfp6_backward_precision='fp8' for a hybrid MXFP6 run." + ) + if self.mxfp6_backward_precision not in MXFP6_BACKWARD_PRECISIONS: + raise ValueError( + f"Unknown mxfp6_backward_precision '{self.mxfp6_backward_precision}'. " + f"Choose from: {list(MXFP6_BACKWARD_PRECISIONS)}." + ) + if self.mxfp6_backward_precision != "mxfp6" and self.fp6 is None: + raise ValueError( + f"mxfp6_backward_precision='{self.mxfp6_backward_precision}' requires fp6 " + "to be set (e.g. fp6: mxfp6); with no MXFP6 linears it has no effect." + ) + if self.mxfp6_fused_wgrad_accum and self.fp6 is None: + raise ValueError( + "mxfp6_fused_wgrad_accum=True requires fp6 to be set (e.g. fp6: mxfp6); " + "with no MXFP6 linears it has no effect." + ) + 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..3259f9f25 100644 --- a/primus/backends/megatron/core/models/diffusion/flux/attention.py +++ b/primus/backends/megatron/core/models/diffusion/flux/attention.py @@ -292,12 +292,24 @@ 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 `.to(value.dtype)` is required, not defensive. Under torch.compile with + # emulate_precision_casts off, inductor may leave the norm's output in its fp32 + # accumulation dtype, and only Q and K go through a norm. V then still carries + # the intended dtype, so it is the reference. Without this, attention receives + # fp32 Q/K against bf16 V and Turbo -- whose dense flash-attention backends all + # require fp16/bf16 -- rejects the call as "No compatible backend found for + # FlashAttnDenseDispatcher", naming shapes but never mentioning dtype. Casting + # here rather than at the attention call keeps it inside the compiled region, + # where it fuses into the norm's epilogue instead of costing an extra pass over + # Q and K. This matches the reference implementation, whose QKNorm.forward in + # backends/diffusion/models/flux/layers.py likewise returns `q.to(v), k.to(v)`. 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 +334,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 @@ -513,6 +525,24 @@ def __init__( tp_comm_buffer_name="proj", ) + def get_query_key_value_tensors(self, *args, **kwargs): + """ + Derive Q, K, V, realigning Q/K onto V's dtype. + + Megatron's implementation returns the QK-norm's output as-is, which under + torch.compile can be its fp32 accumulation dtype. See the note in + JointSelfAttention.get_query_key_value_tensors for why that breaks + attention and why V is the reference. The joint blocks project QKV + themselves and cast inline; the single blocks reuse Megatron's projection, + so the cast goes here, still inside the compiled region. + """ + out = super().get_query_key_value_tensors(*args, **kwargs) + if len(out) < 3: + # split_qkv=False: (mixed_qkv, split_arg_list), no norm applied yet. + return out + query, key, value, *rest = out + return (query.to(value.dtype), key.to(value.dtype), value, *rest) + def forward( self, hidden_states: Tensor, diff --git a/primus/backends/megatron/core/models/diffusion/flux/layer_spec.py b/primus/backends/megatron/core/models/diffusion/flux/layer_spec.py index 9a1932b09..901655a04 100644 --- a/primus/backends/megatron/core/models/diffusion/flux/layer_spec.py +++ b/primus/backends/megatron/core/models/diffusion/flux/layer_spec.py @@ -81,6 +81,15 @@ except ImportError: PrimusTurboMXFP4LocalSpecProvider = None +# MXFP6 in its own guard for the same reason, and additionally because it needs an +# aiter carrying the A6W6 kernels (PR #4859), which is newer than the pinned release. +try: + from primus.backends.megatron.core.extensions.primus_turbo_local_spec import ( + PrimusTurboMXFP6LocalSpecProvider, + ) +except ImportError: + PrimusTurboMXFP6LocalSpecProvider = None + class MMDiTLayer(TransformerLayer): """ @@ -400,6 +409,18 @@ def __call__(self, *args, **kwargs): return super(MegatronModule, self).__call__(*args, **kwargs) +def _mlp_module_for(backend: BackendSpecProvider) -> type: + """Which MLP class this backend wants. + + A backend that can fold the MLP's bias-add + activation into its own quantizer says so + by exposing ``mlp_module``; the MXFP6 local spec is the only one that does today. + Everything else, including TE, gets Megatron's ``MLP``. Duck-typed rather than added to + ``BackendSpecProvider`` so this stays out of the vendored Megatron tree. + """ + getter = getattr(backend, "mlp_module", None) + return getter() if getter is not None else MLP + + def get_flux_single_transformer_spec_for_backend( backend: BackendSpecProvider, ) -> ModuleSpec: @@ -430,7 +451,7 @@ def get_flux_single_transformer_spec_for_backend( ), ), mlp=ModuleSpec( - module=MLP, + module=_mlp_module_for(backend), submodules=MLPSubmodules( linear_fc1=backend.column_parallel_linear(), linear_fc2=backend.row_parallel_linear(), @@ -473,7 +494,7 @@ def get_flux_double_transformer_spec_for_backend( ), ), mlp=ModuleSpec( - module=MLP, + module=_mlp_module_for(backend), submodules=MLPSubmodules( linear_fc1=backend.column_parallel_linear(), linear_fc2=backend.row_parallel_linear(), @@ -536,6 +557,25 @@ def get_flux_layer_spec( backend = PrimusTurboMXFP4LocalSpecProvider() # Resolve sensitive layer backend + sensitive_precision = getattr(config, "sensitive_layer_precision", "bf16") + if sensitive_precision == "tw_fp8": + sensitive_backend = PrimusTurboFloat8LocalSpecProvider() + elif sensitive_precision == "bf16": + sensitive_backend = PrimusTurboLocalSpecProvider() + elif getattr(config, "fp6", None) is not None: + # Unlike fp4, an unavailable MXFP6 provider is fatal rather than a + # silent fall-through to bf16: fp6 is Primus-owned and set explicitly, + # so quietly training in a different precision than asked for would be + # worse than failing. The usual cause is an aiter without A6W6. + if PrimusTurboMXFP6LocalSpecProvider is None: + raise RuntimeError( + f"config.fp6={config.fp6!r} was requested but " + "PrimusTurboMXFP6LocalSpecProvider could not be imported. MXFP6 " + "needs Primus-Turbo with gemm_fp6 and an aiter carrying the A6W6 " + "kernels from https://github.com/ROCm/aiter/pull/4859." + ) + backend = PrimusTurboMXFP6LocalSpecProvider() + sensitive_precision = getattr(config, "sensitive_layer_precision", "bf16") if sensitive_precision == "tw_fp8": sensitive_backend = PrimusTurboFloat8LocalSpecProvider() diff --git a/primus/backends/megatron/data/energon_dataset_provider.py b/primus/backends/megatron/data/energon_dataset_provider.py index cf3bdf704..b512a5fc6 100644 --- a/primus/backends/megatron/data/energon_dataset_provider.py +++ b/primus/backends/megatron/data/energon_dataset_provider.py @@ -27,11 +27,13 @@ get_savable_loader, get_train_dataset, get_val_datasets, + log_exception, ) from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper from primus.backends.megatron.data.dataset_provider import DatasetProvider from primus.backends.megatron.training.eval_budget import ( + get_eval_micro_batch_size, EvalCoverageError, assert_mlperf_timestep_source, assert_val_worker_divisibility, @@ -128,7 +130,6 @@ def create_dataloaders( virtual_epoch_length=getattr(args, "virtual_epoch_length", 1_000_000_000), max_samples_per_sequence=getattr(args, "max_samples_per_sequence", 100), shuffle_buffer_size=getattr(args, "shuffle_buffer_size", None), - handler=lambda *args: None, # Error handler (print errors but continue) ) # Wrap in savable loader for checkpointing support @@ -177,10 +178,9 @@ def create_dataloaders( log_rank_0("Creating validation dataloaders...") val_datasets = get_val_datasets( data_path, - batch_size=args.micro_batch_size, + batch_size=get_eval_micro_batch_size(args), task_encoder=task_encoder, worker_config=val_worker_config, - handler=lambda *args: None, ) # Limit validation datasets to eval_iters * num_microbatches @@ -276,11 +276,21 @@ def _create_worker_config(self, args, num_workers: Optional[int] = None) -> Work if num_workers is None: num_workers = getattr(args, "num_workers", 4) + # Error handling belongs here, not on get_train_dataset/get_val_datasets: those + # pop a handler kwarg and warn that they are ignoring it. log_exception prints the + # traceback, the shard provenance and the offending sample, then continues, so one + # bad sample cannot end a run. It is also the current default, so pinning it + # changes nothing today and only guards against the default being changed. + # + # Whatever goes here must be picklable by name: WorkerConfig is sent to the + # dataloader workers, which run under forkserver rather than fork, so a lambda + # would fail at loader construction. return WorkerConfig( rank=rank, world_size=world_size, num_workers=num_workers, data_parallel_group=data_parallel_group, + global_error_handler=log_exception, ) def _get_data_path(self, args) -> str: diff --git a/primus/backends/megatron/data/energon_picklable.py b/primus/backends/megatron/data/energon_picklable.py new file mode 100644 index 000000000..55e83480c --- /dev/null +++ b/primus/backends/megatron/data/energon_picklable.py @@ -0,0 +1,98 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Picklable stand-ins for the lambdas and closures Energon builds inline. + +Energon 7.3.2 builds its part filters and sample loaders as lambdas in class +bodies and closures inside ``__init__``. That costs nothing under ``fork``, +which shares memory and never serialises the dataset, but ``forkserver`` and +``spawn`` have to pickle it, and pickle resolves functions by qualified name -- +which finds neither a lambda in a class body nor a closure over locals. + +Leaving ``fork`` is not optional once GPU-Direct RDMA is in use: a process that +has registered GPU memory through dmabuf cannot fork, because the workers +inherit device mappings that do not survive the fork and segfault immediately. + +These deliberately live outside ``primus.backends.megatron.patches``. They are +referenced by name from inside pickled datasets, so a dataloader worker imports +this module while unpickling one; importing the patches package instead would +eagerly pull in every ``*_patches`` module for nothing. + +``energon_fork_patches`` is what installs them. +""" + +from typing import Any + +from megatron.energon.flavors.webdataset.field_access import field_access + + +def accept_all_parts(_: str) -> bool: + """Replaces ``lambda _: True``, CrudeWebdataset's default ``part_filter``.""" + return True + + +def identity_sample_loader(sample): + """Replaces ``lambda sample: sample``, the loader CrudeWebdataset passes up.""" + return sample + + +class PartInSet: + """Replaces ``lambda part: part in parts``.""" + + def __init__(self, parts): + self.parts = parts + + def __call__(self, part): + return part in self.parts + + +class FieldMapSampleLoader: + """Replaces the ``field_map`` sample loader closure.""" + + def __init__(self, fields): + self.fields = fields + + def __call__(self, sample): + return {k: field_access(sample, v) for k, v in self.fields.items()} + + +class KeyedSampleLoader: + """Replaces the outer sample loader closure. + + Holds the factory rather than a copy of its subflavors. The closure read + ``self.subflavors`` at call time and the attribute is assigned after this + wrapper is built, so copying the value here would capture it empty. + """ + + def __init__(self, inner, factory): + self.inner = inner + self.factory = factory + + def __call__(self, sample): + return { + "__key__": sample["__key__"], + **self.inner(sample), + "__restore_key__": sample["__restore_key__"], + "__subflavors__": self.factory.subflavors, + "__sources__": sample["__sources__"], + } + + +def worker_config_getstate(self) -> Any: + """``WorkerConfig.__getstate__`` that drops what cannot be pickled. + + A dataloader worker runs no collectives and does not write another + process's debug file, so neither member is needed worker-side, and + ``global_rank()`` already falls back to the local rank when no group is set. + """ + state = object.__getstate__(self) + drop = {"data_parallel_group": None, "_worker_debug_file": None} + # A slotted dataclass serialises as (dict_state, slot_state). + if isinstance(state, tuple) and len(state) == 2 and isinstance(state[1], dict): + return (state[0], {**state[1], **drop}) + if isinstance(state, dict): + return {**state, **drop} + return state diff --git a/primus/backends/megatron/diffusion_trainer.py b/primus/backends/megatron/diffusion_trainer.py index ffa71efc1..b59e687f0 100644 --- a/primus/backends/megatron/diffusion_trainer.py +++ b/primus/backends/megatron/diffusion_trainer.py @@ -296,7 +296,17 @@ def val_loss_func(output_tensor, non_loss_data=False): sample_count = torch.tensor( loss_per_sample.numel(), dtype=loss_sum.dtype, device=loss_sum.device ) - return loss_sum, {"loss": (loss_sum.detach(), sample_count.detach())} + # CLONE, NOT JUST DETACH. Megatron's forward_step treats the first + # element of this pair as the tensor to backpropagate and rescales it + # IN PLACE before storing the dict below -- `output_tensor *= + # cp_group_size`, then `output_tensor /= num_microbatches`. A detached + # view shares that storage, so the reported loss gets rescaled with it + # and what reaches the caller is the true loss divided by the number of + # microbatches. That is invisible at one microbatch per rank per step + # and halves the reported validation loss at two, which under + # mlperf_mode trips the convergence gate at roughly half the samples it + # should. The training path below clones for the same reason. + return loss_sum, {"loss": (loss_sum.detach().clone(), sample_count.detach())} return noise_pred, val_loss_func diff --git a/primus/backends/megatron/flux_pretrain_trainer.py b/primus/backends/megatron/flux_pretrain_trainer.py index 1b41ba896..a6c127ecb 100644 --- a/primus/backends/megatron/flux_pretrain_trainer.py +++ b/primus/backends/megatron/flux_pretrain_trainer.py @@ -204,6 +204,7 @@ def _init_cfg_dropout(self, params): f"CFG dropout: loaded real empty encodings from {encodings_dir}, " f"t5={self.empty_t5_encodings.shape}, clip={self.empty_clip_encodings.shape}" ) + self._make_empty_encodings_resident(params) elif getattr(params, "mock_data", False): image_size = getattr(getattr(params, "mock_dataset", None), "params", None) image_size = getattr(image_size, "image_size", 256) if image_size is not None else 256 @@ -214,6 +215,7 @@ def _init_cfg_dropout(self, params): self.empty_t5_encodings = torch.randn(t5_seq_len, 1, context_dim) self.empty_clip_encodings = torch.randn(vec_in_dim) log_rank_0("CFG dropout: using torch.randn() empty encodings (mock_data mode)") + self._make_empty_encodings_resident(params) else: data_path = getattr(params, "data_path", "") if isinstance(data_path, list): @@ -232,6 +234,32 @@ def _init_cfg_dropout(self, params): log_rank_0(f"CFG dropout prob: {self.cfg_dropout_prob}") + def _make_empty_encodings_resident(self, params): + """Move the empty encodings onto the device in their final dtype. + + These are plain trainer attributes rather than registered buffers, so nothing + else ever moves them, and the CFG dropout branch in forward_step would + otherwise copy them host-to-device on every step. Under CUDA graph capture that + copy is fatal: an unpinned host-to-device copy inside the captured region + raises "Cannot copy between CPU and CUDA tensors during CUDA graph capture". + Landing them here in the compute dtype makes the `.to(...)` in forward_step an + identity. + """ + if getattr(params, "bf16", False): + dtype = torch.bfloat16 + elif getattr(params, "fp16", False): + dtype = torch.float16 + else: + dtype = None + + device = torch.cuda.current_device() + self.empty_t5_encodings = self.empty_t5_encodings.to(device=device, dtype=dtype) + self.empty_clip_encodings = self.empty_clip_encodings.to(device=device, dtype=dtype) + log_rank_0( + f"CFG dropout: empty encodings resident on {self.empty_t5_encodings.device} " + f"as {self.empty_t5_encodings.dtype}" + ) + @staticmethod def _discover_empty_encodings(params) -> "str | None": """Return the first valid empty_encodings directory, or None.""" @@ -547,6 +575,17 @@ def _build_flux_config_from_yaml(self): } ) + # FP6/MXFP6 settings. No recipe field to validate: E2M3 with per-1x32 block + # scales along the contraction axis is the only MXFP6 configuration the A6W6 + # kernels implement, so there is nothing for the YAML to choose between. + config_params.update( + { + "fp6": getattr(params, "fp6", None), + "mxfp6_backward_precision": getattr(params, "mxfp6_backward_precision", "mxfp6"), + "mxfp6_fused_wgrad_accum": getattr(params, "mxfp6_fused_wgrad_accum", False), + } + ) + # Sensitive layer configuration config_params.update( { @@ -742,6 +781,9 @@ def _log_flux_config(self, config, args): "fp4_recipe", "mxfp4_backward_precision", "mxfp4_gradient_stochastic_rounding", + "fp6", + "mxfp6_backward_precision", + "mxfp6_fused_wgrad_accum", "sensitive_layers_enabled", "sensitive_layers_start", "sensitive_layers_end", diff --git a/primus/backends/megatron/patches/_source_patch_utils.py b/primus/backends/megatron/patches/_source_patch_utils.py index 2ef19b5c2..e67da8217 100644 --- a/primus/backends/megatron/patches/_source_patch_utils.py +++ b/primus/backends/megatron/patches/_source_patch_utils.py @@ -20,7 +20,7 @@ import inspect import textwrap -from typing import Any, Callable +from typing import Any, Callable, List, Tuple def patch_method_source( @@ -45,6 +45,35 @@ def patch_method_source( AssertionError: If ``ori_code`` is not found in the method's source (e.g. upstream Megatron-LM changed the function unexpectedly). """ + return patch_method_source_multi(cls, method_name, [(ori_code, new_code)]) + + +def patch_method_source_multi( + cls: Any, + method_name: str, + replacements: List[Tuple[str, str]], +) -> Callable: + """Rewrite several fragments of ``cls.``'s source in one pass. + + Calling :func:`patch_method_source` twice on the same method does not work: + the function it leaves behind was produced by ``exec``, so its code object + has no source file and ``inspect.getsource`` raises on the second call. A + method needing more than one edit therefore has to make all of them before + the single recompile, which is what this does. + + Args: + cls: The class owning the method. + method_name: Name of the method to rewrite. + replacements: ``(anchor, replacement)`` pairs, applied in order. Each + anchor must appear in the source, so a pair that upstream has + already made itself is an error rather than a silent no-op. + + Returns: + The newly compiled function that was set on ``cls``. + + Raises: + AssertionError: If any anchor is not found (e.g. upstream changed). + """ # IMPORTANT: replace on the *raw* (non-dedented) source -- ori_code/new_code # are written using the upstream file's absolute column indentation (i.e. # what you see reading the file directly). Dedent must happen AFTER the @@ -53,11 +82,13 @@ def patch_method_source( # post-dedent indentation is a common source of IndentationError. original = getattr(cls, method_name) source = inspect.getsource(original) - assert ori_code in source, ( - f"[SourcePatch] Anchor not found in {cls.__name__}.{method_name}; " - f"upstream source may have changed. Anchor: {ori_code!r}" - ) - modified_source = textwrap.dedent(source.replace(ori_code, new_code)) + for ori_code, new_code in replacements: + assert ori_code in source, ( + f"[SourcePatch] Anchor not found in {cls.__name__}.{method_name}; " + f"upstream source may have changed. Anchor: {ori_code!r}" + ) + source = source.replace(ori_code, new_code) + modified_source = textwrap.dedent(source) # IMPORTANT: exec'ing `modified_source` as a bare top-level `def` (not # nested in a class body) silently loses the implicit `__class__` closure diff --git a/primus/backends/megatron/patches/args/eval_samples_patches.py b/primus/backends/megatron/patches/args/eval_samples_patches.py index 0763ae68e..b7ae853ee 100644 --- a/primus/backends/megatron/patches/args/eval_samples_patches.py +++ b/primus/backends/megatron/patches/args/eval_samples_patches.py @@ -17,6 +17,7 @@ """ from primus.backends.megatron.training.eval_budget import ( + get_eval_global_batch_size, assert_val_worker_divisibility, get_val_num_workers, read_energon_split_sample_count, @@ -31,7 +32,12 @@ # *after* the build_args phase. So this patch cannot read them off args yet and # must take them from the module config, or it would see every one as unset and # quietly leave eval_iters alone. -PRIMUS_ONLY_EVAL_KEYS = ("eval_samples", "val_num_workers") +PRIMUS_ONLY_EVAL_KEYS = ( + "eval_samples", + "val_num_workers", + "eval_global_batch_size", + "eval_micro_batch_size", +) def _hydrate_primus_only_keys(args, module_config): @@ -82,10 +88,16 @@ def patch_eval_samples(ctx: PatchContext): # for every module, so this almost always overrides something, and a # silent override is how the budget drifted from the intent before. replaced = "" if previous in (None, derived) else f", was {previous}" + eval_batch_size = get_eval_global_batch_size(args) + batch = ( + f"global_batch_size={eval_batch_size}" + if eval_batch_size == args.global_batch_size + else f"eval_global_batch_size={eval_batch_size} " + f"(training global_batch_size={args.global_batch_size})" + ) log_kv_rank_0( "[Patch:megatron.args.eval_samples] -eval_iters", - f"{derived} (from eval_samples={args.eval_samples}, " - f"global_batch_size={args.global_batch_size}{replaced})", + f"{derived} (from eval_samples={args.eval_samples}, {batch}{replaced})", ) return @@ -94,9 +106,9 @@ def patch_eval_samples(ctx: PatchContext): # opt into eval_samples get the coverage guarantee. eval_iters = getattr(args, "eval_iters", 0) or 0 if eval_iters > 0: - assert_val_worker_divisibility(args, eval_iters * args.global_batch_size) + eval_samples = eval_iters * get_eval_global_batch_size(args) + assert_val_worker_divisibility(args, eval_samples) log_kv_rank_0( "[Patch:megatron.args.eval_samples] -val_num_workers", - f"{get_val_num_workers(args)} (coverage verified for " - f"{eval_iters * args.global_batch_size} samples)", + f"{get_val_num_workers(args)} (coverage verified for {eval_samples} samples)", ) diff --git a/primus/backends/megatron/patches/dataloader_patch.py b/primus/backends/megatron/patches/dataloader_patch.py index 774cda1d3..4521ff02f 100644 --- a/primus/backends/megatron/patches/dataloader_patch.py +++ b/primus/backends/megatron/patches/dataloader_patch.py @@ -14,7 +14,17 @@ This patch monkey-patches ``DataLoader.__init__`` during the ``setup`` phase and injects ``multiprocessing_context=`` when the caller has -``num_workers > 0`` and did not pass a context of its own. +``num_workers > 0`` and either did not pass a context of its own or asked +for ``fork``. + +Overriding an explicit ``fork`` matters because the callers that most need +this do not leave the choice open: Megatron-Energon hardcodes +``multiprocessing_context = "fork"`` in both of its dataloader classes, so +only injecting into callers that passed nothing would skip exactly the +dataloader that segfaults. A caller that deliberately asked for ``spawn`` +or ``forkserver`` is left alone, and so is one that passed the context +positionally, where replacing it would collide with the positional +argument. Config (YAML module param, mirrors PyTorch's DataLoader argument): multiprocessing_context: "forkserver" | "spawn" | "fork" | null @@ -55,6 +65,26 @@ def _preload_forkserver_torch() -> None: pass +def _caller_context_is_fork(value) -> bool: + """True when the caller's ``multiprocessing_context`` resolves to fork. + + ``None`` counts: it asks for the default, and the default on Linux is fork, + which is the case this patch exists to fix. + """ + if value is None: + return True + if isinstance(value, str): + return value == "fork" + get_start_method = getattr(value, "get_start_method", None) + if callable(get_start_method): + try: + return get_start_method() == "fork" + except Exception: + return False + # Some other context object we cannot interpret; assume it was deliberate. + return False + + def _install_dataloader_monkeypatch(mp_context) -> None: """Patch ``DataLoader.__init__`` to inject ``multiprocessing_context`` when the caller has ``num_workers > 0`` and did not set one.""" @@ -72,12 +102,19 @@ def _install_dataloader_monkeypatch(mp_context) -> None: def patched_init(self, *args, **kwargs): # Resolve args against the real signature to stay version-agnostic. bound = sig.bind_partial(self, *args, **kwargs) - if ( - int(bound.arguments.get("num_workers", 0)) > 0 - and "multiprocessing_context" not in bound.arguments - ): - kwargs["multiprocessing_context"] = mp_context - log_rank_0(f"Setting DataLoader multiprocessing_context='{mp_context}'.") + if int(bound.arguments.get("num_workers", 0)) > 0: + if "multiprocessing_context" not in bound.arguments: + kwargs["multiprocessing_context"] = mp_context + log_rank_0(f"Setting DataLoader multiprocessing_context='{mp_context}'.") + elif "multiprocessing_context" in kwargs and _caller_context_is_fork( + kwargs["multiprocessing_context"] + ): + requested = kwargs["multiprocessing_context"] + kwargs["multiprocessing_context"] = mp_context + log_rank_0( + f"Overriding DataLoader multiprocessing_context={requested!r} " + f"with '{mp_context}'." + ) return original_init(self, *args, **kwargs) setattr(patched_init, _PATCHED_ATTR, True) @@ -92,7 +129,8 @@ def patched_init(self, *args, **kwargs): description=( "Set DataLoader.multiprocessing_context from the " "'multiprocessing_context' module param to avoid SIGSEGV caused by " - "fork()-hostile native state (RDMA MRs, HIP runtime, IPC handles)." + "fork()-hostile native state (RDMA MRs, HIP runtime, IPC handles). " + "Also overrides callers that hardcode fork, such as Energon." ), condition=_enabled, ) diff --git a/primus/backends/megatron/patches/energon_fork_patches.py b/primus/backends/megatron/patches/energon_fork_patches.py new file mode 100644 index 000000000..a4cf9d116 --- /dev/null +++ b/primus/backends/megatron/patches/energon_fork_patches.py @@ -0,0 +1,225 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Make the Energon dataloader survive a non-fork start method. + +GPU-Direct RDMA registers GPU memory through dmabuf, and a process holding such +a registration cannot ``fork``: the workers inherit device mappings that do not +survive the fork and segfault immediately, inside ``os.fork`` itself. Energon +7.3.2 hardcodes ``fork`` in both of its dataloader classes, so enabling GDR +crashed as soon as the dataloader built its workers. + +``dataloader_patch`` handles the start method itself, by overriding the context +Energon asks for. What is left is that ``fork`` shared the dataset by memory +while ``forkserver`` and ``spawn`` must pickle it, which Energon's dataset is +not able to do: it holds a process group, a debug file handle, two lambdas +defined in a class body and three closures built inside ``__init__``, none of +which pickle resolves by qualified name. + +This patch installs picklable equivalents from +``primus.backends.megatron.data.energon_picklable``. Two are plain rebinds; the +closures are built in the middle of ``__init__`` where nothing can reach them +from outside, so those need the source-rewrite helper. + +Enabled by the same ``dataloader_mp_context`` module param that drives +``dataloader_patch``, since a fork-free start method is the only reason to want +any of this. + +An installed tree that already carries the equivalent source changes -- the +vendored ``energon-7.3.2-no-fork.patch`` in tiger-training-internal applies them +to site-packages -- is detected and left alone, so the two mechanisms can +coexist while one is being retired. A tree that has neither the expected +lambdas nor the replacements is a version mismatch and fails loudly here rather +than segfaulting later. +""" + +import importlib.util + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + +from ._source_patch_utils import patch_method_source_multi + +_PARAM_NAME = "dataloader_mp_context" +_PATCHED_ATTR = "_primus_energon_no_fork_patched" +_LOG = "[Patch:megatron.energon_no_fork]" + + +def _mp_context(ctx: PatchContext): + """Return the configured start method, or None.""" + try: + args = get_args(ctx) + except AssertionError: + return None + return getattr(args, _PARAM_NAME, None) + + +def _enabled(ctx: PatchContext) -> bool: + """Only for an Energon run that is actually leaving fork behind.""" + mp_context = _mp_context(ctx) + if mp_context is None or mp_context == "fork": + return False + return importlib.util.find_spec("megatron.energon") is not None + + +def _rewrite_or_skip(label: str, cls, replacements, applied_marker: str) -> None: + """Rewrite ``cls.__init__``, unless the installed tree already has it. + + Raises: + AssertionError: If the source shows neither the expected lambdas nor the + replacements, or shows a mix of both. + """ + import inspect + + source = inspect.getsource(cls.__init__) + present = [anchor for anchor, _ in replacements if anchor in source] + + if len(present) == len(replacements): + patch_method_source_multi(cls, "__init__", replacements) + log_rank_0(f"{_LOG} rewrote {label}.__init__ ({len(replacements)} fragments)") + return + + if not present and applied_marker in source: + log_rank_0(f"{_LOG} {label}.__init__ already carries the fix on disk; leaving it") + return + + raise AssertionError( + f"{_LOG} {label}.__init__ matched {len(present)} of {len(replacements)} " + f"expected fragments and {'does' if applied_marker in source else 'does not'} " + f"contain {applied_marker!r}. The installed Energon differs from 7.3.2; " + f"refusing to guess, because getting this wrong segfaults the workers." + ) + + +def _patch_crude(picklable) -> None: + """CrudeWebdataset: an identity sample loader and a default part filter.""" + from megatron.energon.flavors import crude + + crude._identity_sample_loader = picklable.identity_sample_loader + _rewrite_or_skip( + "CrudeWebdataset", + crude.CrudeWebdataset, + [ + ( + " sample_loader=lambda sample: sample,\n", + " sample_loader=_identity_sample_loader,\n", + ) + ], + "_identity_sample_loader", + ) + + # part_filter's default is a lambda in the signature, so it is a value rather + # than code: swap the value instead of rewriting the source. + init = crude.CrudeWebdataset.__init__ + kwdefaults = init.__kwdefaults__ + assert kwdefaults is not None and "part_filter" in kwdefaults, ( + f"{_LOG} CrudeWebdataset.__init__ has no keyword-only part_filter default; " + f"installed Energon differs from 7.3.2." + ) + init.__kwdefaults__ = {**kwdefaults, "part_filter": picklable.accept_all_parts} + log_rank_0(f"{_LOG} replaced the CrudeWebdataset part_filter default") + + +def _patch_generic_webdataset(picklable) -> None: + """DefaultGenericWebdatasetFactory: three closures built inside __init__.""" + from megatron.energon.flavors.webdataset import default_generic_webdataset as dgw + + dgw._PartInSet = picklable.PartInSet + dgw._FieldMapSampleLoader = picklable.FieldMapSampleLoader + dgw._KeyedSampleLoader = picklable.KeyedSampleLoader + + # Both part_filter closures are the same line at different depths, so each + # anchor carries the line above it to stay unique. + _rewrite_or_skip( + "DefaultGenericWebdatasetFactory", + dgw.DefaultGenericWebdatasetFactory, + [ + ( + " parts = set(part_filter)\n" + " part_filter = lambda part: part in parts\n", + " parts = set(part_filter)\n" + " part_filter = _PartInSet(parts)\n", + ), + ( + " self._sample_loader = lambda sample: {\n" + " k: field_access(sample, v) for k, v in fields.items()\n" + " }\n", + " self._sample_loader = _FieldMapSampleLoader(fields)\n", + ), + ( + " parts = set(access[0] for options in fields.values() for access in options)\n" + " part_filter = lambda part: part in parts\n", + " parts = set(access[0] for options in fields.values() for access in options)\n" + " part_filter = _PartInSet(parts)\n", + ), + ( + " self._sample_loader = lambda sample: {\n" + ' "__key__": sample["__key__"],\n' + " **inner_sample_loader(sample),\n" + ' "__restore_key__": sample["__restore_key__"],\n' + ' "__subflavors__": self.subflavors,\n' + ' "__sources__": sample["__sources__"],\n' + " }\n", + " self._sample_loader = _KeyedSampleLoader(inner_sample_loader, self)\n", + ), + ], + "_PartInSet", + ) + + +def _pin_start_method(mp_context: str) -> None: + """Pin the process-wide start method as well as the dataloader's. + + Passing a context to the dataloader is not sufficient on its own. Anything + built earlier under the default context stays bound to it, and handing such + an object to a worker in another context fails with "A SemLock created in a + fork context is being shared with a process in a spawn context". Pinning the + default here, before any dataloader exists, is what the deployment used a + site-packages ``.pth`` hook for. + """ + import multiprocessing + + current = multiprocessing.get_start_method(allow_none=True) + if current == mp_context: + return + multiprocessing.set_start_method(mp_context, force=True) + log_rank_0(f"{_LOG} process start method {current!r} -> {mp_context!r}") + + +@register_patch( + "megatron.energon_no_fork", + backend="megatron", + phase="before_train", + priority=50, # Must run before the dataloader is built. + description=( + "Make Energon's dataset picklable so its dataloader can run under " + "forkserver/spawn, which GPU-Direct RDMA requires because a process " + "holding a dmabuf registration cannot fork." + ), + condition=_enabled, +) +def patch_energon_no_fork(ctx: PatchContext) -> None: + from megatron.energon.worker import WorkerConfig + + from primus.backends.megatron.data import energon_picklable + + if getattr(WorkerConfig, _PATCHED_ATTR, False): + log_rank_0(f"{_LOG} already applied; skipping") + return + + mp_context = _mp_context(ctx) + _pin_start_method(mp_context) + + # WorkerConfig holds a process group and a debug file handle, neither of + # which pickles and neither of which a worker needs. + WorkerConfig.__getstate__ = energon_picklable.worker_config_getstate + log_rank_0(f"{_LOG} installed WorkerConfig.__getstate__") + + _patch_crude(energon_picklable) + _patch_generic_webdataset(energon_picklable) + + setattr(WorkerConfig, _PATCHED_ATTR, True) + log_rank_0(f"{_LOG} done; Energon dataset is picklable for '{mp_context}'") diff --git a/primus/backends/megatron/patches/mlperf_warmup_patches.py b/primus/backends/megatron/patches/mlperf_warmup_patches.py index a9fa5a1ee..10305c539 100644 --- a/primus/backends/megatron/patches/mlperf_warmup_patches.py +++ b/primus/backends/megatron/patches/mlperf_warmup_patches.py @@ -222,6 +222,46 @@ def _build_synthetic_iterator(primus_args): return MegatronDataloaderWrapper(mock_loader) +def _reset_ddp_grad_ready_calibration(models): + """Put each DDP bucket group back into the uncalibrated state warmup found it in. + + Megatron's gradient buffers learn, on their first batch, how many times each + parameter registers a ready gradient; from the second batch on they issue the + reduce-scatter only when that golden count is reached again. Warmup steps are + batches like any other, so they consume the calibration: the golden counts end up + describing a synthetic step rather than the first real one. When the real steps + then register a different number of times -- a different microbatch count is + enough -- the bucket either fires early, and the next registration finds a + collective already in flight, or never reaches the golden count at all. + + Any outstanding collective is drained first: the handle belongs to a synthetic step + whose gradients are about to be discarded, and leaving it in flight would hand the + first real step a bucket that is busy for reasons it cannot see. + """ + drained = groups_reset = 0 + for m in models: + groups = list(getattr(m, "bucket_groups", [])) + list( + getattr(m, "expert_parallel_bucket_groups", []) + ) + for group in groups: + if not hasattr(group, "is_first_batch"): + continue + handle = getattr(group, "grad_reduce_handle", None) + if handle is not None: + handle.wait() + group.grad_reduce_handle = None + drained += 1 + group.is_first_batch = True + group.golden_per_param_grad_ready_counts = {} + group.per_param_grad_ready_counts = {} + groups_reset += 1 + _log( + f"Reset DDP grad-ready calibration on {groups_reset} bucket groups " + f"({drained} outstanding collectives drained)" + ) + return groups_reset + + def _run_warmup_and_restore( *, warmup_steps, @@ -365,6 +405,9 @@ def _run_warmup_and_restore( except TypeError: optimizer.zero_grad() + # ---- 11b. Undo the DDP grad-ready calibration the warmup steps consumed ---- + _reset_ddp_grad_ready_calibration(models) + # ---- 12. Reset counters ---- megatron_args.consumed_train_samples = 0 megatron_args.skipped_train_samples = 0 diff --git a/primus/backends/megatron/training/eval_budget.py b/primus/backends/megatron/training/eval_budget.py index 1e83f2235..5bd69d660 100644 --- a/primus/backends/megatron/training/eval_budget.py +++ b/primus/backends/megatron/training/eval_budget.py @@ -35,6 +35,8 @@ "assert_mlperf_timestep_source", "assert_val_worker_divisibility", "get_data_parallel_size", + "get_eval_global_batch_size", + "get_eval_micro_batch_size", "get_eval_num_microbatches", "get_val_num_workers", "read_energon_split_sample_count", @@ -101,15 +103,41 @@ def get_data_parallel_size(args) -> int: return world_size // divisor +def get_eval_micro_batch_size(args) -> int: + """Microbatch width to evaluate at, defaulting to the training width. + + Narrowing this is metric-neutral for a sample-weighted loss: the batches an + evaluation is cut into change, the per-sample timestep pairing does not. + Keep it a multiple of ``NUM_VALIDATION_TIMESTEPS`` if the run injects + equidistant timesteps rather than reading them from the dataset, or the + evaluation will never see the higher timesteps. + """ + return getattr(args, "eval_micro_batch_size", None) or args.micro_batch_size + + +def get_eval_global_batch_size(args) -> int: + """Samples per evaluation iteration, defaulting to the training global batch. + + Exists because the training global batch cannot always cover a validation + split in whole iterations, and is not free to change: it has a floor of + ``data_parallel_size * micro_batch_size``, since every rank must run at + least one microbatch. Four nodes at MBS=64 sit exactly on that floor at + 2048, and 2048 does not divide the 29696-sample MLPerf Flux split + (2^10 * 29), so no training-side batch size can read it exactly. Setting + the evaluation's own batch decouples the two. + """ + return getattr(args, "eval_global_batch_size", None) or args.global_batch_size + + def get_eval_num_microbatches(args) -> int: """Microbatches per evaluation iteration. - Uses the same global batch as training so that ``eval_iters`` counts in - global batches, matching Megatron's convention. + Counts in evaluation global batches, which equal the training global batch + unless the recipe overrides them. """ dp_size = get_data_parallel_size(args) - micro_batch_size = args.micro_batch_size - global_batch_size = args.global_batch_size + micro_batch_size = get_eval_micro_batch_size(args) + global_batch_size = get_eval_global_batch_size(args) samples_per_microbatch = micro_batch_size * dp_size if samples_per_microbatch <= 0: @@ -163,7 +191,7 @@ def assert_val_worker_divisibility(args, eval_samples: int) -> None: this from dividing by zero at the default worker count of 0. """ dp_size = get_data_parallel_size(args) - micro_batch_size = args.micro_batch_size + micro_batch_size = get_eval_micro_batch_size(args) val_num_workers = get_val_num_workers(args) divisor = dp_size * max(1, val_num_workers) * micro_batch_size @@ -288,7 +316,7 @@ def resolve_eval_iters(args) -> Optional[int]: if eval_samples <= 0: raise EvalCoverageError(f"eval_samples must be positive, got {eval_samples}.") - global_batch_size = args.global_batch_size + global_batch_size = get_eval_global_batch_size(args) if eval_samples % global_batch_size != 0: raise EvalCoverageError( f"eval_samples ({eval_samples}) is not divisible by global_batch_size " diff --git a/primus/backends/megatron/training/evaluator.py b/primus/backends/megatron/training/evaluator.py index 2df9dc24f..34e00a9a0 100644 --- a/primus/backends/megatron/training/evaluator.py +++ b/primus/backends/megatron/training/evaluator.py @@ -15,7 +15,11 @@ from megatron.training import ft_integration, get_args, get_timers from megatron.training.utils import is_last_rank -from primus.backends.megatron.training.eval_budget import get_eval_num_microbatches +from primus.backends.megatron.training.eval_budget import ( + get_eval_global_batch_size, + get_eval_micro_batch_size, + get_eval_num_microbatches, +) from primus.backends.megatron.training.eval_session import begin_eval_session from primus.backends.megatron.training.global_vars import get_train_start_time from primus.backends.megatron.training.utils import is_pipeline_stage_containing_loss @@ -192,8 +196,12 @@ def primus_evaluate( total_loss_numerators = {} total_loss_denominators = {} - # make validation batch size independent from training batch size - eval_batch_size = args.global_batch_size + # Validation batch size, independent of the training one where the recipe + # says so: the training global batch has a floor of dp_size x + # micro_batch_size and cannot always divide a validation split, and a split + # it cannot divide would otherwise be evaluated only in part. + eval_batch_size = get_eval_global_batch_size(args) + eval_micro_batch_size = get_eval_micro_batch_size(args) # Shared with the dataloader provider so the loop and the dataset it reads # from cannot disagree about how large an evaluation is. eval_num_microbatches = get_eval_num_microbatches(args) @@ -228,7 +236,7 @@ def primus_evaluate( model=model, num_microbatches=eval_num_microbatches, seq_length=args.seq_length, - micro_batch_size=args.micro_batch_size, + micro_batch_size=eval_micro_batch_size, decoder_seq_length=args.decoder_seq_length, forward_only=True, ) diff --git a/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_mxfp6_local_spec.py b/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_mxfp6_local_spec.py new file mode 100644 index 000000000..01bd60bfc --- /dev/null +++ b/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_mxfp6_local_spec.py @@ -0,0 +1,245 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Integration tests for Flux model with MXFP6 local spec. + +Tests that a real (small) Flux model constructs, uses the MXFP6 linear types, and +produces valid output and gradients under PrimusTurboMXFP6LocalSpecProvider. + +The MXFP4 sibling of this file uses batch_size=2, which MXFP6 cannot: MXFP6 needs +M, N and K all multiples of 256, and at IMG_SIZE_TINY the image stream is only 64 +tokens, so batch_size=2 gives M=128. batch_size=4 lifts every stream over the bar -- +image 4*64=256, text 4*128=512, and the concatenated single-block stream 4*192=768. +Unlike MXFP4 there is also no backend pinning fixture, because MXFP6 has a single +backend (AITER A6W6) and no preshuffle contract to satisfy. +""" + +import os +from unittest import mock + +import pytest +import torch + +from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6ColumnParallelLinear, + MXFP6RowParallelLinear, +) +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +from primus.backends.megatron.core.models.diffusion.flux.model import Flux +from primus.backends.megatron.core.models.diffusion.flux.utils import ( + generate_image_position_ids, + pack_latents, +) +from tests.unit_tests.backends.megatron.conftest import requires_mxfp6 +from tests.unit_tests.backends.megatron.diffusion.constants import ( + CLIP_L_EMBEDDING_DIM, + IMG_SIZE_TINY, + T5_XXL_EMBEDDING_DIM, + TEXT_SEQ_LEN_MEDIUM, + VAE_LATENT_CHANNELS, +) +from tests.utils import PrimusUT + +# Chosen so every MXFP6 GEMM sees M % 256 == 0; see the module docstring. +BATCH_SIZE = 4 + +SPEC_LINEAR_NAMES = { + "linear_qkv", + "added_linear_qkv", + "linear_proj", + "linear_fc1", + "linear_fc2", +} + + +class TestFluxMXFP6LocalSpec(PrimusUT): + """Integration tests for Flux with the MXFP6 local spec provider.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + @pytest.fixture(autouse=True) + def default_fused_mlp_mode(self): + """Pin the mode to the default instead of inheriting it from the shell. + + FluxConfig's default activation is not the one the fused prologue implements, so + these models are built on the fallback path. The submission container exports + PRIMUS_MXFP6_FUSED_MLP=on, which makes that fallback an error. + """ + with mock.patch.dict(os.environ): + os.environ.pop("PRIMUS_MXFP6_FUSED_MLP", None) + yield + + def _make_mxfp6_config(self, **overrides): + defaults = dict(transformer_impl="local", fp6="mxfp6") + defaults.update(overrides) + return FluxConfig.flux_535m(**defaults) + + def _make_inputs(self, batch_size=BATCH_SIZE): + height, width = IMG_SIZE_TINY, IMG_SIZE_TINY + channels = VAE_LATENT_CHANNELS + txt_seq_len = TEXT_SEQ_LEN_MEDIUM + + img = torch.randn(batch_size, channels, height, width, dtype=torch.bfloat16).cuda() + txt = torch.randn(batch_size, txt_seq_len, T5_XXL_EMBEDDING_DIM, dtype=torch.bfloat16).cuda() + y = torch.randn(batch_size, CLIP_L_EMBEDDING_DIM, dtype=torch.bfloat16).cuda() + timesteps = torch.rand(batch_size, dtype=torch.bfloat16).cuda() + + packed_img = pack_latents(img).transpose(0, 1) + txt_t = txt.transpose(0, 1) + + img_ids = generate_image_position_ids(batch_size, height, width, device="cuda") + txt_ids = torch.zeros(batch_size, txt_seq_len, 3).cuda() + + return packed_img, txt_t, y, timesteps, img_ids, txt_ids + + @staticmethod + def _activate_zero_init(model): + """Break Flux's zero-init so the model is actually live. + + Flux deliberately zero-inits ``proj_out`` and the last linear of every + ``adaLN_modulation`` (AdaLN-Zero), which makes a freshly built model output + exactly zero and gives every transformer linear an exactly zero gradient. A + forward/backward test on an unmodified model is therefore vacuous: it passes + whether or not MXFP6 computes anything, which is why the assertions below run + against an activated model and check for non-zero rather than non-None. + """ + with torch.no_grad(): + for name, param in model.named_parameters(): + if "adaLN_modulation" in name and param.dim() == 2: + torch.nn.init.normal_(param, std=0.02) + torch.nn.init.normal_(model.proj_out.weight, std=0.02) + + def _mxfp6_linears(self, model): + return [ + (name, module) + for name, module in model.named_modules() + if isinstance(module, (MXFP6ColumnParallelLinear, MXFP6RowParallelLinear)) + ] + + def _assert_grads_nonzero(self, model, mode): + linears = self._mxfp6_linears(model) + assert linears, "No MXFP6 linears in the model" + + dead = [ + name for name, module in linears if module.weight.grad is None or not module.weight.grad.any() + ] + assert not dead, f"MXFP6 linears with missing or all-zero grad in {mode} mode: {dead}" + + @requires_mxfp6 + def test_flux_535m_mxfp6_constructs(self): + model = Flux(self._make_mxfp6_config()) + assert isinstance(model, Flux) + + @requires_mxfp6 + def test_flux_535m_mxfp6_linear_types(self): + """Every spec-provided linear must be an MXFP6 variant, not a bf16 fallback.""" + from megatron.core.tensor_parallel.layers import ( + ColumnParallelLinear, + RowParallelLinear, + ) + + model = Flux(self._make_mxfp6_config()) + + found_any = False + for name, module in model.named_modules(): + leaf_name = name.rsplit(".", 1)[-1] if "." in name else name + if leaf_name not in SPEC_LINEAR_NAMES: + continue + found_any = True + if isinstance(module, ColumnParallelLinear): + assert isinstance( + module, MXFP6ColumnParallelLinear + ), f"{name}: expected MXFP6ColumnParallelLinear, got {type(module).__name__}" + if isinstance(module, RowParallelLinear): + assert isinstance( + module, MXFP6RowParallelLinear + ), f"{name}: expected MXFP6RowParallelLinear, got {type(module).__name__}" + + assert found_any, "No spec-provided linears found in model" + + @requires_mxfp6 + def test_flux_535m_is_inert_at_init(self): + """Documents why the tests below activate the model first. + + If Flux ever stops zero-initialising its output head, this test fails and the + ``_activate_zero_init`` workaround can be reconsidered. + """ + model = Flux(self._make_mxfp6_config()).cuda().to(torch.bfloat16) + model.train() + + output = model(*self._make_inputs()) + assert not output.any(), "Flux is no longer zero-init at the output head" + + output.sum().backward() + assert all( + module.weight.grad is None or not module.weight.grad.any() + for _, module in self._mxfp6_linears(model) + ), "gradients are non-zero at init, so the activation step may be unnecessary" + + @requires_mxfp6 + def test_flux_535m_mxfp6_forward_backward(self): + model = Flux(self._make_mxfp6_config()).cuda().to(torch.bfloat16) + model.train() + self._activate_zero_init(model) + + output = model(*self._make_inputs()) + + assert len(output.shape) == 3 + assert output.shape[1] == BATCH_SIZE + assert output.any(), "Output is all zero even after activation" + assert not torch.isnan(output).any(), "Output contains NaN" + assert not torch.isinf(output).any(), "Output contains Inf" + + output.sum().backward() + self._assert_grads_nonzero(model, "pure MXFP6") + + @requires_mxfp6 + def test_flux_535m_mxfp6_hybrid_forward_backward(self): + """Hybrid backward (MXFP6 fwd / FP8 bwd) works through the full model.""" + config = self._make_mxfp6_config(mxfp6_backward_precision="fp8") + model = Flux(config).cuda().to(torch.bfloat16) + model.train() + self._activate_zero_init(model) + + output = model(*self._make_inputs()) + + assert output.any(), "Hybrid output is all zero even after activation" + assert not torch.isnan(output).any(), "Hybrid output contains NaN" + assert not torch.isinf(output).any(), "Hybrid output contains Inf" + + output.sum().backward() + self._assert_grads_nonzero(model, "hybrid") + + @requires_mxfp6 + def test_flux_535m_mxfp6_differs_from_bf16_but_tracks_it(self): + """MXFP6 must actually quantize, yet stay close to the bf16 model. + + Guards the two failure modes a construction-only test misses: silently + falling back to bf16 (outputs identical) and a broken GEMM (outputs + uncorrelated). + """ + inputs = self._make_inputs() + + def build(**overrides): + torch.manual_seed(1234) + model = ( + Flux(FluxConfig.flux_535m(transformer_impl="local", **overrides)).cuda().to(torch.bfloat16) + ) + torch.manual_seed(4321) + self._activate_zero_init(model) + return model + + with torch.no_grad(): + mxfp6_out = build(fp6="mxfp6")(*inputs).float() + bf16_out = build()(*inputs).float() + + assert bf16_out.any(), "bf16 reference is all zero, so the comparison is vacuous" + assert not torch.equal( + mxfp6_out, bf16_out + ), "MXFP6 output is bit-identical to bf16, so the MXFP6 path did not run" + + cos = torch.nn.functional.cosine_similarity(mxfp6_out.flatten(), bf16_out.flatten(), dim=0).item() + assert cos > 0.99, f"MXFP6 output cosine similarity to bf16 is only {cos:.5f}" diff --git a/tests/unit_tests/backends/megatron/conftest.py b/tests/unit_tests/backends/megatron/conftest.py index 8c807ac95..83e85e159 100644 --- a/tests/unit_tests/backends/megatron/conftest.py +++ b/tests/unit_tests/backends/megatron/conftest.py @@ -60,6 +60,28 @@ def _is_mxfp4_supported(): ) +def _is_mxfp6_supported(): + if not torch.cuda.is_available(): + return False + try: + from primus_turbo.pytorch.kernels.quantization.mxfp6_pack import ( + check_mxfp6_support, + ) + + supported, _ = check_mxfp6_support() + return supported + except ImportError: + return False + + +# Stricter than requires_mxfp4: as well as gfx950 this needs an aiter carrying the A6W6 +# kernels, which check_mxfp6_support() probes for by attribute rather than by version. +requires_mxfp6 = pytest.mark.skipif( + not _is_mxfp6_supported(), + reason="Requires gfx950+ (MI355X) and an aiter with the A6W6 kernels", +) + + @pytest.fixture(scope="function") def init_parallel_state(): """ diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py b/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py index 487af4193..fb54284d7 100644 --- a/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py @@ -341,3 +341,56 @@ def test_forward_step_loss_func_with_non_loss_data(self, monkeypatch: pytest.Mon # Should return output_tensor directly assert result is output + + def test_validation_report_survives_megatrons_in_place_rescale(self, monkeypatch: pytest.MonkeyPatch): + """The reported validation loss must not move when Megatron rescales the + tensor it backpropagates. + + megatron.core.pipeline_parallel.schedules.forward_step takes the first + element of the pair the loss function returns and rescales it in place -- + ``output_tensor *= cp_group_size``, then ``output_tensor /= + num_microbatches`` -- and only then stores the reported dict. While that + dict held a detached *view* of the same tensor, the report was rescaled + along with it, so the reported validation loss was the true loss divided + by the microbatch count: right at one microbatch per rank per step, half + the true value at two. Under mlperf_mode that halves the number the + convergence gate is compared against, so a run claimed to converge at + about half the samples it really needed. + """ + import torch + + trainer = _build_diffusion_trainer(monkeypatch) + trainer._scheduler = Mock() + trainer.runtime_state = Mock() + trainer.runtime_state.update_metrics = Mock() + + # Four samples, so a summed loss and a per-sample mean cannot be confused. + noise_pred = torch.full((4, 1), 3.0) + clean_latents = torch.zeros(4, 1) + noise = torch.ones(4, 1) + + monkeypatch.setattr( + "primus.backends.megatron.training.diffusion.forward_step.flux_forward_step_func", + lambda *args, **kwargs: (noise_pred, clean_latents, noise, None, {}, True), + ) + + model = Mock() + model.training = False + _, val_loss_func = trainer.forward_step(Mock(), model) + + loss_sum, reported = val_loss_func(noise_pred) + + # target = noise - clean_latents = 1, so each element contributes + # (3 - 1) ** 2 = 4, and the sum over four samples is 16. + assert reported["loss"][0].item() == pytest.approx(16.0) + assert reported["loss"][1].item() == pytest.approx(4.0) + + # Exactly what forward_step does to the tensor it backpropagates. + loss_sum *= 1 # cp_group_size, 1 without context parallelism + loss_sum /= 2 # num_microbatches, 2 at micro batch 32 and GBS 1024 on 16 ranks + + # The rescale has to have happened, or this asserts nothing. + assert loss_sum.item() == pytest.approx(8.0) + + assert reported["loss"][0].item() == pytest.approx(16.0) + assert reported["loss"][1].item() == pytest.approx(4.0) diff --git a/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp6_local.py b/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp6_local.py new file mode 100644 index 000000000..76cf8e122 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp6_local.py @@ -0,0 +1,1108 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for compile-friendly MXFP6 linear layers (primus_turbo_mxfp6_local). + +Covers cross-validation against Primus-Turbo's FP6GemmMXFunction reference, +gradient accuracy against BF16 truth, torch.compile graph breaks, the hybrid +MXFP6-forward / FP8-backward mode, the 256-alignment contract, and init guards. + +Every shape here keeps M, N and K multiples of 256. That is not test tidiness: +MXFP6's backward GEMMs use K as an output dimension, so a 128-aligned K that a +single forward would accept fails in backward. See the module docstring of +``primus_turbo_mxfp6_local``. +""" + +import functools +import os +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from tests.unit_tests.backends.megatron.conftest import requires_mxfp6 +from tests.utils import PrimusUT + +# M, N, K all multiples of 256, the MXFP6 training contract. +M, N, K = 256, 512, 256 + + +def _init_method(): + return functools.partial(torch.nn.init.xavier_uniform_) + + +def _make_mxfp6_config(**overrides): + """A BaseDiffusionConfig for the linear layers. + + BaseDiffusionConfig rather than a plain TransformerConfig because ``fp6`` and + ``mxfp6_backward_precision`` are Primus-owned fields that only exist on the + diffusion config, and it is a TransformerConfig subclass so the Megatron linears + accept it unchanged. + """ + from primus.backends.megatron.core.models.diffusion.common.config import ( + BaseDiffusionConfig, + ) + + defaults = dict( + hidden_size=K, + num_attention_heads=8, + num_layers=1, + params_dtype=torch.bfloat16, + fp6="mxfp6", + ) + defaults.update(overrides) + return BaseDiffusionConfig(**defaults) + + +def _snr_db(got, want): + signal = (want**2).mean() + noise = ((got.float() - want) ** 2).mean() + return (10 * torch.log10(signal / noise)).item() + + +@pytest.fixture +def megatron_global_args(monkeypatch): + """The subset of Megatron's global args the parallel linears read at init. + + ``ColumnParallelLinear``/``RowParallelLinear`` reach for ``get_args()`` on + construction, which the unit-test harness never populates. + """ + dummy_args = SimpleNamespace( + rank=0, + world_size=1, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + offload=False, + offload_ops=[], + patch_primus_pipeline=False, + pp_algorithm=None, + patch_zero_bubble=False, + enable_zero_bubble=False, + rampup_batch_size=None, + global_batch_size=1, + micro_batch_size=1, + data_parallel_size=1, + decrease_batch_size_if_needed=False, + ) + import megatron.training.global_vars as gvars + + monkeypatch.setattr(gvars, "_GLOBAL_ARGS", dummy_args) + + +def _pure_args(fuse_wgrad_accum=False, bias=None): + """Trailing MXFP6LinearFunction args for the pure-MXFP6 path, bias first.""" + return (bias, False, None, 0, 0, fuse_wgrad_accum) + + +def _hybrid_args(bias=None): + """Trailing MXFP6LinearFunction args for the MXFP6-fwd / FP8-bwd path, bias first.""" + from primus_turbo.pytorch.core.backend import BackendType + from primus_turbo.pytorch.core.low_precision import ScalingGranularity, float8_e5m2 + + return ( + bias, + True, + float8_e5m2, + ScalingGranularity.TENSORWISE.value, + BackendType.HIPBLASLT.value, + False, + ) + + +# --------------------------------------------------------------------------- +# Cross-validation against Primus-Turbo's FP6GemmMXFunction reference +# --------------------------------------------------------------------------- + + +class TestMXFP6CrossValidation(PrimusUT): + """Verify MXFP6LinearFunction matches Primus-Turbo's canonical FP6 autograd op. + + Both paths quantize with the same dual packer and call the same GEMM, so the + results should be bit-identical. This is the test that would catch a swapped + operand or a row/column blob mix-up, which SNR-vs-BF16 assertions can miss -- + a wrong-but-consistent pairing still produces a plausible-looking number. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + @requires_mxfp6 + def test_forward_matches_reference_fp6gemm(self): + from primus_turbo.pytorch.core.low_precision import Float6QuantConfig + from primus_turbo.pytorch.ops.gemm_fp6 import FP6GemmMXFunction + + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6LinearFunction, + ) + + torch.manual_seed(42) + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + + our_output = MXFP6LinearFunction.apply(x, w, *_pure_args())[0] + ref_output = FP6GemmMXFunction.apply(x.clone(), w.clone(), x.dtype, Float6QuantConfig()) + + assert torch.equal(our_output, ref_output), ( + "Forward output differs from the Primus-Turbo reference. Max abs diff: " + f"{(our_output - ref_output).abs().max().item():.6e}" + ) + + @requires_mxfp6 + def test_backward_matches_reference_fp6gemm(self): + from primus_turbo.pytorch.core.low_precision import Float6QuantConfig + from primus_turbo.pytorch.ops.gemm_fp6 import FP6GemmMXFunction + + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6LinearFunction, + ) + + torch.manual_seed(42) + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + x_ref = x.detach().clone().requires_grad_(True) + w_ref = w.detach().clone().requires_grad_(True) + + our_output = MXFP6LinearFunction.apply(x, w, *_pure_args())[0] + grad_out = torch.randn_like(our_output) + our_output.backward(grad_out) + + ref_output = FP6GemmMXFunction.apply(x_ref, w_ref, x_ref.dtype, Float6QuantConfig()) + ref_output.backward(grad_out.clone()) + + # Unlike MXFP4, both gradients are bit-identical: MXFP6 has no recipe knobs, so + # there is no way for the two paths to pick different quantizations. + assert torch.equal(x.grad, x_ref.grad), ( + "grad_input differs from reference. Max abs diff: " + f"{(x.grad - x_ref.grad).abs().max().item():.6e}" + ) + assert torch.equal(w.grad, w_ref.grad), ( + "grad_weight differs from reference. Max abs diff: " + f"{(w.grad - w_ref.grad).abs().max().item():.6e}" + ) + + +# --------------------------------------------------------------------------- +# Accuracy against BF16 truth +# --------------------------------------------------------------------------- + + +class TestMXFP6Accuracy(PrimusUT): + """Check all three GEMM directions against the unquantized result. + + The 24 dB floor is set well under the ~28 dB measured on these shapes, but well + above the ~18 dB an MXFP4 path scores, so it fails if MXFP6 silently degrades to + 4-bit-grade error. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + @requires_mxfp6 + def test_forward_and_backward_snr(self): + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6LinearFunction, + ) + + torch.manual_seed(42) + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + + output = MXFP6LinearFunction.apply(x, w, *_pure_args())[0] + grad_out = torch.randn_like(output) + output.backward(grad_out) + + xf, wf, gf = x.detach().float(), w.detach().float(), grad_out.detach().float() + for name, got, want in ( + ("forward", output, xf @ wf.T), + ("grad_input", x.grad, gf @ wf), + ("grad_weight", w.grad, gf.T @ xf), + ): + snr = _snr_db(got, want) + assert snr > 24, f"{name} SNR {snr:.1f} dB vs BF16 is below the 24 dB floor" + + @requires_mxfp6 + def test_accepts_3d_input(self): + """The transformer passes [batch, seq, hidden]; the flatten must round-trip.""" + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6LinearFunction, + ) + + torch.manual_seed(42) + x = torch.randn(2, M // 2, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + + output = MXFP6LinearFunction.apply(x, w, *_pure_args())[0] + assert output.shape == (2, M // 2, N) + + output.backward(torch.randn_like(output)) + assert x.grad.shape == x.shape + assert w.grad.shape == w.shape + + @requires_mxfp6 + def test_hybrid_fp8_backward_runs(self): + """Hybrid mode keeps the MXFP6 forward bit-exact and produces usable grads. + + The FP8 backward is a different numerical path, so it only gets a loose floor; + the point of the test is that the saved-tensor bookkeeping differs between the + two modes (BF16 activations vs packed blobs) and that switch must not break. + """ + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6LinearFunction, + ) + + torch.manual_seed(42) + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + x_pure = x.detach().clone().requires_grad_(True) + w_pure = w.detach().clone().requires_grad_(True) + + out_hybrid = MXFP6LinearFunction.apply(x, w, *_hybrid_args())[0] + out_pure = MXFP6LinearFunction.apply(x_pure, w_pure, *_pure_args())[0] + + assert torch.equal(out_hybrid, out_pure), "hybrid mode changed the MXFP6 forward" + + grad_out = torch.randn_like(out_hybrid) + out_hybrid.backward(grad_out) + + xf, wf, gf = x.detach().float(), w.detach().float(), grad_out.detach().float() + assert x.grad.shape == x.shape + assert w.grad.shape == w.shape + for name, got, want in ( + ("grad_input", x.grad, gf @ wf), + ("grad_weight", w.grad, gf.T @ xf), + ): + snr = _snr_db(got, want) + assert snr > 15, f"hybrid {name} SNR {snr:.1f} dB vs BF16 is below the 15 dB floor" + + +# --------------------------------------------------------------------------- +# torch.compile +# --------------------------------------------------------------------------- + + +class TestMXFP6Compile(PrimusUT): + """The whole reason this module exists is to trace cleanly, so guard that.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + # PrimusUT is a unittest.TestCase, which silently ignores pytest.mark.parametrize, + # so the two modes are spelled out rather than parametrized. + def _assert_no_graph_break(self, args): + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6LinearFunction, + ) + + torch._dynamo.reset() + + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + + explanation = torch._dynamo.explain(MXFP6LinearFunction.apply)(x, w, *args) + + assert explanation.graph_break_count == 0, ( + f"Expected 0 graph breaks, got {explanation.graph_break_count}. " + f"Reasons: {explanation.break_reasons}" + ) + + @requires_mxfp6 + def test_no_graph_break_pure(self): + self._assert_no_graph_break(_pure_args()) + + @requires_mxfp6 + def test_no_graph_break_hybrid(self): + self._assert_no_graph_break(_hybrid_args()) + + + @requires_mxfp6 + def test_compiled_forward_matches_eager(self): + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6LinearFunction, + ) + + torch._dynamo.reset() + torch.manual_seed(42) + + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + + eager_out = MXFP6LinearFunction.apply(x, w, *_pure_args())[0] + compiled_out = torch.compile(MXFP6LinearFunction.apply)(x, w, *_pure_args())[0] + + assert torch.equal(eager_out, compiled_out), ( + "Compiled output differs from eager. Max abs diff: " + f"{(eager_out - compiled_out).abs().max().item():.6e}" + ) + + +# --------------------------------------------------------------------------- +# Module instantiation and init guards +# --------------------------------------------------------------------------- + + +class TestMXFP6LinearModules(PrimusUT): + """Instantiate the real Megatron parallel linears and take a training step.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state, megatron_global_args): + pass + + def _column_linear(self, **config_overrides): + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6ColumnParallelLinear, + ) + + return MXFP6ColumnParallelLinear( + input_size=K, + output_size=N, + config=_make_mxfp6_config(**config_overrides), + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + + def _row_linear(self, **config_overrides): + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6RowParallelLinear, + ) + + return MXFP6RowParallelLinear( + input_size=N, + output_size=K, + config=_make_mxfp6_config(**config_overrides), + init_method=_init_method(), + bias=False, + input_is_parallel=True, + skip_bias_add=False, + is_expert=False, + ) + + @requires_mxfp6 + def test_column_parallel_training_step(self): + layer = self._column_linear().cuda() + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + + output, _ = layer(x) + assert output.shape == (M, N) + + output.sum().backward() + assert layer.weight.grad is not None + assert torch.isfinite(layer.weight.grad).all() + assert torch.isfinite(x.grad).all() + + @requires_mxfp6 + def test_row_parallel_training_step(self): + layer = self._row_linear().cuda() + x = torch.randn(M, N, dtype=torch.bfloat16, device="cuda", requires_grad=True) + + output, _ = layer(x) + assert output.shape == (M, K) + + output.sum().backward() + assert layer.weight.grad is not None + assert torch.isfinite(layer.weight.grad).all() + + @requires_mxfp6 + def test_training_loop_decreases_loss(self): + """The gradients must actually point downhill over a run of steps. + + The step is applied to an FP32 master copy of the weight, which is what + Megatron's optimizer does and what makes this test meaningful. Stepping the + BF16 parameter in place instead is vacuous: with this loss the update is ~7e-6 + relative to the weight while BF16 resolves ~8e-3, so every step rounds away and + the loss stays bit-identical regardless of whether the gradients are correct. + """ + torch.manual_seed(42) + layer = self._column_linear().cuda() + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + target = torch.randn(M, N, dtype=torch.bfloat16, device="cuda") + + master = layer.weight.detach().float().clone().requires_grad_(True) + opt = torch.optim.Adam([master], lr=3e-3) + + losses = [] + for _ in range(8): + layer.weight.grad = None + output, _ = layer(x) + loss = ((output.float() - target.float()) ** 2).mean() + loss.backward() + master.grad = layer.weight.grad.float() + opt.step() + with torch.no_grad(): + layer.weight.copy_(master) + losses.append(loss.item()) + + assert all( + b < a for a, b in zip(losses, losses[1:]) + ), f"loss did not decrease monotonically: {[round(v, 5) for v in losses]}" + + @requires_mxfp6 + def test_hybrid_backward_selected_from_config(self): + layer = self._column_linear(mxfp6_backward_precision="fp8").cuda() + assert layer._backward_is_fp8 is True + + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + output, _ = layer(x) + output.sum().backward() + assert torch.isfinite(layer.weight.grad).all() + + @requires_mxfp6 + def test_pure_backward_is_the_default(self): + assert self._column_linear()._backward_is_fp8 is False + + @requires_mxfp6 + def test_bias_is_applied_in_the_gemm_epilogue(self): + """Bias now rides in the A6W6 store epilogue rather than a separate pass; check it lands. + + The epilogue reads the bias out of the kernarg slot the fp6 kernel used to ignore, so + a bias that fails to reach the kernel shows up as no shift at all rather than as an + error. That is what this pins down. + """ + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6ColumnParallelLinear, + ) + + torch.manual_seed(42) + layer = MXFP6ColumnParallelLinear( + input_size=K, + output_size=N, + config=_make_mxfp6_config(), + init_method=_init_method(), + bias=True, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ).cuda() + + with torch.no_grad(): + layer.bias.fill_(1.0) + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + + with_bias, _ = layer(x) + with torch.no_grad(): + layer.bias.zero_() + without_bias, _ = layer(x) + + # Both outputs are BF16, whose spacing around these magnitudes (|out| up to ~5) + # is about 0.03, so the difference of the two roundings cannot recover 1.0 + # exactly. The tolerance is that spacing, not an accuracy claim about MXFP6. It is + # if anything looser than needed now: the epilogue rounds once where the separate + # add rounded twice. + delta = (with_bias.float() - without_bias.float()).abs() + assert torch.allclose( + delta, torch.ones_like(delta), atol=0.05 + ), f"bias shift deviates from 1.0 by up to {(delta - 1).abs().max().item():.4f}" + + @requires_mxfp6 + def test_epilogue_bias_matches_adding_it_afterwards(self): + """The epilogue must agree with the separate add it replaced, to within one rounding. + + These cannot be compared bitwise, and the reason is worth stating because a naive + equality check here fails loudly on ~2% of elements. The separate add rounded the GEMM + result to bf16 and then added, rounding twice; the epilogue adds to the fp32 accumulator + and rounds once. Where bias nearly cancels the accumulator, the amount the old path + discarded is many ulps *of the result*, so the bound has to be one ulp of the pre-add + magnitude rather than of the answer. The epilogue is the more accurate of the two. + """ + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6LinearFunction, + ) + + torch.manual_seed(7) + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(N, dtype=torch.bfloat16, device="cuda") + + fused = MXFP6LinearFunction.apply(x, w, *_pure_args(bias=b))[0] + unbiased = MXFP6LinearFunction.apply(x, w, *_pure_args())[0] + separate = unbiased + b + + # The bias has to actually arrive: a dropped pointer looks like no shift, not an error. + assert not torch.equal(fused, unbiased), "the epilogue bias never reached the kernel" + + def ulp(t): + return torch.exp2(torch.floor(torch.log2(t.float().abs().clamp_min(1e-30))) - 7.0) + + diff = (fused.float() - separate.float()).abs() + bound = ulp(unbiased) + ulp(separate) + over = int((diff > bound).sum()) + assert over == 0, ( + f"{over} of {diff.numel()} elements exceed the double-rounding bound, " + f"worst {diff.max().item():.3e}" + ) + + @requires_mxfp6 + def test_bias_gradient_comes_from_the_packer_column_sums(self): + """The bias gradient rides along with the backward's existing quantization. + + The reference is the arrangement this replaced: the caller added the bias after + ``.apply()`` returned, so autograd owned the reduction over ``grad_output``. Both + reduce the same bf16 tensor, but the packer accumulates across tiles in fp32, so + this is checked as a reduction and not for bit equality -- in practice the two + agree exactly at these shapes, and at worst by a bf16 ulp. + """ + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6LinearFunction, + ) + + torch.manual_seed(0) + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(N, dtype=torch.bfloat16, device="cuda") + upstream = torch.randn(M, N, dtype=torch.bfloat16, device="cuda") + + xg, wg, bg = (t.clone().requires_grad_() for t in (x, w, b)) + MXFP6LinearFunction.apply(xg, wg, *_pure_args(bias=bg))[0].backward(upstream) + + xr, wr, br = (t.clone().requires_grad_() for t in (x, w, b)) + out_ref = MXFP6LinearFunction.apply(xr, wr, *_pure_args())[0] + br + out_ref.backward(upstream.clone()) + + assert bg.grad.shape == br.grad.shape == (N,) + snr = _snr_db(bg.grad, br.grad.float()) + assert snr > 35, f"bias grad diverges from the autograd reduction: {snr:.1f} dB" + + # The other two gradients must be untouched by where the bias was added. + assert torch.equal(xg.grad, xr.grad), "grad_input changed" + assert torch.equal(wg.grad, wr.grad), "grad_weight changed" + + @requires_mxfp6 + def test_bias_gradient_is_skipped_when_bias_needs_no_grad(self): + """``want_col_sum=False`` when the bias is frozen, so the packer does no extra work.""" + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6LinearFunction, + ) + + torch.manual_seed(0) + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + frozen = torch.randn(N, dtype=torch.bfloat16, device="cuda") + + MXFP6LinearFunction.apply(x, w, *_pure_args(bias=frozen))[0].sum().backward() + + assert frozen.grad is None + assert x.grad is not None and w.grad is not None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_column_parallel_rejects_tp_gt_1(self): + with pytest.raises(ValueError, match="tensor_model_parallel_size=1"): + self._column_linear(tensor_model_parallel_size=2) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_row_parallel_rejects_tp_gt_1(self): + with pytest.raises(ValueError, match="tensor_model_parallel_size=1"): + self._row_linear(tensor_model_parallel_size=2) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_rejects_gradient_accumulation_fusion(self): + """Megatron's knob stays rejected: it would also move the plain linears.""" + with pytest.raises(ValueError, match="mxfp6_fused_wgrad_accum=True instead"): + self._column_linear(gradient_accumulation_fusion=True) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_fused_wgrad_accum_is_opt_in(self): + assert self._column_linear()._fuse_wgrad_accum is False + assert self._column_linear(mxfp6_fused_wgrad_accum=True)._fuse_wgrad_accum is True + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_fused_wgrad_accum_rejects_fp8_backward(self): + """The FP8 backward forms its wgrad with a GEMM that has no out-variant.""" + with pytest.raises(ValueError, match="mxfp6_backward_precision"): + self._column_linear( + mxfp6_fused_wgrad_accum=True, mxfp6_backward_precision="fp8" + ) + + @requires_mxfp6 + def test_fused_wgrad_lands_in_main_grad(self): + """The fused wgrad must equal the unfused one and bypass the DDP hook's add. + + Megatron's hook keys off ``grad_added_to_main_grad`` to skip its ``add_``, so a + wgrad that landed in ``main_grad`` without setting the flag would be counted + twice, and one that set the flag without landing would be dropped. + """ + torch.manual_seed(42) + layer = self._column_linear(mxfp6_fused_wgrad_accum=True).cuda() + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + + # Stand in for what Megatron DDP attaches: a bf16 buffer pre-filled with garbage + # a beta=0 store has to obliterate. + layer.weight.main_grad = torch.full_like(layer.weight, 7.0) + layer(x)[0].sum().backward() + + reference = self._column_linear().cuda() + with torch.no_grad(): + reference.weight.copy_(layer.weight) + reference(x)[0].sum().backward() + + assert layer.weight.grad_added_to_main_grad is True + assert torch.equal(layer.weight.main_grad, reference.weight.grad) + + @requires_mxfp6 + def test_fused_wgrad_does_not_break_the_compiled_region(self): + """The fusion must not cost the compiled region, which is what actually pays here. + + Traced through the module with a real ``main_grad`` rather than through + ``MXFP6LinearFunction.apply`` on a bare tensor: dynamo decides whether to inline an + autograd Function by tracing its backward too, and the backward's ``main_grad`` + lookup is only there to trace when the attribute exists. Getting this wrong is not + a small loss. Breaking the region splits one compiled block into hundreds of + fragments, and on the MBS=32 Flux 12B arm that cost 42.7 ms of eager elementwise + work per 512 images against the ~16 ms of ``add_`` the fusion removes. + """ + torch._dynamo.reset() + layer = self._column_linear(mxfp6_fused_wgrad_accum=True).cuda() + layer.weight.main_grad = torch.zeros_like(layer.weight) + # requires_grad is load-bearing: with no autograd graph to build, dynamo never + # traces the Function's backward and reports zero breaks whatever is in there. + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + + explanation = torch._dynamo.explain(lambda inp: layer(inp)[0])(x) + + assert explanation.graph_break_count == 0, ( + f"Expected 0 graph breaks, got {explanation.graph_break_count}. " + f"Reasons: {explanation.break_reasons}" + ) + + @requires_mxfp6 + def test_fused_wgrad_needs_a_bf16_main_grad(self): + layer = self._column_linear(mxfp6_fused_wgrad_accum=True).cuda() + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + + # Matched on message rather than type: the autograd engine is free to rewrap an + # exception raised inside a Function's backward. + with pytest.raises((RuntimeError, TypeError), match="weight.main_grad"): + layer(x)[0].sum().backward() + + layer.weight.main_grad = torch.zeros_like(layer.weight, dtype=torch.float32) + with pytest.raises((RuntimeError, TypeError), match="main_grads_dtype"): + layer(x)[0].sum().backward() + + @requires_mxfp6 + def test_rejects_unaligned_k(self): + """K only 128-aligned is accepted by a lone forward but fails in backward. + + Enforced in Primus-Turbo's gemm_fp6 entry point, so the assertion fires at the + call rather than part-way through the backward pass. Asserted here because it is + the constraint most likely to bite someone configuring a model. + """ + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6LinearFunction, + ) + + x = torch.randn(M, 128, dtype=torch.bfloat16, device="cuda") + w = torch.randn(N, 128, dtype=torch.bfloat16, device="cuda") + + with pytest.raises((AssertionError, RuntimeError, ValueError)): + MXFP6LinearFunction.apply(x, w, *_pure_args())[0].sum().backward() + + +# --------------------------------------------------------------------------- +# BaseDiffusionConfig validation +# --------------------------------------------------------------------------- + + +class TestMXFP6ConfigValidation: + """The fp6 field is Primus-owned, so these cross-checks are the only guard. + + Megatron validates fp4-vs-fp8 itself but has never heard of fp6, so without these + a config asking for both would just silently get one of them. + """ + + @staticmethod + def _config(**overrides): + from primus.backends.megatron.core.models.diffusion.common.config import ( + BaseDiffusionConfig, + ) + + defaults = dict(hidden_size=K, num_attention_heads=8, num_layers=1) + defaults.update(overrides) + return BaseDiffusionConfig(**defaults) + + def test_defaults_leave_fp6_off(self): + config = self._config() + assert config.fp6 is None + assert config.mxfp6_backward_precision == "mxfp6" + + def test_accepts_mxfp6(self): + assert self._config(fp6="mxfp6").fp6 == "mxfp6" + + def test_accepts_hybrid_backward(self): + config = self._config(fp6="mxfp6", mxfp6_backward_precision="fp8") + assert config.mxfp6_backward_precision == "fp8" + + def test_rejects_unknown_fp6_format(self): + with pytest.raises(ValueError, match="Unknown fp6"): + self._config(fp6="mxfp6_e3m2") + + def test_rejects_fp6_with_fp4(self): + with pytest.raises(ValueError, match="cannot both be set"): + self._config(fp6="mxfp6", fp4="mxfp4") + + def test_rejects_fp6_with_fp8(self): + with pytest.raises(ValueError, match="cannot both be set"): + self._config(fp6="mxfp6", fp8="e4m3") + + def test_rejects_unknown_backward_precision(self): + with pytest.raises(ValueError, match="Unknown mxfp6_backward_precision"): + self._config(fp6="mxfp6", mxfp6_backward_precision="mxfp4") + + def test_rejects_backward_precision_without_fp6(self): + with pytest.raises(ValueError, match="requires fp6"): + self._config(mxfp6_backward_precision="fp8") + + def test_rejects_fp6_with_mxfp4_to_fp8_switch(self): + """The switch would build a zero-layer plan and silently never fire. + + The MXFP4 -> FP8 switch is a separate change. Until it lands the field does not + exist and there is no combination to reject, so this asserts nothing rather than + failing on the constructor. + """ + import dataclasses + + from primus.backends.megatron.core.models.diffusion.common.config import ( + BaseDiffusionConfig, + ) + + if not any(f.name == "mxfp4_to_fp8_switch_iter" for f in dataclasses.fields(BaseDiffusionConfig)): + pytest.skip("mxfp4_to_fp8_switch_iter not present; nothing to cross-check") + + with pytest.raises(ValueError, match="mxfp4_to_fp8_switch_iter"): + self._config(fp6="mxfp6", mxfp4_to_fp8_switch_iter=100) + + +class TestMXFP6RecipeConfig: + """The shipped Flux 12B MXFP6 recipe must actually carry the fp6 fields through. + + ``fp6`` is not a Megatron argument, so if the YAML key were dropped anywhere in the + resolution path the run would silently train in BF16 rather than fail. + """ + + RECIPE = ( + "examples/megatron/configs/MI355X/diffusion/" + "flux_12b_ddp_energon_schnell_resample_local_spec_mxfp6.yaml" + ) + + @staticmethod + def _load(rel_path): + import pathlib + + from primus.core.config.primus_config import load_primus_config + from primus.core.utils import file_utils + + root = pathlib.Path(__file__).resolve().parents[4] + with pytest.MonkeyPatch.context() as mp: + # PrimusConfig.__init__ mkdir's the workspace; keep the load side-effect + # free, matching tests/unit_tests/configs/test_example_configs.py. + mp.setattr(file_utils, "create_path_if_not_exists", lambda *a, **k: None) + return load_primus_config(root / rel_path, None) + + def test_recipe_carries_fp6_fields(self): + cfg = self._load(self.RECIPE) + pre_trainer = next(m for m in cfg.modules if m.name == "pre_trainer") + params = pre_trainer.params + + assert getattr(params, "fp6", None) == "mxfp6" + assert getattr(params, "mxfp6_backward_precision", None) == "mxfp6" + # fp6 is mutually exclusive with both of these. + assert getattr(params, "fp4", None) is None + assert getattr(params, "fp8", None) is None + # Megatron's knob would also move the plain linears; the MXFP6 linears reject it. + assert getattr(params, "gradient_accumulation_fusion", True) is False + # The MXFP6-only fused wgrad store has no beta=1 epilogue, so it is correct at one + # microbatch per step only. This is the general recipe, whose global batch a user + # is free to raise, so it stays off here. + assert getattr(params, "mxfp6_fused_wgrad_accum", False) is False + assert getattr(params, "transformer_impl", None) == "local" + + +# --------------------------------------------------------------------------- +# Spec provider wiring +# --------------------------------------------------------------------------- + + +class TestMXFP6SpecProvider: + """Verify the provider hands back MXFP6 linears and that fp6 selects it.""" + + @requires_mxfp6 + def test_provider_returns_mxfp6_linears(self): + from primus.backends.megatron.core.extensions.primus_turbo_local_spec import ( + PrimusTurboMXFP6LocalSpecProvider, + ) + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6ColumnParallelLinear, + MXFP6RowParallelLinear, + ) + + provider = PrimusTurboMXFP6LocalSpecProvider() + assert provider.column_parallel_linear() is MXFP6ColumnParallelLinear + assert provider.row_parallel_linear() is MXFP6RowParallelLinear + + @requires_mxfp6 + def test_flux_layer_spec_selects_mxfp6_backend(self): + """fp6 on the config must route get_flux_layer_spec to the MXFP6 linears.""" + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6ColumnParallelLinear, + ) + from primus.backends.megatron.core.models.diffusion.flux.layer_spec import ( + get_flux_layer_spec, + ) + + config = SimpleNamespace( + transformer_impl="local", + fp4=None, + fp6="mxfp6", + fp8=None, + num_joint_layers=1, + num_single_layers=1, + sensitive_layers_enabled=False, + sensitive_layer_precision="bf16", + ) + + spec = get_flux_layer_spec(config) + submodules = spec.layer_specs[0].submodules + rendered = str(submodules) + assert ( + MXFP6ColumnParallelLinear.__name__ in rendered + ), "get_flux_layer_spec did not select the MXFP6 linears for fp6='mxfp6'" + + @requires_mxfp6 + def test_provider_returns_fused_mlp(self): + from primus.backends.megatron.core.extensions.primus_turbo_local_spec import ( + PrimusTurboMXFP6LocalSpecProvider, + ) + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6FusedMLP, + ) + + assert PrimusTurboMXFP6LocalSpecProvider().mlp_module() is MXFP6FusedMLP + + @requires_mxfp6 + def test_flux_layer_spec_uses_fused_mlp_for_fp6(self): + """The MXFP6 spec must actually reach the MLP, not just the linears. + + The two Flux block factories build their ``mlp`` ModuleSpec independently, so this + is the check that neither was missed -- a spec still naming Megatron's ``MLP`` would + train correctly and simply never fuse anything. + """ + from megatron.core.transformer.mlp import MLP + + from primus.backends.megatron.core.models.diffusion.flux.layer_spec import ( + get_flux_layer_spec, + ) + + config = SimpleNamespace( + transformer_impl="local", + fp4=None, + fp6="mxfp6", + fp8=None, + num_joint_layers=1, + num_single_layers=1, + sensitive_layers_enabled=False, + sensitive_layer_precision="bf16", + ) + + spec = get_flux_layer_spec(config) + # One joint block and one single block, both of which own an MLP. + assert len(spec.layer_specs) == 2 + for layer in spec.layer_specs: + mlp_module = layer.submodules.mlp.module + assert mlp_module is not MLP, "layer spec still uses the unfused MLP" + assert mlp_module.__name__ == "MXFP6FusedMLP" + + +# --------------------------------------------------------------------------- +# Fused MLP epilogue +# --------------------------------------------------------------------------- + + +def _make_fused_mlp(**config_overrides): + """Build an MXFP6FusedMLP with the Flux MLP configuration.""" + from megatron.core.transformer.mlp import MLPSubmodules + + from primus.backends.megatron.core.extensions.primus_turbo_mxfp6_local import ( + MXFP6ColumnParallelLinear, + MXFP6FusedMLP, + MXFP6RowParallelLinear, + ) + + defaults = dict( + ffn_hidden_size=512, + add_bias_linear=True, + gated_linear_unit=False, + bias_activation_fusion=False, + activation_func=functools.partial(torch.nn.functional.gelu, approximate="tanh"), + ) + defaults.update(config_overrides) + config = _make_mxfp6_config(**defaults) + submodules = MLPSubmodules(linear_fc1=MXFP6ColumnParallelLinear, linear_fc2=MXFP6RowParallelLinear) + return MXFP6FusedMLP(config, submodules).to("cuda:0") + + +class TestMXFP6FusedMLP(PrimusUT): + """The fused MLP must be numerically the same module as the one it replaces. + + The fusion removes traffic, not arithmetic: the packed operands the GEMMs consume are the + same ones either way, down to a rounding of the activation's tanh, so forward and both + weight gradients should agree to well within MXFP6's own quantization error rather than + merely correlate. The bias gradient is looser and is checked separately, because it comes + from a reduction the fusion had to reorder. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state, megatron_global_args): + pass + + @pytest.fixture(autouse=True) + def default_fused_mlp_mode(self): + """Pin the mode to the default instead of inheriting it from the shell. + + The submission container exports PRIMUS_MXFP6_FUSED_MLP=on, which is what makes a + fallback an error there -- so the two fallback tests below would fail inside the + very container the fused MLP is meant to run in. + """ + with mock.patch.dict(os.environ): + os.environ.pop("PRIMUS_MXFP6_FUSED_MLP", None) + yield + + @requires_mxfp6 + def test_forward_and_grads_match_stock_mlp(self): + torch.manual_seed(0) + fused = _make_fused_mlp() + assert fused._fused_epilogue, "fused path unexpectedly disabled" + + # Same module, same weights, epilogue not fused: the reference is MLP.forward. + x = torch.randn((M, 1, K), dtype=torch.bfloat16, device="cuda:0", requires_grad=True) + x_ref = x.detach().clone().requires_grad_() + + out, out_bias = fused(x) + # MLP.forward via the base class is exactly the path being replaced. + from megatron.core.transformer.mlp import MLP + + ref_out, ref_bias = MLP.forward(fused, x_ref) + + assert out.shape == ref_out.shape + assert out_bias is ref_bias # both hand back linear_fc2.bias unadded + + snr = _snr_db(out, ref_out.float()) + assert snr > 40, f"forward diverges from the unfused MLP: {snr:.1f} dB" + + grad = torch.randn_like(out) + out.backward(grad) + ref_out.backward(grad.clone()) + + for name, got, want in ( + ("input", x.grad, x_ref.grad), + ("fc1.weight", fused.linear_fc1.weight.grad, None), + ): + if want is None: + continue + snr = _snr_db(got, want.float()) + assert snr > 40, f"{name} grad diverges: {snr:.1f} dB" + + @requires_mxfp6 + def test_fused_wgrad_lands_in_both_main_grads(self): + """The fused MLP owns two weights, so both have to reach the right main_grad. + + Worth its own test rather than trusting the linear's: the two wgrads are formed + from different operands in the same backward, so swapping them produces gradients + that are individually plausible and jointly wrong. + """ + torch.manual_seed(0) + fused = _make_fused_mlp(mxfp6_fused_wgrad_accum=True) + assert fused._fused_epilogue + + reference = _make_fused_mlp() + with torch.no_grad(): + for name, param in fused.named_parameters(): + reference.get_parameter(name).copy_(param) + + # Stand in for what Megatron DDP attaches, pre-filled with garbage a beta=0 + # store has to obliterate. + for linear in (fused.linear_fc1, fused.linear_fc2): + linear.weight.main_grad = torch.full_like(linear.weight, 7.0) + + x = torch.randn((M, 1, K), dtype=torch.bfloat16, device="cuda:0", requires_grad=True) + x_ref = x.detach().clone().requires_grad_() + + out, _ = fused(x) + grad = torch.randn_like(out) + out.backward(grad) + + ref_out, _ = reference(x_ref) + ref_out.backward(grad.clone()) + + for name in ("linear_fc1", "linear_fc2"): + got = getattr(fused, name).weight + want = getattr(reference, name).weight + assert got.grad_added_to_main_grad is True, f"{name} did not claim main_grad" + assert torch.equal(got.main_grad, want.grad), f"{name} wgrad differs" + + @requires_mxfp6 + def test_bias_gradient_matches_eager_reduction(self): + """fc1's bias gradient comes from the packer's side output, not a separate sum. + + This is the one quantity the fusion does not reproduce bit-for-bit: the tensor it + would be reduced from no longer reaches HBM, so the sum is taken over LDS tiles in + fp32 and finished across tiles. That is a different -- and more accurate -- + summation order than a single pass over a bf16 tensor, so it is checked as a + reduction rather than for equality. + """ + torch.manual_seed(0) + fused = _make_fused_mlp() + assert fused._fused_epilogue + + from megatron.core.transformer.mlp import MLP + + x = torch.randn((M, 1, K), dtype=torch.bfloat16, device="cuda:0", requires_grad=True) + x_ref = x.detach().clone().requires_grad_() + + out, _ = fused(x) + grad = torch.randn_like(out) + out.backward(grad) + got = fused.linear_fc1.bias.grad.detach().clone() + + fused.zero_grad(set_to_none=True) + ref_out, _ = MLP.forward(fused, x_ref) + ref_out.backward(grad.clone()) + want = fused.linear_fc1.bias.grad + + assert got.shape == want.shape + snr = _snr_db(got, want.float()) + assert snr > 35, f"fc1 bias grad diverges: {snr:.1f} dB" + + @requires_mxfp6 + def test_falls_back_when_activation_is_not_tanh_gelu(self): + """An activation the prologue does not implement must disable the fusion. + + FluxConfig's *default* activation is a hand-written tanh GELU with a different + association than ATen's, so this is the realistic misconfiguration, not a synthetic + one. Silently fusing it would change numerics with nothing to flag it. + """ + with pytest.warns(UserWarning, match="fused MLP epilogue disabled"): + fused = _make_fused_mlp(activation_func=torch.nn.functional.silu) + assert not fused._fused_epilogue + + @requires_mxfp6 + def test_falls_back_when_backward_is_fp8(self): + """FP8 backward keeps the activation live, which the fusion removes.""" + with pytest.warns(UserWarning, match="mxfp6_backward_precision"): + fused = _make_fused_mlp(mxfp6_backward_precision="fp8") + assert not fused._fused_epilogue + + @requires_mxfp6 + def test_env_kill_switch_disables_and_requires(self): + with mock.patch.dict(os.environ, {"PRIMUS_MXFP6_FUSED_MLP": "off"}): + assert not _make_fused_mlp()._fused_epilogue + + with mock.patch.dict(os.environ, {"PRIMUS_MXFP6_FUSED_MLP": "on"}): + assert _make_fused_mlp()._fused_epilogue + # "on" turns an unusable configuration into an error rather than a fallback. + with pytest.raises(RuntimeError, match="fused MLP is unusable"): + _make_fused_mlp(activation_func=torch.nn.functional.silu)