feat(model): add Gemma-4 E4B support (layer spec, checkpoint loader, parity check)

[DOGEUNNKIM]

What does this PR do?

• Adds full Gemma-4 E4B (3.8B dense text) support to Megatron-Bridge,
  consolidating all Gemma4-specific code that was previously scattered in
  Megatron-LM core (see MCore PR: [Feature-request] Add Megatron Core support geglu for Gemma 4 dense model (e.g. E4B) NVIDIA/Megatron-LM#5090).
• Introduces Gemma4DenseProvider — a Bridge-native model provider that works
  against clean Megatron-Core (no Gemma4-specific CLI args or
  TransformerConfig fields required in MCore).
• Separates the text-only path (Gemma4Bridge / Gemma4ForCausalLM) from the
  VL path (Gemma4VLBridge / Gemma4ForConditionalGeneration); gemma_vl
  imports from gemma, not vice versa.
• Provides HF → Megatron checkpoint conversion, logit parity verification, and
  training launch scripts.

Changelog

src/megatron/bridge/models/gemma/modeling_gemma4.py

Core Gemma-4 architecture:

• Gemma4TransformerLayer — 4-norm residual structure, Phase 4 PLE residual
  block (gelu(gate(h)) × per_layer_input → proj → norm → residual), per-layer
  scalar, optional MoE block.
• Gemma4SelfAttention — GQA with per-head q_norm/k_norm,
  heterogeneous head_dim (sliding=256, global=512), shared-KV forwarding.
• Gemma4DenseRotaryEmbedding — Dual-theta RoPE: sliding layers use
  θ=10 000, global layers use θ=1 000 000 with partial rotation (factor=0.25).
• wire_gemma4_kv_sharing() — Post-construction wiring of shared-KV source
  references for the last num_kv_shared_layers layers.

src/megatron/bridge/models/gemma/gemma4_provider.py

• Gemma4DenseProvider — All-in-one dataclass provider for clean MCore:
  builds GPTModel, attaches dual RoPE, PLE modules, and wires KV sharing.
  activation_func uses fast_gelu (registered as "gelu_pytorch_tanh" in
  ACTIVATION_FUNC_MAP) so megatron_to_hf_activation can match it by
  identity — avoids "Unsupported activation function" on HF export.

src/megatron/bridge/models/gemma_vl/modeling_gemma4_vl.py

• Gemma4VLModel — Wraps HF vision/audio towers with Megatron language
  model.
  • _compute_attention_mask returns a (B, 1, S, S) bool tensor
    (True = blocked) instead of a dict, matching unit-test expectations and
    downstream pipeline usage.
  • forward() returns logits directly in inference mode
    (labels=None, loss_mask=None) to avoid tuple AttributeError in the
    pipeline forward pass.

src/megatron/bridge/models/gemma/gemma4_bridge.py / src/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py

• Gemma4Bridge handles Gemma4ForCausalLM; includes a load-state alias hook
  that maps self_attention_sliding/global.* checkpoint keys to
  self_attention.* model keys, silencing the strict-load warning on export.
• Gemma4VLBridge handles Gemma4ForConditionalGeneration; respects
  GEMMA4_CONVERSION_MODE=text|vl|audio to select the appropriate provider
  and parameter mappings.

examples/models/gemma/gemma4/conversion.sh

HF → Megatron import, Megatron → HF export, and round-trip parity check.
Sets GEMMA4_CONVERSION_MODE=text to use Gemma4DenseProvider (GPTModel)
for gemma-4-E4B-it which is otherwise detected as
Gemma4ForConditionalGeneration.

examples/models/gemma/gemma4/inference.sh

Text generation from HF or Megatron checkpoint using
hf_to_megatron_generate_text.py. Sets GEMMA4_CONVERSION_MODE=text; prompts
use the Gemma 4 IT chat template for the instruction-tuned model.

examples/models/gemma/gemma4/parity_check_e4b.py

Distributed logit parity check (TP=1/2). Expected max|diff| < 3.0 (bf16).

examples/models/gemma/gemma4/README.md

Usage guide: requirements, workspace setup, conversion, inference, parity
checks, pretraining, and architecture notes.

Relation to MCore PR

This PR is the Bridge half of the split requested in
Megatron-LM #5090.
The only MCore dependency retained is the generic per_layer_inputs hook
already threaded through TransformerBlock.forward().

Validation

Validated locally with:

• HF → Megatron import ✅
• Megatron → HF export ✅
• Round-trip weight verification ✅
• Text,vision,audio modality logit parity check ✅
• Unit tests: all 227 tests passing ✅

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
• Reviewer: Does the PR have correct import guards for all optional libraries?

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[DOGEUNNKIM]

@yaoyu-33 Thanks for the review. I updated the PR based on the feedback.

• Kept Gemma4ForCausalLM registration in Gemma4Bridge for pure text use cases.
• Added examples/models/gemma/gemma4/conversion.sh and inference.sh.
• Refactored pure LLM code out of models/gemma_vl.

The shared/text Gemma4 implementation now lives under models/gemma:

• modeling_gemma4.py
• gemma4_provider.py
• gemma4_bridge.py

models/gemma_vl now only keeps multimodal-specific wrappers, VL providers, and VL/audio mappings, while reusing the text provider/modeling code from models/gemma.

I also fixed the Gemma4 Dense export-load warning caused by self_attention_sliding/global.* checkpoint keys versus self_attention.* model keys using a Gemma4-only load-state alias hook.

Local validation:

• HF -> Megatron import: passed
• Megatron -> HF export: passed
• round-trip weight verification: passed
• text ,vision, audio logit parity: passed

[yaoyu-33]

Structured review comments from local review.

[DOGEUNNKIM]

@yaoyu-33 Thanks for the thorough review! I've addressed all the comments in [cc94737]:

Copyright headers:

Added NVIDIA 2026 Apache 2.0 headers to all new non-test Python files (gemma4_bridge.py, gemma4_provider.py, modeling_gemma4_vl.py, gemma4_vl_bridge.py, gemma4.py, parity_check_e4b.py).

Hardcoded paths:

Replaced /mnt/nvme0/kdg6245 defaults in slurm_pretrain.sh with ${VAR:?'Error: set VAR to a writable directory'} guards — the script now fails explicitly if the variable is unset rather than silently using a personal path.

####torchrun → uv:
Updated TORCHRUN_BIN default to uv run python -m torch.distributed.run per repo convention.

Recipe unit test:

Added tests/unit_tests/recipes/test_gemma4_recipe.py that verifies 23 critical model config fields stay in sync between gemma4_e4b_pretrain_config() and Gemma4Bridge._build_dense_provider(), catching silent drift at test time.

[gautham-kollu]

@DOGEUNNKIM gentle reminder of the merge conflicts

[DOGEUNNKIM]

@DOGEUNNKIM gentle reminder of the merge conflicts

@gautham-kollu Thanks for the reminder. I resolved the merge conflicts and pushed the updated changes.

[ko3n1g]

/ok to test 2715f5d

[DOGEUNNKIM]

@weijiac0619

Hi, all CI checks are passing now.

Could you please review this again when you have a chance?

Thank you.

[weijiac0619]

sure. taking a look now

[weijiac0619]

Hi @DOGEUNNKIM , LGTM. Could you resolve the conflicts?

[DOGEUNNKIM]

Hi @DOGEUNNKIM , LGTM. Could you resolve the conflicts?

@weijiac0619
The conflicts have been resolved by @yaoyu-33.
Could you please rerun the tests when you have a chance?

Thank you!