Add Gemma 4 E2B / E4B (text) support to MaxText

[gagika]

Description

Adds Gemma 4 small variants — E2B and E4B (text-only) — to MaxText.

These are the smaller members of the Gemma 4 family. They share the
broader Gemma 4 attention / norm structure but introduce two new features
that drive their parameter efficiency:

• Per-Layer Embeddings (PLE). Each decoder layer consumes a per-layer
  slice of an extra embedding tensor injected by a new Gemma4SmallPLE
  block. Controlled by hidden_size_per_layer_input /
  vocab_size_per_layer_input.
• KV sharing. The last num_kv_shared_layers decoder layers reuse
  K / V from the most recent non-shared layer of the same attention type
  (sliding↔sliding, full↔full). E2B additionally widens the MLP on those
  shared layers (use_double_wide_mlp: true) to compensate for the
  missing parameters.

Both features carry per-layer state that is not expressible inside
nn.scan, so a new GEMMA4_SMALL DecoderBlockType is added with its
own non-scanned execution path (Decoder._apply_gemma4_small_layers).
The model validator enforces scan_layers=False for these variants.

What's included

• New model file src/maxtext/models/gemma4_small.py (PLE + attention
  with optional KV sharing + decoder layer).
• New configs configs/models/gemma4-e2b.yml and gemma4-e4b.yml.
• HF round-trip: hf_model_configs.py, hf_shape.py,
  param_mapping.py updated to handle PLE params, KV-shared layers, and
  the (optional) double-wide MLP.
• TFLOP/MFU accounting: calculate_gemma4_small_tflops_training_per_device.
• Config plumbing: DecoderBlockType.GEMMA4_SMALL, four new
  Attention fields in configs/types.py, base.yml defaults, and
  validation that rejects scan_layers=true / use_multimodal=true for
  E2B / E4B.

Out of scope

• Multimodal. E2B / E4B ship a vision tower in their HF configs, but
  MaxText support for the gemma4-small vision encoder (clipped linears
  in particular) is not in this PR. use_multimodal=true is rejected by
  the validator with a clear error.
• Scanned layers. Per-layer KV sharing isn't expressible under
  nn.scan; rejected by the validator.

Tests

• New unit tests:
  • tests/unit/gemma4_small_test.py — attention-pattern dispatch,
    layer-type tuples, KV donor/shared-layer mapping for both variants.
  • tests/unit/flop_calculation_test.py::test_calculate_gemma4_small_tflops_* —
    closed-form TFLOP accounting matching the layer/donor structure.
  • tests/unit/configs_test.py — E2B / E4B yml configs are loaded by
    the existing config-instantiation sweep.
• End-to-end forward-pass logit checks:
  • tests/end_to_end/tpu/gemma4/e2b/{convert_gemma4,convert_gemma4_pt}.sh
  • tests/end_to_end/tpu/gemma4/e4b/{convert_gemma4,convert_gemma4_pt}.sh
  • Each converts the HF checkpoint with to_maxtext, then runs
    forward_pass_logit_checker against the HF model with
    --max_kl_div=0.03. This is the recommended smoke test after
    touching the model code, param map, or either YAML.

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.

[github-actions]

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

[codecov]

Codecov Report

❌ Patch coverage is 27.11864% with 344 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/models/gemma4_small.py	35.29%	108 Missing and 2 partials ⚠️
...xtext/checkpoint_conversion/utils/param_mapping.py	1.96%	100 Missing ⚠️
...rc/maxtext/checkpoint_conversion/utils/hf_shape.py	0.00%	59 Missing ⚠️
src/maxtext/layers/decoders.py	3.03%	30 Missing and 2 partials ⚠️
src/maxtext/layers/attentions.py	28.57%	22 Missing and 8 partials ⚠️
src/maxtext/multimodal/processor.py	0.00%	6 Missing and 1 partial ⚠️
src/maxtext/utils/maxtext_utils.py	93.87%	1 Missing and 2 partials ⚠️
...xt/checkpoint_conversion/utils/hf_model_configs.py	75.00%	2 Missing ⚠️
src/maxtext/layers/encoders.py	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 I'm sorry @gagika, but I was unable to process your request. Please see the logs for more details.

[aireenmei]

Thanks for the rapid implementation! I wonder if you have test results from forward_pass_logit_checker? Also do you add some unit tests for comparison with torch on the new modules such as Gemma4SmallPLE, Gemma4SmallAttention, Gemma4SmallDecoderLayer? https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/unit/gemma4_layers_test.py This was added recently.

[shuningjin]

Thanks! I agree with @aireenmei that it would be good if we could reuse rope/attention, along with component-wise unit tests (potentially as follow-up). Some minor comments.

[github-actions]

🤖 Hi @shuningjin, 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]

🤖 I'm sorry @shuningjin, but I was unable to process your request. Please see the logs for more details.

[gagika]

Thanks for the rapid implementation! I wonder if you have test results from forward_pass_logit_checker? Also do you add some unit tests for comparison with torch on the new modules such as Gemma4SmallPLE, Gemma4SmallAttention, Gemma4SmallDecoderLayer? https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/unit/gemma4_layers_test.py This was added recently.

Added those unit test, PTL

for forward logits test, yes, I have done for both models.
I did again after addressing PR feedback:

https://paste.googleplex.com/5790253452492800
https://paste.googleplex.com/6349200758538240

[github-actions]

🤖 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]

🤖 I'm sorry @gagika, but I was unable to process your request. Please see the logs for more details.