Skip to content

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

Open
RexBearIU wants to merge 2 commits into
mainfrom
jackyf/nnx-lora-native-demo
Open

docs(nnx): add native LoRA Gemma4 and Qwen3 tutorial notebooks#4417
RexBearIU wants to merge 2 commits 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.

These tutorials cover both:

  1. Pathway A: SFT Fine-Tuning with native LoRA.
  2. Pathway B: Pre-training with QLoRA, including checkpoint saving and restoration.

IMPORTANT NOTE: This pull request is blocked by PR #4346 (feat(nnx): support native Flax NNX PEFT/LoRA training loop). It targets main directly to keep its commit history clean, but should only be merged after PR #4346 has been successfully integrated into main.

BUGS: #4400

Tests

Both tutorial notebooks have been extensively verified and executed cell-by-cell under standard TPU VM environments using papermill with zero errors:

  1. Qwen3 Tutorial Notebook (qwen3_native_lora_demo.ipynb)

    • Compiles, runs, saves checkpoints, and restores checkpoints for pre-training.
    • Validated with 100% cell execution success.
  2. Gemma4 Tutorial Notebook (gemma4_native_lora_demo.ipynb)

    • Configured with scan_layers=False to handle the specific multi-query KV sharing architectural constraints of Gemma4-e2b.
    • Compiles, runs, saves checkpoints, and restores checkpoints.
    • Validated with 100% cell execution success.

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 6 times, most recently from 09394fa to 356da98 Compare July 10, 2026 17:54
@RexBearIU RexBearIU requested a review from jacoguzo as a code owner July 10, 2026 17:54
@RexBearIU RexBearIU force-pushed the jackyf/nnx-lora-native-demo branch 2 times, most recently from 9fc481f to a8bb4ad Compare July 10, 2026 18:13
@RexBearIU RexBearIU force-pushed the jackyf/nnx-lora-native-demo branch from a8bb4ad to 2180f1b Compare July 13, 2026 07:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant