Skip to content

docs(nnx): add native LoRA Gemma4 and Qwen3 notebooks and update tutorial guides - #4417

Merged
copybara-service[bot] merged 1 commit into
mainfrom
jackyf/nnx-lora-native-demo
Aug 11, 2026
Merged

copybara-service[bot] merged 1 commit into
mainfrom
jackyf/nnx-lora-native-demo

Conversation

@RexBearIU

@RexBearIU RexBearIU commented Jul 10, 2026

Copy link
Copy Markdown
Collaborator

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.ipynb from PR #4708.

These tutorials cover both standard PEFT workflows:

  1. Pathway A: Native LoRA Supervised Fine-Tuning (SFT) starting from converted Hugging Face checkpoints.
  2. Pathway B: Native LoRA Pre-training & Checkpoint Restore with QLoRA.

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 set MODEL_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 using papermill with zero errors (100% clean exit code 0).

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 Bot commented Jul 10, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@RexBearIU
RexBearIU force-pushed the jackyf/nnx-lora-native-demo branch from a8bb4ad to 2180f1b Compare July 13, 2026 07:24
@RexBearIU
RexBearIU force-pushed the jackyf/nnx-lora-native-demo branch from d6ddb5b to b728e39 Compare July 15, 2026 09:22
@RexBearIU RexBearIU changed the title docs(nnx): add native LoRA Gemma4 and Qwen3 tutorial notebooks docs(nnx): add native LoRA Gemma4 and Qwen3 notebooks and update tutorial guides Jul 15, 2026
@RexBearIU
RexBearIU force-pushed the jackyf/nnx-lora-native-demo branch 2 times, most recently from 8867214 to bd49cd6 Compare July 21, 2026 05:03
@RexBearIU

Copy link
Copy Markdown
Collaborator Author

Rebased branch onto main and fixed Sphinx documentation build warnings regarding notebook links.

@RexBearIU
RexBearIU force-pushed the jackyf/nnx-lora-native-demo branch 2 times, most recently from 1b14194 to 5c2e04d Compare August 4, 2026 02:58
@RexBearIU
RexBearIU force-pushed the jackyf/nnx-lora-native-demo branch 2 times, most recently from 0b4ed26 to 83a529c Compare August 4, 2026 07:25
Comment thread src/maxtext/examples/gemma4_native_lora_demo.ipynb Outdated
Comment thread src/maxtext/examples/qwen3_native_lora_demo.ipynb Outdated
Comment thread src/maxtext/examples/gemma4_native_lora_demo.ipynb Outdated
@RexBearIU
RexBearIU force-pushed the jackyf/nnx-lora-native-demo branch from 83a529c to 6404d52 Compare August 7, 2026 08:04
@RexBearIU
RexBearIU requested a review from shuningjin as a code owner August 7, 2026 08:04
Comment thread src/maxtext/examples/native_lora_demo.ipynb
Comment thread docs/tutorials/posttraining/native_lora.md
Comment thread docs/tutorials/posttraining/native_lora.md Outdated
Comment thread docs/tutorials/posttraining/native_lora.md Outdated
@RexBearIU
RexBearIU force-pushed the jackyf/nnx-lora-native-demo branch 4 times, most recently from d522c0d to 2d50e47 Compare August 10, 2026 10:29
```sh
python3 -m maxtext.trainers.pre_train.train \
run_name="native_qlora_pretrain_demo" \
model_name="gemma4-e2b" \

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should be ${MODEL?}?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, thank you. Updated model_name to use ${MODEL?} for consistency with the SFT command above.

@RexBearIU
RexBearIU force-pushed the jackyf/nnx-lora-native-demo branch from 2d50e47 to 1cc2c0b Compare August 11, 2026 07:58
@copybara-service
copybara-service Bot merged commit 89c1a23 into main Aug 11, 2026
45 checks passed
@copybara-service
copybara-service Bot deleted the jackyf/nnx-lora-native-demo branch August 11, 2026 16:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants