Skip to content

Restore Engram defaults in RLConfig - #4621

Merged
copybara-service[bot] merged 1 commit into
AI-Hypercomputer:mainfrom
csgoogle:gemma4-post-training
Jul 27, 2026
Merged

Restore Engram defaults in RLConfig#4621
copybara-service[bot] merged 1 commit into
AI-Hypercomputer:mainfrom
csgoogle:gemma4-post-training

Conversation

@csgoogle

@csgoogle csgoogle commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

Summary

  • add the existing Engram configuration mixin to RLConfig
  • add a focused regression test for the disabled Engram defaults

Root cause

The dedicated RLConfig schema omitted Engram when RL configuration was split from MaxTextConfig. Shared decoder code reads engram_layers even when Engram is disabled, so RL models need those fields to satisfy the decoder configuration contract.

This change supplies the existing disabled defaults (engram_layers=[] and engram_max_ngram_size=3). It does not enable Engram or change Gemma 4 weights, architecture, or numerics.

Tests

  • JAX_PLATFORMS=cpu pytest -q tests/post_training/unit/train_rl_test.py -k test_rl_config_includes_decoder_engram_defaults — passed
  • git diff --check — passed

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.

@gemini-code-assist

Copy link
Copy Markdown

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@csgoogle
csgoogle force-pushed the gemma4-post-training branch from d70fdb0 to 23aa9f8 Compare July 27, 2026 08:28
@csgoogle csgoogle changed the title Verify Gemma post-training with Tunix Restore Engram defaults in RLConfig Jul 27, 2026
@csgoogle
csgoogle marked this pull request as ready for review July 27, 2026 08:29
@gemini-code-assist

Copy link
Copy Markdown

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@gemini-code-assist

Copy link
Copy Markdown

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@codecov

codecov Bot commented Jul 27, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@copybara-service
copybara-service Bot merged commit 40aaa08 into AI-Hypercomputer:main Jul 27, 2026
57 of 63 checks passed
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