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.

Latest Updates (July 9, 2026)

1. Rebased and Squashed: Cleanly rebased the nnx-lora-support branch onto the latest main branch, resolving merge conflicts in src/maxtext/utils/train_utils.py while preserving important # pyrefly: ignore comments. All commits have been squashed into exactly 1 clean, production-ready commit: feat(nnx): support native Flax NNX PEFT/LoRA training loop.
2. Traversal Fix for Flax NNX State: Resolved an AttributeError: No attribute 'flat' in State crash that occurred during sft_native trainer initialization. Modern nnx.State structures do not support .flat(); this traversal was updated to use the robust and compliant zip(state.flat_state().paths, state.flat_state().leaves) API pattern.
3. Recovery Helper Renamed: Renamed restore_custom_types_from_state to the more descriptive restore_qlora_base_weights in src/maxtext/utils/lora_utils.py and src/maxtext/layers/linears.py to ensure type-restoration is specific and self-explanatory.
4. Local Verification: Verified all unit and integration tests successfully (19 passed under CPU backend).
5. Remote Synced: Successfully force-pushed the squashed and updated branch to origin/nnx-lora-support.

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):

  • Added test_pure_nnx_setup_with_lora_enabled which asserts that setup_train_loop correctly instantiates a model with LoRA adapters injected and the optimizer configured with wrt=nnx.LoRAParam.
  • Added test_train_step_updates_only_lora_weights which 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.
  • Updated the tiny NNX testing config (_tiny_nnx_pyconfig) to use attention="dot_product" so that CPU backend runs are fully supported without relying on TPU/GPU Pallas attention kernels.
  • Run command:

    JAX_PLATFORMS='cpu' PYTHONPATH=src pytest tests/integration/setup_train_loop_nnx_test.py -k SetupTrainLoopNNXLoraTest

• Unit tests (tests/unit/lora_utils_nnx_test.py):

  • Added test_sharding_extracts_only_lora_params to verify that sharding.maybe_update_params_sharding_with_opt extracts only nnx.LoRAParam under LoRA configuration.
  • Run command:

    JAX_PLATFORMS='cpu' PYTHONPATH=src pytest tests/unit/lora_utils_nnx_test.py

---

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

Files with missing lines	Patch %	Lines
src/maxtext/utils/lora_utils.py	74.41%	7 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	71.42%	2 Missing and 2 partials ⚠️
src/maxtext/utils/sharding.py	75.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!