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

[RexBearIU]

Description

This PR implements native parameter-efficient fine-tuning (PEFT) and LoRA training support within the Flax NNX training loop inside MaxText (train.py and train_utils.py).

Previously, training inside train.py with Flax NNX was limited to full-parameter fine-tuning. This was due to:

1. setup_train_loop inside train_utils.py initializing nnx.Optimizer with wrt=nnx.Param, which allocated optimizer states for the entire base model.
2. train_step and diff_wrapper inside train.py hardcoding model splits using nnx.Param.
3. Sharding extraction inside sharding.py hardcoding nnx.Param extraction, which failed to cleanly separate base parameters from LoRA parameters.

This PR addresses those limitations through the following design:

• Dynamic Parameter Split/Filter Target: We define wrt dynamically based on configuration (nnx.LoRAParam if config.lora.enable_lora is True, otherwise nnx.Param). This is used consistently across optimizer initialization, training step splits/merges, and sharding parameter extraction.
• LoRA Adapter Injection & Restoration: In setup_train_loop, Qwix LoRA adapters are dynamically injected when config.lora.enable_lora is True. On fresh training runs (no previous checkpoint step), pre-trained adapters can be loaded from lora.lora_restore_path.
• Mesh Context Tracing Safety: Updated apply_lora_to_model inside lora_utils.py to only invoke jax.set_mesh(mesh) if tracing is not currently active (jax_core.trace_state_clean()), avoiding compilation errors during eager/eval paths under NNX.
• Strict Checkpointing Validation: _raise_on_weight_mismatch recursively validates that the restored checkpoint matches the target model structure, fully supporting standard model and optimizer state restorations.

Tests

The changes have been thoroughly tested on CPU and verified with extensive multi-device TPU benchmarks.

1. Local Verification Tests

• Integration tests (tests/integration/setup_train_loop_nnx_test.py):
  • test_pure_nnx_setup_with_lora_enabled asserts that setup_train_loop correctly instantiates a model with LoRA adapters injected and the optimizer configured with wrt=nnx.LoRAParam.
  • test_train_step_updates_only_lora_weights runs a full forward and backward pass for a single training step and asserts that only the adapter parameters (nnx.LoRAParam) are modified, while the base weights (nnx.Param) remain unchanged.
• Unit tests:
  • tests/unit/checkpointing_nnx_load_test.py covers full structural checkpointer loading logic.

---

2. TPU E2E Verification & Performance Comparison (8x TPU v6 Lite)

To thoroughly validate native Flax NNX QLoRA (int8) support in multi-device TPU environments, we executed end-to-end 20-step benchmarks across all three trainers (pre_train, sft_native, sft_custom) under both scan_layers=True and scan_layers=False configurations.

All performance metrics (step times, tokens/s/device) are averaged across steps 10 to 15, and training losses (including MoE load balancing losses, where applicable) are reported at step 15.

A. Llama 3.1 8B QLoRA Performance Comparison

Trainer	Scan Layers	Loss (Step 15)	Step Time (avg)	Tokens/s/dev (avg)
pre_train	True	5.9930	0.0065s	19,356.8
pre_train	False	5.9900	0.0130s	10,039.2
sft_native	True	5.9930	0.0065s	19,604.7
sft_native	False	5.9900	0.0127s	10,017.4
sft_custom	True	5.9990	0.0147s	8,764.6
sft_custom	False	6.0230	0.0298s	4,342.0

Technical Architecture Insights (Llama 3.1 8B):

• Scan vs. Non-Scan Performance: Enabling sequential layer scanning (scan_layers=True) results in a ~2x performance increase in step times and throughput. This is due to compiler optimization through sequential loop unrolling and reduced JIT compilation overhead on TPUs.
• Tiling & Sharding Divisibility: Running QLoRA on Llama 3.1 with a default lora_tile_size=32 when embedding dimensions are small (or on small FSDP meshes) can cause IndivisibleError since the tiled dimension doesn't divide evenly across the mesh devices. Setting lora.lora_tile_size=16 guarantees a tile count (128 / 16 = 8 tiles) that perfectly divides across the 8-way FSDP mesh.

B. Gemma 4 26B (Mixture of Experts) QLoRA Performance Comparison

Trainer	Scan Layers	Loss (Step 15)	MoE LB Loss	Step Time (avg)	Tokens/s/dev (avg)
pre_train	True	6.2810	0.0080	0.0073s	16,814.3
pre_train	False	6.2810	0.0080	0.0080s	16,383.8
sft_native	True	6.2810	0.0080	0.0075s	16,997.4
sft_native	False	6.2810	0.0080	0.0080s	16,386.5
sft_custom	True	6.2900	0.0080	0.0283s	4,539.8
sft_custom	False	6.2830	0.0080	0.0307s	4,189.1

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 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 68.86792% with 33 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/lora_utils.py	73.33%	8 Missing and 4 partials ⚠️
src/maxtext/utils/train_utils.py	46.15%	5 Missing and 2 partials ⚠️
src/maxtext/layers/linears.py	45.45%	4 Missing and 2 partials ⚠️
src/maxtext/utils/maxtext_utils.py	72.72%	1 Missing and 5 partials ⚠️
src/maxtext/common/checkpointing.py	75.00%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[SurbhiJainUSC]

nit: getattr(getattr(config, "lora", None), "enable_lora", False)

[SurbhiJainUSC]

Make sure this is uniform across all files to ensure safe configuration check.

[RexBearIU]

Done! Made the config check uniform across all relevant modified files in the change list using the safer nested getattr pattern: getattr(getattr(config, "lora", None), "enable_lora", False).