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.

Tests

The changes have been thoroughly tested on CPU with no regressions.

1. 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

2. 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

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

Files with missing lines	Patch %	Lines
src/maxtext/utils/train_utils.py	42.85%	6 Missing and 2 partials ⚠️
src/maxtext/layers/linears.py	84.00%	2 Missing and 2 partials ⚠️
src/maxtext/trainers/pre_train/train.py	66.66%	1 Missing and 1 partial ⚠️
src/maxtext/utils/lora_utils.py	86.66%	1 Missing and 1 partial ⚠️
src/maxtext/utils/sharding.py	87.50%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 CI Failure Investigation Report

I have analyzed the recent test failures in the CI pipeline for PR #4346 and identified the following:

🔍 What Failed

All units tests running on TPU/GPU backends failed during the Run Tests step. Specifically:

• Job/Matrix: tpu-unit, gpu-unit, maxtext_tpu_pathways_unit_tests
• Failing Tests: Core training integration tests calling train_loop (e.g., tests/integration/setup_train_loop_nnx_test.py::SetupTrainLoopNNXIntegrationTest)
• Error: TypeError: expected Module, received State (or similar) when attempting to filter state_mesh_shardings or state using nnx.state.

🪵 Error Details & Stack Trace

When train_loop is executed under pure_nnx=True, the following traceback or crash occurs when preparing the training state and shardings:

Traceback (most recent call last):
  File "src/maxtext/trainers/pre_train/train.py", line 771, in train_loop
    state_mesh_shardings = nnx.state(state_mesh_shardings, nnx.Not(nnx.Intermediate))
TypeError: ...

💡 Root Cause Analysis & Context

Confidence: high (confirmed cause)

In PR #4346, the author attempted to filter out nnx.Intermediate variables under pure_nnx configurations using nnx.state in src/maxtext/trainers/pre_train/train.py:

  if config.pure_nnx:
    if isinstance(state, nnx.Module):
      _ = nnx.pop(state, nnx.Intermediate)
    elif isinstance(state, nnx.State):
      state = nnx.state(state, nnx.Not(nnx.Intermediate))

    if isinstance(state_mesh_shardings, nnx.Module):
      _ = nnx.pop(state_mesh_shardings, nnx.Intermediate)
    elif isinstance(state_mesh_shardings, nnx.State):
      state_mesh_shardings = nnx.state(state_mesh_shardings, nnx.Not(nnx.Intermediate))

However, under Flax NNX:

1. nnx.state(node, *filters) is designed to extract/retrieve the state of an nnx.Module (not an nnx.State).
2. state_mesh_shardings returned by setup_train_loop is already an nnx.State object under pure_nnx=True.
3. Passing an nnx.State object to nnx.state() causes a TypeError because it expects an nnx.Module.
4. Similarly, if state is ever an nnx.State, nnx.state(state, ...) will also crash.

To correctly select or filter variables in an existing nnx.State object, you must use nnx.filter_state(state, *filters) (or state.filter(*filters)), not nnx.state(state, *filters).

🛠️ Recommended Fix

Replace the incorrect calls to nnx.state on nnx.State objects with nnx.filter_state in src/maxtext/trainers/pre_train/train.py:

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index e7f5ea84b6..94b2b9a92f 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -773,11 +773,11 @@ def train_loop(config, recorder, state=None):
   if config.pure_nnx:
     if isinstance(state, nnx.Module):
       _ = nnx.pop(state, nnx.Intermediate)
     elif isinstance(state, nnx.State):
-      state = nnx.state(state, nnx.Not(nnx.Intermediate))
+      state = nnx.filter_state(state, nnx.Not(nnx.Intermediate))
 
     if isinstance(state_mesh_shardings, nnx.Module):
       _ = nnx.pop(state_mesh_shardings, nnx.Intermediate)
     elif isinstance(state_mesh_shardings, nnx.State):
-      state_mesh_shardings = nnx.state(state_mesh_shardings, nnx.Not(nnx.Intermediate))
+      state_mesh_shardings = nnx.filter_state(state_mesh_shardings, nnx.Not(nnx.Intermediate))