Feat: integrate NNX LoRA support via Qwix with unified configuration

[RexBearIU]

Description

Overview
This pull request introduces native LoRA support in MaxText by leveraging the NNX model definition and the Qwix library. It enables a seamless workflow for applying LoRA adapters during
training and provides utilities for bidirectional checkpoint conversion with the HuggingFace ecosystem.

Key Changes

• Core NNX Integration:
  • Refactored NNXDecoder layer application logic to support nnx.scan with dynamic graph initialization, ensuring compatibility with Qwix's parameter materialization.
• SFT Pipeline Enhancements:
  • Integrated apply_lora_to_model and restore_lora_from_path into the SFT trainer.
  • Added dummy input preparation to materialize LoRA parameters before trainer initialization.
• Bidirectional Conversion Scripts:
  • hf_lora_to_maxtext.py: Converts HuggingFace PEFT adapters to MaxText checkpoint format. Updated to 2026 copyright and cleaned up comments.
  • maxtext_to_hf_lora.py: Converts MaxText LoRA checkpoints back to HuggingFace format. Updated to use max_logging and 2026 copyright.
• Configuration & Type System:
  • Added lora_module_path auto-detection logic for popular models (Llama, etc.) via lora_module_path.yml.
  • Updated types.py with specific LoRA/QLoRA fields.
• Current Limitations:
  • QLoRA flags (lora_weight_qtype, lora_tile_size) are included in the configuration but explicitly marked as TODO / Not Working for this initial release.

Tests

The Qwix-based LoRA implementation was validated through a new unit test suite and verified via a comprehensive tutorial.

1. Unit Tests
   Implemented tests/unit/lora_utils_test.py to ensure structural correctness and trainer compatibility. Key areas covered:

• Model Transformation: Verified that apply_lora_to_model correctly injects nnx.LoRAParam into the model state.
• Layer Scanning: Confirmed the implementation works with both scan_layers=True and scan_layers=False by handling the resulting differences in the nnx module path tree.
• Trainer Compatibility: Validated that tunix.sft.peft_trainer.PeftTrainer correctly identifies the LoRA parameters for optimization, ensuring only adapter weights are trained.
• Path Matching: Tested the regex logic for auto-detecting LoRA target modules across different model architectures (e.g., Llama).

Command to run unit tests:

1 # From the maxtext root directory
2 export PYTHONPATH=$PYTHONPATH:$(pwd)/src:$(pwd)
3 python3 tests/unit/lora_utils_test.py

2. Documentation

• Added docs/tutorials/posttraining/lora.md, which provides a step-by-step guide for running LoRA fine-tuning, including environment setup and checkpoint conversion. This tutorial serves as the reference for end-to-end functional verification.

Logit test result

https://paste.googleplex.com/6233928391327744

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

Files with missing lines	Patch %	Lines
src/maxtext/checkpoint_conversion/to_maxtext.py	0.00%	234 Missing ⚠️
...rc/maxtext/checkpoint_conversion/to_huggingface.py	0.00%	63 Missing ⚠️
src/maxtext/utils/lora_utils.py	73.84%	26 Missing and 8 partials ⚠️
src/maxtext/utils/sharding.py	73.61%	10 Missing and 9 partials ⚠️
src/maxtext/checkpoint_conversion/utils/utils.py	0.00%	17 Missing ⚠️
src/maxtext/layers/nnx_decoders.py	82.14%	4 Missing and 1 partial ⚠️
src/maxtext/utils/maxtext_utils.py	33.33%	3 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[bvandermoon]

Thank you @RexBearIU. Left some comments but this is generally looking good to me