feat(maxtext): Maxtext & MaxDiffusion refactor

.github/workflows/benchmark.yaml

@@ -463,43 +463,43 @@ jobs:
           GLOO_SOCKET_IFNAME: "eth0"
         run: |
           ./runner/primus-cli direct --log_file "${BENCHMARK_LOG_DIR}/maxtext/deepseek_v2_16B.log" \
-          -- train pretrain --config examples/maxtext/configs/MI300X/deepseek_v2_16B-pretrain.yaml
+          -- train pretrain --config examples/maxtext/configs/MI300X/deepseek_v2_16B-bf16-pretrain.yaml
       - name: "[2/7][MaxText] llama2_70B"
         timeout-minutes: 120
         continue-on-error: true
         run: |
           ./runner/primus-cli direct --log_file "${BENCHMARK_LOG_DIR}/maxtext/llama2_70B.log" \
-          -- train pretrain --config examples/maxtext/configs/MI300X/llama2_70B-pretrain.yaml
+          -- train pretrain --config examples/maxtext/configs/MI300X/llama2_70B-bf16-pretrain.yaml
       - name: "[3/7][MaxText] llama2_7B"
         timeout-minutes: 120
         continue-on-error: true
         run: |
           ./runner/primus-cli direct --log_file "${BENCHMARK_LOG_DIR}/maxtext/llama2_7B.log" \
-          -- train pretrain --config examples/maxtext/configs/MI300X/llama2_7B-pretrain.yaml
+          -- train pretrain --config examples/maxtext/configs/MI300X/llama2_7B-bf16-pretrain.yaml
       - name: "[4/7][MaxText] llama3.3_70B"
         timeout-minutes: 120
         continue-on-error: true
         run: |
           ./runner/primus-cli direct --log_file "${BENCHMARK_LOG_DIR}/maxtext/llama3.3_70B.log" \
-          -- train pretrain --config examples/maxtext/configs/MI300X/llama3.3_70B-pretrain.yaml
+          -- train pretrain --config examples/maxtext/configs/MI300X/llama3.3_70B-bf16-pretrain.yaml
       - name: "[5/7][MaxText] llama3_70B"
         timeout-minutes: 120
         continue-on-error: true
         run: |
           ./runner/primus-cli direct --log_file "${BENCHMARK_LOG_DIR}/maxtext/llama3_70B.log" \
-          -- train pretrain --config examples/maxtext/configs/MI300X/llama3_70B-pretrain.yaml
+          -- train pretrain --config examples/maxtext/configs/MI300X/llama3_70B-bf16-pretrain.yaml
       - name: "[6/7][MaxText] llama3_8B"
         timeout-minutes: 120
         continue-on-error: true
         run: |
           ./runner/primus-cli direct --log_file "${BENCHMARK_LOG_DIR}/maxtext/llama3_8B.log" \
-          -- train pretrain --config examples/maxtext/configs/MI300X/llama3_8B-pretrain.yaml
+          -- train pretrain --config examples/maxtext/configs/MI300X/llama3_8B-bf16-pretrain.yaml
       - name: "[7/7][MaxText] mixtral_8x7B"
         timeout-minutes: 120
         continue-on-error: true
         run: |
           ./runner/primus-cli direct --log_file "${BENCHMARK_LOG_DIR}/maxtext/mixtral_8x7B.log" \
-          -- train pretrain --config examples/maxtext/configs/MI300X/mixtral_8x7B-pretrain.yaml
+          -- train pretrain --config examples/maxtext/configs/MI300X/mixtral_8x7B-bf16-pretrain.yaml
       - name: Generate Summary Report
         run: |
           echo "Generate Summary Report"

.gitignore

@@ -36,4 +36,5 @@ pp_simulation_result
 /nohup.*
 /core.*.gpu
 .triton_cache_shared/
+.jax_cache/
 .cursor/

.gitmodules

@@ -8,7 +8,7 @@
 [submodule "third_party/maxtext"]
 	path = third_party/maxtext
 	url = https://github.com/ROCm/maxtext.git
-	branch = release/v26.5
+	branch = release/v26.6
 [submodule "third_party/Emerging-Optimizers"]
 	path = third_party/Emerging-Optimizers
 	url = https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git
@@ -23,6 +23,9 @@
 [submodule "third_party/HummingbirdXT"]
 	path = third_party/HummingbirdXT
 	url = https://github.com/AMD-AGI/HummingbirdXT.git
+[submodule "third_party/maxdiffusion"]
+	path = third_party/maxdiffusion
+	url = https://github.com/AI-Hypercomputer/maxdiffusion.git
 [submodule "third_party/Automodel"]
 	path = third_party/Automodel
 	url = https://github.com/NVIDIA-NeMo/Automodel.git

docs/02-user-guide/pretraining.md

@@ -15,6 +15,7 @@ The following table describes the four backend types supported by Primus and the
 | Megatron-LM | `framework: megatron` | Large-scale transformer pretraining with Megatron-style parallelism (TP/PP/EP). |
 | TorchTitan | `framework: torchtitan` | PyTorch-native scaled training (FSDP / tensor / pipeline / expert parallelism per config). |
 | MaxText (JAX) | `framework: maxtext` | JAX/MaxText single- and multi-node runs; parallelism via MaxText `ici_*` / `dcn_*` settings. |
+| MaxDiffusion (JAX) | `framework: maxdiffusion` | JAX/MaxDiffusion diffusion pretraining (WAN 2.1, FLUX.1-dev). Source is vendored as the `third_party/maxdiffusion` submodule; deps/patches installed by `examples/maxdiffusion/setup_maxdiffusion_env.sh`. |
 | Megatron Bridge | `framework: megatron_bridge` | Bridge-oriented workflows (configure like other backends; see parameter reference). |
 
 > Several setup steps apply to **all** backends (mock vs. real data, Hugging Face tokens, scaling to multiple nodes, and HipBLASLt autotuning). After you read the backend section that applies to you, see [Common patterns](#common-patterns) below.
@@ -197,7 +198,7 @@ Some configurations omit an explicit `parallelism:` block; in that case the defa
 
 ```bash
 ./runner/primus-cli container -- train pretrain \
-  --config examples/maxtext/configs/MI300X/llama2_7B-pretrain.yaml
+  --config examples/maxtext/configs/MI300X/llama2_7B-bf16-pretrain.yaml
 ```
 
 ### JAX-specific requirements
@@ -212,18 +213,87 @@ pip install -r requirements-jax.txt
 
 | File | Key parallelism (`ici_*` intra-node, `dcn_*` inter-node) |
 | --- | --- |
-| `deepseek_v2_16B-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
-| `grok1-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
-| `llama2_70B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
-| `llama2_7B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
-| `llama3.3_70B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
-| `llama3_70B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
-| `llama3_8B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
-| `mixtral_8x7B-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
-| `qwen3_14B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
-| `qwen3_30B_A3B-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
-
-The `llama2_7B-pretrain.yaml` example also sets `dataset_type: "synthetic"` and `hf_access_token: ${HF_TOKEN:""}` for gated Hugging Face assets when you switch to real data.
+| `deepseek_v2_16B-bf16-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `grok1-nanoo_fp8-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `llama2_70B-bf16-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `llama2_7B-bf16-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `llama3.3_70B-bf16-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `llama3_70B-bf16-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `llama3_8B-bf16-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `mixtral_8x7B-bf16-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `qwen3_14B-bf16-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `qwen3_30B_A3B-bf16-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+
+The `llama2_7B-bf16-pretrain.yaml` example also sets `dataset_type: "synthetic"` and `hf_access_token: ${HF_TOKEN:""}` for gated Hugging Face assets when you switch to real data.
+
+> **fp8 MoE (v26.6):** fp8 Mixture-of-Experts configs must set `pure_nnx_decoder: false` in their overrides; otherwise they crash at step 1 under the v26.6 pure-NNX decoder default. See [MaxText parameters → Precision and quantization](../03-configuration-reference/maxtext-parameters.md#7-precision-and-quantization). Dense fp8 and bf16 configs are unaffected.
+
+---
+
+## MaxDiffusion (JAX) pretraining
+
+The MaxDiffusion backend runs JAX diffusion pretraining (WAN 2.1, FLUX.1-dev). Environment setup depends on your image:
+
+| Image has `maxdiffusion` installed? | What happens |
+| --- | --- |
+| **Yes** (e.g. MAD `primus_maxdiffusion` image, unified docker) | `setup_maxdiffusion_env.sh` detects it and is a **no-op**. Set `PRIMUS_SKIP_PIP=1` to skip calling it entirely. |
+| **No** (e.g. bare `rocm/jax-training:maxtext-*` image) | The script installs everything from the Primus checkout: torch (ROCm wheels), deps, editable submodule, and patches. Requires `third_party/maxdiffusion` submodule to be initialized. |
+
+The relevant pieces:
+
+- **Source** is vendored as the `third_party/maxdiffusion` submodule.
+- **Dependencies** live in `requirements-maxdiffusion.txt` (kept separate from `requirements-jax.txt` so the MaxDiffusion pins never affect MaxText runs).
+- **Install + patches** are applied by `examples/maxdiffusion/setup_maxdiffusion_env.sh` (idempotent): torch/torchvision (ROCm wheels), the requirements above, an editable install of the vendored submodule, and four source patches (Flax-T5 clip rename, TensorFlow-preload-before-TransformerEngine, Shardy-on, and the TransformerEngine empty context-parallel-axis fix).
+
+### Prerequisites
+
+Initialize the vendored submodule (a plain clone will not populate it):
+
+```bash
+git submodule update --init third_party/maxdiffusion
+```
+
+Run on a JAX base image (for example `rocm/jax-training`) or a bare-metal JAX environment, and export `HF_TOKEN` for gated Hugging Face assets.
+
+### Quick start (run from a bare Primus checkout)
+
+Use `run_pretrain.sh` with `BACKEND=MaxDiffusion`. When `PRIMUS_SKIP_PIP` is unset, the launcher runs `setup_maxdiffusion_env.sh` for you (installs the stack + applies the patches), sets `NVTE_FRAMEWORK=jax` and `MAXDIFFUSION_PATH`, then launches:
+
+```bash
+BACKEND=MaxDiffusion \
+EXP=examples/maxdiffusion/configs/MI355X/wan2.1_1.3b-pretrain.yaml \
+  bash ./examples/run_pretrain.sh
+```
+
+To run the environment setup once by itself (e.g. to warm an image or a shared venv), invoke the script directly, then launch with `PRIMUS_SKIP_PIP=1`:
+
+```bash
+bash examples/maxdiffusion/setup_maxdiffusion_env.sh
+PRIMUS_SKIP_PIP=1 BACKEND=MaxDiffusion \
+EXP=examples/maxdiffusion/configs/MI355X/flux_dev-pretrain.yaml \
+  bash ./examples/run_pretrain.sh
+```
+
+### Quick start (container mode)
+
+`primus-cli` bootstraps the same environment: the `train/pretrain/maxdiffusion` prepare hooks run `setup_maxdiffusion_env.sh` before training and select the plain-python launcher (JAX drives every GPU from one process, so `torchrun` is never used).
+
+```bash
+./primus-cli container -- train pretrain \
+  --config examples/maxdiffusion/configs/MI300X/wan2.1_1.3b-pretrain.yaml --max_train_steps 10
+```
+
+> Container launches start from a clean image each time, so the setup runs on every launch. Wheels are cached under `$DATA_PATH/pip_cache` inside the mounted checkout, so only the first run pays for downloads. Set `PRIMUS_SKIP_PIP=1` to skip the step entirely on images that already ship the stack.
+
+> Step counts use `--max_train_steps` (the MaxDiffusion field name). `--steps` belongs to MaxText and is silently ignored here.
+
+### Example configurations under `examples/maxdiffusion/configs/MI355X/`
+
+| File | Model | Status on MI355X (gfx950) |
+| --- | --- | --- |
+| `flux_dev-pretrain.yaml` | FLUX.1-dev | ✅ validated |
+| `wan2.1_1.3b-pretrain.yaml` | WAN 2.1 1.3B | ✅ validated |
+| `wan2.1_14b-pretrain.yaml` | WAN 2.1 14B | ✅ validated (requires `RCCL_WARP_SPEED_AUTO=0`, set in config) |
 
 ---
 

docs/03-configuration-reference/maxtext-parameters.md

@@ -111,6 +111,7 @@ These are Primus overlay defaults. MaxText also loads upstream `base.yml` at run
 | `checkpoint_is_quantized` | `false`           | Set `true` when loading an AQT-quantized checkpoint.                  |
 | `logits_dot_in_fp32`      | `false`           | Compute logits matmul in `float32` for numerical stability.           |
 
+**fp8 + MoE (v26.6):** MoE fp8 configs (e.g. `mixtral_8x7B-fp8`, `qwen3_30B_A3B-fp8`) must set `pure_nnx_decoder: false`. v26.6 defaults to `true`, which invokes the legacy Linen `Fp8Einsum` without a binding scope. See [Advanced](#9-advanced). Dense fp8 configs are unaffected.
 
 ---
 
@@ -123,7 +124,7 @@ From `model_base.yaml` and per-model files such as `llama3_8B.yaml`.
 | ----------------------- | ------------------------------------------------------------------- | ----------------------------------------------------------------------- |
 | `model_name`            | `"default"` in `model_base`; e.g. `"llama3-8b"` in `llama3_8B.yaml` | Selects MaxText’s bundled model YAML when present.                      |
 | `override_model_config` | `true`                                                              | When `true`, CLI / kwargs override values from the loaded model config. |
-| `attention`             | `"cudnn_flash_te"`                                                  | Attention implementation (Primus default favors TE flash on AMD GPUs).  |
+| `attention`             | `"autoselected"`                                                    | Attention implementation. **gemma4 requires `cudnn_flash_te`**: its local layers use sliding-window attention, and only `cudnn_flash_te` passes the window to the kernel (as `window_size`). `autoselected` resolves to a pallas kernel that builds a plain causal mask and silently drops the window. |
 | `use_iota_embed`        | `true`                                                              | Use iota-based embedding for performance on accelerator backends.       |
 | `tokenizer_path`        | e.g. `"meta-llama/Meta-Llama-3-8B"`                                 | Hugging Face tokenizer id or local path.                                |
 
@@ -133,9 +134,10 @@ From `model_base.yaml` and per-model files such as `llama3_8B.yaml`.
 ## 9. Advanced
 
 
-| Parameter | Default | Description                                                           |
-| --------- | ------- | --------------------------------------------------------------------- |
-| `shardy`  | `false` | Enable Shardy-related integration in MaxText when building shardings. |
+| Parameter           | Default | Description                                                                                                                                                                                                                            |
+| ------------------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `shardy`            | `false` | Enable Shardy-related integration in MaxText when building shardings.                                                                                                                                                                 |
+| `pure_nnx_decoder`  | `true`  | Run the decoder as pure Flax NNX. As of v26.6 the base config defaults this to `true`. **fp8 MoE configs must override this to `false`** — the legacy per-module Linen `Fp8Einsum` (used by the MoE sparse-matmul quant path) needs an active Linen binding scope, and the pure-NNX decoder invokes it unbound, crashing at the first step. See the note in [Precision and quantization](#7-precision-and-quantization). |
 
 
 ---

docs_deprecated/cli/PRIMUS-CLI-GUIDE.md

@@ -286,7 +286,7 @@ If you are using AINIC in your cluster, you can use the `runner/use_ainic.yaml`
 
 Here is an example of using the AINIC configuration file to run a training job:
 ```bash
-./primus-cli --config runner/use_ainic.yaml slurm srun -N 2 -- train pretrain --config examples/maxtext/configs/MI355X/llama2_7B-pretrain.yaml
+./primus-cli --config runner/use_ainic.yaml slurm srun -N 2 -- train pretrain --config examples/maxtext/configs/MI355X/llama2_7B-bf16-pretrain.yaml
 ```
 
 ### Configuration Priority

examples/README.md

@@ -144,6 +144,12 @@ export EXP=examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml
 export EXP=examples/torchtitan/configs/MI300X/llama3.1_8B-pretrain.yaml
 ./primus-cli direct -- train pretrain --config "$EXP"
 
+# MaxDiffusion (JAX) directly from Primus (on a JAX base image, e.g. rocm/jax-training).
+# Requires the vendored submodule: git submodule update --init third_party/maxdiffusion
+# run_pretrain.sh runs examples/maxdiffusion/setup_maxdiffusion_env.sh to install deps + patches,
+# then launches. See docs/02-user-guide/pretraining.md ("MaxDiffusion (JAX) pretraining").
+BACKEND=MaxDiffusion EXP=examples/maxdiffusion/configs/MI355X/wan2.1_1.3b-pretrain.yaml bash ./examples/run_pretrain.sh
+
 ```
 
 ---