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

docs/guides/models/gemma/gemma4.md

@@ -0,0 +1,148 @@
+# Gemma 4
+
+NeMo RL supports post-training the Gemma 4 family with the AutoModel training
+backend and vLLM generation. The reference configurations cover text-only DAPO
+for the E2B, 12B, 26B-A4B, and 31B variants and multimodal GRPO for the E4B
+variant.
+
+> [!IMPORTANT]
+> **Status: Functionally Ready.** The listed configurations provide short-run
+> functional and CI coverage. They are starting points for post-training, not a
+> claim of long-run convergence on every model and parallel layout.
+
+## Support Matrix
+
+| Model | Task | Training backend | Training parallelism | Generation backend | Status |
+| --- | --- | --- | --- | --- | --- |
+| `google/gemma-4-E2B-it` | Text-only DAPO | AutoModel (FSDP2) | CP1 or CP2 | vLLM | Functionally Ready |
+| `google/gemma-4-E4B-it` | VLM GRPO | AutoModel (FSDP2) | CP1 | vLLM TP4 | Functionally Ready |
+| `google/gemma-4-12B-it` | Text-only DAPO | AutoModel (FSDP2) | CP1 | vLLM TP2 | Functionally Ready |
+| `google/gemma-4-26B-A4B-it` | Text-only DAPO | AutoModel (FSDP2) | EP32 with CP1; EP16 with CP2 | vLLM TP4 | Functionally Ready |
+| `google/gemma-4-31B-it` | Text-only DAPO | AutoModel (FSDP2) | CP1 or CP2 | vLLM TP4 | Functionally Ready |
+
+## Reference Recipes
+
+Recipe YAML files under `examples/configs/recipes/` are the source of truth.
+
+| Model | Layout | Nodes and GPUs | Recipe |
+| --- | --- | --- | --- |
+| E2B | FSDP2, CP1 | 1n8g | [`dapo-gemma4-e2b-it-1n8g-fsdp2-automodel.yaml`](../../../../examples/configs/recipes/llm/dapo-gemma4-e2b-it-1n8g-fsdp2-automodel.yaml) |
+| E2B | FSDP2, CP2 | 1n8g | [`dapo-gemma4-e2b-it-1n8g-fsdp2cp2-automodel.yaml`](../../../../examples/configs/recipes/llm/dapo-gemma4-e2b-it-1n8g-fsdp2cp2-automodel.yaml) |
+| E4B | VLM FSDP2, CP1 | 1n8g | [`vlm_grpo-gemma4-e4b-geo3k-1n8g-automodel.yaml`](../../../../examples/configs/recipes/vlm/vlm_grpo-gemma4-e4b-geo3k-1n8g-automodel.yaml) |
+| 12B | FSDP2, CP1 | 2n8g | [`dapo-gemma4-12b-it-2n8g-fsdp2-automodel.yaml`](../../../../examples/configs/recipes/llm/dapo-gemma4-12b-it-2n8g-fsdp2-automodel.yaml) |
+| 26B-A4B | FSDP2, EP32, CP1 | 4n8g | [`dapo-gemma4-26ba4b-it-4n8g-fsdp2-automodel.yaml`](../../../../examples/configs/recipes/llm/dapo-gemma4-26ba4b-it-4n8g-fsdp2-automodel.yaml) |
+| 26B-A4B | FSDP2, EP16, CP2 | 4n8g | [`dapo-gemma4-26ba4b-it-4n8g-fsdp2ep16cp2-automodel.yaml`](../../../../examples/configs/recipes/llm/dapo-gemma4-26ba4b-it-4n8g-fsdp2ep16cp2-automodel.yaml) |
+| 31B | FSDP2, CP1 | 4n8g | [`dapo-gemma4-31b-it-4n8g-fsdp2-automodel.yaml`](../../../../examples/configs/recipes/llm/dapo-gemma4-31b-it-4n8g-fsdp2-automodel.yaml) |
+| 31B | FSDP2, CP2 | 4n8g | [`dapo-gemma4-31b-it-4n8g-fsdp2cp2-automodel.yaml`](../../../../examples/configs/recipes/llm/dapo-gemma4-31b-it-4n8g-fsdp2cp2-automodel.yaml) |
+
+## Run a Recipe
+
+From an allocation matching the recipe, launch the standard GRPO entry point.
+For example, run the 12B configuration with:
+
+```bash
+uv run examples/run_grpo.py \
+  --config examples/configs/recipes/llm/dapo-gemma4-12b-it-2n8g-fsdp2-automodel.yaml
+```
+
+See the [GRPO guide](../../grpo.md) for algorithm and launch details.
+
+## Context Parallel
+
+Context Parallel support uses the refactored AutoModel CP interface introduced
+by [PR #3498](https://github.com/NVIDIA-NeMo/RL/pull/3498). The E2B and 31B CP2
+recipes are thin overrides of their CP1 parents that change
+`policy.dtensor_cfg.context_parallel_size` to `2`. The 26B-A4B CP2 recipe also
+changes expert parallelism from EP32 to EP16 so the model-parallel product fits
+on the same four-node allocation. Other model, optimizer, sequence-length, and
+generation settings remain identical to their CP1 parents.
+
+Context Parallel currently applies to the text-only E2B, 26B-A4B, and 31B
+recipes backed by AutoModel's Gemma 4 model-owned attention. Do not enable it
+for the E4B VLM recipe or the 12B unified checkpoint. The text-only recipes
+freeze the vision and audio towers, disable sequence packing, and configure
+vLLM with `language_model_only: true` where the checkpoint requires it.
+
+Training and generation parallelism are independent. CP partitions training
+sequences; the `tensor_parallel_size` under `policy.generation.vllm_cfg`
+controls vLLM. The AutoModel model-parallel product must divide the training
+world size. The 26B-A4B CP2 recipe uses EP16 × CP2 across 32 GPUs, provided by
+four nodes with eight GPUs per node.
+
+## Validation Curves
+
+The following plots show raw step-level metrics captured from completed
+validation runs. Validation accuracy is plotted only at validation steps; the
+other metrics are plotted at every recorded training step.
+
+### E2B Context Parallel Parity
+
+The E2B CP1 and CP2 runs use the same 100-step training configuration except
+for `policy.dtensor_cfg.context_parallel_size`. Their trajectories remain close
+across the six metrics: the mean absolute difference in validation accuracy is
+0.0083 across six aligned validation points, while the mean absolute difference
+in generation KL error is 1.6e-5 across 100 aligned training points.
+
+![Gemma 4 E2B CP1 versus CP2 curves for train/loss, train/mean_gen_tokens_per_sample, train/reward, validation/accuracy, train/gen_kl_error, and train/approx_entropy](../../../assets/gemma4/gemma4-e2b-cp1-vs-cp2-100steps.png)
+
+### 26B-A4B Context Parallel Parity
+
+The 26B-A4B pair keeps the four-node, 32-GPU training world size fixed while
+changing the model-parallel layout from EP32 × CP1 to EP16 × CP2. Both runs
+complete 100 steps. The mean absolute difference in validation accuracy is
+0.0114 across 20 aligned validation points, while the mean absolute difference
+in generation KL error is 5.6e-5 across 100 aligned training points. The final
+validation accuracies are 0.741 for EP32 × CP1 and 0.751 for EP16 × CP2.
+
+![Gemma 4 26B-A4B EP32 CP1 versus EP16 CP2 curves for train/loss, train/mean_gen_tokens_per_sample, train/reward, validation/accuracy, train/gen_kl_error, and train/approx_entropy](../../../assets/gemma4/gemma4-26ba4b-cp1-vs-cp2-100steps.png)
+
+### 31B Context Parallel Parity
+
+The 31B CP1 and CP2 runs use the same four-node training configuration except
+for `policy.dtensor_cfg.context_parallel_size`. Both runs complete 100 steps.
+The mean absolute difference in validation accuracy is 0.0099 across 20 aligned
+validation points, while the mean absolute difference in generation KL error is
+2.2e-5 across 100 aligned training points. The final validation accuracies are
+0.823 for CP1 and 0.818 for CP2.
+
+![Gemma 4 31B CP1 versus CP2 curves for train/loss, train/mean_gen_tokens_per_sample, train/reward, validation/accuracy, train/gen_kl_error, and train/approx_entropy](../../../assets/gemma4/gemma4-31b-cp1-vs-cp2-100steps.png)
+
+### 12B Long Run
+
+The 12B CP1 run completed 200 steps. Validation accuracy increases from 0.540
+at step 5 to 0.617 at step 200 and reaches a maximum of 0.631. Training loss
+decreases from 0.166 at step 1 to 0.054 at step 200, while generation KL error
+stays below 6.9e-4. The recorded history contains 199 training-metric rows and
+39 validation rows; step 125 has no values for the requested metrics, so the
+plot does not interpolate that missing row.
+
+![Gemma 4 12B CP1 curves for train/loss, train/mean_gen_tokens_per_sample, train/reward, validation/accuracy, train/gen_kl_error, and train/approx_entropy](../../../assets/gemma4/gemma4-12b-cp1-200steps.png)
+
+## 12B Unified Checkpoint
+
+The 12B checkpoint reports `model_type: gemma4_unified` and architecture
+`Gemma4UnifiedForConditionalGeneration`. NeMo RL routes it through the
+image-text AutoModel class so the complete checkpoint can load, while the
+reference recipe trains only the language path:
+
+- The vision and audio towers are frozen.
+- vLLM generation uses `language_model_only: true` and TP2.
+- Weight refits omit the frozen vision/audio tensors because vLLM's text-only
+  unified model uses encoder-free multimodal stubs with a different layout.
+- vLLM tokenizer initialization is enabled for the unified architecture.
+- Activation checkpointing and optimizer offload during log-probability
+  computation are enabled.
+- Sequence packing and Liger kernels are disabled.
+- The total sequence length is 4,096 tokens, with up to 3,072 generated tokens.
+
+## Limitations and Tracking
+
+- Context Parallel is not supported by the Gemma 4 VLM recipe.
+- The 12B `gemma4_unified` checkpoint currently supports CP1 only; unlike the
+  other text variants, it does not use AutoModel's Gemma 4 model-owned CP
+  attention implementation.
+- The recipes disable sequence packing; validate any packing change separately.
+- The functional status does not imply long-run convergence for every variant.
+- Gemma 4 CP support is tracked by
+  [#2914](https://github.com/NVIDIA-NeMo/RL/issues/2914), and 12B support by
+  [#2913](https://github.com/NVIDIA-NeMo/RL/issues/2913).

docs/guides/models/gemma/index.md

@@ -0,0 +1,20 @@
+# Gemma
+
+This is the landing page for Gemma model guidance in NeMo RL. It links to
+version-specific pages covering reference recipes, supported parallel layouts,
+and model-specific configuration requirements.
+
+For the full list of supported models, see
+[Model Support](../../../about/model-support.md).
+
+## Version Guides
+
+- **[Gemma 4](gemma4.md)** — text-only DAPO recipes for the E2B, 12B, 26B-A4B,
+  and 31B variants, including AutoModel Context Parallel recipes, plus an E4B
+  VLM GRPO recipe.
+
+```{toctree}
+:hidden:
+
+gemma4.md
+```

docs/guides/models/index.md

@@ -12,6 +12,8 @@ For the full list of supported models, see
 
 - **[GLM](glm/index.md)** — GLM-5.1 and GLM-5.2 GRPO recipes on the Megatron
   backend, colocated and non-colocated with vLLM.
+- **[Gemma](gemma/index.md)** — Gemma 4 text-only DAPO and VLM GRPO recipes on
+  AutoModel, including Context Parallel layouts for the text-only variants.
 - **[Nemotron](nemotron/index.md)** — post-training recipes for Nemotron 3
   Nano, Nano Omni, Super, Ultra, and Nemotron 3.5 Lightning, spanning the
   Megatron and AutoModel backends.
@@ -27,6 +29,7 @@ migrated into this hub as their guidance grows.
 :hidden:
 
 glm/index.md
+gemma/index.md
 nemotron/index.md
 qwen/index.md
 ```

examples/configs/recipes/llm/dapo-gemma4-12b-it-2n8g-fsdp2-automodel.yaml

@@ -0,0 +1,98 @@
+defaults: ../../grpo_math_1B.yaml
+grpo:
+  batch_multiplier: 3
+  val_period: 5
+  max_val_samples: 960
+  val_batch_size: 960
+  use_leave_one_out_baseline: false
+  use_dynamic_sampling: true
+  reward_scaling:
+    enabled: true
+    target_min: -1.0
+  reward_shaping:
+    enabled: true
+    overlong_buffer_length: 384
+    max_response_length: 3072
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_type: tis
+  truncated_importance_sampling_ratio: 2
+  ratio_clip_max: 0.28
+  ratio_clip_c: 10
+checkpointing:
+  checkpoint_dir: results/dapo-gemma4-12b-it-2n8g-fsdp2-automodel
+  save_period: 5
+policy:
+  model_name: google/gemma-4-12B-it
+  tokenizer:
+    name: google/gemma-4-12B-it
+  train_micro_batch_size: 1
+  logprob_batch_size: 1
+  max_total_sequence_length: 4096
+  logprob_chunk_size: 4096
+  offload_optimizer_for_logprob: true
+  optimizer:
+    name: transformer_engine.pytorch.optimizers.fused_adam.FusedAdam
+    kwargs:
+      lr: 1.0e-06
+      weight_decay: 0.1
+      master_weights: true
+      store_param_remainders: true
+      exp_avg_dtype: torch.bfloat16
+      exp_avg_sq_dtype: torch.bfloat16
+  scheduler:
+  - name: torch.optim.lr_scheduler.LinearLR
+    kwargs:
+      start_factor: 1.0e-08
+      end_factor: 1.0
+      total_iters: 10
+  - name: torch.optim.lr_scheduler.ConstantLR
+    kwargs:
+      factor: 1.0
+      total_iters: 10000000000
+  - milestones:
+    - 10
+  dtensor_cfg:
+    activation_checkpointing: true
+    automodel_kwargs:
+      use_liger_kernel: false
+      use_sdpa_patching: false
+      freeze_config:
+        freeze_vision_tower: true
+        freeze_audio_tower: true
+        freeze_language_model: false
+  sequence_packing:
+    enabled: false
+  dynamic_batching:
+    enabled: true
+  make_sequence_length_divisible_by: 1
+  generation:
+    max_new_tokens: 3072
+    vllm_cfg:
+      tensor_parallel_size: 2
+      gpu_memory_utilization: 0.5
+    vllm_kwargs:
+      language_model_only: true
+data:
+  max_input_seq_length: 2048
+  train:
+    dataset_name: DAPOMath17K
+  validation:
+    dataset_name: DAPOMathAIME2024
+  default:
+    prompt_file: null
+env:
+  math:
+    num_workers: 16
+    math_verify_impl: dapo_math_verify
+logger:
+  log_dir: logs/dapo-gemma4-12b-it-2n8g-fsdp2-automodel
+  wandb_enabled: true
+  tensorboard_enabled: true
+  wandb:
+    project: nemorl-gemma4
+    name: dapo-gemma4-12b-it-2n8g-fsdp2-automodel
+cluster:
+  gpus_per_node: 8
+  num_nodes: 2

examples/configs/recipes/llm/dapo-gemma4-26ba4b-it-4n8g-fsdp2ep16cp2-automodel.yaml

@@ -0,0 +1,11 @@
+defaults: ./dapo-gemma4-26ba4b-it-4n8g-fsdp2-automodel.yaml
+policy:
+  dtensor_cfg:
+    expert_parallel_size: 16
+    context_parallel_size: 2
+checkpointing:
+  checkpoint_dir: results/dapo-gemma4-26ba4b-it-4n8g-fsdp2ep16cp2-automodel
+logger:
+  log_dir: logs/dapo-gemma4-26ba4b-it-4n8g-fsdp2ep16cp2-automodel
+  wandb:
+    name: dapo-gemma4-26ba4b-it-4n8g-fsdp2ep16cp2-automodel

examples/configs/recipes/llm/dapo-gemma4-31b-it-4n8g-fsdp2cp2-automodel.yaml

@@ -0,0 +1,10 @@
+defaults: ./dapo-gemma4-31b-it-4n8g-fsdp2-automodel.yaml
+policy:
+  dtensor_cfg:
+    context_parallel_size: 2
+checkpointing:
+  checkpoint_dir: results/dapo-gemma4-31b-it-4n8g-fsdp2cp2-automodel
+logger:
+  log_dir: logs/dapo-gemma4-31b-it-4n8g-fsdp2cp2-automodel
+  wandb:
+    name: dapo-gemma4-31b-it-4n8g-fsdp2cp2-automodel

examples/configs/recipes/llm/dapo-gemma4-e2b-it-1n8g-fsdp2cp2-automodel.yaml

@@ -0,0 +1,10 @@
+defaults: ./dapo-gemma4-e2b-it-1n8g-fsdp2-automodel.yaml
+policy:
+  dtensor_cfg:
+    context_parallel_size: 2
+checkpointing:
+  checkpoint_dir: results/dapo-gemma4-e2b-it-1n8g-fsdp2cp2-automodel
+logger:
+  log_dir: logs/dapo-gemma4-e2b-it-1n8g-fsdp2cp2-automodel
+  wandb:
+    name: dapo-gemma4-e2b-it-1n8g-fsdp2cp2-automodel

examples/configs/recipes/vlm/vlm_grpo-gemma4-e4b-geo3k-1n8g-automodel.yaml

@@ -24,6 +24,8 @@ policy:
       exp_avg_sq_dtype: torch.bfloat16
   dtensor_cfg:
     activation_checkpointing: true
+    env_vars:
+      PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True
     automodel_kwargs:
       backend:
         _target_: nemo_automodel.components.models.common.utils.BackendConfig

nemo_rl/models/automodel/setup.py

@@ -585,6 +585,15 @@ def setup_model_and_optimizer(
                 "Please refer to https://github.com/NVIDIA/NeMo-RL/blob/main/docs/model-quirks.md#context-parallel-with-fsdp2 for more details."
             )
 
+        if model_config.model_type == "gemma4_unified":
+            raise AssertionError(
+                "Context parallel is not supported for the Gemma 4 unified "
+                "checkpoint (model_type='gemma4_unified'). Its global-attention "
+                "GQA uses head_dim=512, for which no CP SDPA kernel is available. "
+                "Set policy.dtensor_cfg.context_parallel_size = 1. See "
+                "docs/guides/models/gemma/gemma4.md."
+            )
+
         if tp_size > 1 and sequence_parallel_enabled:
             raise AssertionError(
                 "It's a known issue that context parallel can't be used together with sequence parallel in DTensor worker. "

nemo_rl/models/generation/vllm/vllm_backend.py

@@ -61,6 +61,11 @@
 UnsupportedNativeRefitTransport = Literal["checkpoint_engine", "sparse_delta"]
 WeightUpdateFinalizer = Callable[[], None]
 
+_GEMMA4_UNIFIED_MULTIMODAL_WEIGHT_MARKERS = (
+    "model.embed_vision.",
+    "model.embed_audio.",
+)
+
 
 def _format_refit_key_error(label: str, keys: set[str]) -> str:
     """Format a bounded refit-key diagnostic."""
@@ -798,6 +803,29 @@ def _load_weights(self, weights):
             for idx, (key, weight) in enumerate(weights):
                 weights[idx] = (fix_gemma3_vision_weight_name(key), weight)
 
+        model_config = self.model_runner.vllm_config.model_config
+        multimodal_config = getattr(model_config, "multimodal_config", None)
+        if any("Gemma4Unified" in arch for arch in model_config.architectures) and (
+            multimodal_config is not None and multimodal_config.language_model_only
+        ):
+            # HF loads the full unified checkpoint, while text-only vLLM uses
+            # encoder-free multimodal stubs with a different parameter layout.
+            # The recipe freezes these towers and never invokes them, so only
+            # refit the language path and leave the unused vLLM stubs untouched.
+            num_weights = len(weights)
+            weights = [
+                (key, weight)
+                for key, weight in weights
+                if not key.startswith(_GEMMA4_UNIFIED_MULTIMODAL_WEIGHT_MARKERS)
+            ]
+            num_dropped = num_weights - len(weights)
+            if num_dropped:
+                logger.info(
+                    "Gemma4 Unified text-only refit dropped %d frozen "
+                    "vision/audio weights",
+                    num_dropped,
+                )
+
         policy_weights, draft_weights = self._split_policy_and_draft_weights(weights)
         self._load_hf_weights(policy_weights)
         # Eagle3 draft weights are exported with the `draft.` prefix.