Commit 606baf8

committed

NNX: add sharding tools, Linen<->NNX converter, and checkpoint comparison utility

- modify print_shardings_params to support NNX (maxtext_utils.py) - add --pure_nnx flag to run_sharding_dump.py - add bidirectional Linen<->NNX checkpoint conversion utility (linen_nnx_converter.py) - add checkpoint comparison utility for Linen vs NNX validation (compare_linen_nnx_checkpoint.py)

1 parent 813868a commit 606baf8Copy full SHA for 606baf8

6 files changed

src/maxtext
- checkpoint_conversion
  - compare_linen_nnx_checkpoint.py
  - linen_nnx_converter.py
- utils
  - maxtext_utils.py
tests
- unit
  - compare_linen_nnx_checkpoint_test.py
  - linen_nnx_converter_test.py
- utils
  - run_sharding_dump.py

Comments

(0)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Commit 606baf8

File tree

0 commit comments