[NNX] NNX migration prep (4/N): sharding tools, Linen↔NNX checkpoint utilities, and post-training fixes

[ecnal-cienet]

NNX Migration Route Map

1. ✅ Add NNX scaffolding: pure_nnx flag, init_state_fn, TrainStateNNX, NNX utils. Linen workflow unchanged. (PR #3427)
2. ✅ NNX sharding utilities: get_abstract_state_nnx, get_named_sharding_nnx, set_named_sharding_nnx, get_partition_spec_nnx, get_mesh_from_config. (PR #3470)
3. ✅ NNX fully supported end-to-end: TrainStateNNX, model creation, gradient accumulation, checkpointing, and training loop dispatch. (PR #3500)
4. 🔄 [This PR] NNX sharding diagnostics, bidirectional Linen↔NNX checkpoint conversion utilities, and post-training fixes.
5. ❌ Enable NNX by default; fix unit and integration test failures.
6. ❌ Remove Linen-specific code paths and NNX compatibility flags.

Description

This PR combines two closely related changes — NNX tooling additions and post-training bugfixes — into a single reviewable unit.

Part 1: Sharding diagnostics and Linen↔NNX checkpoint utilities

File	Change
src/maxtext/utils/maxtext_utils.py	Extend print_shardings_params to support NNX state trees in addition to Linen param dicts.
src/maxtext/tools/run_sharding_dump.py	Add --pure_nnx flag to select the NNX sharding path.
src/maxtext/utils/linen_nnx_converter.py (new)	Bidirectional Linen↔NNX checkpoint conversion utility: convert a Linen checkpoint to NNX format and vice versa.
src/maxtext/tools/compare_linen_nnx_checkpoint.py (new)	Checkpoint comparison tool for validating numerical equivalence between Linen and NNX checkpoints.

Part 2: Post-training bug fixes

File	Fix
src/maxtext/models/models.py	Transformer.__call__ was passing multimodal_input=MultimodalInput(...) to NNXDecoder, which only accepts individual fields (image_embeddings, image_masks, audio_embeddings, audio_masks, bidirectional_mask). Fixed by unpacking the object at the call site.
src/maxtext/optimizers/optimizers.py	adam_pax called learning_rate_fn(count) unconditionally, failing when optax.inject_hyperparams (used by the distillation trainer) passes a pre-evaluated scalar instead of a callable schedule. Fixed with a callable() guard.
src/maxtext/trainers/post_train/sft/train_sft.py	Tunix's default PeftTrainer._train_step nests nnx.value_and_grad inside nnx.jit, causing Flax NNX to assign conflicting outer_index values and raising ValueError: The graph structure of a node added to cached_partial was mutated. Fixed by adding MaxTextPeftTrainer subclass that overrides create_train_step_fn to use jax.value_and_grad with an explicit nnx.split/nnx.merge pattern (matching MaxText's pre-training NNX train step).
src/maxtext/trainers/post_train/distillation/train_distill.py	Same nested NNX transform issue in MaxTextDistillationTrainer._train_step. Additionally, teacher inference is now run outside value_and_grad (it's frozen via stop_gradient anyway) to avoid tracing it unnecessarily. Fixed with the same jax.value_and_grad + explicit split/merge pattern.
src/maxtext/trainers/post_train/rl/train_rl.py	Two RL-specific fixes: (1) JAX 0.9+ with_sharding_constraint asserts instead of reshards under Explicit mesh axes — patched with a try/except AssertionError fallback to jax.sharding.reshard. (2) tpu_inference initializes weights as float32 by default; during weight sync, tunix._apply_dtype_cast was upcasting incoming bfloat16 MaxText weights to float32, causing a dtype mismatch in the ragged paged attention kernel — patched to skip bfloat16→float32 upcasts so synced weights stay bfloat16. Also passes dtype explicitly to vLLM init.
src/maxtext/utils/model_creation_utils.py	Guard against None checkpoint metadata in create_nnx_model with a descriptive error message when the checkpoint directory is empty or the save did not complete.

Tests

• Sharding dump validated on V6e-8 with --pure_nnx flag.
• Linen↔NNX checkpoint conversion validated via compare_linen_nnx_checkpoint.py on gemma2-2b.
• Post-training fixes validated during SFT and distillation trial runs on V6e-8.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 65.82734% with 380 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/trainers/pre_train/train.py	32.35%	94 Missing and 21 partials ⚠️
...ckpoint_conversion/compare_linen_nnx_checkpoint.py	65.03%	97 Missing and 10 partials ⚠️
src/maxtext/utils/sharding.py	5.55%	50 Missing and 1 partial ⚠️
...xtext/checkpoint_conversion/linen_nnx_converter.py	87.82%	22 Missing and 16 partials ⚠️
src/maxtext/utils/train_utils.py	27.58%	19 Missing and 2 partials ⚠️
src/maxtext/utils/maxtext_utils.py	84.84%	8 Missing and 7 partials ⚠️
src/maxtext/utils/gradient_accumulation.py	27.77%	8 Missing and 5 partials ⚠️
src/maxtext/trainers/post_train/rl/train_rl.py	62.50%	6 Missing ⚠️
src/maxtext/common/checkpointing.py	68.75%	4 Missing and 1 partial ⚠️
src/maxtext/utils/model_creation_utils.py	89.13%	3 Missing and 2 partials ⚠️
... and 2 more

📢 Thoughts on this report? Let us know!