Repository navigation

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

Merged

wenxie-amd merged 52 commits into

mainfrom

jax-maxtext-v26.6

Aug 21, 2026

Merged

feat(maxtext): Maxtext & MaxDiffusion refactor#925
wenxie-amd merged 52 commits into
mainfrom
jax-maxtext-v26.6

Commits on Jul 27, 2026
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
- MI355X: raise gemma4 per-device batch size to the measured ceiling
  yeandy
  committed
  3cc97e2
- maxdiffusion: skip env setup when the stack is already installed
  yeandy
  committed
  873f6af
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
- Add XLA_PYTHON_CLIENT_MEM_FRACTION=0.96 to 70B MI300X configs
  yeandy
  committed
  258e27a
- Fix 70B MI300X configs: move XLA_PYTHON_CLIENT_MEM_FRACTION to top-level env block
  yeandy
  committed
  1fa44e1
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

Commits

Commits on Jul 27, 2026

Commits on Jul 30, 2026

Commits on Jul 31, 2026

Commits on Aug 1, 2026

Commits on Aug 2, 2026

Commits on Aug 3, 2026

Commits on Aug 4, 2026

Commits on Aug 5, 2026

Commits on Aug 7, 2026

Commits on Aug 10, 2026

Commits on Aug 11, 2026

Commits on Aug 12, 2026

Commits on Aug 13, 2026

Commits on Aug 14, 2026

Commits on Aug 16, 2026

Commits on Aug 18, 2026

Commits on Aug 21, 2026