Skip to content

Commit 262ad18

Browse files
tdenetdene
committed
Address reviewer comments
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
1 parent ead709a commit 262ad18

8 files changed

Lines changed: 110 additions & 20 deletions

File tree

‎examples/configs/grpo_math_1B.yaml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -284,13 +284,14 @@ policy:
284284
stop_token_ids: null
285285
stop_strings: null
286286
mcore_generation_config:
287-
buffer_size_gb: 10 # Total buffer size is 2x this value (active requests + paused requests)
287+
buffer_size_gb: 10
288288
num_cuda_graphs: 4 # Number of CUDA graphs to pre-compile for different batch sizes
289289
block_size_tokens: 256 # Size of each KV cache block in tokens (affects memory granularity)
290290
use_cuda_graphs_for_non_decode_steps: true # Enable CUDA graphs for prefill/context processing
291291
enable_chunked_prefill: true # Split long prefills into chunks for better memory management
292292
unified_memory_level: 0 # Unified memory usage level (0=disabled, higher values enable more aggressive paging)
293293
max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens
294+
refit_backend: "nvshmem" # Copy-service backend for non-colocated megatron weight refit. Options: "gloo" or "nvshmem".
294295
vllm_cfg:
295296
async_engine: false
296297
precision: ${policy.precision}

‎examples/configs/grpo_math_1B_megatron.yaml‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -104,14 +104,10 @@ policy:
104104
moe_enable_deepep: false
105105
moe_token_dispatcher_type: "alltoall"
106106
moe_shared_expert_overlap: false
107-
#gives ~20% training perf speedup with sequence packing
107+
#gives ~20% training perf speedup with sequence packing
108108
apply_rope_fusion: True
109-
moe_pad_experts_for_cuda_graph_inference: false
110-
cuda_graph_impl: "local"
111-
cuda_graph_scope: "full_iteration_inference"
112109
use_te_rng_tracker: true
113110
inference_rng_tracker: true
114-
batch_invariant_mode: false
115111

116112
peft:
117113
enabled: false
@@ -150,7 +146,6 @@ policy:
150146
clip_grad: ${policy.max_grad_norm}
151147

152148
scheduler:
153-
override_opt_param_scheduler: true
154149
start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
155150
end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
156151
weight_decay_incr_style: "constant"
@@ -173,6 +168,8 @@ policy:
173168
top_p: 1.0
174169
top_k: null
175170
mcore_generation_config:
171+
cuda_graph_impl: "local"
172+
cuda_graph_scope: "full_iteration_inference"
176173
buffer_size_gb: 10 # Total buffer size is 2x this value (active requests + paused requests)
177174
num_cuda_graphs: 4 # Number of CUDA graphs to pre-compile for different batch sizes
178175
block_size_tokens: 256 # Size of each KV cache block in tokens (affects memory granularity)
@@ -182,7 +179,7 @@ policy:
182179
enable_chunked_prefill: false
183180
kv_cache_management_mode: "persist" # Can be "persist", "offload", or "recompute"
184181
static_kv_memory_pointers: false # Relevant only for offload and recompute modes
185-
materialize_only_last_token_logits: false
182+
materialize_only_last_token_logits: true
186183
async_engine: false
187184

188185
vllm_cfg:

‎examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation.yaml‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
defaults: ../../grpo_math_1B.yaml
2+
grpo:
3+
num_prompts_per_step: 2
4+
num_generations_per_prompt: 8
5+
checkpointing:
6+
checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_generation
7+
policy:
8+
model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16
9+
tokenizer:
10+
name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
11+
train_global_batch_size: 16
12+
train_micro_batch_size: 1
13+
logprob_batch_size: 1
14+
max_total_sequence_length: 2048
15+
dtensor_cfg:
16+
enabled: false
17+
megatron_cfg:
18+
enabled: true
19+
bias_activation_fusion: false
20+
tensor_model_parallel_size: 2
21+
expert_model_parallel_size: 8
22+
sequence_parallel: true
23+
sequence_packing:
24+
enabled: false
25+
generation:
26+
backend: megatron
27+
mcore_generation_config:
28+
transformer_impl: "inference_optimized"
29+
cuda_graph_impl: "local"
30+
cuda_graph_scope: "full_iteration_inference"
31+
activation_checkpointing: false
32+
tensor_model_parallel_size: 2
33+
expert_model_parallel_size: 8
34+
sequence_parallel: true
35+
async_engine: true
36+
colocated:
37+
enabled: false
38+
resources:
39+
gpus_per_node: 8
40+
num_nodes: 1
41+
logger:
42+
wandb_enabled: true
43+
tensorboard_enabled: true
44+
wandb:
45+
project: nemo-rl
46+
name: grpo-nanov3-30BA3B-2n8g-megatron_generation
47+
cluster:
48+
gpus_per_node: 8
49+
num_nodes: 2

examples/nemo_gym/grpo_nanov3_megatron.yaml renamed to examples/nemo_gym/grpo_nanov3_megatron_generation.yaml

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,24 +7,20 @@ grpo:
77

88
policy:
99
megatron_cfg:
10-
distributed_data_parallel_config:
11-
overlap_grad_reduce: false
12-
converter_type: "LlamaForCausalLM"
13-
cuda_graph_impl: "local"
14-
cuda_graph_scope: "full_iteration"
1510
use_te_rng_tracker: true
1611
batch_invariant_mode: false
17-
# Needed to set this to false because it causes "full" recompute_granularity
18-
# which isn't compatible with "full_iteration_inference" CUDA graph scope.
19-
activation_checkpointing: false
2012

2113
generation:
2214
backend: "megatron"
2315

2416
mcore_generation_config:
17+
# Scoped to generation workers only; training keeps its own defaults
18+
# via the base yaml. The mcore_generation_config block is merged into
19+
# megatron_cfg for inference workers before model setup runs.
2520
transformer_impl: "inference_optimized"
21+
cuda_graph_impl: "local"
2622
cuda_graph_scope: "full_iteration_inference"
27-
moe_pad_experts_for_cuda_graph_inference: false
23+
activation_checkpointing: false
2824
mamba_inference_ssm_states_dtype: "float32"
2925
inference_moe_token_dispatcher_type: "nccl" # Fall back to NCCL for now
3026
inference_grouped_gemm_backend: "vllm"

‎nemo_rl/algorithms/grpo.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -787,9 +787,9 @@ def initialize_generation_with_policy(
787787
if backend == "megatron":
788788
# Non-colocated Megatron: use init_refit_collective with the
789789
# configured refit backend (gloo or nvshmem).
790-
refit_backend = policy_config.get("megatron_cfg", {}).get(
791-
"refit_backend", "gloo"
792-
)
790+
refit_backend = policy_config["generation"]["mcore_generation_config"][
791+
"refit_backend"
792+
]
793793
futures_train = policy.init_refit_collective(
794794
ip, port, world_size,
795795
rank_offset=0,

‎nemo_rl/models/megatron/config.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,9 @@ class MegatronGenerationConfig(TypedDict):
4141
# Can cause OOM if set too high so should be tuned with buffer_size_gb if OOMing. If set too
4242
# low, then will only do 512 tokens at a time, which can be slow.
4343
max_tokens: int
44+
# Copy-service backend used to transfer weights from training to inference
45+
# workers via the refit collective. Options: "gloo" or "nvshmem".
46+
refit_backend: str
4447

4548

4649
## returned from validate_and_set_config

‎tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation.sh‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
#!/bin/bash
2+
SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
3+
source $SCRIPT_DIR/common.env
4+
5+
# ===== BEGIN CONFIG =====
6+
NUM_NODES=2
7+
GPUS_PER_NODE=8
8+
STEPS_PER_RUN=10
9+
MAX_STEPS=10
10+
NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up
11+
NUM_MINUTES=60
12+
# ===== END CONFIG =====
13+
14+
exit_if_max_steps_reached
15+
16+
# Run the experiment
17+
cd $PROJECT_ROOT
18+
uv run examples/run_grpo.py \
19+
--config $CONFIG_PATH \
20+
grpo.max_num_steps=$MAX_STEPS \
21+
policy.generation.backend=megatron \
22+
logger.log_dir=$LOG_DIR \
23+
logger.wandb_enabled=True \
24+
logger.wandb.project=nemo-rl \
25+
logger.wandb.name=$EXP_NAME \
26+
logger.monitor_gpus=True \
27+
logger.tensorboard_enabled=True \
28+
checkpointing.enabled=True \
29+
checkpointing.checkpoint_dir=$CKPT_DIR \
30+
$@ \
31+
2>&1 | tee $RUN_LOG
32+
33+
# Convert tensorboard logs to json
34+
uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
35+
36+
# Only run metrics if the target step is reached
37+
if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then
38+
uv run tests/check_metrics.py $JSON_METRICS \
39+
'max(data["train/reward"]) > 0.0'
40+
41+
# Clean up checkpoint directory after successful run to save space.
42+
rm -rf "$CKPT_DIR"
43+
fi

‎tests/test_suites/nightly.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@ tests/test_suites/llm/grpo-nano-v2-12b-2n8g-fsdp2tp1.sh
8484
tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-fsdp2.sh
8585
tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-fsdp2-lora.sh
8686
tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-lora.sh
87+
tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation.sh
8788

8889
# Lora
8990
tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora.sh

0 commit comments

Comments
 (0)