Enable native-TE MXFP4 training for Flux on MI355X

[amital-amd]

Make MXFP4 Flux training work end to end on the pure TransformerEngine path, with no dependency on the Primus-Turbo autocast layer. Previously MXFP4 could only be driven through Primus-Turbo; running the TE spec with enable_primus_turbo=false either fell back to BF16 GEMMs or ignored the configured attention backend, so a "TE MXFP4" run was silently neither.

This change adds the switch that selects the native-TE path and fixes the two quantization/attention gaps that were undermining it, then ships a validated MI355X recipe that exercises the whole path.

Route FP4 autocast through TransformerEngine

• FluxConfig.fp4_use_native_te_autocast (config.py) opts a run into TE's own fp8_autocast with the MXFP4BlockScaling recipe (TE -> AITER a4w4).
• get_fp4_context (fp4_utils.py) treats the flag as an override: when set it takes the TE-native branch even if Primus-Turbo is otherwise enabled, so enable_primus_turbo=false yields a genuine pure-TE MXFP4 run.
• flux_pretrain_trainer.py plumbs the option from params into FluxConfig.

Serve the right quantization context for FP4 runs

• Flux.get_fp8_context() (model.py) previously always delegated to the FP8 utility, so an FP4-only run received a context that issued no FP4 GEMMs and silently degraded to BF16. It now returns the FP8 context for FP8, the FP4 context for FP4, and nullcontext otherwise, with FP8 taking precedence when both are set (mirroring stock Megatron's transformer block). The local spec still returns nullcontext, since local handles quantization per module at init and must stay free of global state for torch.compile.

Honor the configured attention backend

• _build_flux_config_from_yaml() (flux_pretrain_trainer.py) now coerces the YAML attention_backend string into the AttnBackend enum before it reaches FluxConfig. As a plain string it matched no branch in _set_attention_backend() and silently fell back to FlashAttention, so a recipe requesting fused CK attention never got it. An omitted or null/blank value maps to AttnBackend.auto; an unknown value raises a ValueError listing the valid names.

Ship the MI355X recipe

• Add the Flux 12B Schnell config: Megatron DDP + distributed optimizer, TE spec, fused CK attention, and MXFP4 via the TE-native autocast, with Primus-Turbo and torch.compile disabled (TE's FP4 autocast toggles process-global state that breaks graph capture).

Add unit tests covering get_fp8_context routing (FP8/FP4/precedence/local), the get_fp4_context native-TE-vs-Turbo branch selection, and attention_backend coercion (string, default, null, blank, whitespace, invalid, case-sensitivity).

[Copilot]

Pull request overview

Enable end-to-end MXFP4 Flux training on MI355X using TransformerEngine’s native autocast path (no dependency on Primus-Turbo autocast), while also ensuring the requested attention backend is honored and quantization contexts are routed correctly.

Changes:

• Add fp4_use_native_te_autocast plumbing/config to force TE-native FP4 autocast selection in get_fp4_context.
• Fix Flux quantization context routing so FP4 runs use FP4 context (and FP8 takes precedence when both are set).
• Coerce YAML attention_backend strings to AttnBackend (with robust handling for null/blank/invalid), and add a validated MI355X training recipe plus unit tests.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py	Adds unit tests for YAML attention_backend coercion and error handling.
tests/unit_tests/backends/megatron/diffusion/test_fp4_context.py	Adds unit tests verifying TE-native vs Primus-Turbo FP4 autocast branch selection.
tests/unit_tests/backends/megatron/diffusion/test_flux_fp8_context.py	Extends tests to cover Flux FP8/FP4 context routing and precedence behavior.
primus/backends/megatron/flux_pretrain_trainer.py	Coerces attention_backend strings to AttnBackend and plumbs fp4_use_native_te_autocast into FluxConfig.
primus/backends/megatron/core/models/diffusion/flux/model.py	Updates get_fp8_context() to return FP8, FP4, or nullcontext depending on config and spec.
primus/backends/megatron/core/models/diffusion/flux/config.py	Introduces fp4_use_native_te_autocast configuration flag (TE-spec FP4 autocast backend selection).
primus/backends/megatron/core/fp4_utils.py	Changes Turbo-vs-TE-native FP4 autocast selection logic to honor fp4_use_native_te_autocast.
examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_mxfp4.yaml	Adds MI355X recipe exercising TE-spec + fused attention + MXFP4 via TE-native autocast.