feat(maxtext): Maxtext & MaxDiffusion refactor - #925
Merged
Conversation
Adds per-precision MaxText experiment configs (fp8, nanoo_fp8) derived from the existing bf16 configs by setting `quantization` in overrides, for 9 models across MI300X and MI355X. Upstream MaxText previously shipped only bf16 (`<model>-pretrain.yaml`); this fills the fp8/nanoo_fp8 gap so quantized runs can be launched by config path.
yeandy
requested review from
Xiaoming-AMD,
limou102 and
wenxie-amd
as code owners
July 27, 2026 21:22
…MI355X fp8) MI300X uses NANOO FP8 and MI355X uses FP8, so drop the unused plain-fp8 configs on MI300X and the unused nanoo_fp8 configs on MI355X. Each device keeps bf16 plus its supported quantization.
Adds the fp8 quantized variant for Llama 3.1 405B on MI355X (bf16 already existed); MI355X uses FP8.
…mem_fraction 0.96 Add a single, launcher-agnostic hook for per-config environment overrides: TrainRuntime._apply_config_env() reads a top-level `env:` mapping from the experiment YAML and exports it to os.environ right after config load, before runtime env / distributed init / any `import jax`. The block lives outside `modules:`, so it is never swept into module params nor forwarded to MaxText. Per-config `env:` wins over image-baked defaults. MaxTextAdapter.prepare_backend() applies arch defaults via os.environ.setdefault (gfx950 -> RCCL_WARP_SPEED_AUTO=0) so they never override a per-config `env:`. Precedence (highest wins): per-config env: > arch defaults > image bake. Use it to run mixtral_8x7B at per_device_batch_size=11 on a single MI355X: set XLA_PYTHON_CLIENT_MEM_FRACTION=0.96 (image default 0.9 is too low; bs11 peaks ~270-274 GB / 288 GB) and moe_dispatch_no_expert_sharding=true. Applied to both the bf16 (mixtral_8x7B-pretrain) and fp8 (mixtral_8x7B-fp8-pretrain) MI355X configs.
Run Google's MaxDiffusion WAN/FLUX trainers through primus/cli, mirroring the MaxText backend so JAX diffusion follows the same launch pattern (distinct from the in-tree PyTorch `diffusion` backend). - primus/backends/maxdiffusion: adapter, argument_builder, pretrain trainer (pyconfig.initialize -> WanTrainer/FluxTrainer inside transformer_engine_context) - primus/configs/models|modules/maxdiffusion: native maxdiffusion configs as model presets + module preset - examples/maxdiffusion: MI300X + MI355X wrapper configs (top-level env:) and prepare.py hook - examples/run_pretrain.sh: generalize the JAX-launch checks (MaxText -> MaxText | MaxDiffusion) so JAX backends launch without torchrun
…pec (Option C) Make Primus the single source of truth for all MaxText/JAX environment. A declarative registry (env_spec.py) is applied by a reusable mechanism (env_registry.py) through an opt-in BackendAdapter hook, replacing env exports scattered across run_pretrain.sh, the maxtext prepare hook, and ad-hoc adapter logic. - env_registry.py (new): EnvVar dataclass, gpu arch detection (gfx950/gfx942), and XLA_FLAGS merge where managed knobs win over image-baked flags (e.g. xla_gpu_autotune_level, which fixes the fp8 MoE NaN) while preserving unmanaged ones. - env_spec.py (new): MaxText spec incl. arch-gated RCCL_WARP_SPEED_AUTO (gfx950) and HSA_NO_SCRATCH_RECLAIM (gfx942), plus multi-node JAX coordinator derivation. - backend_adapter.py: opt-in env_defaults() (empty by default, so other backends are unchanged) and apply_env_defaults(), called in prepare_backend before run_setup. - maxtext_adapter.py: supplies maxtext_env_defaults(); drops the old prepare_backend override. - run_pretrain.sh / maxtext prepare hook: stop emitting perf/arch env; Primus owns it. Validated across run_pretrain.sh, primus-cli, and MAD: 18/18 models on both the rock ci and rocm/jax-training:maxtext-v26.5 images.
Run Google's MaxDiffusion WAN/FLUX trainers through primus/cli, mirroring the MaxText backend so JAX diffusion follows the same launch pattern (distinct from the in-tree PyTorch `diffusion` backend). - primus/backends/maxdiffusion: adapter, argument_builder, pretrain trainer (pyconfig.initialize -> WanTrainer/FluxTrainer inside transformer_engine_context) - primus/configs/models|modules/maxdiffusion: native maxdiffusion configs as model presets + module preset - examples/maxdiffusion: MI300X + MI355X wrapper configs (top-level env:) and prepare.py hook - examples/run_pretrain.sh: generalize the JAX-launch checks (MaxText -> MaxText | MaxDiffusion) so JAX backends launch without torchrun
…ext Option C) Bring the MaxDiffusion backend in line with the jax-maxtext-v26.6 env refactor: replace the ad-hoc prepare_backend override / _apply_arch_env_defaults / _is_gfx950 rocminfo logic with the shared env_registry mechanism. - env_spec.py (new): declarative maxdiffusion_env_defaults() owning the single arch-gated knob (RCCL_WARP_SPEED_AUTO=0 on gfx950). All other JAX/XLA/NVTE tuning stays in each wrapper config's top-level env: block. - maxdiffusion_adapter.py: implement env_defaults() (applied by the base adapter before JAX/XLA import); drop the custom prepare_backend + arch-detection code.
…every step The MI355X wan2.1_1.3b / wan2.1_14b / flux_dev configs were missing the ROCm/HIP + RCCL correctness vars that the verified-good base_wan_*/base_flux_dev env scripts set (NCCL_PROTO=Simple, HSA_NO_SCRATCH_RECLAIM=1, GPU_MAX_HW_QUEUES=2, HIP_FORCE_DEV_KERNARG=1, RCCL_MSCCL_ENABLE=0, NCCL_MAX_NCHANNELS=112, NCCL_IB_TC/SL, NCCL_DEBUG, LIBTPU_INIT_ARGS, JAX/MIOpen cache dirs). Dropping them caused an RCCL clique-init hang (GPUs 0%, CPU busy-wait) during model construction. Restoring the full 30-var env block lets wan2.1_1.3b / wan2.1_14b train to completion on gfx950 (exclusive-node allocation still required to avoid GPU co-tenancy RCCL deadlock). Also set log_period: 1 so per-step throughput lines are emitted at the 20-step benchmark length (image default log_period=100 logged nothing, leaving the perf CSV empty). env_spec.py now returns [] so the per-config env: block is the single source of truth and never injects RCCL_WARP_SPEED_AUTO, which the known-good baseline does not set.
…f capture
Bind each MI355X config's metrics_file to ${PERF_METRICS_FILE} (set by
scripts/jax-maxdiffusion/run.sh). MaxDiffusion's per-step stdout line
("completed step: ..., TFLOP/s/device: ...") does not survive the Primus
launcher's stdout handling, so the perf CSV came back empty. With metrics_file
set, the trainer writes per-step JSON (perf/step_time_seconds,
perf/per_device_tflops_per_sec, ...) via max_utils.write_metrics_locally, which
the extractor parses instead. Empty when PERF_METRICS_FILE is unset (no-op).
The MI355X WAN/FLUX runs OOM'd on the first step trying to allocate a ~221 GiB attention matrix, even though the effective config (attention: cudnn_flash_te, ici_fsdp=8/context=1, 720x1280x85, all NVTE_* + XLA_FLAGS) was byte-identical to the known-good standalone launch. Root cause: MaxDiffusionPretrainTrainer.run() imported the trainer class (which pulls in the TransformerEngine attention layers) BEFORE entering transformer_engine_context(). That context installs TE's global_shard_guard MeshResource (fsdp_resource="fsdp", cp_resource="context", ...), which is what shards the attention across the mesh. Upstream train_wan.py / train_flux.py import the trainer INSIDE the guard (train() does the import within `with transformer_engine_context()`). Importing it outside left the attention unsharded, so TE's CK fused-attention kernel was unavailable for the full per-device sequence and fell back to a materialized S^2 matmul. Move the trainer import inside the guard (parity with upstream). Verified on wan2.1_1.3b: no OOM, 19 steps, 4.47s/step, 913.8 TFLOP/s/gpu. Also point JAX_COMPILATION_CACHE_DIR/MIOPEN_CUSTOM_CACHE_DIR at the bind-mounted /myworkspace/.jax_cache (NFS) instead of the ephemeral in-container /app/.cache, so the ~50-min cold compile (autotune_level=5 on the long-sequence fused-attention graph) is paid once and reused across runs instead of recompiling every run and risking a launcher timeout.
origin/jax-maxdiffusion-v26.6 (dcf9725) was an earlier standalone push of the MaxDiffusion backend, based directly on jax-maxtext-v26.6 without the declarative "Option C" env refactor. This local branch supersedes it: it carries the Option C env spec, the full WAN/FLUX backend, per-step metrics capture, and the TE mesh-guard OOM fix -- verified end-to-end on MI355X (wan2.1_1.3b 913.8 TFLOP/s/gpu, wan2.1_14b ~795, flux_dev 760.8). Reconcile with `-s ours` so the branch fast-forwards cleanly onto origin while keeping this verified tree intact (origin's diverged content remains recoverable from history / the bak/ tag).
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
| except Exception: # noqa: BLE001 - detection must never abort a run | ||
| pass | ||
|
|
||
| _ARCH_CACHE = arch |
Add Primus model configs (gemma4_31B, gemma4_26B -> MaxText gemma4-31b/26b, bundled tokenizer.gemma3) and MI355X + MI300X example configs (bf16 + fp8) for pretraining on synthetic data. Discovered automatically by MAD's jax-maxtext generator (tags: maxtext/jax/<arch>/<precision>). Known limitation (gfx950 / jax_the_rock_ci_45156b7 + jaxlib 0.10.0): gemma4 does not yet run on this JAX-ROCm stack -- the efficient attention kernels (cudnn_flash_te default, and flash) crash natively during XLA compile for gemma4's head_dim=256 + bidirectional sliding-window attention (fails even at per_device_batch_size=1, so not a memory/batch-size issue), and dot_product cannot compile in reasonable time at seq len 8192. Configs are committed ready for when gemma4 attention support lands; excluded from the active MAD sweep.
The MAD v26.6 sweep confirmed gemma4 trains on maxtext release/v26.6 via the reference dot_product attention (31B dense ~395 TFLOP/s/gpu, 26B MoE ~117 TFLOP/s/gpu, seq 4096). The fused/flash kernels still native-crash on gemma4's head_dim=256 + bidirectional sliding-window on the gfx950 ROCm/jaxlib 0.10.0 stack. Sets attention=dot_product, max_target_length=4096, per_device_batch_size=1 on the MI355X bf16 pretrain configs, and updates the gemma4 model-config notes to reflect the validated status (previously marked as not-running/excluded-from-sweep, which the v26.6 sweep disproved).
The v26.6 base.yml flipped pure_nnx_decoder false->true. The legacy per-module linen Fp8Einsum used by the MoE quant path requires an active Linen binding scope; under the pure-NNX decoder it is invoked unbound and crashes at the first step. deepseek_v2_16B-fp8, mixtral_8x7B-fp8 and qwen3_30B_A3B-fp8 all regressed vs v26.5 for this reason. Override pure_nnx_decoder: false in the three fp8-MoE MI355X configs to run the decoder on the bridged path (v26.5 behavior). Validated on gfx950/MI355X with base ce580ea_20260804: deepseek ~456 TFLOP/s and qwen3_30B_A3B ~328 TFLOP/s train cleanly; mixtral no longer crashes.
Make Primus self-contained for the MaxDiffusion (JAX) backend so it can run from a bare checkout on the MaxText JAX base image, without the separate MAD primus_maxdiffusion image. Its env stays isolated from MaxText's. - third_party/maxdiffusion submodule (pin 68e06965) - requirements-maxdiffusion.txt: MaxDiffusion-only deps (transformers==4.57.3, scikit-image, imageio-ffmpeg, loguru, wandb); torch/torchvision installed separately from the ROCm wheel index - examples/maxdiffusion/setup_maxdiffusion_env.sh: idempotent installer that brings up the stack (torch + requirements + editable submodule) and applies the four historical patches (Flax-T5 clip rename, TF-preload-before-TE, Shardy-on, TE empty-CP-axis) previously baked into the MAD Dockerfile - run_pretrain.sh: BACKEND=maxdiffusion runs the setup script (unless PRIMUS_SKIP_PIP=1) and exports NVTE_FRAMEWORK=jax + MAXDIFFUSION_PATH; the MaxText path is unchanged (guarded by _is_maxdiffusion_backend) Also bump third_party/maxtext a7c6c7e (release/v26.5) -> 59a87dca (release/v26.6), the validated release commit, so the branch pin matches the release it is named for.
- Document running MaxDiffusion directly from a bare Primus checkout (vendored submodule + setup_maxdiffusion_env.sh via run_pretrain.sh), including MI355X config status, in pretraining.md and examples/README.md. - Document the fp8-MoE requirement to set pure_nnx_decoder: false under the v26.6 pure-NNX decoder default (parameters ref + pretraining note). - Document that gemma4 on gfx950 (MI355X) must use attention: dot_product (head_dim=256 + sliding-window has no fused/flash kernel coverage).
yeandy
changed the title
MaxDiffusion's GPU init (max_utils.initialize_jax_for_gpu) only calls jax.distributed.initialize() when JAX_COORDINATOR_IP is set, keyed on num_processes=NNODES / process_id=NODE_RANK (one process per node). The launcher exported NNODES/NODE_RANK/MASTER_ADDR but not JAX_COORDINATOR_IP/PORT, so a multi-node MaxDiffusion job had each node initialize as a standalone single-node run. Export JAX_COORDINATOR_IP/PORT (from MASTER_ADDR/PORT) in the maxdiffusion branch, mirroring the MaxText env_spec, so 2+ node runs rendezvous.
The 14B graph makes RCCL instantiate 448 collective channels, over WarpSpeed's
256-channel ceiling. RCCL warns on every rank ("WarpSpeed does not support more
than 256 channels ... To avoid hang, run with RCCL_WARP_SPEED_AUTO=0") and the
launcher dies at comm-init (exit 1) right after re-enabling GC. Set
RCCL_WARP_SPEED_AUTO=0 in this config's env block (per-config, per the
MaxDiffusion env-policy in backends/maxdiffusion/env_spec.py), mirroring the
MaxText gfx950 default. wan2.1_1.3b/flux stay under 256 channels and are
untouched.
The MI355X fp8-MoE configs were pinned in 6038de6 but their MI300X counterparts were missed, so mixtral_8x7B-nanoo_fp8 crashed before step 1 with: AttributeError: "Fp8Einsum" object has no attribute "input_scale" v26.6 flipped the base.yml default to pure_nnx_decoder: true, which runs the decoder in pure NNX and invokes the legacy Linen Fp8Einsum (the MoE sparse-matmul quant path) without a Linen binding scope. Applied to all three MoE + fp8 configs on MI300X. Dense fp8 and bf16 configs are unaffected and left alone. Co-authored-by: Cursor <cursoragent@cursor.com>
The MI300X configs set JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES=all but no cache directory, so the cache was never initialized: Not writing persistent cache entry with key 'jit_train_step-...' since cache is disabled/not initialized Every run therefore recompiled the whole graph. Measured on a 10-step wan2.1_1.3b run: ~41 min of 46 min wall clock was compilation. Paths are relative to the Primus root, which is the working directory on every launch path and is bind-mounted in container mode, so the cache survives the container. The MI355X configs use absolute /myworkspace paths, which only exist in the MAD image. Deliberately no CK_JIT_CACHE_DIR: TransformerEngine builds its CK kernels via ck_jit_compile.sh, spawned with a working directory that is not the Primus root, so a relative path there resolves elsewhere and the build fails with "mktemp: ... No such file or directory", taking the run down. Recorded as a warning comment in the 1.3B config. Co-authored-by: Cursor <cursoragent@cursor.com>
A crashing run exited silently. init_worker_logger rebinds builtins.print
to debug_rank_all, which discarded the file= kwarg and logged everything
at DEBUG. traceback.print_exc() emits via print(..., file=sys.stderr), so
the fatal stack landed at DEBUG in debug.log while the console (INFO) and
error.log showed nothing -- error.log stayed 0 bytes on a failed run.
Affects every backend: the patch site is core runtime with no backend
branching.
Promote to ERROR only for stderr writes made while an exception is being
handled. That is precisely the traceback case and it excludes backends
that use stderr for ordinary chatter outside any handler (e.g.
fused_residual_rmsnorm logs "disabled (set PRIMUS_FUSED_RESIDUAL_NORM=1
...)" on every Megatron run), which would otherwise put a spurious ERROR
on the console of healthy runs.
Also wire up --debug, which was parsed and never read: its help promised
"verbose error output" while the console stayed pinned at the
module_base.yaml default. It now publishes PRIMUS_LOG_LEVEL=DEBUG, which
the runtime logger honours as a verbosity *floor* -- never an override,
because runner/lib/common.sh exports PRIMUS_LOG_LEVEL unconditionally as
${PRIMUS_LOG_LEVEL:-INFO}, and treating that always-present default as an
override would silently disable the stderr_sink_level / sink_level config
knobs. WARN/FATAL are aliased to loguru's WARNING/CRITICAL, and unknown
values fall back to the config instead of raising.
Verified against real sinks: the traceback reaches the console and
error.log at every level, and informational stderr stays hidden at the
default verbosity.
Co-authored-by: Cursor <cursoragent@cursor.com>
MaxDiffusion's max_logging.log is a bare print, and Primus rebinds builtins.print to a DEBUG-level logger call. Every per-step "completed step: N, ... loss: ..." line was therefore emitted at DEBUG and dropped by the console sink (INFO), surviving only in debug.log -- so a healthy run looked like it produced no output at all. MaxText already ships the equivalent patch; MaxDiffusion had no patches package, which its own runs reported unnoticed as "[Patch] Applied 0/0 patches for maxdiffusion/setup". Rebinding the module attribute is sufficient because every MaxDiffusion call site is module-qualified. log_rank_0 is used rather than MaxText's primus_logger.info: it resolves the sink at call time instead of capturing _logger at import, gates on rank 0, and attributes each line to the calling trainer rather than to max_logging itself. The rebinding happens before the confirmation is logged, so the patch cannot be skipped by a failure in the very logging path it repairs. Co-authored-by: Cursor <cursoragent@cursor.com>
gemma4 uses bidirectional sliding-window attention in its local layers, and with dataset_type: synthetic the data is unpacked, which MaxText's efficient attention path rejects outright: AssertionError: Sliding window attention requires context parallelism with load-balanced ring strategy and packing enabled. fc4dc19 fixed this for MI355X but only touched the MI355X configs. Mirror it here: dot_product avoids the cudnn path so the restriction never applies, and max_target_length drops to 4096 because dot_product compilation is very heavy at 8192. Verified on MI300X / gfx942, 5 steps each: gemma4_31B (Dense) 3.22 s/step, 245 TFLOP/s/device gemma4_26B (MoE) 1.28 s/step, 80.4 TFLOP/s/device (MI355X reports ~395 and ~117 on gfx950 for comparison.) Co-authored-by: Cursor <cursoragent@cursor.com>
- gemma4_26B-fp8: max_target_length 8192 -> 4096 and set attention: dot_product to match the bf16 config (fused/flash native-crashes on gemma4 head_dim=256 + sliding-window). - qwen3_30B_A3B (bf16 + fp8): max_target_length 8192 -> 4096, mirroring the MI300X seq-length change.
Derive the precision token from each config's own quantization setting rather than its previous filename, so a name can no longer disagree with what the config actually runs: no quantization key means bf16, "fp8" means fp8, and "nanoo_fp8" means nanoo_fp8. This corrects three files that were previously misdescribed. Both MI300X gemma4 configs set nanoo_fp8 while being named -fp8-, and grok1 carried no precision in its name at all despite quantizing. The MI355X -fp8- configs keep their token because they genuinely use OCP fp8, which is a different format from nanoo_fp8. grok1 also switches from fp8 to nanoo_fp8, the one behavior change here; every other config is a pure rename. Updates the 28 references in the benchmark workflow, the maxtext trainer tests, and the docs. Co-authored-by: Cursor <cursoragent@cursor.com>
- gemma4 (26B bf16/fp8, 31B bf16/fp8): attention dot_product -> cudnn_flash_te. The fused path used to abort at startup because maxtext routed synthetic batches through the context-parallel branch, which rejects AttentionType.LOCAL_SLIDING; maxtext now gives synthetic data its own mask-free causal branch, so the window reaches the kernel as window_size. - The pin stays explicit rather than inheriting autoselected: on GPU that resolves to a pallas kernel that builds a plain causal mask, which would silently drop the window on gemma4's local layers. - gemma4_31B (bf16): max_target_length 4096 -> 8192, which had only been lowered to keep dot_product's O(S^2) step tractable. - gemma4_26B-fp8: pure_nnx_decoder false, as on the other MI355X fp8-MoE configs. This config still does not run, but it now reaches the same upstream flax tracer leak MI300X hits instead of failing in attention. - docs: drop the claim that gemma4 must use dot_product on MI355X. Validated on MI355X (8x gfx950) through both primus-cli and MAD, agreeing to within 0.1%: 26B bf16 115, 31B bf16 281, 31B fp8 341 TFLOP/s/device. Needs the matching third_party/maxtext bump, which lands separately.
Picks up the synthetic-data mask-free causal branch the MI355X gemma4 configs depend on. On the previous pin, synthetic batches took the context-parallel path, which rejects AttentionType.LOCAL_SLIDING, so gemma4 aborted at startup under cudnn_flash_te. Tree-identical to the commit the MI355X sweep validated against, so this changes no code relative to what was tested. It also moves off 59a87dca, which was left unreferenced when release/v26.6 was rewritten upstream.
The two arch config dirs drifted in both directions: the MI355X hardening in 91552d6/30f5648 never reached MI300X, and the MI300X collective-timeout fix in 53790aa never reached MI355X or the MI300X wan configs. - Raise both first_collective_call timeouts in all six configs. Every config runs autotune_level=5, so step 1 spends >10 min in per-device autotune and the 8 device threads reach the first kAllGather minutes apart; XLA's 20s/40s rendezvous defaults then abort a process where every thread is still making progress (SIGABRT, exit 134). Only MI300X/flux_dev carried the raised values. - Port the 11-var ROCm/HIP + RCCL/NCCL block to the MI300X configs. Dropping five of these caused an RCCL init hang on MI355X (GPUs 0%, CPU busy-wait) and MI300X carried none of them. Nothing else supplies them: the MaxDiffusion adapter contributes no env (env_spec.py returns []) and MAD's run.sh exports no tuning, so the per-config env: block really is the only source. - Bind metrics_file to PERF_METRICS_FILE and set log_period: 1 on MI300X. run.sh always exports PERF_METRICS_FILE and passes --metrics-file to the extractor, but only a config that binds metrics_file makes the trainer write it, and the per-step stdout line does not survive the Primus launcher -- so MI300X perf capture had no working source. - Unify the compile-cache dirs on the relative ./.jax_cache/<model> form, which resolves against the Primus root (the working directory on every launch path, bind-mounted in container mode), instead of /myworkspace, which only exists in the MAD image. RCCL_WARP_SPEED_AUTO stays MI355X-only, since WarpSpeed is gfx950. The MI300X comments claiming run.sh applies it were stale and are corrected here.
Mirrors e29d78b, which did the same for MI355X. The assertion that forced dot_product here ("Sliding window attention requires context parallelism with load-balanced ring strategy and packing enabled") was a maxtext bug: synthetic batches were routed through the context-parallelism branch, which rejects sliding-window attention. That is fixed in the pinned release/v26.6 at 2ec83add and the fix is not arch-specific, so both arches now take the same path -- cudnn_flash_te is the only setting that passes the local layers' sliding window to the kernel, as window_size. The global layers' head_dim=512 is past TE's fused limit and still falls back to unfused attention. Validated on gfx950; not yet re-run on gfx942, which each config records. Also rewrites the gemma4_31B-nanoo_fp8 status note. Its kAllGather rendezvous abort after a ~10 min compile matches the first-collective-call watchdog analysis in examples/maxdiffusion/configs/MI300X/flux_dev-pretrain.yaml, not the missing NCCL env it blamed: maxtext already applies GPU_MAX_HW_QUEUES, HIP_FORCE_DEV_KERNARG and HSA_FORCE_FINE_GRAIN_PCIE on every arch (plus HSA_NO_SCRATCH_RECLAIM on gfx942) from primus/backends/maxtext/env_spec.py, so only NCCL_PROTO and RCCL_MSCCL_ENABLE are absent there.
| from pathlib import Path | ||
| from typing import Any, List, Optional | ||
|
|
||
| import primus.backends.maxdiffusion.patches # noqa: F401 # Register patches |
yeandy
changed the title
The shipped batch sizes were leaving most of the card idle on 26B-bf16. Walking each model up until it OOMs on this 288GB host gives 26B-bf16 20 (was 1), 31B-bf16 3 (was 1), and 31B-fp8 3 (was 2); one rung higher OOMs in every case. Measured over 50 steps through MAD: gemma4_26B-bf16 114.5 -> 204.5 TFLOP/s (+78.6%) gemma4_31B-bf16 277.9 -> 300.9 TFLOP/s (+8.3%) gemma4_31B-fp8 341.4 -> 348.5 TFLOP/s (+2.1%) Note this changes the recipe, not just the speed: per-device batch multiplies by the 8 devices, so 26B-bf16's global batch goes from 8 to 160. These three models also take 64 minutes instead of 36, since a step now does 20x or 3x the work. gemma4_26B-fp8 keeps batch 2 because it does not train at any batch size (flax fp8-MoE tracer leak, upstream). The ceilings are specific to MI355X's 288GB, so the MI300X configs are untouched. The comment explaining the cudnn_flash_te pin is dropped here and recorded with the sweep results instead.
On an image that bakes maxdiffusion there is nothing to install, and often no submodule to install from, so the script failed its source check on a perfectly good environment. Detect an importable maxdiffusion first and exit early. PRIMUS_SKIP_SETUP_CHECK=1 forces the install path. PRIMUS_SKIP_PIP=1 already short-circuits this script from the launcher, so this only matters when it is invoked directly.
# Conflicts: # .gitmodules # docs/02-user-guide/training-recipes.md # examples/maxtext/configs/MI300X/qwen3_30B_A3B-bf16-pretrain.yaml # primus/_thirdparty.lock
cudnn_flash_te is now validated; remove verbose history of the earlier dot_product workaround. Condense STATUS blocks and model YAML headers.
Add a table in pretraining.md and a header comment in the script itself explaining the three scenarios (image has it, image doesn't, submodule missing).
Fp8Einsum.setup() allocates Flax variables (amax_history, scales) inside jax.lax.scan when gemma4's scanned local layers invoke RoutedMoE.dense_matmul, causing UnexpectedTracerError. The new patch makes get_einsum() return plain jnp.einsum for FP8 MoE matmuls when scan_layers is true, so expert matmuls run in bf16 while dense layers still use FP8 via dot_general_cls. Validated on MI355X: 148.7 TFLOP/s/device, 5940 tok/s/device (gemma4_26B).
mixtral_8x7b and qwen3-30b-a3b explicitly set scan_layers: false in their upstream model configs, so Fp8Einsum.setup() runs outside any scan transform and doesn't leak. gemma4 inherits the default scan_layers: true and uses its own _scan_local_layers, which is why the patch is needed only for gemma4.
The Dockerfile blanks LD_LIBRARY_PATH to prevent TE 2.17's broken hipblaslt Tensile kernels from being loaded on gfx950. However, ensure_rocm_ld_library_path() re-introduces _rocm_sdk_devel/lib at runtime, undoing that fix and causing SIGSEGV during training. Gate the call behind a BACKEND check so JAX backends (MaxText, MaxDiffusion) keep LD_LIBRARY_PATH empty and resolve ROCm libs via wheel RPATH as intended.
Prevents OOM on 70B models by allowing JAX to use 96% of device memory.
…vel env block - Move XLA_PYTHON_CLIENT_MEM_FRACTION from overrides (wrong) to top-level env: section (.97) for llama2_70B, llama2_70B-nanoo_fp8, llama3_70B, llama3.3_70B MI300X configs - Add status comment to qwen3_30B_A3B-fp8 noting possible NaN from FP8 dynamic range overflow in MoE expert layers
wenxie-amd
approved these changes
Aug 21, 2026
GeneDer
added a commit
that referenced
this pull request
Aug 22, 2026
Adds per-precision MaxText experiment configs (fp8, nanoo_fp8) derived from the existing bf16 configs by setting `quantization` in overrides, for 9 models across MI300X and MI355X. Upstream MaxText previously shipped only bf16 (`<model>-pretrain.yaml`); this fills the fp8/nanoo_fp8 gap so quantized runs can be launched by config path. --------- Co-authored-by: Andy <14128880+yeandy@users.noreply.github.com> Co-authored-by: Andy Ye <andyye12@crs-m2m-cpu-spur-009.us-east2-a.compute.internal> Co-authored-by: Fuyuan Jing <Fuyuan.Jing@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds per-precision MaxText experiment configs (fp8, nanoo_fp8) derived from the existing bf16 configs by setting
quantizationin overrides, for 9 models across MI300X and MI355X. Upstream MaxText previously shipped only bf16 (<model>-pretrain.yaml); this fills the fp8/nanoo_fp8 gap so quantized runs can be launched by config path.