NNX train#3442

Draft

charlesli640 wants to merge 11 commits intoAI-Hypercomputer:mainfrom

CIeNET-International:charlesli/nnx_train

Commits on Mar 30, 2026

Add nnx_train
Charles Li
committed
Combine loss_fn for both train and evel
Charles Li
committed
Support gradient_accumulation and align to latest train.py
Charles Li
committed
Support muon
Charles Li
committed
NNX train_compile
Charles Li
committed
Fix unit test failures
Charles Li
committed
Chagne the config base.yml to set nnx flags to True
Charles Li
committed
Using jax.lax.scan instead of nnx.vmap to decrease memory usage
Charles Li
committed
Update MTP

hsuan-lun-chiang
authored and
Charles Li
committed
Fix out_sharding issue
Charles Li
committed

Commits on Mar 31, 2026

Add checkpoint converter
Charles Li
committed