feat(maxtext): Maxtext & MaxDiffusion refactor - #925

Merged

wenxie-amd merged 52 commits into

mainfrom

jax-maxtext-v26.6

Aug 21, 2026

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

Commits on Aug 7, 2026

make primus-cli work for maxdiffusion
amd-fuyuajin
committed
096530e

Commits on Aug 10, 2026

Fix fp8-MoE crash on MI300X by pinning pure_nnx_decoder: false

amd-fuyuajin
and
cursoragent
committed
1dafc1a
maxdiffusion: persist XLA/MIOpen compile caches on MI300X

amd-fuyuajin
and
cursoragent
committed
90e5e9b
fix(logging): surface fatal tracebacks and make --debug work

amd-fuyuajin
and
cursoragent
committed
cc8d826
feat(maxdiffusion): route max_logging through the Primus logger

amd-fuyuajin
and
cursoragent
committed
f059e02
Fix gemma4 bf16 pretrain on MI300X with dot_product attention

amd-fuyuajin
and
cursoragent
committed
9bcac19
gemma4 fp8 on MI300X: use nanoo_fp8, record both as non-working

amd-fuyuajin
and
cursoragent
committed
ea72b1a
change sqeuence lenght of Qwen3-30B-A3B on MI300X to 4096
amd-fuyuajin
committed
94e1df9
MI355X: 4k seq length for gemma4_26B fp8 and qwen3_30B_A3B (bf16+fp8)
yeandy
committed
52fa494

Commits on Aug 11, 2026

change sequence length of gemma4_31B to 8192
amd-fuyuajin
committed
ceb1777
update MI300X gemma4_31B config file
amd-fuyuajin
committed
5468531
maxtext: name MI300X/MI355X configs <model>-<precision>-pretrain.yaml

amd-fuyuajin
and
cursoragent
committed
6ce4fc9
fix the config file for flux model on mi300x
amd-fuyuajin
committed
53790aa
MI355X: use cudnn_flash_te for gemma4
yeandy
committed
e29d78b
third_party/maxtext: bump to 2ec83add (release/v26.6)
yeandy
committed
ba052cf
fix(maxdiffusion): make the MI300X and MI355X configs consistent
yeandy
committed
57ad153
MI300X: use cudnn_flash_te for gemma4
yeandy
committed
eea17f8
primus/_thirdparty.lock: regenerate from submodule pins
yeandy
committed
966908b

Commits on Aug 12, 2026

Commits on Aug 13, 2026

Commits on Aug 14, 2026

Commits on Aug 16, 2026

fix: skip LD_LIBRARY_PATH injection for JAX backends on gfx950
yeandy
committed
9e9e031

Commits on Aug 18, 2026

Commits on Aug 21, 2026

Merge branch 'main' into jax-maxtext-v26.6
yeandy
committed
ed6d3f3
Merge branch 'main' into jax-maxtext-v26.6
yeandy
committed
371cf9b