docs(nnx): add native LoRA Gemma4 and Qwen3 notebooks and update tutorial guides - #4417
Merged
Merged
Conversation
RexBearIU
force-pushed
the
jackyf/nnx-lora-native-demo
branch
from
July 10, 2026 10:29
1e932e9 to
0ba215e
Compare
4 tasks
RexBearIU
force-pushed
the
jackyf/nnx-lora-native-demo
branch
from
July 10, 2026 11:38
0ba215e to
95fb143
Compare
RexBearIU
force-pushed
the
nnx-lora-support
branch
2 times, most recently
from
July 10, 2026 11:54
054531c to
7fa9df0
Compare
RexBearIU
force-pushed
the
jackyf/nnx-lora-native-demo
branch
from
July 10, 2026 12:15
95fb143 to
af4da59
Compare
RexBearIU
marked this pull request as ready for review
July 10, 2026 12:21
RexBearIU
requested review from
A9isha,
NicoGrande,
NuojCheng,
RissyRan,
SurbhiJainUSC,
abhinavclemson,
aireenmei,
bvandermoon,
darisoy,
dipannita08,
gagika,
gobbleturk,
hengtaoguo,
igorts-git,
jiangjy1982,
khatwanimohit,
richjames0,
shralex,
suexu1025,
vipannalla and
xibinliu
as code owners
July 10, 2026 12:21
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
RexBearIU
force-pushed
the
jackyf/nnx-lora-native-demo
branch
from
July 13, 2026 07:24
a8bb4ad to
2180f1b
Compare
RexBearIU
force-pushed
the
jackyf/nnx-lora-native-demo
branch
from
July 15, 2026 09:22
d6ddb5b to
b728e39
Compare
RexBearIU
force-pushed
the
jackyf/nnx-lora-native-demo
branch
2 times, most recently
from
July 21, 2026 05:03
8867214 to
bd49cd6
Compare
Collaborator
Author
|
Rebased branch onto main and fixed Sphinx documentation build warnings regarding notebook links. |
RexBearIU
force-pushed
the
jackyf/nnx-lora-native-demo
branch
2 times, most recently
from
August 4, 2026 02:58
1b14194 to
5c2e04d
Compare
4 tasks
RexBearIU
force-pushed
the
jackyf/nnx-lora-native-demo
branch
2 times, most recently
from
August 4, 2026 07:25
0b4ed26 to
83a529c
Compare
shralex
approved these changes
Aug 6, 2026
RexBearIU
force-pushed
the
jackyf/nnx-lora-native-demo
branch
from
August 7, 2026 08:04
83a529c to
6404d52
Compare
RexBearIU
force-pushed
the
jackyf/nnx-lora-native-demo
branch
4 times, most recently
from
August 10, 2026 10:29
d522c0d to
2d50e47
Compare
| ```sh | ||
| python3 -m maxtext.trainers.pre_train.train \ | ||
| run_name="native_qlora_pretrain_demo" \ | ||
| model_name="gemma4-e2b" \ |
Collaborator
There was a problem hiding this comment.
Should be ${MODEL?}?
Collaborator
Author
There was a problem hiding this comment.
Good catch, thank you. Updated model_name to use ${MODEL?} for consistency with the SFT command above.
SurbhiJainUSC
approved these changes
Aug 10, 2026
RexBearIU
force-pushed
the
jackyf/nnx-lora-native-demo
branch
from
August 11, 2026 07:58
2d50e47 to
1cc2c0b
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds native Gemma4 and Qwen3 (0.6B) LoRA and QLoRA tutorial notebooks for Flax NNX, along with accompanying post-training setup documentation. Also includes a follow-up fix for
lora_llama3_demo.ipynbfrom PR #4708.These tutorials cover both standard PEFT workflows:
Summary of Changes:
src/maxtext/examples/gemma4_native_lora_demo.ipynb: Gemma4 Native LoRA and QLoRA demo notebook.src/maxtext/examples/qwen3_native_lora_demo.ipynb: Qwen3 Native LoRA and QLoRA demo notebook.src/maxtext/examples/lora_llama3_demo.ipynb: Follow-up fix from PR Fix JAX 0.11.0 sharding compatibility and SFT trainer signature for Tunix integration #4708 to setMODEL_NAME = "llama3.1-8b-Instruct", ensuring checkpoint conversion downloads the Instruct-aligned weights (meta-llama/Llama-3.1-8B-Instruct).docs/tutorials/posttraining/native_lora.md: Guide for native Flax NNX LoRA.docs/guides/run_python_notebook.md&docs/tutorials/post_training_index.md: Updated documentation indices.Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.
Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.
Tests
All tutorial notebooks (
qwen3_native_lora_demo.ipynb,gemma4_native_lora_demo.ipynb,lora_llama3_demo.ipynb) have been extensively verified and executed cell-by-cell under standard TPU VM environments usingpapermillwith zero errors (100% clean exit code 0).Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.