Scale Gemma 3 RL on GKE with NeMo RL, Custom ComputeClasses, and 8-NIC RDMA

[Baichuan_Sun]

Scale Gemma 3 reinforcement learning on GKE with NVIDIA NeMo RL, Custom ComputeClasses, and 8-NIC RDMA

Co-authors: Baichuan Sun and Injae Kwak, Google Cloud
@Baichuan_Sun @Injae_Kwak

TL;DR: This guide demonstrates how to apply Group Relative Policy Optimization (GRPO) reinforcement learning to google/gemma-3-1b-it using NVIDIA NeMo RL (v0.5.0) and KubeRay on Google Kubernetes Engine (GKE) to improve domain-specific mathematical reasoning accuracy while maximizing GPU capacity obtainability through GKE Custom ComputeClasses:

• Domain task accuracy (MATH-500 Pass@1): Post-training google/gemma-3-1b-it with synchronous on-policy GRPO for 30 steps lifts zero-shot MATH-500 Pass@1 accuracy from 40.20% (201 / 500) to 46.40% (232 / 500) – a +6.20 percentage-point (+15.4% relative) gain (and +12.11 percentage points on in-loop validation before Step-40 KL length contraction).
• Dedicated Gemma 3 architectural support (NVIDIA NeMo RL v0.5.0): Standard RL post-training pipelines tuned around Llama and Qwen defaults often fail at startup or suffer silent policy collapse on Gemma 3 due to four model-specific quirks – vLLM’s non-persistent normalizer scaling buffer during weight refits, the 1B (Gemma3ForCausalLM) versus 4B/12B/27B (Gemma3ForConditionalGeneration) VLM architecture split with 4 layernorms per block and dual global/local RoPE, a strict alternating user/model Jinja template that omits <eos>, and 256-dimensional attention heads (head_dim=256) with non-null sliding-window attn_bias. NeMo RL (v0.5.0) provides built-in ModelFlag and DTensor parallelization support for the first three and clean vllm_kwargs / dynamic-batching hooks for the fourth.
• Accelerator capacity obtainability (Custom ComputeClasses): A declarative GKE ComputeClass (gpu-rl-ccc) diversifies accelerator provisioning across scale-to-zero Dynamic Workload Scheduler (DWS) Flex-Start (b200-dws-ccc / h200-dws-ccc) and fallback Spot GPU node pools (b200-spot-ccc / h200-spot-ccc) on 8-GPU NVIDIA B200 (a4-highgpu-8g) and H200 (a3-ultragpu-8g) nodes.
• Unified GKE infrastructure and ML workload: A single GKE Standard cluster combines a 3-VPC 8-NIC RoCE GPUDirect-RDMA (gIB) network fabric and Cloud Storage FUSE (30Gi head ephemeral cache) with colocated 1D PyTorch FSDP2 training, vLLM V1 (block_size: 32, load_format: "auto") rollouts, and CPU-consolidated PyTorch Distributed Checkpoints (DCP).

To build multi-step reasoning agents without paying frontier-model latency and token costs on every intermediate turn, you can post-train compact open-weight models (1B to 8B parameters) using reinforcement learning with verifiable rewards. Post-training Google’s open-weight Gemma 3 family with Group Relative Policy Optimization (GRPO) sharpens domain-specific mathematical, code, and tool-planning accuracy while keeping production serving within a single-GPU footprint.

Unlike supervised fine-tuning (SFT), on-policy reinforcement learning alternates continuously between high-throughput autoregressive generation, deterministic reward evaluation, sharded backward passes, and collective weight synchronization. Running that loop reliably in Kubernetes stresses GPU capacity provisioning, multi-NIC RDMA collectives, storage sidecar caching, attention kernel geometry, and distributed checkpoint consolidation at the same time.

This article provides a self-contained, end-to-end blueprint for post-training google/gemma-3-1b-it (extensible to 4B and 12B variants) on a single GKE cluster across 8-GPU NVIDIA B200 (a4-highgpu-8g) and NVIDIA H200 (a3-ultragpu-8g) nodes, organized across two parts of the same cluster stack:

• GKE infrastructure: A private multi-network GKE Standard cluster spanning three VPCs (management, north-south gVNIC, and an 8-subnet RoCE RDMA fabric), paired scale-to-zero Dynamic Workload Scheduler (DWS) Flex-Start and Spot GPU node pools, a declarative ComputeClass, and Google Cloud’s nccl-plugin-gib RDMA installer DaemonSet.
• ML workload (KubeRay and NeMo RL): Self-contained KubeRay (Ray 2.49.2) and NVIDIA NeMo RL (v0.5.0) manifests, Gemma 3 vLLM V1 kernel overrides, Cloud Storage FUSE sidecar cache tuning, and offline PyTorch Distributed Checkpoint (DCP) consolidation commands that lifted zero-shot MATH-500 Pass@1 accuracy from 40.20% to 46.40% (+6.20 percentage points, or +15.4% relative) in 30 GRPO steps.

Architect the GKE infrastructure and ML workload

Decoupling cluster-level network, storage, and node-pool provisioning (GKE infrastructure) from the Ray-orchestrated RL training loop (ML workload) lets platform teams reuse the same RDMA-enabled GKE cluster across multiple post-training workloads without rebuilding node pools or VPCs. Figure 1 illustrates how the GKE infrastructure and the KubeRay ML workload run together on a single cluster.

Figure 1. GKE reference architecture pairing 3-VPC RoCE RDMA and Custom ComputeClasses (GKE infrastructure) with a KubeRay colocated synchronous GRPO loop (ML workload) on a single cluster.

Choose 1D PyTorch FSDP2 and vLLM V1 over 3D parallelism for 1B to 8B models

NVIDIA NeMo RL supports two distributed training backbones: 3D parallelism through Megatron-Core (megatron_cfg) and 1D sharded data parallelism through native PyTorch Distributed Tensor / FSDP2 (dtensor_cfg._v2: true). Matching the parallelism backend to the parameter regime eliminates format-conversion toil without sacrificing throughput.

Architectural dimension	PyTorch FSDP2 and vLLM V1 (dtensor_cfg)	Megatron-Core 3D parallelism (megatron_cfg)
Target model scale	Compact and mid-size models (1B to ~30B parameters) where a single layer fits cleanly in per-GPU HBM	Frontier-scale models (70B+ parameters) or 128K+ context windows exceeding single-node HBM
Checkpoint and weight format	Native Hugging Face weights and PyTorch Distributed Checkpoints (DCP), with zero .nemo conversion	Requires explicit import and export between Hugging Face .safetensors and Megatron .nemo layouts
Parallelism configuration	Automatic 1D parameter, gradient, and AdamW state sharding across GPUs (tensor_parallel_size: 1)	Explicit tuning of Tensor Parallel (TP), Pipeline Parallel (PP), and Sequence Parallel geometries
Rollout and weight refit path	Direct collective broadcast from FSDP2 shards into colocated vLLM V1 PagedAttention workers	Weight resharding across 3D pipeline and tensor ranks into inference buffers

For Gemma 3 models in the 1B to 8B range on 8-GPU B200 (180 GB HBM3e per GPU) or H200 (141 GB HBM3e per GPU) nodes, setting dtensor_cfg.enabled: true and megatron_cfg.enabled: false keeps the entire workflow native to Hugging Face tokenizers and chat templates while saturating PagedAttention rollout workers.

Start with synchronous on-policy rollouts before enabling async execution

When monitoring GPU compute utilization on the Ray Dashboard during a colocated GRPO run (colocated.enabled: true, async_grpo.enabled: false), you observe periodic sawtooth oscillations rather than a flat 100% line. That profile reflects the four sequential phases of strict on-policy reinforcement learning:

1. Autoregressive rollout (vLLM V1): PyTorch training workers pause while vLLM V1 PagedAttention workers generate 16 candidate reasoning trajectories for each of the 32 prompts in the step (512 sequences total).
2. Deterministic outcome verification (SymPy): Rule-based math workers parse final \boxed{} expressions and compute leave-one-out group-relative advantages across each prompt’s 16 rollouts.
3. Policy gradient update (PyTorch FSDP2): Training actors compute token-level log-probabilities, apply the k3 KL-divergence penalty against the reference policy, run loss.backward(), and step AdamW in bfloat16.
4. Collective weight broadcast: Updated policy parameters broadcast directly over GPU memory and RDMA into the colocated vLLM V1 workers before the next rollout step begins.

NeMo RL also supports decoupled asynchronous rollouts (async_grpo.enabled: true), where vLLM generates trajectories continuously from slightly older policy weights while FSDP2 updates run concurrently. Although asynchronous execution smooths GPU utilization, off-policy trajectory staleness requires importance-sampling ratio clipping and can destabilize mathematical reasoning runs early on. Always establish a converging synchronous baseline (async_grpo.enabled: false) before experimenting with asynchronous rollout buffers.

Provision the GKE RDMA and ComputeClass foundation

Before deploying KubeRay, configure a GKE Standard cluster with multi-networking (datapath_provider = "ADVANCED_DATAPATH", enable_multi_networking = true), managed KubeRay and Cloud Storage FUSE CSI addons, Workload Identity Federation for GKE, and Container-Optimized OS image streaming (gcfs_config).

Build the 3-VPC GPUDirect-RDMA network and scale-to-zero GPU node pools

To isolate east-west GPU collectives from north-south dataset and control-plane traffic, attach three distinct VPC networks across every 8-GPU node:

• Management VPC (nic0 / eth0): Hosts the private GKE control plane, CPU system pool (default-pool), and Cloud NAT egress.
• Secondary gVNIC VPC (nic1 / gvnic-1): Registers a GKENetworkParamSet (deviceMode: "NetDevice") and Network (type: "Device") for high-throughput north-south storage and API traffic without contending with RDMA rings.
• RoCE RDMA HPC VPC (nic2 through nic9 / rdma-0 through rdma-7): Uses the zonal RoCE network profile (projects/PROJECT_ID/global/networkProfiles/ZONE-vpc-roce) with eight /24 subnets mapped 1:1 to eight GKENetworkParamSet (deviceMode: "RDMA") and Network (type: "Device") Kubernetes custom resources.

On top of those networks, provision two scale-to-zero (initial_node_count = 0, min_node_count = 0) GPU node pools that share a cloud.google.com/compute-class: gpu-rl-ccc label and NO_SCHEDULE taint (local_ssd_count = 32):

1. Primary DWS Flex-Start node pool (b200-dws-ccc or h200-dws-ccc): Configured with flex_start = true, spot = false, reservation_affinity = "NO_RESERVATION", and max_run_duration = "604800s".
2. Fallback Spot node pool (b200-spot-ccc or h200-spot-ccc): Configured with spot = true and local_ssd_count = 32 for immediate cost-optimized fallback capacity.

Configure Custom ComputeClass priority fallback and the gIB NCCL DaemonSet

With both scale-to-zero node pools registered, apply the following ComputeClass manifest (ccc.yaml) so GKE attempts DWS Flex-Start provisioning first and falls back automatically to Spot capacity:

# SPDX-License-Identifier: Apache-2.0
apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: gpu-rl-ccc
spec:
  activeMigration:
    optimizeRulePriority: false
  priorities:
  - machineType: a4-highgpu-8g
    nodepools: ["b200-dws-ccc"]
    gpu:
      type: "nvidia-b200"
      count: 8
      driverVersion: "latest"
    storage:
      localSSDCount: 32
    flexStart:
      enabled: true
  - machineType: a4-highgpu-8g
    nodepools: ["b200-spot-ccc"]
    gpu:
      type: "nvidia-b200"
      count: 8
      driverVersion: "latest"
    storage:
      localSSDCount: 32
    spot: true

(To target 8-GPU NVIDIA H200 nodes instead, set machineType: a3-ultragpu-8g, gpu.type: "nvidia-h200-141gb", and nodepools to ["h200-dws-ccc"] and ["h200-spot-ccc"].)

Next, deploy the nccl-rdma-installer DaemonSet in kube-system using us-docker.pkg.dev/gce-ai-infra/gpudirect-gib/nccl-plugin-gib:v1.1.0 (scoped by nodeAffinity to cloud.google.com/gke-accelerator values nvidia-b200 and nvidia-h200-141gb) so every autoscaled GPU node installs Google Cloud’s gIB NCCL tuner and RDMA libraries onto /home/kubernetes/bin/gib and /home/kubernetes/bin/nvidia/lib64 before worker pods start.

Deploy the KubeRay cluster manifest for NeMo RL v0.5.0

Running nvcr.io/nvidia/nemo-rl:v0.5.0 (Python 3.12) on KubeRay requires three pod-level patterns in your RayCluster manifest:

1. Raise the head-pod Cloud Storage FUSE sidecar ephemeral cache from "0" to "30Gi": Redirecting HF_HOME to a Cloud Storage FUSE mount (/data/huggingface) keeps container images lean, but compiling multi-gigabyte Parquet splits (OpenMathInstruct-2) stages chunks inside the CSI sidecar’s local cache volume. Leaving gke-gcsfuse/ephemeral-storage-limit at "0" bounds the sidecar to the head container’s 9Gi ephemeral-storage limit, triggering kubelet pod eviction around 11% into dataset split generation. Setting gke-gcsfuse/ephemeral-storage-limit: "30Gi" on headGroupSpec (while omitting aggressive parallel file-cache mount flags) eliminates head-pod evictions within standard CPU node quotas.
2. Stage Ray into pod-local emptyDir and proxy KubeRay’s startup command with eval "$*": Installing ray[default]==2.49.2 into pod-local emptyDir (/tmp/ray/packages) avoids metadata bottlenecks on network FUSE storage. Wrapping source /usr/local/gib/scripts/set_nccl_env.sh && eval "$*" in the worker command field ensures that sourcing RDMA environment variables does not overwrite KubeRay’s dynamically injected ray start arguments.
3. Bind all RDMA interfaces with gIB NCCL environment tuning: Each 8-GPU worker attaches the management, gVNIC, and eight rdma-0 through rdma-7 interfaces, sets NCCL_CROSS_NIC=1, NCCL_ALGO=Ring,Tree, NCCL_PROTO=Simple, and NCCL_TUNER_PLUGIN=libnccl-tuner.so, enables NRL_VLLM_USE_V1=1, and mounts a 250Gi in-memory /dev/shm volume.

Save the following manifest as ray-cluster-b200-nemo.yaml and apply it with kubectl apply -f ray-cluster-b200-nemo.yaml:

# SPDX-License-Identifier: Apache-2.0
apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-cluster-b200-nemo
  namespace: default
spec:
  enableInTreeAutoscaling: true
  rayVersion: "2.49.2"
  headGroupSpec:
    rayStartParams:
      dashboard-host: "0.0.0.0"
      num-cpus: "12"
    template:
      metadata:
        annotations:
          gke-gcsfuse/volumes: "true"
          gke-gcsfuse/cpu-limit: "0"
          gke-gcsfuse/memory-limit: "0"
          gke-gcsfuse/ephemeral-storage-limit: "30Gi"
      spec:
        serviceAccountName: WORKLOAD_IDENTITY_SA
        restartPolicy: Never
        nodeSelector:
          cloud.google.com/gke-nodepool: "default-pool"
        initContainers:
        - name: install-ray
          image: nvcr.io/nvidia/nemo-rl:v0.5.0
          command: ["/bin/bash", "-c", "pip install -t /tmp/ray/packages \"ray[default]==2.49.2\""]
          volumeMounts:
          - name: ray-packages
            mountPath: /tmp/ray
        containers:
        - name: ray-head
          image: nvcr.io/nvidia/nemo-rl:v0.5.0
          env:
          - name: PYTHONPATH
            value: "/tmp/ray/packages"
          - name: RAY_enable_autoscaler_v2
            value: "1"
          - name: ORCHESTRATOR
            value: "ray"
          command: ["/bin/bash", "-c", "export PATH=/tmp/ray/packages/bin:$PATH; eval \"$*\""]
          resources:
            requests:
              cpu: "12"
              memory: "32G"
              ephemeral-storage: "9Gi"
            limits:
              cpu: "12"
              memory: "32G"
              ephemeral-storage: "9Gi"
          volumeMounts:
          - name: ray-packages
            mountPath: /tmp/ray
          - name: gcs-fuse-csi-eph
            mountPath: /data
        volumes:
        - name: ray-packages
          emptyDir: {}
        - name: gcs-fuse-csi-eph
          csi:
            driver: gcsfuse.csi.storage.gke.io
            volumeAttributes:
              bucketName: GCS_BUCKET_NAME
              gcsfuseMetadataPrefetchOnMount: "true"
              mountOptions: "implicit-dirs,write:enable-streaming-writes:true,read_ahead_kb=1024,file-system:kernel-list-cache-ttl-secs:-1"
  workerGroupSpecs:
  - groupName: gpu-group
    replicas: 2
    minReplicas: 1
    maxReplicas: 4
    rayStartParams:
      num-cpus: "220"
    template:
      metadata:
        annotations:
          networking.gke.io/default-interface: "eth0"
          networking.gke.io/interfaces: |
            [
              {"interfaceName":"eth0","network":"default"},
              {"interfaceName":"eth1","network":"gvnic-1"},
              {"interfaceName":"eth2","network":"rdma-0"},
              {"interfaceName":"eth3","network":"rdma-1"},
              {"interfaceName":"eth4","network":"rdma-2"},
              {"interfaceName":"eth5","network":"rdma-3"},
              {"interfaceName":"eth6","network":"rdma-4"},
              {"interfaceName":"eth7","network":"rdma-5"},
              {"interfaceName":"eth8","network":"rdma-6"},
              {"interfaceName":"eth9","network":"rdma-7"}
            ]
          gke-gcsfuse/volumes: "true"
          gke-gcsfuse/cpu-limit: "0"
          gke-gcsfuse/memory-limit: "0"
          gke-gcsfuse/ephemeral-storage-limit: "0"
      spec:
        serviceAccountName: WORKLOAD_IDENTITY_SA
        restartPolicy: Never
        nodeSelector:
          cloud.google.com/compute-class: "gpu-rl-ccc"
        tolerations:
        - key: "cloud.google.com/compute-class"
          operator: "Equal"
          value: "gpu-rl-ccc"
          effect: "NoSchedule"
        - key: "nvidia.com/gpu"
          operator: "Exists"
          effect: "NoSchedule"
        initContainers:
        - name: install-ray
          image: nvcr.io/nvidia/nemo-rl:v0.5.0
          command: ["/bin/bash", "-c", "pip install -t /tmp/ray/packages \"ray[default]==2.49.2\""]
          volumeMounts:
          - name: ray-tmp-storage
            mountPath: /tmp
        containers:
        - name: ray-worker
          image: nvcr.io/nvidia/nemo-rl:v0.5.0
          env:
          - name: PYTHONPATH
            value: "/tmp/ray/packages"
          - name: LD_LIBRARY_PATH
            value: "/usr/local/nvidia/lib64:/usr/local/gib/lib64"
          - name: NCCL_CROSS_NIC
            value: "1"
          - name: NCCL_ALGO
            value: "Ring,Tree"
          - name: NCCL_PROTO
            value: "Simple"
          - name: NCCL_MIN_NCHANNELS
            value: "4"
          - name: NCCL_TUNER_PLUGIN
            value: "libnccl-tuner.so"
          - name: NRL_VLLM_USE_V1
            value: "1"
          command: ["/bin/bash", "-c", "export PATH=/tmp/ray/packages/bin:$PATH; source /usr/local/gib/scripts/set_nccl_env.sh && eval \"$*\""]
          resources:
            requests:
              cpu: "220"
              memory: "2800Gi"
              nvidia.com/gpu: "8"
              ephemeral-storage: "1000Gi"
            limits:
              cpu: "220"
              memory: "2800Gi"
              nvidia.com/gpu: "8"
              ephemeral-storage: "1000Gi"
          volumeMounts:
          - name: nvidia
            mountPath: /usr/local/nvidia
          - name: gib
            mountPath: /usr/local/gib
          - name: shared-memory
            mountPath: /dev/shm
          - name: ray-tmp-storage
            mountPath: /tmp
          - name: gcs-fuse-csi-eph
            mountPath: /data
        volumes:
        - name: gib
          hostPath:
            path: /home/kubernetes/bin/gib
        - name: nvidia
          hostPath:
            path: /home/kubernetes/bin/nvidia
        - name: shared-memory
          emptyDir:
            medium: "Memory"
            sizeLimit: 250Gi
        - name: ray-tmp-storage
          emptyDir: {}
        - name: gcs-fuse-csi-eph
          csi:
            driver: gcsfuse.csi.storage.gke.io
            volumeAttributes:
              bucketName: GCS_BUCKET_NAME
              gcsfuseMetadataPrefetchOnMount: "true"
              mountOptions: "implicit-dirs,write:enable-streaming-writes:true,read_ahead_kb=1024,file-system:kernel-list-cache-ttl-secs:-1"

Replace the following:

• WORKLOAD_IDENTITY_SA: the Kubernetes service account bound through Workload Identity Federation for GKE to roles/storage.objectUser on your Cloud Storage bucket.
• GCS_BUCKET_NAME: the Cloud Storage bucket used for Hugging Face dataset caches, TensorBoard event logs, and checkpoints.

Once the head and worker pods reach Running, port-forwarding svc/ray-cluster-b200-nemo-head-svc on port 8265 exposes live GPU compute and HBM telemetry across the cluster (Figure 2). Each a4-highgpu-8g worker node registers 220 CPUs, 2.73 TB of system memory, and eight NVIDIA B200 GPUs (183,359 MiB HBM3e per GPU), hitting 87% to 95% active GPU utilization during rollout and training bursts.

Figure 2. Ray Dashboard cluster view showing the ray-head pod and an 8-GPU NVIDIA B200 worker node (183,359 MiB HBM3e per GPU, 2.73 TB host RAM) at 87% to 95% GPU compute utilization.

Configure NeMo RL and vLLM V1 for Gemma 3 architectural quirks

Most reinforcement learning harnesses are built around Llama and Qwen conventions: a uniform ForCausalLM class across all text model sizes, two block-level RMSNorms per decoder layer (input_layernorm and post_attention_layernorm), a single rotary_emb module, 64- or 128-dimensional attention heads (head_dim <= 128), and chat templates with a dedicated system turn that append <eos> tokens automatically.

Gemma 3 diverges from those conventions in four ways that cause generic RL pipelines to either crash during worker initialization or suffer silent policy collapse on Step 1. Three of these quirks are handled automatically in NVIDIA NeMo RL (v0.5.0) source code, while the fourth requires explicit YAML configuration in your experiment manifest:

Gemma 3 architectural quirk	Failure mode in generic RL harnesses	How NVIDIA NeMo RL (v0.5.0) and this blueprint handle it
1. Non-persistent vLLM normalizer buffer during weight refits	In Hugging Face transformers, Gemma3TextScaledWordEmbedding stores \sqrt{\text{hidden\_size}} as a Python float (embed_scale), whereas vLLM’s gemma3.py registers normalizer as a torch.Tensor buffer (persistent=False). When vLLM initializes with load_format="dummy", its dummy weight loader overwrites normalizer with dummy values; because normalizer does not exist in the Hugging Face .safetensors state dict, FSDP-to-vLLM collective weight refits never overwrite it, corrupting rollout token embeddings on Step 1.	Automatic (ModelFlag.VLLM_LOAD_FORMAT_AUTO): nemo_rl/models/huggingface/common.py and vllm_worker.py detect gemma2, gemma3, and gemma3_text architectures and force load_format="auto" so normalizer is initialized from real model config/weights at engine startup.
2. Split 1B (Gemma3ForCausalLM) vs. 4B/12B/27B (Gemma3ForConditionalGeneration) architecture, 4 block-level RMSNorms + QK-norm, and dual RoPE	Gemma 3 1B is text-only (Gemma3ForCausalLM), whereas 4B, 12B, and 27B wrap the language model inside Gemma3ForConditionalGeneration (AutoModelForImageTextToText) alongside a SiglipVisionEmbeddings tower. Every decoder layer uses four block-level Gemma3RMSNorm layers (input, post_attention, pre_feedforward, and post_feedforward) plus q_norm/k_norm with (1 + weight) unit-offset scaling, and dual global (\theta = 10^6) and local sliding-window (\theta = 10^4) RoPE modules (rotary_emb and rotary_emb_local). Generic FSDP auto-wrappers fail on unactivated vision blocks, miss the extra norms, or crash vLLM when skip_tokenizer_init=True.	Automatic (_parallelize_gemma3 and AUTOMODEL_FACTORY): nemo_rl/models/dtensor/parallelize.py provides a dedicated Gemma 3 parallelization plan across both Gemma3ForCausalLM and Gemma3ForConditionalGeneration (including RotaryEmbedParallel on both RoPE modules and all four block-level layernorms), while vllm_worker.py forces skip_tokenizer_init=False for Gemma3ForConditionalGeneration.
3. Strict alternating user/model Jinja chat template without a native system turn or trailing <eos>	Gemma 3’s Jinja template has no <start_of_turn>system token (prepending messages[0]['role'] == 'system' into the first user turn), raises TemplateError: Conversation roles must alternate user/assistant/user/assistant/... on consecutive user turns or mid-conversation system roles (crashing harnesses that probe system prompts with synthetic consecutive user messages), and ends turns with <end_of_turn>\n (id=106) rather than tokenizer.eos_token <eos> (id=1).	Automatic (llm_message_utils.py and examples/prompts/cot.txt): nemo_rl/data/llm_message_utils.py formats prefix slices without synthetic consecutive-user probing and explicitly checks not stripped_message_chunk.endswith(tokenizer.eos_token) to handle missing <eos> tokens on Gemma templates, while cot.txt embeds the chain-of-thought instruction (\boxed{}) directly inside the user turn.
4. Decoupled 256-dimensional attention heads (1B/4B/12B head_dim=256) and non-null sliding-window attn_bias	Gemma 3 1B, 4B, and 12B explicitly set head_dim = 256 (even on 1B, where hidden_size=1152 and num_attention_heads=4; 27B uses head_dim=128), and all sizes use 5:1 local-to-global sliding-window attention masks where attn_bias is non-null. On 1B/4B/12B, vLLM’s default 16-token KV-cache block size triggers a FlashInfer assertion (block_size 16 head size 256), while flash_attention_2 sequence packing or PyTorch Context Parallelism fails on non-null attn_bias.	Explicit YAML override (vllm_kwargs.block_size: 32 and dynamic_batching: true): Because vllm_worker.py selectively reconstructs vllm_cfg keys, pass block_size: 32 inside policy.generation.vllm_kwargs and pair dynamic_batching.enabled: true with sequence_packing.enabled: false.

To apply these settings in nvcr.io/nvidia/nemo-rl:v0.5.0, create the experiment override file grpo-gemma3-1b-it-1n8g-fsdp2tp1-b200.yaml (extending NeMo RL’s built-in examples/configs/grpo_math_1b.yaml base config):

# SPDX-License-Identifier: Apache-2.0
defaults: /opt/nemo-rl/examples/configs/grpo_math_1b.yaml

grpo:
  num_prompts_per_step: 32
  num_generations_per_prompt: 16
  max_num_steps: 50
  val_period: 10
  val_at_start: true
  val_at_end: true
  async_grpo:
    enabled: false

loss_fn:
  reference_policy_kl_penalty: 0.01
  reference_policy_kl_type: "k3"

checkpointing:
  enabled: true
  checkpoint_dir: /data/nemo-rl-results/grpo-gemma3-1b-it-1n8g-fsdp2tp1-b200
  metric_name: "val:accuracy"
  higher_is_better: true
  model_save_format: "torch_save"
  save_consolidated: false
  save_period: 10
  keep_top_k: 5

policy:
  model_name: google/gemma-3-1b-it
  precision: "bfloat16"
  train_global_batch_size: 512
  train_micro_batch_size: 4
  max_total_sequence_length: 2048
  dtensor_cfg:
    _v2: true
    enabled: true
    tensor_parallel_size: 1
  megatron_cfg:
    enabled: false
  tokenizer:
    name: google/gemma-3-1b-it
    chat_template_kwargs:
      enable_thinking: true
  dynamic_batching:
    enabled: true
  sequence_packing:
    enabled: false
  make_sequence_length_divisible_by: 1
  generation:
    backend: "vllm"
    max_new_tokens: 1024
    colocated:
      enabled: true
    vllm_cfg:
      async_engine: false
      max_model_len: 2048
      enforce_eager: true
      load_format: "auto"
      gpu_memory_utilization: 0.6
    vllm_kwargs:
      block_size: 32

data:
  max_input_seq_length: 1024
  dataset_name: OpenMathInstruct-2
  split_validation_size: 0.05
  prompt_file: "examples/prompts/cot.txt"
  processor: "math_hf_data_processor"
  env_name: "math"

env:
  math:
    num_workers: 8
    math_verify_impl: "hf_math_verify"

logger:
  log_dir: /data/nemo-rl-logs/grpo-gemma3-1b-it-1n8g-fsdp2tp1-b200
  tensorboard_enabled: true
  wandb_enabled: false

cluster:
  gpus_per_node: 8
  num_nodes: 1

Submit the GRPO job without a local Ray installation

Because nvcr.io/nvidia/nemo-rl:v0.5.0 already bundles the compiled virtual environment at /opt/nemo_rl_venv, you do not need a matching Python or Ray environment on your workstation. Copy the configuration into the head pod and invoke ray job submit directly inside the cluster:

# SPDX-License-Identifier: Apache-2.0
RAY_HEAD_POD=$(kubectl get pods \
  -l ray.io/cluster=ray-cluster-b200-nemo,ray.io/node-type=head \
  -o jsonpath='{.items[0].metadata.name}')

kubectl cp grpo-gemma3-1b-it-1n8g-fsdp2tp1-b200.yaml \
  "${RAY_HEAD_POD}:/workspace/grpo-gemma3-1b-it-1n8g-fsdp2tp1-b200.yaml"

kubectl exec "${RAY_HEAD_POD}" -c ray-head -- bash -c "
  export PATH=/tmp/ray/packages/bin:\$PATH
  ray job submit --working-dir /workspace --no-wait -- bash -c '
    export HF_TOKEN=${HF_TOKEN}
    export HF_HOME=/data/huggingface
    cd /opt/nemo-rl
    /opt/nemo_rl_venv/bin/python examples/run_grpo_math.py \
      --config /workspace/grpo-gemma3-1b-it-1n8g-fsdp2tp1-b200.yaml
  '
"

Replace the following:

• HF_TOKEN: a Hugging Face access token with read permission for the gated google/gemma-3-1b-it repository.

Decouple PyTorch DCP checkpointing from Hugging Face export and evaluation

Distributed FSDP2 training and standalone Hugging Face / vLLM serving optimize for opposite I/O patterns. During training, writing native sharded PyTorch Distributed Checkpoints (model_save_format: "torch_save", save_consolidated: false) flushes per-rank shards in parallel without blocking accelerator compute or risking unconcatenated vocabulary shards ([32768, 2048] per rank versus Gemma 3’s full 262144 vocabulary size).

When a run completes or saves a milestone step (such as step_30), convert the sharded DCP directory into a consolidated Hugging Face .safetensors directory on CPU using NeMo RL’s built-in examples/converters/convert_dcp_to_hf.py utility, and stage the merged weights onto each worker node’s Local SSD (/tmp/hf_merged_model_eval) before launching vLLM evaluation. Staging onto node-local NVMe avoids mmap page-fault stalls over network FUSE mounts when eight vLLM workers map model weights concurrently:

# SPDX-License-Identifier: Apache-2.0
STEP_DIR="/data/nemo-rl-results/grpo-gemma3-1b-it-1n8g-fsdp2tp1-b200/step_30"
HF_OUT_DIR="${STEP_DIR}/hf_consolidated"
LOCAL_EVAL_DIR="/tmp/hf_merged_model_eval"

# 1. Convert sharded PyTorch DCP weights to consolidated Hugging Face safetensors on CPU
kubectl exec "${RAY_HEAD_POD}" -c ray-head -- bash -c "
  cd /opt/nemo-rl
  /opt/nemo_rl_venv/bin/python examples/converters/convert_dcp_to_hf.py \
    --config ${STEP_DIR}/config.yaml \
    --dcp-ckpt-path ${STEP_DIR}/policy/weights/ \
    --hf-ckpt-path ${HF_OUT_DIR}
"

# 2. Stage consolidated weights onto each GPU worker's Local SSD for mmap-safe vLLM loading
for POD in $(kubectl get pods -l ray.io/cluster=ray-cluster-b200-nemo,ray.io/node-type=worker -o jsonpath='{.items[*].metadata.name}'); do
  kubectl exec "${POD}" -c ray-worker -- bash -c "
    mkdir -p ${LOCAL_EVAL_DIR} && cp -r ${HF_OUT_DIR}/* ${LOCAL_EVAL_DIR}/
  "
done

# 3. Run deterministic zero-shot Pass@1 evaluation on MATH-500
kubectl exec "${RAY_HEAD_POD}" -c ray-head -- bash -c "
  export PATH=/tmp/ray/packages/bin:\$PATH
  ray job submit --working-dir /workspace -- bash -c '
    export HF_HOME=/data/huggingface
    cd /opt/nemo-rl
    /opt/nemo_rl_venv/bin/python examples/run_eval.py \
      --config examples/configs/evals/math_eval.yaml \
      generation.model_name=${LOCAL_EVAL_DIR} \
      generation.temperature=0.0 \
      generation.vllm_cfg.max_model_len=8192 \
      generation.vllm_kwargs.block_size=32 \
      data.dataset_name=math500 \
      eval.num_tests_per_prompt=1 \
      cluster.num_nodes=1 \
      cluster.gpus_per_node=8
  '
"

(Note: During multi-GPU vLLM startup, PyTorch’s Inductor graph cache can log transient UnpicklingError or CompiledFxGraph warnings if multiple workers populate the default /root/.cache directory simultaneously. PyTorch catches those cache misses automatically and compiles in memory without affecting accuracy or throughput.)

Analyze training dynamics: the Step-30 KL rebound and MATH-500 gains

Streaming /data/nemo-rl-logs directly through Cloud Storage FUSE lets you attach TensorBoard (port 6006) to the head pod and inspect reward, generation length, and KL-divergence dynamics in real time. Figure 3 captures five scalar curves across a 50-step GRPO run on Gemma-3-1b-it.

Figure 3. TensorBoard metrics (validation/accuracy, train/kl_penalty, train/mean_gen_tokens, train/reward, and validation/avg_length) showing the exploration peak at Step 30 and the subsequent KL penalty rebound.

Understand the Step-30 KL rebound

Inspecting train/kl_penalty alongside train/mean_gen_tokens and validation/accuracy across the 256-sample in-loop validation split highlights why periodic checkpointing by validation metric (metric_name: "val:accuracy") matters in RL post-training:

Training step	In-loop validation accuracy (256 samples)	Mean validation length	Validation wall time (8x B200)	Observed policy regime
Step 10	30.47% (78 / 256)	722.7 tokens	30.49s	Early exploration with verbose Markdown scaffolding (kl_penalty < 0.005)
Step 20	34.38% (88 / 256)	675.8 tokens	29.71s	Sheds rigid list headers and explores algebraic variable substitutions (y = 1/x)
Step 30 (Peak)	42.58% (109 / 256)	584.3 tokens	30.41s	Concise algebraic reasoning peak (+12.11 pp over Step 10) right before KL acceleration
Step 40	39.06% (100 / 256)	506.1 tokens	~30.0s	Rising k3 KL penalty (> 0.06) triggers sharp sequence-length contraction and hesitation
Step 50	33.59% (86 / 256)	658.8 tokens	29.61s	Post-rebound length recovery (kl_penalty ~ 0.095), prior to long-horizon stabilization

Inspect qualitative reasoning evolution from console traces

Because math_verify_impl: "hf_math_verify" evaluates symbolic equivalence inside \boxed{} through SymPy without penalizing omitted <think>...</think> XML tags, Gemma-3-1b-it develops structured derivations organically. Comparing the raw evaluation console traces at Step 10 versus Step 30 on the same factoring prompt (Factor \frac{1}{x^2} + \frac{2}{x} + 1, Figure 4) shows the policy progressing from wordy four-bullet prose (722.7 tokens average) to a tight, common-denominator derivation (\frac{1+2x+x^2}{x^2} = \frac{(x+1)^2}{x^2}, 584.3 tokens average) while jumping +12.11 percentage points in validation accuracy.

Figure 4. In-loop validation console traces comparing Step 10 (30.47% accuracy, 722.7 tokens) and Step 30 (42.58% accuracy, 584.3 tokens) on the same algebraic factoring problem.

Examining Step 40 and Step 50 (Figure 5) reveals the exact mechanics of the KL rebound and a subtle symbolic-verifier edge case. At Step 40, rising KL divergence penalties compress average response length down to 506.1 tokens and cause the policy to loop across repeated rewrites before emitting the unfactored input expression \boxed{\frac{1}{x^2} + \frac{2}{x} + 1}. Because hf_math_verify tests algebraic equivalence in SymPy ((x+1)^2 / x^2 == 1/x^2 + 2/x + 1), that unfactored identity still receives Reward: 1.0000 on factoring prompts even as overall validation accuracy drops on non-identity problems. For tasks where structural form matters (such as factoring or simplification), pair SymPy equivalence with a lightweight structural form check in NeMo RL’s reward registry.

Figure 5. In-loop validation console traces at Step 40 (39.06% accuracy, 506.1 tokens – showing KL-induced rewrite hesitation and SymPy identity equivalence) and Step 50 (33.59% accuracy, 658.8 tokens).

Zero-shot MATH-500 benchmark results

Evaluating the consolidated step_30 checkpoint against the base google/gemma-3-1b-it model on the held-out 500-problem MATH-500 benchmark (temperature=0.0, max_new_tokens=8192, seed=42) confirms the generalization gain from 30 on-policy GRPO steps:

Model checkpoint	Benchmark	Decoding configuration	Zero-shot Pass@1	Gain over baseline
Vanilla google/gemma-3-1b-it	MATH-500	Greedy (temperature=0.0, max_model_len=8192)	40.20% (201 / 500)	Baseline
GRPO post-trained Gemma-3-1b-it (step_30)	MATH-500	Greedy (temperature=0.0, max_model_len=8192)	46.40% (232 / 500)	+6.20 percentage points (+15.4% relative)

Scale beyond 8B and next steps on GKE

Once this single-node 8-GPU baseline is converging on your domain verifier, you can extend the same GKE and KubeRay foundation along three axes:

• Multi-node 27B+ post-training over 8-NIC RDMA: For larger models such as Gemma-3-27B where parameter, optimizer, and KV-cache footprints span multiple 8-GPU B200 or H200 nodes (workerGroupSpecs.replicas >= 2), switch from 1D FSDP2 to Megatron-Core 3D parallelism (megatron_cfg.enabled: true) over the attached rdma-0..rdma-7 GPUDirect-RDMA interfaces.
• Custom and neural reward verifiers: For tasks without closed-form symbolic answers (such as code refactoring or multi-turn tool use), register custom Python verifiers in NeMo RL’s environment registry or host a dedicated LLM-as-a-Judge / Process Reward Model (PRM) on an independent KubeRay vLLM actor group.
• Test-time compute scaling (Best-of-N): At evaluation and serving time, pair the GRPO-aligned checkpoint with parallel Best-of-N sampling in vLLM (temperature > 0) and majority-vote aggregation over \boxed{} candidates to trade cheap batch tokens on compact Gemma 3 models for substantially higher pass rates.

To deploy and adapt this architecture on your own GKE cluster, explore the official documentation:

• Compute and scheduling: Custom ComputeClasses on GKE and Dynamic Workload Scheduler on GKE
• Networking and storage: GPUDirect-RDMA network profiles on GKE and Cloud Storage FUSE CSI driver on GKE
• Orchestration and RL framework: KubeRay on GKE overview, AI and ML infrastructure on GKE, NVIDIA NeMo RL documentation, and the NVIDIA GTC 2026 session on scaling RL pipelines on GKE (S82359)