feat(gemma4): add CP recipes and 12B support

[sharonyu-115]

What does this PR do?

Adds AutoModel Context Parallel recipes and functional coverage for the Gemma 4 text variants, and adds text-only DAPO support for the Gemma 4 12B unified checkpoint.

• Adds CP2 recipe overrides for google/gemma-4-E2B-it, google/gemma-4-26B-A4B-it, and google/gemma-4-31B-it.
• Uses EP16 x CP2 on 4 nodes for the 26B-A4B CP2 recipe.
• Adds a 2-node AutoModel/FSDP2 recipe for google/gemma-4-12B-it.
• Routes the 12B gemma4_unified architecture through the image-text AutoModel class and ensures vLLM initializes its tokenizer.
• Makes text-only vLLM refits omit only the exact frozen vision/audio state-dict prefixes whose layout differs from the full unified checkpoint.
• Enforces the diagnosed CP1-only constraint for gemma4_unified at runtime.
• Adds functional drivers and Gemma 4 model-hub documentation with self-contained E2B CP parity and 12B training curves.

The 12B unified checkpoint is intentionally CP1-only. Its global-attention GQA uses head dimension 512, for which no CP SDPA kernel is available. E2B, 26B-A4B, and 31B use the refactored model-owned CP interface from #3498.

Issues

Closes #2913

Closes #2914

Usage

Run the 12B text-only recipe:

uv run examples/run_grpo.py \
  --config examples/configs/recipes/llm/dapo-gemma4-12b-it-2n8g-fsdp2-automodel.yaml

Run the E2B CP2 recipe:

uv run examples/run_grpo.py \
  --config examples/configs/recipes/llm/dapo-gemma4-e2b-it-1n8g-fsdp2cp2-automodel.yaml

Run the 4-node 26B-A4B EP16 x CP2 recipe:

uv run examples/run_grpo.py \
  --config examples/configs/recipes/llm/dapo-gemma4-26ba4b-it-4n8g-fsdp2ep16cp2-automodel.yaml

Validation

Local and static checks:

• Recipe minimization and the new functional driver's TEST_DRYRUN=1 config resolution: passed.
• Nightly suite dry-run: 4,192 GPU-hours, below the 4,197-hour cap.
• Tokenizer architecture helper: 5 parametrized cases passed, including None, empty, causal-only, conditional, and ordered multi-match inputs.
• Gemma 4 unified CP runtime guard: 1 focused test passed.
• Gemma 4 unified refit filtering: 2 focused cases passed (text-only filtering and full-model preservation).
• Ruff format, Ruff lint, shell/config checks, and git diff --check: passed.

GPU functional evidence:

• E2B CP1: Slurm job 17503334, W&B run ID e2c1a830, completed 100 steps.
• E2B CP2: Slurm job 17503337, W&B run ID e2c2a830, completed 100 steps.
  • Validation-accuracy mean absolute difference: 0.0083 over six aligned validation points.
  • Generation-KL mean absolute difference: 1.6e-5 over 100 aligned training steps.
• 12B CP1: W&B run ID g412b200, completed 200 steps across the resumed Slurm segments (16900585, 16995049, 17503331, 17503333).
  • Validation accuracy: 0.540 at step 5 to 0.617 at step 200; maximum 0.631.
  • Training loss: 0.166 at step 1 to 0.054 at step 200.
  • Generation KL remained below 6.9e-4.
• AutoModel CP contract suite: Slurm job 16896034; all 12 applicable tests passed.
• 26B-A4B EP16 x CP2 review smoke: submitted on the batch partition as Slurm job 17620431, W&B run ID g26e16c2 in nv-welcome/nemo-rl-gemma4. This is pending and is not claimed as passed.

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 run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• Rebased onto upstream main at 9b25508a3.
• Validation used the refreshed AutoModel container pinned for these Gemma 4 runs.
• W&B run IDs are recorded without private dashboard links; the documentation keeps the plots and metric summaries self-contained.

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

[sharonyu-115]

/ok to test 3a52dbf

[sharonyu-115]

/ok to test e1c593b

[sharonyu-115]

/ok to test bebe687

[sharonyu-115]

/ok to test d05ccf7

[sharonyu-115]

/ok to test 2773a4d

[sharonyu-115]

/ok to test 5223aee

[sharonyu-115]

Hi @yuki-97 would you help review this PR when you get a chance? Thank you!

[sharonyu-115]

/ok to test 31e8be4

[yuki-97]

Overall LGTM. The CP recipes and the 12B gemma4_unified path both look solid, and the CP1-vs-CP2 parity evidence is stronger than most model-enablement PRs ship — a 0.0083 validation-accuracy MAD and a 1.6e-5 generation-KL MAD over 100 aligned steps is exactly the right way to argue a new parallelism axis is a no-op. Thanks also for the self-contained curves in the guide.

Scope: Gemma 4 CP recipes, the 12B gemma4_unified path, and the test-suite/doc changes that ship with them.

Checked:

• vLLM refit path — both refit entry points traced, and the new filter's premise verified against pinned vLLM 0.25.1 source.
• Test drivers — every newly asserted metric key traced to its producer, and each changed threshold compared against its pre-PR form.
• Docs vs code — every support-matrix cell and recipe link checked against the YAML it names.
• Guards, config keys, copyright, docs index, and the nightly budget bump against the repo's own conventions.

Remaining comments are one correctness item on the reload-API refit path, plus four nits on vllm_backend.py and its unit test. Nothing blocking beyond the first.

[sharonyu-115]

/ok to test a304926

[sharonyu-115]

/ok to test a304926

[sharonyu-115]

/ok to test 6c778f4

[copy-pr-bot]

/ok to test 6c778f41cfe76e9667397f9450075578c9aeeb77

@sharonyu-115, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[sharonyu-115]

/ok to test 6c778f4

[sharonyu-115]

/ok to test 6c778f4

[sharonyu-115]

Verification update for a30492675, which addresses the review feedback:

• The Gemma 4 unified text-only architecture check is now shared by both the legacy load and native reload paths. Multimodal weights are filtered lazily, negative cases remain unfiltered, and the drop summary is logged only once.
• Focused Gemma 4 vLLM tests: 7 passed. The full PR workflow also completed successfully, including all three vLLM unit-test shards.
• Current-head 12B nightly, Slurm 18089008: 20/20 steps, exit 0. All configured metric assertions passed:
  • final train/loss: 0.101327
  • median train/token_mult_prob_error: 1.007972
  • mean train/gen_kl_error: 0.000403992
  • mean train/reward: 0.251566
  • mean train/filtered_reward: -0.001452
  • mean/final validation/accuracy: 0.553385 / 0.548958
• Current-head non-colocated native-reload run, Slurm 18089203: 2/2 steps, exit 0, using two policy nodes plus one dedicated 8-GPU vLLM node. Two native reload cycles completed in 0.34 s and 0.42 s, and the second rollout completed after the updated weights were loaded:
  • final train/loss: 0.156278
  • median train/token_mult_prob_error: 1.042639
  • final train/gen_kl_error: 0.000410787
  • final train/reward: 0.270365
  • no OOM, NCCL error, traceback, or Ray task failure was observed.

Both runtime jobs used current head 6c778f41c, whose only commit after a30492675 is the empty signed CI-retry commit.

[sharonyu-115]

Hi @yuki-97 Thank you very much for your review and comments. I have addressed all the comments and tested good. Please see the details above. Could you help give another round of review?

[yuki-97]

LGTM, thanks @sharonyu-115