feat: Add Gemma3 LoRA SFT integration and programmatic weight mapping…

[RexBearIU]

Description

This PR adds end-to-end integration and weight-mapping support for Gemma3 LoRA SFT and vLLM decoding. Previously, standalone offline decoding ( vllm_decode.py ) and the online Tunix server ( adapter.py ) lacked the capability
to
map and load LoRA-adapted weights specifically for Gemma3, causing failures during parameter restoration and serving.

Key Changes:

• Gemma3 Weight Mapping: Added gemma3.py mapping to translate Gemma3 layers programmatically between MaxText and HuggingFace/vLLM structures.
• Generalized Restore API: Updated restore_lora_from_path to handle raw nnx.Module models (used in vLLM adapters) in addition to trainer states.
• vLLM Integration: Hooked both online and offline vLLM decode paths to restore LoRA parameters when lora.enable_lora=True .

Tests

1. Unit Tests

python3 -m pytest tests/post_training/unit/lora_utils_test.py

• Result: 10 passed

2. End-to-End Tests

Tested the end-to-end flow using the newly converted checkpoint:

1. Convert Checkpoint:
   export HF_TOKEN=
   export RUN_ID=$(date +%Y-%m-%d-%H-%M)
   bash tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_mt.sh $RUN_ID

2. Run LoRA & Decode (starts from the newly converted checkpoint above):
   bash tests/end_to_end/tpu/gemma3/4b/test_gemma3_lora.sh $RUN_ID

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

Files with missing lines	Patch %	Lines
src/maxtext/integration/tunix/utils.py	0.00%	28 Missing ⚠️
src/maxtext/checkpoint_conversion/utils/utils.py	10.00%	18 Missing ⚠️
...maxtext/integration/tunix/weight_mapping/gemma3.py	60.00%	8 Missing ⚠️
...t/integration/vllm/maxtext_vllm_adapter/adapter.py	0.00%	5 Missing ⚠️
src/maxtext/utils/lora_utils.py	54.54%	3 Missing and 2 partials ⚠️
...rc/maxtext/checkpoint_conversion/to_huggingface.py	0.00%	2 Missing ⚠️
...xtext/integration/tunix/weight_mapping/__init__.py	33.33%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[RexBearIU]

close by split pr to #4066 #4067 #4068