feat(nnx): support native Flax NNX PEFT/LoRA training loop

[RexBearIU]

Description

This PR implements full native parameter-efficient fine-tuning (PEFT), LoRA, and QLoRA training support and complete checkpoint lifecycle resume/restoration within Flax NNX training in MaxText (train.py, train_sft.py, maxtext_utils.py, checkpointing.py, train_state_nnx.py, lora_utils.py).

Key Implementation Highlights:

1. Dynamic Parameter Filter Target:

   • wrt parameter filtering is dynamically selected based on configuration (nnx.LoRAParam if config.lora.enable_lora is True, otherwise nnx.Param), ensuring optimizer states and parameter updates affect only LoRA adapters when LoRA is enabled.

2. Checkpoint Resume & Clean Utility Integration:

   • Re-integrated LoRA/NNX checkpoint helpers (_norm_path_key, quantize_params_safe, _unwrap_restored_params, _update_nnx_state_from_pure_dict) directly into src/maxtext/common/checkpointing.py while strictly preserving baseline pre-training and native SFT checkpointing code paths.
   • Resolved pjit sharding mismatches during checkpoint resume by re-evaluating state_mesh_shardings after parameter restoration and LoRA/QLoRA module injection in setup_initial_state.
   • Updated split_for_checkpoint in train_state_nnx.py to route both nnx.Param and nnx.LoRAParam into the params collection, ensuring Orbax saves and restores complete PyTree metadata for LoRA adapter runs.
   • Normalized layer key matching (0 <-> "layers_0") across structural variations in layer naming.

3. Dynamic Quantization & PyTree Sanitization:

   • Added automatic dynamic quantization (quantize_params_safe) in checkpointing.py so unquantized base checkpoints are dynamically quantized when loading into QLoRA models.
   • Updated l2norm_pytree in max_utils.py to safely skip non-inexact leaves (e.g., packed void or QArray quantized arrays and WithAux wrappers) by checking jnp.issubdtype(leaf.dtype, jnp.inexact).
   • Added mesh axis shape sanitization in sharding.py (get_leaf_pspec) for non-divisible dimensions (e.g. Gemma 4 / Qwen 3 head dimensions vs TPU mesh shapes).

4. Multi-Architecture Support:

   • Fully verified across gemma4-26b (MoE) and qwen3-4b for both scan_layers=True and scan_layers=False modes across Base Training, Resume Training, and Standalone Restore.

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned. This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

1. Comprehensive Unit Test Suite

Ran full test suite (checkpointing_test.py, lora_utils_test.py, model_creation_utils_test.py) with 100% passing results (103/103 tests passed):

PYTHONPATH=src python -m pytest tests/unit/checkpointing_test.py tests/post_training/unit/lora_utils_test.py tests/unit/model_creation_utils_test.py

2. Pre-Commit Verification

Ran pre-commit hooks across all changed files with 100% passing results:

git diff --name-only --diff-filter=d origin/main...HEAD | xargs pre-commit run --files

3. Multi-Device TPU E2E Verification (gemma4-26b & qwen3-4b)

Verified full 4-step E2E checkpoint lifecycle for gemma4-26b and qwen3-4b across both scan_layers=False and scan_layers=True:

• Step 1: Generating base-only checkpoint (PASSED)
• Step 2: Training with LoRA starting from base checkpoint (PASSED)
• Step 3: Resuming training under same run name (PASSED)
• Step 4: Verifying standalone restore of LoRA checkpoint (PASSED)

Model	Trainer	Scan Layers	Loss Step 2	Loss Step 3	Standalone Restore	Status
gemma4-26b	pre_train/train.py	False	0.0084	0.0084	PASSED	PASSED
gemma4-26b	pre_train/train.py	True	0.0084	0.0084	PASSED	PASSED

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 details above.
• 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 15.56503% with 396 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/common/checkpointing.py	6.54%	257 Missing ⚠️
src/maxtext/utils/maxtext_utils.py	22.22%	48 Missing and 1 partial ⚠️
src/maxtext/utils/lora_utils.py	40.98%	28 Missing and 8 partials ⚠️
src/maxtext/utils/model_creation_utils.py	6.89%	27 Missing ⚠️
src/maxtext/utils/train_utils.py	6.25%	15 Missing ⚠️
src/maxtext/utils/sharding.py	25.00%	5 Missing and 1 partial ⚠️
src/maxtext/common/train_state_nnx.py	0.00%	4 Missing ⚠️
src/maxtext/utils/max_utils.py	80.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!