feat(maxtext): Maxtext & MaxDiffusion refactor + FP8 separate configs - #925

Open

yeandy wants to merge 19 commits into

mainfrom

jax-maxtext-v26.6

Commits on Jul 27, 2026

feat(maxtext): add fp8 and nanoo_fp8 pretrain configs for MI300X/MI355X
yeandy
committed
74d7a2c
feat(maxtext): scope quantized configs per device (MI300X nanoo_fp8, MI355X fp8)
yeandy
committed
b5b43b6
feat(maxtext): add MI355X fp8 config for llama3.1_405B
yeandy
committed
7310eda

Commits on Jul 30, 2026

feat(maxtext): per-config top-level env: loader; mixtral 8x7B bs11 @ mem_fraction 0.96
yeandy
committed
1ad7522

Commits on Jul 31, 2026

feat(maxdiffusion): add Primus JAX MaxDiffusion backend (WAN/FLUX)
yeandy
committed
dcf9725

Commits on Aug 1, 2026

feat(maxtext): consolidate env management into a declarative Primus spec (Option C)
Andy Ye
committed
427702c
feat(maxdiffusion): add Primus JAX MaxDiffusion backend (WAN/FLUX)

yeandy
authored and
Andy Ye
committed
d61b61c
feat(maxdiffusion): adopt declarative env_defaults() (align with maxtext Option C)
Andy Ye
committed
40dfae8

Commits on Aug 2, 2026

fix(maxdiffusion): restore full ROCm/RCCL env in MI355X configs; log every step
Andy Ye
committed
91552d6
feat(maxdiffusion): write per-step metrics to a file for reliable perf capture
Andy Ye
committed
30f5648
fix(maxdiffusion): import trainer inside TE mesh-guard to fix WAN OOM
Andy Ye
committed
8ba08d2
Merge origin/jax-maxdiffusion-v26.6 (superseded backend) using ours
Andy Ye
committed
d7264f0

Commits on Aug 3, 2026

Add Gemma 4 31B (Dense) & 26B (MoE) MaxText pretrain configs
Andy Ye
committed
8b4e441

Commits on Aug 4, 2026

Use validated dot_product attention for gemma4 MI355X pretrain configs
Andy Ye
committed
fc4dc19

Commits on Aug 5, 2026