Update NemoRL on GKE scripts. (#2070)

Author: yuyantingzero

ai-ml/nemo-rl-on-gke/nemoRL/gemma3-27b-it/gemma3-27b-gsm8k.sh

@@ -46,12 +46,12 @@ export WANDB_API_KEY=$WANDB_API_KEY
 export HF_TOKEN=$HF_TOKEN
 export HF_HOME=/opt/nemo-rl/
 
-###-----Example to launch Gemma3-27B on 3 nodes (24 GPUs)----------
+###-----Example to launch Gemma3-27B on 2 nodes (16 GPUs)----------
 uv run python examples/run_grpo_math.py \
   --config examples/configs/recipes/llm/grpo-gemma3-27b-it-8n4g-fsdp2tp4-actckpt-long.yaml \
   cluster.num_nodes=2 \
   cluster.gpus_per_node=8 \
-  grpo.max_num_steps=300 \
+  grpo.max_num_steps=10 \
   checkpointing.checkpoint_dir=/data/nemo_rl_gemma3_27b_3_17 \
   data.dataset_name=ResponseDataset \
   +data.train_data_path=openai/gsm8k \
@@ -65,7 +65,7 @@ uv run python examples/run_grpo_math.py \
   logger.wandb_enabled=True \
   logger.wandb.name='nemo_rl_gemma3_27b_3_17' \
   grpo.num_prompts_per_step=16 \
-  grpo.num_generations_per_prompt=64 \
+  grpo.num_generations_per_prompt=32 \
   policy.generation.colocated.enabled=False \
   policy.generation.colocated.resources.num_nodes=1 \
   policy.generation.colocated.resources.gpus_per_node=8 \

ai-ml/nemo-rl-on-gke/nemoRL/lustre/lustre-pv.yaml

@@ -19,7 +19,7 @@ metadata:
 spec:
   storageClassName: lustre-rwx-500mbps-per-tib
   capacity:
-    storage: 9000Gi
+    storage: 18000Gi
   accessModes:
     - ReadWriteMany
   persistentVolumeReclaimPolicy: Retain
@@ -29,7 +29,7 @@ spec:
     name: lustre-pvc
   csi:
     driver: lustre.csi.storage.gke.io
-    volumeHandle: "northam-ce-mlai-tpu/asia-northeast1-b/pmotgi-lustre"
+    volumeHandle: "${PROJECT_ID}/${NODE_ZONE}/${LUSTRE_NAME}"
     volumeAttributes:
-      ip: 10.247.224.2
-      filesystem: lustrefs
+      ip: ${LUSTRE_IP}
+      filesystem: lustrefs

ai-ml/nemo-rl-on-gke/nemoRL/lustre/lustre-pvc.yaml

@@ -23,4 +23,4 @@ spec:
   volumeName: lustre-pv
   resources:
     requests:
-      storage: 1000Gi
+      storage: 18000Gi

ai-ml/nemo-rl-on-gke/nemoRL/templates/fluent-bit-config.yaml.j2

@@ -1,4 +1,4 @@
-{# Copyright 2026 Google LLC
+# Copyright 2026 Google LLC
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -10,7 +10,7 @@
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
-# limitations under the License. #}
+# limitations under the License. #
 
 {{- if .Values.configMap.fluentbit }}
 apiVersion: v1

ai-ml/nemo-rl-on-gke/nemoRL/templates/raycluster-cluster.yaml.j2

@@ -1,4 +1,4 @@
-{# Copyright 2026 Google LLC
+# Copyright 2026 Google LLC
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -10,7 +10,7 @@
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
-# limitations under the License. #}
+# limitations under the License. #
 
 apiVersion: ray.io/v1
 kind: RayCluster

ai-ml/nemo-rl-on-gke/nemoRL/values.yaml

@@ -60,6 +60,8 @@ head:
   containerEnv:
   - name: RAY_GROUP
     value: "head"
+  nodeSelector:
+    cloud.google.com/gke-nodepool: default-pool
   resources:
     limits:
       cpu: "64"