NNX Migration: all commits#3075

Draft

xibinliu wants to merge 16 commits intomainfrom

Commits on Mar 24, 2026

Fix unit test errors

Charles Li
authored and
ecnal-cienet
committed
NNX: exclude Intermediate variables from scan state to fix qk_clip leaf count mismatch
ecnal-cienet
committed
NNX migration preparation: pure_nnx flag and init_state_fn

xibinliu
authored and
ecnal-cienet
committed
NNX migration: NNX utils

xibinliu
authored and
ecnal-cienet
committed
NNX TrainState and unit tests

xibinliu
authored and
ecnal-cienet
committed
NNX migration: add NNX support to muon_utils and refactor model_creation_utils

xibinliu
authored and
ecnal-cienet
committed
NNX train: add nnx model and trainstate support

xibinliu
authored and
ecnal-cienet
committed
NNX: add maybe_update_params_sharding_with_opt_nnx

xibinliu
authored and
ecnal-cienet
committed
NNX: support nnx model in the gradient accumulation

xibinliu
authored and
ecnal-cienet
committed
NNX: update train/eval step sharding signatures to omit rng for pure_nnx

xibinliu
authored and
ecnal-cienet
committed
NNX: fix checkpointing in the training loop
ecnal-cienet
committed
NNX: add bidirectional Linen<->NNX checkpoint conversion utility
ecnal-cienet
committed
NNX migration: modify the print_shardings_params to support NNX

xibinliu
authored and
ecnal-cienet
committed
NNX: add checkpoint comparison utility for Linen vs NNX validation
ecnal-cienet
committed
NNX: add --pure_nnx flag to run_sharding_dump.py
ecnal-cienet
committed
[WIP] NNX: fix model and test compatibility issues
ecnal-cienet
committed