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.
Implementing a thought channel bypass in input_pipeline_utils.py to prevent validation failures when the generation prompt includes the <|channel>thought block.
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.

github-actions · 2026-06-02T13:34:33Z

  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


🟡 Disabling NNX graph caching for MoE models is a critical fix for metadata synchronization. However, this change is currently only applied to the SFT trainer. Other post-training trainers, such as DPO, RL, and Distillation, may also be affected when training MoE models and should ideally receive a similar update for consistency and to avoid potential issues.

github-actions · 2026-06-02T13:34:33Z

 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)))"
 olmo3: "decoder/layers/.*(attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))"


🟢 The regex for `gemma4` includes `MoeBlock_0` within the MLP group. Since the `.*` wildcard already allows for matching nested paths, `MoeBlock_0` may be redundant unless there are specific parameters directly under it that require targeting. Additionally, please ensure that `layers.*/.*` is sufficiently precise to avoid matching unintended components in deeper architectures.

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.)

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?

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

Are there any tests for Gemma4 Lora?

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

feat: Gemma4 LoRA Extension

6a64bd0

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

gagika reviewed Jun 2, 2026

View reviewed changes

Conversation

RexBearIU commented May 22, 2026

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!

github-actions Bot Jun 2, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jun 2, 2026

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov Bot commented May 22, 2026 •

edited

Loading