Extend LoRA for Gemma4 by RexBearIU · Pull Request #3969 · AI-Hypercomputer/maxtext

RexBearIU · 2026-05-22T07:31:35Z

Description

This PR extends the recent LoRA support to accurately target and process Gemma 4 architectures (including MoE).

Gemma 4 introduces complex nested structures (like scanned_blocks and layers_remainder) and unique chat template behaviors (such as the <|channel>thought block) that are incompatible with standard LoRA targeting and data
processing. Furthermore, MoE models require dynamic metadata synchronization during forward passes which is broken by aggressive NNX graph caching.

This PR addresses these challenges by:

Adding accurate regex mapping for Gemma 4 standard and MoE LoRA targets in lora_module_path.yml.
Dynamically disabling NNX graph caching in train_sft.py specifically for MoE models (where experts > 1) to allow necessary metadata synchronization.

Tests

Added unit tests for the Gemma 4 tokenizer bypass in tests/post_training/unit/sft_data_processing_test.py (test_tokenizer_gemma4_thought_channel_bypass).
Verified caching behavior changes by running Gemma-4 MoE LoRA tuning on TPU.

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 · 2026-05-22T07:37:22Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

SurbhiJainUSC · 2026-05-27T16:25:44Z

  def test_tokenizer_wo_generation_prompt(self):
    verify_chat_template_generation_prompt_logic(self.llama2_tokenizer)

+  def test_tokenizer_gemma4_thought_channel_bypass(self):


This test expects to not fail with TemplateError or ValueError. Can you add an assertion for this so that it is readable what this test actually verifies?

I updated verify_chat_template_generation_prompt_logic to return True on success, and wrapped all three test cases in explicit self.assertTrue() assertions.

This cleanly verifies that validation succeeds and keeps all tests uniform. All tests pass successfully!"

SurbhiJainUSC · 2026-05-28T16:08:38Z

  actual_prefix_in_full_turn = full_turn_ids[len(prompt_wo_gen_ids) : len(prompt_wo_gen_ids) + len(assistant_prefix)]

  if actual_prefix_in_full_turn != assistant_prefix:
+    # Allow the generation prompt to include a thought channel block (e.g., for Gemma 4).


This logic looks like a hacky approach to support Gemma4. I am working on a generalized logic to support any model that requires specific prefix shifting. I will send out the PR soon.

Agreed, it was definitely a workaround. Thanks for taking the lead on a generalized solution! I'll track #4010 and we can use that approach instead.

I've rebased this PR on top of #4010

github-actions · 2026-06-02T13:31:56Z

🤖 Hi @gagika, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

github-actions

## 📋 Review Summary

This Pull Request successfully extends LoRA support for Gemma 4 architectures and addresses a critical issue with NNX graph caching in MoE models. However, there is a significant discrepancy between the PR description and the actual changes, as several mentioned files and unit tests are missing from the diff.

🔍 General Feedback

Missing Implementation: The PR description mentions a "thought channel bypass" in input_pipeline_utils.py and new unit tests in tests/post_training/unit/sft_data_processing_test.py, but these files are not included in the PR. Please ensure all intended changes are staged and pushed.
Consistency across Trainers: The dynamic disabling of NNX graph caching is a great addition for MoE stability; consider applying this same logic to DPO, RL, and Distillation trainers to ensure consistent behavior across the post-training suite.
LoRA Targeting: The regex for Gemma 4 LoRA targeting is comprehensive but should be monitored to ensure it doesn't become overly broad as the architecture evolves.

gagika · 2026-06-02T13:48:11Z

 deepseek2: "decoder/(dense_layers|moe_stack)/self_attention/(query|out|wkv_a|wkv_b)|decoder/(dense_layers|moe_stack)/(mlp|shared_experts)/(wi_0|wi_1|wo)"
 gemma2: "decoder/layers/(self_attention_local|self_attention_global)/(query|key|value|out)|decoder/layers/(mlp_local|mlp_global)/(wi_0|wi_1|wo)"
 gemma3: "decoder/layers/.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo|gate|up|down))"
+gemma4: "decoder/(scanned_blocks|layers_remainder)/layers.*/.*(self_attention/(query|key|value|out)|mlp/.*(MoeBlock_0|wi_0|wi_1|wo|shared_experts/(wi_0|wi_1|wo)))"


The pattern also adapts the MoE router: probably by accident. The MoeBlock_0 term matches everything inside the routed-MoE block, including the router/gate that decides which expert each token goes to. You normally don't want LoRA on the router. The actual expert weights (wi_0, wi_1, wo) are already matched by the other terms in the pattern, so you can just delete MoeBlock_0|. Experts and shared experts still get LoRA; the router no longer does. (Checked against the parameter tree.)

Thank you for the excellent catch! Yes, including MoeBlock_0 was a mistake that dragged in the routing gate.

I have removed MoeBlock_0| from the gemma4 pattern in lora_module_path.yml. This successfully excludes the router/gate from LoRA adaptation while still perfectly capturing the expert MLP weights (wi_0, wi_1, wo)
and the shared experts.

gagika · 2026-06-02T13:50:01Z

  with mesh, nn_partitioning.axis_rules(mt_config.logical_axis_rules):
+    # Disable NNX graph caching for MoE models (where experts > 1) to allow
+    # necessary dynamic metadata synchronization during forward passes (e.g., in jax.lax.scan).
+    enable_nnx_cache = getattr(mt_config, "num_experts", 1) <= 1


For MaxText flags just directly use dot notation:
mt_config.num_experts

Also why do you need to disable it for MoE, what happens if you don't?

Simplified Flag Access: Updated to use direct dot notation (mt_config.num_experts) as suggested.

Why we disable caching for MoE:
If we run training with cache_nnx_graph=True (which uses nnx.cached_partial under the hood to cache graph topology), MoE training fails during step execution with the following traceback:

File "/home/jackyf_google_com/5-4-qlora/lib/python3.12/site-packages/flax/nnx/graph.py", line 1716, in unflatten raise ValueError( ValueError: The graph structure of a node added to cached_partial was mutated inside the transformation, this is not allowed. Node: Transformer( ... )

nnx.cached_partial expects a static topological structure ( graphdef ) across JIT boundaries. However, Mixture-of-Experts (MoE) models perform dynamic token routing and scan-indexing during the forward pass. These operations introduce dynamic tracer metadata that mutates the internal structure of the Transformer node during transformation. This violates the static structural invariants of cached_partial , triggering the traceback. Disabling the graph cache allows NNX to split/merge definitions dynamically at each step, preventing this error.

gagika · 2026-06-02T13:56:04Z

Are there any tests for Gemma4 Lora?

We ran the end-to-end LoRA training loop for Gemma 4 successfully without any issues.4 successfully without any issues. log

RexBearIU changed the title ~~Jackyf/gemma4 lora~~ Extend LoRA for Gemma4 May 22, 2026

RexBearIU mentioned this pull request May 22, 2026

feat: QLoRA support for Dense/MoE models across Pathways and McJAX #3702

Closed

4 tasks

RexBearIU force-pushed the jackyf/gemma4-lora branch from 2bc8632 to ab61640 Compare May 22, 2026 07:38

SurbhiJainUSC reviewed May 27, 2026

View reviewed changes

RexBearIU force-pushed the jackyf/gemma4-lora branch 2 times, most recently from 61626bd to ef50ff7 Compare May 28, 2026 08:41

SurbhiJainUSC reviewed May 28, 2026

View reviewed changes

RexBearIU force-pushed the jackyf/gemma4-lora branch from ef50ff7 to 5fd616b Compare May 29, 2026 07:48

RexBearIU force-pushed the jackyf/gemma4-lora branch from 5fd616b to 6a64bd0 Compare June 1, 2026 06:59

SurbhiJainUSC approved these changes Jun 1, 2026

View reviewed changes

gagika added the gemini-review label Jun 2, 2026

github-actions Bot reviewed Jun 2, 2026

View reviewed changes

Comment thread src/maxtext/trainers/post_train/sft/train_sft.py Outdated

Comment thread src/maxtext/configs/post_train/lora_module_path.yml

gagika reviewed Jun 2, 2026

View reviewed changes

RexBearIU force-pushed the jackyf/gemma4-lora branch 2 times, most recently from 4c2d52f to 491b37c Compare June 3, 2026 15:43

feat: Gemma4 LoRA Extension

b52137c

RexBearIU force-pushed the jackyf/gemma4-lora branch from 491b37c to b52137c Compare June 3, 2026 16:24

Conversation

RexBearIU commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Checklist

Uh oh!

codecov Bot commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jun 2, 2026

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

🔍 General Feedback

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

RexBearIU Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

RexBearIU commented May 22, 2026 •

edited

Loading

codecov Bot commented May 22, 2026 •

edited

Loading

RexBearIU Jun 3, 2026 •

edited

Loading