From e263f489c121c93adee40731b519837df7a9c7c6 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Fri, 4 Sep 2026 23:13:55 -0700 Subject: [PATCH 01/17] test: add Nano 4B Gym training E2E Signed-off-by: yaoyu-33 --- .../L1_Functional_Tests_Gym_Training.sh | 28 ++++ .../grpo_nano4b_gym_training_e2e.sh | 58 +++++++++ .../grpo_nano4b_gym_training_e2e.yaml | 122 ++++++++++++++++++ tests/functional/nano4b_math_e2e.jsonl | 2 + 4 files changed, 210 insertions(+) create mode 100755 tests/functional/L1_Functional_Tests_Gym_Training.sh create mode 100755 tests/functional/grpo_nano4b_gym_training_e2e.sh create mode 100644 tests/functional/grpo_nano4b_gym_training_e2e.yaml create mode 100644 tests/functional/nano4b_math_e2e.jsonl diff --git a/tests/functional/L1_Functional_Tests_Gym_Training.sh b/tests/functional/L1_Functional_Tests_Gym_Training.sh new file mode 100755 index 00000000000..c984ee80d0b --- /dev/null +++ b/tests/functional/L1_Functional_Tests_Gym_Training.sh @@ -0,0 +1,28 @@ +#!/bin/bash + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# 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. + +set -euo pipefail + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") + +cd "${PROJECT_ROOT}" + +if [[ "${FAST:-0}" == "1" ]]; then + echo "FAST: Skipping Nano 4B Gym training E2E" +else + time uv run --no-sync bash ./tests/functional/grpo_nano4b_gym_training_e2e.sh +fi diff --git a/tests/functional/grpo_nano4b_gym_training_e2e.sh b/tests/functional/grpo_nano4b_gym_training_e2e.sh new file mode 100755 index 00000000000..21ea6ae3f5f --- /dev/null +++ b/tests/functional/grpo_nano4b_gym_training_e2e.sh @@ -0,0 +1,58 @@ +#!/bin/bash + +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# 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. + +set -euo pipefail + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") +EXP_NAME=$(basename "$0" .sh) +EXP_DIR="${SCRIPT_DIR}/${EXP_NAME}" +LOG_DIR="${EXP_DIR}/logs" +JSON_METRICS="${EXP_DIR}/metrics.json" +RUN_LOG="${EXP_DIR}/run.log" +CONFIG_PATH="${SCRIPT_DIR}/grpo_nano4b_gym_training_e2e.yaml" + +export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}" + +rm -rf "${EXP_DIR}" +mkdir -p "${LOG_DIR}" + +cd "${PROJECT_ROOT}" + +uv run coverage run -a --data-file="${PROJECT_ROOT}/tests/.coverage" --source="${PROJECT_ROOT}/nemo_rl" \ + "${PROJECT_ROOT}/examples/nemo_gym/run_grpo_nemo_gym.py" \ + --config "${CONFIG_PATH}" \ + logger.log_dir="${LOG_DIR}" \ + "$@" \ + 2>&1 | tee "${RUN_LOG}" + +grep -Fq "Running synchronous GRPO training" "${RUN_LOG}" + +uv run tests/json_dump_tb_logs.py "${LOG_DIR}" --output_path "${JSON_METRICS}" + +# The fixture intentionally contains one accepted and one rejected rollout. In +# addition to testing both verifier outcomes, this gives Reinforce++ a non-zero +# advantage so grad_norm proves that an optimizer step was actually exercised. +uv run tests/check_metrics.py "${JSON_METRICS}" \ + 'len(data["train/loss"]) == 1' \ + 'all_finite(data["train/loss"])' \ + 'all_finite(data["train/grad_norm"])' \ + 'min(data["train/grad_norm"]) > 0' \ + 'data["train/total_reward/min"]["1"] == 0' \ + 'data["train/total_reward/max"]["1"] == 1' \ + 'data["train/total_reward/mean"]["1"] == 0.5' \ + 'all_finite(data["train/token_mult_prob_error"])' \ + 'data["timing/train/generation"]["1"] > 0' diff --git a/tests/functional/grpo_nano4b_gym_training_e2e.yaml b/tests/functional/grpo_nano4b_gym_training_e2e.yaml new file mode 100644 index 00000000000..e1733980933 --- /dev/null +++ b/tests/functional/grpo_nano4b_gym_training_e2e.yaml @@ -0,0 +1,122 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# 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. + +defaults: ../../examples/configs/grpo_math_1B.yaml + +grpo: + num_prompts_per_step: 2 + num_generations_per_prompt: 1 + max_num_steps: 1 + normalize_rewards: false + use_leave_one_out_baseline: false + val_period: -1 + val_at_start: false + val_at_end: false + overlong_filtering: false + skip_reference_policy_logprobs_calculation: true + adv_estimator: + name: reinforce_plus_plus + normalize_rewards: false + use_leave_one_out_baseline: false + minus_baseline: false + +loss_fn: + reference_policy_kl_penalty: 0.0 + use_importance_sampling_correction: false + use_kl_in_reward: false + +checkpointing: + enabled: false + +policy: + model_name: nvidia/Llama-3.1-Nemotron-Nano-4B-v1.1 + tokenizer: + name: ${policy.model_name} + train_global_batch_size: 2 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 512 + dtensor_cfg: + enabled: true + activation_checkpointing: true + tensor_parallel_size: 1 + megatron_cfg: + enabled: false + dynamic_batching: + enabled: false + sequence_packing: + enabled: false + make_sequence_length_divisible_by: 1 + generation: + max_new_tokens: 64 + temperature: 0.0 + top_p: 1.0 + vllm_cfg: + async_engine: true + tensor_parallel_size: 2 + max_model_len: ${policy.max_total_sequence_length} + gpu_memory_utilization: 0.7 + enforce_eager: true + expose_http_server: true + colocated: + enabled: false + resources: + gpus_per_node: 2 + num_nodes: 1 + +data: + max_input_seq_length: null + shuffle: false + num_workers: 0 + train: + data_path: tests/functional/nano4b_math_e2e.jsonl + validation: null + default: + dataset_name: NemoGymDataset + env_name: nemo_gym + prompt_file: null + system_prompt_file: null + processor: nemo_gym_data_processor + +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + should_mask_flagged_samples: false + nemo_gym: + port_range_low: 5000 + port_range_high: 5999 + is_trajectory_collection: false + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/math_with_judge/configs/math_with_judge.yaml + math_with_judge: + resources_servers: + math_with_judge: + judge_model_server: + name: policy_model + should_use_judge: false + policy_model: + responses_api_models: + vllm_model: + uses_reasoning_parser: false + +logger: + log_dir: tests/functional/grpo_nano4b_gym_training_e2e/logs + wandb_enabled: false + tensorboard_enabled: true + monitor_gpus: true + +cluster: + gpus_per_node: 4 + num_nodes: 1 diff --git a/tests/functional/nano4b_math_e2e.jsonl b/tests/functional/nano4b_math_e2e.jsonl new file mode 100644 index 00000000000..5178df2262e --- /dev/null +++ b/tests/functional/nano4b_math_e2e.jsonl @@ -0,0 +1,2 @@ +{"responses_create_params":{"input":[{"role":"user","content":"Return exactly \\boxed{2}."}],"max_output_tokens":64},"question":"Return exactly \\boxed{2}.","expected_answer":"2","agent_ref":{"type":"responses_api_agents","name":"math_with_judge_simple_agent"}} +{"responses_create_params":{"input":[{"role":"user","content":"Return exactly \\boxed{12}."}],"max_output_tokens":64},"question":"Return exactly \\boxed{12}.","expected_answer":"999999","agent_ref":{"type":"responses_api_agents","name":"math_with_judge_simple_agent"}} From 933b72b1bfbe394d001b03c20c20d8591cdef479 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Fri, 4 Sep 2026 23:26:34 -0700 Subject: [PATCH 02/17] test: exercise post-update Gym rollout Signed-off-by: yaoyu-33 --- tests/functional/grpo_nano4b_gym_training_e2e.sh | 9 ++++++++- tests/functional/grpo_nano4b_gym_training_e2e.yaml | 12 ++++++++++-- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/tests/functional/grpo_nano4b_gym_training_e2e.sh b/tests/functional/grpo_nano4b_gym_training_e2e.sh index 21ea6ae3f5f..260aca7ab58 100755 --- a/tests/functional/grpo_nano4b_gym_training_e2e.sh +++ b/tests/functional/grpo_nano4b_gym_training_e2e.sh @@ -51,8 +51,15 @@ uv run tests/check_metrics.py "${JSON_METRICS}" \ 'all_finite(data["train/loss"])' \ 'all_finite(data["train/grad_norm"])' \ 'min(data["train/grad_norm"]) > 0' \ + 'all_finite(data["train/advantages/min"])' \ + 'all_finite(data["train/advantages/max"])' \ + 'min(data["train/advantages/min"]) < 0' \ + 'max(data["train/advantages/max"]) > 0' \ 'data["train/total_reward/min"]["1"] == 0' \ 'data["train/total_reward/max"]["1"] == 1' \ 'data["train/total_reward/mean"]["1"] == 0.5' \ 'all_finite(data["train/token_mult_prob_error"])' \ - 'data["timing/train/generation"]["1"] > 0' + 'max(data["train/token_mult_prob_error"]) < 1.05' \ + 'data["timing/train/generation"]["1"] > 0' \ + 'data["validation/accuracy"]["1"] == 0.5' \ + 'data["timing/validation/total_validation_time"]["1"] > 0' diff --git a/tests/functional/grpo_nano4b_gym_training_e2e.yaml b/tests/functional/grpo_nano4b_gym_training_e2e.yaml index e1733980933..8a2d9235921 100644 --- a/tests/functional/grpo_nano4b_gym_training_e2e.yaml +++ b/tests/functional/grpo_nano4b_gym_training_e2e.yaml @@ -20,9 +20,10 @@ grpo: max_num_steps: 1 normalize_rewards: false use_leave_one_out_baseline: false + max_val_samples: null val_period: -1 val_at_start: false - val_at_end: false + val_at_end: true overlong_filtering: false skip_reference_policy_logprobs_calculation: true adv_estimator: @@ -41,8 +42,12 @@ checkpointing: policy: model_name: nvidia/Llama-3.1-Nemotron-Nano-4B-v1.1 + hf_config_overrides: + revision: d552708a9d575fa8d4a690b988fd870d65279f98 tokenizer: name: ${policy.model_name} + tokenizer_kwargs: + revision: ${policy.hf_config_overrides.revision} train_global_batch_size: 2 train_micro_batch_size: 1 logprob_batch_size: 1 @@ -69,6 +74,8 @@ policy: gpu_memory_utilization: 0.7 enforce_eager: true expose_http_server: true + vllm_kwargs: + revision: ${policy.hf_config_overrides.revision} colocated: enabled: false resources: @@ -81,7 +88,8 @@ data: num_workers: 0 train: data_path: tests/functional/nano4b_math_e2e.jsonl - validation: null + validation: + data_path: tests/functional/nano4b_math_e2e.jsonl default: dataset_name: NemoGymDataset env_name: nemo_gym From 2a3edcb0c5db8ad9d07a42b65dd18eed5764b219 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Fri, 4 Sep 2026 23:39:01 -0700 Subject: [PATCH 03/17] test: allowlist pinned model revision Signed-off-by: yaoyu-33 --- tests/functional/grpo_nano4b_gym_training_e2e.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/functional/grpo_nano4b_gym_training_e2e.yaml b/tests/functional/grpo_nano4b_gym_training_e2e.yaml index 8a2d9235921..c351b72c2ff 100644 --- a/tests/functional/grpo_nano4b_gym_training_e2e.yaml +++ b/tests/functional/grpo_nano4b_gym_training_e2e.yaml @@ -43,7 +43,7 @@ checkpointing: policy: model_name: nvidia/Llama-3.1-Nemotron-Nano-4B-v1.1 hf_config_overrides: - revision: d552708a9d575fa8d4a690b988fd870d65279f98 + revision: d552708a9d575fa8d4a690b988fd870d65279f98 # pragma: allowlist secret tokenizer: name: ${policy.model_name} tokenizer_kwargs: From 585004a74540b6dc3de0c547606897238cd1be78 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Fri, 4 Sep 2026 23:43:28 -0700 Subject: [PATCH 04/17] test: prove post-update generation refit Signed-off-by: yaoyu-33 --- nemo_rl/algorithms/grpo.py | 1 + nemo_rl/models/generation/vllm/vllm_worker.py | 6 +++++- tests/functional/grpo_nano4b_gym_training_e2e.sh | 1 + 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 7df8f4db53a..20282b69774 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -3638,6 +3638,7 @@ def grpo_train( _refit_buffer_size_gb=refit_buffer_size_gb, kv_scales=kv_scales_cache if sync_kv_scales else None, ) + refit_metrics["post_update_success"] = 1.0 POLICY_GENERATION_STALE = False else: if colocated_inference: diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index ba9ad5fb837..c8240110aab 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -602,7 +602,11 @@ def _load_model(self, bundle_indices, seed): # Override HF config for gpt-oss models to ensure compatibility with megatron # The megatron --> hf export is done in bf16, so we disable quantization - hf_config = AutoConfig.from_pretrained(self.model_name, trust_remote_code=True) + hf_config = AutoConfig.from_pretrained( + self.model_name, + trust_remote_code=True, + revision=vllm_kwargs.get("revision"), + ) self.routed_experts_dtype = resolve_routed_experts_dtype( get_num_routed_experts(hf_config) ) diff --git a/tests/functional/grpo_nano4b_gym_training_e2e.sh b/tests/functional/grpo_nano4b_gym_training_e2e.sh index 260aca7ab58..27e9aaccc69 100755 --- a/tests/functional/grpo_nano4b_gym_training_e2e.sh +++ b/tests/functional/grpo_nano4b_gym_training_e2e.sh @@ -60,6 +60,7 @@ uv run tests/check_metrics.py "${JSON_METRICS}" \ 'data["train/total_reward/mean"]["1"] == 0.5' \ 'all_finite(data["train/token_mult_prob_error"])' \ 'max(data["train/token_mult_prob_error"]) < 1.05' \ + 'data["refit/post_update_success"]["1"] == 1' \ 'data["timing/train/generation"]["1"] > 0' \ 'data["validation/accuracy"]["1"] == 0.5' \ 'data["timing/validation/total_validation_time"]["1"] > 0' From d9619282a8180eb37237a1ea2da69aa50294239c Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 01:33:06 -0700 Subject: [PATCH 05/17] test: fit Nano 4B Gym E2E to H100 runner Signed-off-by: yaoyu-33 --- tests/functional/grpo_nano4b_gym_training_e2e.yaml | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/functional/grpo_nano4b_gym_training_e2e.yaml b/tests/functional/grpo_nano4b_gym_training_e2e.yaml index c351b72c2ff..c3feab0ead5 100644 --- a/tests/functional/grpo_nano4b_gym_training_e2e.yaml +++ b/tests/functional/grpo_nano4b_gym_training_e2e.yaml @@ -69,7 +69,7 @@ policy: top_p: 1.0 vllm_cfg: async_engine: true - tensor_parallel_size: 2 + tensor_parallel_size: 1 max_model_len: ${policy.max_total_sequence_length} gpu_memory_utilization: 0.7 enforce_eager: true @@ -79,7 +79,7 @@ policy: colocated: enabled: false resources: - gpus_per_node: 2 + gpus_per_node: 1 num_nodes: 1 data: @@ -87,7 +87,9 @@ data: shuffle: false num_workers: 0 train: + dataset_name: NemoGymDataset data_path: tests/functional/nano4b_math_e2e.jsonl + split_validation_size: 0 validation: data_path: tests/functional/nano4b_math_e2e.jsonl default: @@ -125,6 +127,8 @@ logger: tensorboard_enabled: true monitor_gpus: true +# The H100 functional runner exposes two GPUs: one stays with non-colocated +# vLLM above and the other runs the FSDP2 policy/optimizer. cluster: - gpus_per_node: 4 + gpus_per_node: 2 num_nodes: 1 From 669c8c785916c609c60e06a6bc693967dbf059cc Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 11:20:21 -0700 Subject: [PATCH 06/17] test: shard Nano 4B E2E training across two GPUs Signed-off-by: yaoyu-33 --- tests/functional/grpo_nano4b_gym_training_e2e.yaml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/functional/grpo_nano4b_gym_training_e2e.yaml b/tests/functional/grpo_nano4b_gym_training_e2e.yaml index c3feab0ead5..7fd4d95d8de 100644 --- a/tests/functional/grpo_nano4b_gym_training_e2e.yaml +++ b/tests/functional/grpo_nano4b_gym_training_e2e.yaml @@ -77,10 +77,10 @@ policy: vllm_kwargs: revision: ${policy.hf_config_overrides.revision} colocated: - enabled: false + enabled: true resources: - gpus_per_node: 1 - num_nodes: 1 + gpus_per_node: null + num_nodes: null data: max_input_seq_length: null @@ -127,8 +127,8 @@ logger: tensorboard_enabled: true monitor_gpus: true -# The H100 functional runner exposes two GPUs: one stays with non-colocated -# vLLM above and the other runs the FSDP2 policy/optimizer. +# The H100 functional runner exposes two GPUs. Colocation lets FSDP2 shard the +# full-parameter policy and optimizer across both while vLLM shares that cluster. cluster: gpus_per_node: 2 num_nodes: 1 From d759166c42ddfb3dfcb3f00e36767264fb5354ed Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 13:07:56 -0700 Subject: [PATCH 07/17] test: make Nano 4B training sampling finite and reproducible Signed-off-by: yaoyu-33 --- tests/functional/grpo_nano4b_gym_training_e2e.yaml | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/functional/grpo_nano4b_gym_training_e2e.yaml b/tests/functional/grpo_nano4b_gym_training_e2e.yaml index 7fd4d95d8de..8491cd5f995 100644 --- a/tests/functional/grpo_nano4b_gym_training_e2e.yaml +++ b/tests/functional/grpo_nano4b_gym_training_e2e.yaml @@ -65,7 +65,8 @@ policy: make_sequence_length_divisible_by: 1 generation: max_new_tokens: 64 - temperature: 0.0 + temperature: 1.0 + val_temperature: 0.0 top_p: 1.0 vllm_cfg: async_engine: true @@ -120,6 +121,10 @@ env: responses_api_models: vllm_model: uses_reasoning_parser: false + # Keep temperature positive for policy logprobs while making sampled + # training rollouts reproducible enough for exact reward goldens. + extra_body: + seed: 42 logger: log_dir: tests/functional/grpo_nano4b_gym_training_e2e/logs From 174f3d37a873f8e4bb282077e3c0e032a335222b Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 15:14:57 -0700 Subject: [PATCH 08/17] fix: align greedy rollout and training logprobs Signed-off-by: yaoyu-33 --- nemo_rl/algorithms/grpo.py | 30 +++++++----- nemo_rl/algorithms/logits_sampling_utils.py | 29 +++++++++++ nemo_rl/models/automodel/train.py | 18 +------ nemo_rl/models/generation/__init__.py | 10 ++++ nemo_rl/models/megatron/train.py | 18 +------ .../policy/workers/dtensor_policy_worker.py | 5 +- .../grpo_nano4b_gym_training_e2e.sh | 2 +- .../grpo_nano4b_gym_training_e2e.yaml | 6 +-- tests/unit/algorithms/test_grpo.py | 49 +++++++++++++++++-- .../algorithms/test_logits_sampling_utils.py | 19 ++++++- .../models/automodel/test_automodel_train.py | 11 +++++ .../models/generation/test_vllm_generation.py | 30 ++++++++++++ tests/unit/models/megatron/test_train.py | 13 +++++ .../unit/models/policy/test_dtensor_worker.py | 17 +++++++ 14 files changed, 199 insertions(+), 58 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 20282b69774..9d71dd78192 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -521,6 +521,17 @@ def shutdown_environments( print(f"Error stopping environment {task_name}: {kill_error}") +def _training_sampling_params_from_generation_config( + generation_config: GenerationConfig, +) -> TrainingSamplingParams: + """Build the effective policy-logprob sampling configuration.""" + return TrainingSamplingParams( + top_k=generation_config["top_k"], + top_p=generation_config["top_p"], + temperature=generation_config["temperature"], + ) + + def setup( master_config: MasterConfig, tokenizer: TokenizerType, @@ -770,10 +781,7 @@ def init_train_dataloader(dataset, suffix: str = ""): # applied. This also keeps prev/reference logprobs (computed via the fused # get_logprobs path) consistent with the actor logprobs. assert not need_top_k_or_top_p_filtering( - TrainingSamplingParams( - top_k=generation_config["top_k"], - top_p=generation_config["top_p"], - ) + _training_sampling_params_from_generation_config(generation_config) ), ( "Linear CE fusion loss is not supported with top-k/top-p training-time " "filtering for GRPO. The fused path computes logprobs from unfiltered " @@ -2532,7 +2540,8 @@ def refit_policy_generation( kv_scales: Optional dictionary of KV cache scales for FP8 quantization. Returns: - Scalar metrics reported by the selected weight synchronizer. + Scalar metrics reported by the selected weight synchronizer, or the + number of generation worker groups that acknowledged a direct refit. """ # Every SGLang deployment reaches its refit through this hook: `setup` # attaches an SGLang synchronizer that owns the whole lifecycle (phase @@ -2562,7 +2571,7 @@ def refit_policy_generation( ) with timer_context: # update weights - update_success = False + acknowledged_updates: list[bool] = [] if colocated_inference: # get model param keys, which is grouped by size if _refit_buffer_size_gb is not None: @@ -2587,7 +2596,7 @@ def refit_policy_generation( # wait for all futures to complete ray.get(futures_train) results = ray.get(futures_inference) - update_success = all(result for result in results if result is not None) + acknowledged_updates = [result is True for result in results] else: # update weights through nccl (vLLM) futures_train = policy.broadcast_weights_for_collective( @@ -2597,10 +2606,10 @@ def refit_policy_generation( # wait for all futures to complete ray.get(futures_train) results = ray.get(futures_inference) - update_success = all(result for result in results if result is not None) + acknowledged_updates = [result is True for result in results] # check if update is successful - if not update_success: + if not acknowledged_updates or not all(acknowledged_updates): error_tag = "cuda-ipc" if colocated_inference else "nccl" error_message = ( "❌ Error: Updating weights for the generation policy failed during refit.\n" @@ -2613,7 +2622,7 @@ def refit_policy_generation( policy.offload_after_refit() policy_generation.prepare_for_generation(tags=["kv_cache"]) - return {} + return {"generation_workers_updated": float(len(acknowledged_updates))} def _initial_policy_generation_stale( @@ -3638,7 +3647,6 @@ def grpo_train( _refit_buffer_size_gb=refit_buffer_size_gb, kv_scales=kv_scales_cache if sync_kv_scales else None, ) - refit_metrics["post_update_success"] = 1.0 POLICY_GENERATION_STALE = False else: if colocated_inference: diff --git a/nemo_rl/algorithms/logits_sampling_utils.py b/nemo_rl/algorithms/logits_sampling_utils.py index 2b1583080aa..418473618e6 100644 --- a/nemo_rl/algorithms/logits_sampling_utils.py +++ b/nemo_rl/algorithms/logits_sampling_utils.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import math from dataclasses import dataclass from typing import Optional @@ -43,6 +44,25 @@ class TrainingSamplingParams: top_p: float = 1.0 temperature: float = 1.0 + def __post_init__(self) -> None: + """Normalize generation parameters to vLLM's logprob semantics. + + Generation backends use ``temperature=0`` to select tokens greedily, + while the selected tokens' policy logprobs still come from the unscaled + model distribution. Top-k and top-p do not constrain greedy selection, + so training-time filtering must be disabled as well. + """ + if not math.isfinite(self.temperature): + raise ValueError(f"temperature must be finite, got {self.temperature}") + if self.temperature < 0.0: + raise ValueError( + f"temperature must be non-negative, got {self.temperature}" + ) + if self.temperature == 0.0: + self.temperature = 1.0 + self.top_k = None + self.top_p = 1.0 + def _need_top_k_filtering(top_k: int | None) -> bool: """Check if top-k filtering is needed.""" @@ -66,6 +86,15 @@ def need_top_k_or_top_p_filtering( return _need_top_k_filtering(top_k) or _need_top_p_filtering(top_p) +def apply_temperature_scaling( + logits: torch.Tensor, sampling_params: Optional[TrainingSamplingParams] +) -> torch.Tensor: + """Apply the effective training temperature to logits in place.""" + if sampling_params is not None and sampling_params.temperature != 1.0: + logits.div_(sampling_params.temperature) + return logits + + @torch.no_grad() def _apply_top_k_only_fn( logits: torch.Tensor, diff --git a/nemo_rl/models/automodel/train.py b/nemo_rl/models/automodel/train.py index f5932477e21..d65e60270ce 100644 --- a/nemo_rl/models/automodel/train.py +++ b/nemo_rl/models/automodel/train.py @@ -44,6 +44,7 @@ from nemo_rl.algorithms.logits_sampling_utils import ( TrainingSamplingParams, + apply_temperature_scaling, apply_top_k_top_p, need_top_k_or_top_p_filtering, ) @@ -229,23 +230,6 @@ def extract_logits( return outputs.logits -def apply_temperature_scaling( - logits: torch.Tensor, sampling_params: Optional[TrainingSamplingParams] -) -> torch.Tensor: - """Apply temperature scaling to logits. - - Args: - logits: Logits tensor to scale - sampling_params: Sampling parameters - - Returns: - torch.Tensor: Temperature-scaled logits - """ - if sampling_params is not None and sampling_params.temperature != 1.0: - logits.div_(sampling_params.temperature) - return logits - - def apply_top_k_top_p_filtering_for_local_logits( logits: torch.Tensor, sampling_params: Optional[TrainingSamplingParams] ) -> torch.Tensor: diff --git a/nemo_rl/models/generation/__init__.py b/nemo_rl/models/generation/__init__.py index 87c2d5de14a..dc8533a4226 100644 --- a/nemo_rl/models/generation/__init__.py +++ b/nemo_rl/models/generation/__init__.py @@ -22,6 +22,7 @@ from nemo_rl.models.generation.vllm.config import VLLM_SPARSE_REFIT_TRANSPORTS TokenizerType = PreTrainedTokenizerBase +_VLLM_MIN_NON_ZERO_TEMPERATURE = 1e-2 def resolve_generation_class( @@ -64,6 +65,15 @@ def configure_generation_config( trains_mtp: bool = False, ) -> GenerationConfig: """Apply specific configurations to generation config.""" + # vLLM clamps tiny positive temperatures before sampling. Normalize the + # shared config at this backend boundary so policy logprob recomputation + # receives the same effective temperature without changing other backends. + if ( + config["backend"] in ("vllm", "dynamo") + and 0.0 < config["temperature"] < _VLLM_MIN_NON_ZERO_TEMPERATURE + ): + config["temperature"] = _VLLM_MIN_NON_ZERO_TEMPERATURE + # tokenizer setting if "_pad_token_id" in config: warnings.warn( diff --git a/nemo_rl/models/megatron/train.py b/nemo_rl/models/megatron/train.py index 4b263383db1..2b0aa8d0710 100644 --- a/nemo_rl/models/megatron/train.py +++ b/nemo_rl/models/megatron/train.py @@ -34,6 +34,7 @@ from nemo_rl.algorithms.logits_sampling_utils import ( TrainingSamplingParams, + apply_temperature_scaling, need_top_k_or_top_p_filtering, ) from nemo_rl.algorithms.loss import ( @@ -201,23 +202,6 @@ def model_forward( return output_tensor -def apply_temperature_scaling( - logits: torch.Tensor, sampling_params: Optional[TrainingSamplingParams] -) -> torch.Tensor: - """Apply temperature scaling to logits. - - Args: - logits: Logits tensor to scale - sampling_params: Sampling parameters - - Returns: - torch.Tensor: Temperature-scaled logits - """ - if sampling_params is not None and sampling_params.temperature != 1.0: - logits.div_(sampling_params.temperature) - return logits - - def forward_with_post_processing_fn( data_iterator: Iterator[ProcessedMicrobatch], model: GPTModel, diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker.py b/nemo_rl/models/policy/workers/dtensor_policy_worker.py index 55d62ab88ea..11d11f78cfd 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker.py @@ -52,6 +52,7 @@ from nemo_rl.algorithms.logits_sampling_utils import ( TrainingSamplingParams, + apply_temperature_scaling, apply_top_k_top_p, need_top_k_or_top_p_filtering, ) @@ -571,9 +572,7 @@ def create_context_parallel_ctx( # based on https://github.com/pytorch/torchtitan/blob/cddd7dc809f36fe0ed51cdaaea0671c084d75442/torchtitan/distributed/utils.py#L178 def _apply_temperature_scaling(self, logits: torch.Tensor) -> torch.Tensor: - if self.sampling_params is not None and self.sampling_params.temperature != 1.0: - logits.div_(self.sampling_params.temperature) - return logits + return apply_temperature_scaling(logits, self.sampling_params) def _apply_top_k_top_p_filtering(self, logits: torch.Tensor) -> torch.Tensor: """Apply top-k and top-p filtering to the logits locally when TP is disabled.""" diff --git a/tests/functional/grpo_nano4b_gym_training_e2e.sh b/tests/functional/grpo_nano4b_gym_training_e2e.sh index 27e9aaccc69..dc8ae176f66 100755 --- a/tests/functional/grpo_nano4b_gym_training_e2e.sh +++ b/tests/functional/grpo_nano4b_gym_training_e2e.sh @@ -60,7 +60,7 @@ uv run tests/check_metrics.py "${JSON_METRICS}" \ 'data["train/total_reward/mean"]["1"] == 0.5' \ 'all_finite(data["train/token_mult_prob_error"])' \ 'max(data["train/token_mult_prob_error"]) < 1.05' \ - 'data["refit/post_update_success"]["1"] == 1' \ + 'data["refit/generation_workers_updated"]["1"] > 0' \ 'data["timing/train/generation"]["1"] > 0' \ 'data["validation/accuracy"]["1"] == 0.5' \ 'data["timing/validation/total_validation_time"]["1"] > 0' diff --git a/tests/functional/grpo_nano4b_gym_training_e2e.yaml b/tests/functional/grpo_nano4b_gym_training_e2e.yaml index 8491cd5f995..9afecdb7537 100644 --- a/tests/functional/grpo_nano4b_gym_training_e2e.yaml +++ b/tests/functional/grpo_nano4b_gym_training_e2e.yaml @@ -65,7 +65,7 @@ policy: make_sequence_length_divisible_by: 1 generation: max_new_tokens: 64 - temperature: 1.0 + temperature: 0.0 val_temperature: 0.0 top_p: 1.0 vllm_cfg: @@ -121,10 +121,6 @@ env: responses_api_models: vllm_model: uses_reasoning_parser: false - # Keep temperature positive for policy logprobs while making sampled - # training rollouts reproducible enough for exact reward goldens. - extra_body: - seed: 42 logger: log_dir: tests/functional/grpo_nano4b_gym_training_e2e/logs diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 901f986b65f..e76324186dc 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -17,7 +17,7 @@ from pathlib import Path from threading import Event from types import SimpleNamespace -from typing import Any +from typing import Any, cast from unittest.mock import MagicMock, patch import pytest @@ -51,6 +51,7 @@ _resolve_message_level_advantage_penalties, _save_async_replay_buffer_checkpoint, _startup_pipeline_ready, + _training_sampling_params_from_generation_config, _validate_multimodal_dedup_capability, _validate_use_kl_in_reward_compat, aggregate_rollout_metrics, @@ -91,7 +92,10 @@ from nemo_rl.experience.rollouts import calculate_rewards from nemo_rl.models.generation import configure_generation_config from nemo_rl.models.generation.dynamo import DynamoConfig -from nemo_rl.models.generation.interfaces import should_use_async_rollouts +from nemo_rl.models.generation.interfaces import ( + GenerationConfig, + should_use_async_rollouts, +) from nemo_rl.models.generation.megatron import MegatronGeneration from nemo_rl.utils.config import load_config, register_omegaconf_resolvers from nemo_rl.utils.timer import Timer @@ -113,6 +117,19 @@ def _mock_policy_generation() -> MagicMock: return policy_generation +def test_training_sampling_params_from_generation_config_normalizes_greedy() -> None: + generation_config = cast( + GenerationConfig, + {"top_k": 5, "top_p": 0.8, "temperature": 0.0}, + ) + + params = _training_sampling_params_from_generation_config(generation_config) + + assert params.temperature == 1.0 + assert params.top_k is None + assert params.top_p == 1.0 + + def test_save_async_replay_buffer_checkpoint(tmp_path): replay_buffer = MagicMock() replay_buffer.save_to_path.remote.return_value = 7 @@ -205,7 +222,7 @@ def test_refit_policy_generation_forwards_kv_scales_on_colocated_ipc( policy_generation.weight_synchronizer = None kv_scales = {"layer.0": 0.5} - refit_policy_generation( + metrics = refit_policy_generation( policy, policy_generation, colocated_inference=True, @@ -217,6 +234,32 @@ def test_refit_policy_generation_forwards_kv_scales_on_colocated_ipc( buffer_size_bytes=1024**3, kv_scales=kv_scales, ) + assert metrics == {"generation_workers_updated": 1.0} + + +@pytest.mark.parametrize("colocated_inference", [True, False]) +@pytest.mark.parametrize( + "acknowledgements", + [[], [False], [True, False], [True, None]], +) +@patch("nemo_rl.algorithms.grpo.ray") +def test_refit_policy_generation_requires_a_worker_acknowledgement( + mock_ray: MagicMock, + acknowledgements: list[bool | None], + colocated_inference: bool, +) -> None: + mock_ray.get.side_effect = [None, acknowledgements] + policy = MagicMock() + policy_generation = MagicMock() + policy_generation.weight_synchronizer = None + + with pytest.raises(RuntimeError, match="Updating weights.*failed"): + refit_policy_generation( + policy, + policy_generation, + colocated_inference=colocated_inference, + _refit_buffer_size_gb=1.0, + ) class TestMaskSampleFilter: diff --git a/tests/unit/algorithms/test_logits_sampling_utils.py b/tests/unit/algorithms/test_logits_sampling_utils.py index cfdee9db82c..04084db538e 100644 --- a/tests/unit/algorithms/test_logits_sampling_utils.py +++ b/tests/unit/algorithms/test_logits_sampling_utils.py @@ -15,7 +15,24 @@ import pytest import torch -from nemo_rl.algorithms.logits_sampling_utils import apply_top_k_top_p +from nemo_rl.algorithms.logits_sampling_utils import ( + TrainingSamplingParams, + apply_top_k_top_p, +) + + +def test_greedy_sampling_uses_unscaled_unfiltered_training_logprobs() -> None: + params = TrainingSamplingParams(top_k=5, top_p=0.8, temperature=0.0) + + assert params.temperature == 1.0 + assert params.top_k is None + assert params.top_p == 1.0 + + +@pytest.mark.parametrize("temperature", [-1.0, float("nan"), float("inf")]) +def test_invalid_temperature_is_rejected(temperature: float) -> None: + with pytest.raises(ValueError, match="temperature must"): + TrainingSamplingParams(temperature=temperature) @pytest.mark.parametrize( diff --git a/tests/unit/models/automodel/test_automodel_train.py b/tests/unit/models/automodel/test_automodel_train.py index 92737b121e1..334c30a788a 100644 --- a/tests/unit/models/automodel/test_automodel_train.py +++ b/tests/unit/models/automodel/test_automodel_train.py @@ -462,6 +462,17 @@ def test_temperature_scaling_with_temperature_one(self): assert torch.equal(result, original_logits) + def test_greedy_temperature_keeps_logits_finite_and_unscaled(self): + """Greedy generation logprobs come from the unscaled distribution.""" + logits = torch.randn(4, 64, 32000) + original_logits = logits.clone() + sampling_params = TrainingSamplingParams(temperature=0.0) + + result = apply_temperature_scaling(logits, sampling_params) + + assert torch.equal(result, original_logits) + assert torch.isfinite(result).all() + def test_temperature_scaling_with_temperature_two(self): """Test that logits are divided by the configured temperature=2.0.""" logits = torch.randn(4, 64, 32000) diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index ec45edf72ac..8c6c7eb019a 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -146,6 +146,36 @@ } +@pytest.mark.parametrize("backend", ["vllm", "dynamo"]) +def test_configure_generation_config_matches_vllm_tiny_temperature_clamp( + backend: str, +): + generation_config = deepcopy(basic_vllm_test_config) + generation_config["backend"] = backend + generation_config["temperature"] = 1e-4 + + configured = configure_generation_config( + generation_config, MagicMock(pad_token_id=0, eos_token_id=1) + ) + + assert configured["temperature"] == 1e-2 + + +@pytest.mark.parametrize("backend", ["sglang", "trtllm", "megatron"]) +def test_configure_generation_config_preserves_other_backend_tiny_temperature( + backend: str, +): + generation_config = deepcopy(basic_vllm_test_config) + generation_config["backend"] = backend + generation_config["temperature"] = 1e-4 + + configured = configure_generation_config( + generation_config, MagicMock(pad_token_id=0, eos_token_id=1) + ) + + assert configured["temperature"] == 1e-4 + + def test_context_capped_max_new_tokens(): assert ( _context_capped_max_new_tokens( diff --git a/tests/unit/models/megatron/test_train.py b/tests/unit/models/megatron/test_train.py index 57ff3504d89..71e36c8a1b5 100644 --- a/tests/unit/models/megatron/test_train.py +++ b/tests/unit/models/megatron/test_train.py @@ -193,6 +193,19 @@ def test_temperature_scaling_with_temperature_one(self): assert torch.allclose(result, original) + def test_greedy_temperature_keeps_logits_finite_and_unscaled(self): + """Greedy generation logprobs come from the unscaled distribution.""" + from nemo_rl.models.megatron.train import apply_temperature_scaling + + logits = torch.randn(2, 10, 100) + original = logits.clone() + sampling_params = TrainingSamplingParams(temperature=0.0) + + result = apply_temperature_scaling(logits, sampling_params) + + assert torch.equal(result, original) + assert torch.isfinite(result).all() + def test_temperature_scaling_with_temperature_two(self): """Test that logits are divided by the configured temperature=2.0.""" from nemo_rl.models.megatron.train import apply_temperature_scaling diff --git a/tests/unit/models/policy/test_dtensor_worker.py b/tests/unit/models/policy/test_dtensor_worker.py index 7ae249d480a..fbda36ec22a 100644 --- a/tests/unit/models/policy/test_dtensor_worker.py +++ b/tests/unit/models/policy/test_dtensor_worker.py @@ -19,6 +19,7 @@ import torch from transformers import AutoModelForCausalLM +from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams from nemo_rl.algorithms.loss import ClippedPGLossConfig, ClippedPGLossFn, NLLLossFn from nemo_rl.algorithms.loss.interfaces import LossFunction from nemo_rl.algorithms.utils import get_tokenizer @@ -43,6 +44,22 @@ def eval(self): self.eval_called = True +def test_dtensor_greedy_temperature_keeps_logits_finite_and_unscaled(): + from nemo_rl.models.policy.workers.dtensor_policy_worker import ( + DTensorPolicyWorkerImpl, + ) + + worker = object.__new__(DTensorPolicyWorkerImpl) + worker.sampling_params = TrainingSamplingParams(temperature=0.0) + logits = torch.randn(2, 10, 100) + original = logits.clone() + + result = DTensorPolicyWorkerImpl._apply_temperature_scaling(worker, logits) + + assert torch.equal(result, original) + assert torch.isfinite(result).all() + + def test_dtensor_prepare_for_training_restores_optimizer(monkeypatch): from nemo_rl.models.policy.workers.dtensor_policy_worker import ( DTensorPolicyWorkerImpl, From 9b3522543f5a6175aa593d7529181d92051ff1bb Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 15:28:35 -0700 Subject: [PATCH 09/17] fix: validate fused sampling and refit acknowledgements Signed-off-by: yaoyu-33 --- nemo_rl/algorithms/grpo.py | 32 ++++++++++++----- .../generation/trtllm/trtllm_worker_async.py | 12 +++---- tests/unit/algorithms/test_grpo.py | 15 ++++++++ .../trtllm/test_trtllm_worker_async.py | 36 +++++++++++++++++-- 4 files changed, 77 insertions(+), 18 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 9d71dd78192..25dcd398b4f 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -532,6 +532,28 @@ def _training_sampling_params_from_generation_config( ) +def _validate_fused_linear_logprobs_sampling( + generation_config: GenerationConfig, +) -> None: + """Reject sampling transforms unsupported by fused Megatron logprobs.""" + sampling_params = _training_sampling_params_from_generation_config( + generation_config + ) + assert sampling_params.temperature == 1.0, ( + "Linear CE fusion loss is not supported with non-unit training-time " + "temperature for GRPO. The fused path computes logprobs before " + "temperature scaling. Set policy.megatron_cfg.use_fused_linear_logprobs=false, " + "or set policy.generation.temperature to 1.0 (or 0.0 for greedy generation)." + ) + assert not need_top_k_or_top_p_filtering(sampling_params), ( + "Linear CE fusion loss is not supported with top-k/top-p training-time " + "filtering for GRPO. The fused path computes logprobs from unfiltered " + "logits. Set policy.megatron_cfg.use_fused_linear_logprobs=false, or " + "disable filtering (policy.generation.top_k=null, " + "policy.generation.top_p=1.0)." + ) + + def setup( master_config: MasterConfig, tokenizer: TokenizerType, @@ -780,15 +802,7 @@ def init_train_dataloader(dataset, suffix: str = ""): # (unfiltered) logits, so top-k/top-p training-time filtering cannot be # applied. This also keeps prev/reference logprobs (computed via the fused # get_logprobs path) consistent with the actor logprobs. - assert not need_top_k_or_top_p_filtering( - _training_sampling_params_from_generation_config(generation_config) - ), ( - "Linear CE fusion loss is not supported with top-k/top-p training-time " - "filtering for GRPO. The fused path computes logprobs from unfiltered " - "logits. Set policy.megatron_cfg.use_fused_linear_logprobs=false, or " - "disable filtering (policy.generation.top_k=null, " - "policy.generation.top_p=1.0)." - ) + _validate_fused_linear_logprobs_sampling(generation_config) loss_fn = ClippedPGLossFn( loss_config, use_fused_linear_logprobs=use_fused_linear_logprobs diff --git a/nemo_rl/models/generation/trtllm/trtllm_worker_async.py b/nemo_rl/models/generation/trtllm/trtllm_worker_async.py index 32839ada63a..784aff382c7 100644 --- a/nemo_rl/models/generation/trtllm/trtllm_worker_async.py +++ b/nemo_rl/models/generation/trtllm/trtllm_worker_async.py @@ -362,10 +362,10 @@ async def update_weights_from_collective_async( "update_weights_from_collective", kwargs={"drain": drain, "recompute_kv": recompute_kv}, ) - worker_result = results[0] if results else True - if not worker_result: + if not results or not all(result is True for result in results): print( - f"Error: TRT-LLM worker failed to update weights. Result: {worker_result}" + "Error: TRT-LLM worker failed to update weights. " + f"Results: {results}" ) return False return True @@ -380,10 +380,10 @@ async def update_weights_via_ipc_zmq_async(self) -> bool: assert self.llm is not None try: results = await self.llm.collective_rpc("update_weights_via_ipc_zmq") - worker_result = results[0] if results else True - if not worker_result: + if not results or not all(result is True for result in results): print( - f"Error: TRT-LLM worker failed to update weights via IPC. Result: {worker_result}" + "Error: TRT-LLM worker failed to update weights via IPC. " + f"Results: {results}" ) return False return True diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index e76324186dc..d29a40ea26d 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -52,6 +52,7 @@ _save_async_replay_buffer_checkpoint, _startup_pipeline_ready, _training_sampling_params_from_generation_config, + _validate_fused_linear_logprobs_sampling, _validate_multimodal_dedup_capability, _validate_use_kl_in_reward_compat, aggregate_rollout_metrics, @@ -128,6 +129,20 @@ def test_training_sampling_params_from_generation_config_normalizes_greedy() -> assert params.temperature == 1.0 assert params.top_k is None assert params.top_p == 1.0 + _validate_fused_linear_logprobs_sampling(generation_config) + + +@pytest.mark.parametrize("temperature", [0.5, 1e-2]) +def test_fused_linear_logprobs_reject_non_unit_temperature( + temperature: float, +) -> None: + generation_config = cast( + GenerationConfig, + {"top_k": None, "top_p": 1.0, "temperature": temperature}, + ) + + with pytest.raises(AssertionError, match="non-unit training-time temperature"): + _validate_fused_linear_logprobs_sampling(generation_config) def test_save_async_replay_buffer_checkpoint(tmp_path): diff --git a/tests/unit/models/generation/trtllm/test_trtllm_worker_async.py b/tests/unit/models/generation/trtllm/test_trtllm_worker_async.py index a674f23f3e8..3159b601278 100644 --- a/tests/unit/models/generation/trtllm/test_trtllm_worker_async.py +++ b/tests/unit/models/generation/trtllm/test_trtllm_worker_async.py @@ -227,8 +227,17 @@ async def test_async_lifecycle_resets_cache_before_sleep_and_resumes_selected_ta @pytest.mark.asyncio -@pytest.mark.parametrize("result", [[True], [], [False]]) -async def test_async_collective_refit_propagates_worker_result(result): +@pytest.mark.parametrize( + ("result", "expected"), + [ + ([True], True), + ([], False), + ([False], False), + ([True, False], False), + ([True, None], False), + ], +) +async def test_async_collective_refit_propagates_worker_result(result, expected): worker = _worker() worker.llm.collective_rpc.return_value = result @@ -237,13 +246,34 @@ async def test_async_collective_refit_propagates_worker_result(result): recompute_kv=True, ) - assert succeeded is (result != [False]) + assert succeeded is expected worker.llm.collective_rpc.assert_awaited_once_with( "update_weights_from_collective", kwargs={"drain": False, "recompute_kv": True}, ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("result", "expected"), + [ + ([True], True), + ([], False), + ([False], False), + ([True, False], False), + ([True, None], False), + ], +) +async def test_async_ipc_refit_propagates_every_worker_result(result, expected): + worker = _worker() + worker.llm.collective_rpc.return_value = result + + succeeded = await worker.update_weights_via_ipc_zmq_async() + + assert succeeded is expected + worker.llm.collective_rpc.assert_awaited_once_with("update_weights_via_ipc_zmq") + + @pytest.mark.asyncio async def test_async_ipc_refit_returns_false_on_worker_exception(): worker = _worker() From ee56dd6075d0ac924e2b41f009b8c8f47898b28d Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 21:53:03 -0700 Subject: [PATCH 10/17] fix: allow generation configs without temperature Signed-off-by: yaoyu-33 --- nemo_rl/models/generation/__init__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/nemo_rl/models/generation/__init__.py b/nemo_rl/models/generation/__init__.py index dc8533a4226..1dacd93f76a 100644 --- a/nemo_rl/models/generation/__init__.py +++ b/nemo_rl/models/generation/__init__.py @@ -68,9 +68,11 @@ def configure_generation_config( # vLLM clamps tiny positive temperatures before sampling. Normalize the # shared config at this backend boundary so policy logprob recomputation # receives the same effective temperature without changing other backends. + temperature = config.get("temperature") if ( config["backend"] in ("vllm", "dynamo") - and 0.0 < config["temperature"] < _VLLM_MIN_NON_ZERO_TEMPERATURE + and temperature is not None + and 0.0 < temperature < _VLLM_MIN_NON_ZERO_TEMPERATURE ): config["temperature"] = _VLLM_MIN_NON_ZERO_TEMPERATURE From 4e6e44ecde057de137ae9cafc9e24d0550a5dce7 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 22:30:36 -0700 Subject: [PATCH 11/17] test: require a positive training learning rate Signed-off-by: yaoyu-33 --- tests/functional/grpo_nano4b_gym_training_e2e.sh | 5 ++++- tests/functional/grpo_nano4b_gym_training_e2e.yaml | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/functional/grpo_nano4b_gym_training_e2e.sh b/tests/functional/grpo_nano4b_gym_training_e2e.sh index dc8ae176f66..f4975a6f5f3 100755 --- a/tests/functional/grpo_nano4b_gym_training_e2e.sh +++ b/tests/functional/grpo_nano4b_gym_training_e2e.sh @@ -45,12 +45,15 @@ uv run tests/json_dump_tb_logs.py "${LOG_DIR}" --output_path "${JSON_METRICS}" # The fixture intentionally contains one accepted and one rejected rollout. In # addition to testing both verifier outcomes, this gives Reinforce++ a non-zero -# advantage so grad_norm proves that an optimizer step was actually exercised. +# advantage so a finite nonzero grad norm and positive learning rate prove that +# a trainable update signal reached the optimizer path. uv run tests/check_metrics.py "${JSON_METRICS}" \ 'len(data["train/loss"]) == 1' \ 'all_finite(data["train/loss"])' \ 'all_finite(data["train/grad_norm"])' \ 'min(data["train/grad_norm"]) > 0' \ + 'all_finite(data["train/lr"])' \ + 'min(data["train/lr"]) > 0' \ 'all_finite(data["train/advantages/min"])' \ 'all_finite(data["train/advantages/max"])' \ 'min(data["train/advantages/min"]) < 0' \ diff --git a/tests/functional/grpo_nano4b_gym_training_e2e.yaml b/tests/functional/grpo_nano4b_gym_training_e2e.yaml index 9afecdb7537..4fd2e1c86f5 100644 --- a/tests/functional/grpo_nano4b_gym_training_e2e.yaml +++ b/tests/functional/grpo_nano4b_gym_training_e2e.yaml @@ -128,7 +128,7 @@ logger: tensorboard_enabled: true monitor_gpus: true -# The H100 functional runner exposes two GPUs. Colocation lets FSDP2 shard the +# The functional test reserves two GPUs. Colocation lets FSDP2 shard the # full-parameter policy and optimizer across both while vLLM shares that cluster. cluster: gpus_per_node: 2 From 579e298692b5f3267fccbf2cff880d953d611416 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 8 Sep 2026 13:30:57 -0700 Subject: [PATCH 12/17] refactor: inline GRPO sampling validation setup Signed-off-by: yaoyu-33 --- nemo_rl/algorithms/grpo.py | 17 ++++------------- tests/unit/algorithms/test_grpo.py | 8 +------- 2 files changed, 5 insertions(+), 20 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 25dcd398b4f..97ee7683f7d 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -521,23 +521,14 @@ def shutdown_environments( print(f"Error stopping environment {task_name}: {kill_error}") -def _training_sampling_params_from_generation_config( - generation_config: GenerationConfig, -) -> TrainingSamplingParams: - """Build the effective policy-logprob sampling configuration.""" - return TrainingSamplingParams( - top_k=generation_config["top_k"], - top_p=generation_config["top_p"], - temperature=generation_config["temperature"], - ) - - def _validate_fused_linear_logprobs_sampling( generation_config: GenerationConfig, ) -> None: """Reject sampling transforms unsupported by fused Megatron logprobs.""" - sampling_params = _training_sampling_params_from_generation_config( - generation_config + sampling_params = TrainingSamplingParams( + top_k=generation_config["top_k"], + top_p=generation_config["top_p"], + temperature=generation_config["temperature"], ) assert sampling_params.temperature == 1.0, ( "Linear CE fusion loss is not supported with non-unit training-time " diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index d29a40ea26d..226e24c9f9f 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -51,7 +51,6 @@ _resolve_message_level_advantage_penalties, _save_async_replay_buffer_checkpoint, _startup_pipeline_ready, - _training_sampling_params_from_generation_config, _validate_fused_linear_logprobs_sampling, _validate_multimodal_dedup_capability, _validate_use_kl_in_reward_compat, @@ -118,17 +117,12 @@ def _mock_policy_generation() -> MagicMock: return policy_generation -def test_training_sampling_params_from_generation_config_normalizes_greedy() -> None: +def test_fused_linear_logprobs_allows_greedy_sampling() -> None: generation_config = cast( GenerationConfig, {"top_k": 5, "top_p": 0.8, "temperature": 0.0}, ) - params = _training_sampling_params_from_generation_config(generation_config) - - assert params.temperature == 1.0 - assert params.top_k is None - assert params.top_p == 1.0 _validate_fused_linear_logprobs_sampling(generation_config) From b22c4e8fb79700bbc221ef0503b02d7d56eb01a4 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 8 Sep 2026 13:36:11 -0700 Subject: [PATCH 13/17] refactor: inline fused sampling validation Signed-off-by: yaoyu-33 --- nemo_rl/algorithms/grpo.py | 83 ++++++++++++------------------ tests/unit/algorithms/test_grpo.py | 45 ++++++++-------- 2 files changed, 53 insertions(+), 75 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 97ee7683f7d..80dec4d0252 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -521,30 +521,6 @@ def shutdown_environments( print(f"Error stopping environment {task_name}: {kill_error}") -def _validate_fused_linear_logprobs_sampling( - generation_config: GenerationConfig, -) -> None: - """Reject sampling transforms unsupported by fused Megatron logprobs.""" - sampling_params = TrainingSamplingParams( - top_k=generation_config["top_k"], - top_p=generation_config["top_p"], - temperature=generation_config["temperature"], - ) - assert sampling_params.temperature == 1.0, ( - "Linear CE fusion loss is not supported with non-unit training-time " - "temperature for GRPO. The fused path computes logprobs before " - "temperature scaling. Set policy.megatron_cfg.use_fused_linear_logprobs=false, " - "or set policy.generation.temperature to 1.0 (or 0.0 for greedy generation)." - ) - assert not need_top_k_or_top_p_filtering(sampling_params), ( - "Linear CE fusion loss is not supported with top-k/top-p training-time " - "filtering for GRPO. The fused path computes logprobs from unfiltered " - "logits. Set policy.megatron_cfg.use_fused_linear_logprobs=false, or " - "disable filtering (policy.generation.top_k=null, " - "policy.generation.top_p=1.0)." - ) - - def setup( master_config: MasterConfig, tokenizer: TokenizerType, @@ -615,6 +591,38 @@ def setup( policy_config["generation"] = generation_config _validate_multimodal_dedup_capability(master_config) + # Fused linear logprobs compute next-token logprobs directly from hidden states + # and therefore cannot apply sampling transforms to materialized logits. + megatron_cfg = policy_config.get("megatron_cfg", {}) + use_fused_linear_logprobs = bool( + megatron_cfg.get("enabled") and megatron_cfg.get("use_fused_linear_logprobs") + ) + if use_fused_linear_logprobs: + assert not policy_config["sequence_packing"]["enabled"], ( + "Linear CE fusion loss is not supported with sequence packing for GRPO. " + "The fused path has not been validated with cu_seqlens-based logprob " + "aggregation. Set policy.megatron_cfg.use_fused_linear_logprobs=false " + "or policy.sequence_packing.enabled=false." + ) + sampling_params = TrainingSamplingParams( + top_k=generation_config["top_k"], + top_p=generation_config["top_p"], + temperature=generation_config["temperature"], + ) + assert sampling_params.temperature == 1.0, ( + "Linear CE fusion loss is not supported with non-unit training-time " + "temperature for GRPO. The fused path computes logprobs before " + "temperature scaling. Set policy.megatron_cfg.use_fused_linear_logprobs=false, " + "or set policy.generation.temperature to 1.0 (or 0.0 for greedy generation)." + ) + assert not need_top_k_or_top_p_filtering(sampling_params), ( + "Linear CE fusion loss is not supported with top-k/top-p training-time " + "filtering for GRPO. The fused path computes logprobs from unfiltered " + "logits. Set policy.megatron_cfg.use_fused_linear_logprobs=false, or " + "disable filtering (policy.generation.top_k=null, " + "policy.generation.top_p=1.0)." + ) + # Validation-only sampling is honored only on the NeMo-Gym vLLM rollout # path; everywhere else validation must sample exactly like training. val_sampling_overridden = ( @@ -768,33 +776,6 @@ def init_train_dataloader(dataset, suffix: str = ""): # ========================== # Loss Function # ========================== - # Fused linear logprobs compute next-token logprobs directly from hidden states - # (chunked over the sequence) and never materialize the full - # [batch, seq_len, vocab_size] logit tensor, which significantly reduces peak - # memory. It is only available on the Megatron backend. - # Both megatron_cfg and use_fused_linear_logprobs are NotRequired, and many - # configs (e.g. nemo_gym, modelopt, non-megatron) omit them -- use .get() with - # a {} fallback to avoid a KeyError. - megatron_cfg = policy_config.get("megatron_cfg", {}) - use_fused_linear_logprobs = bool( - megatron_cfg.get("enabled") and megatron_cfg.get("use_fused_linear_logprobs") - ) - if use_fused_linear_logprobs: - # Sequence packing is not yet validated with the fused path: the fused - # forward rolls labels over the whole (packed) sequence and would mix - # tokens across packed-sequence boundaries. - assert not policy_config["sequence_packing"]["enabled"], ( - "Linear CE fusion loss is not supported with sequence packing for GRPO. " - "The fused path has not been validated with cu_seqlens-based logprob " - "aggregation. Set policy.megatron_cfg.use_fused_linear_logprobs=false " - "or policy.sequence_packing.enabled=false." - ) - # The fused forward gathers the logprob of the realized token from the raw - # (unfiltered) logits, so top-k/top-p training-time filtering cannot be - # applied. This also keeps prev/reference logprobs (computed via the fused - # get_logprobs path) consistent with the actor logprobs. - _validate_fused_linear_logprobs_sampling(generation_config) - loss_fn = ClippedPGLossFn( loss_config, use_fused_linear_logprobs=use_fused_linear_logprobs ) diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 226e24c9f9f..8a54df87795 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -17,7 +17,7 @@ from pathlib import Path from threading import Event from types import SimpleNamespace -from typing import Any, cast +from typing import Any from unittest.mock import MagicMock, patch import pytest @@ -51,7 +51,6 @@ _resolve_message_level_advantage_penalties, _save_async_replay_buffer_checkpoint, _startup_pipeline_ready, - _validate_fused_linear_logprobs_sampling, _validate_multimodal_dedup_capability, _validate_use_kl_in_reward_compat, aggregate_rollout_metrics, @@ -92,10 +91,7 @@ from nemo_rl.experience.rollouts import calculate_rewards from nemo_rl.models.generation import configure_generation_config from nemo_rl.models.generation.dynamo import DynamoConfig -from nemo_rl.models.generation.interfaces import ( - GenerationConfig, - should_use_async_rollouts, -) +from nemo_rl.models.generation.interfaces import should_use_async_rollouts from nemo_rl.models.generation.megatron import MegatronGeneration from nemo_rl.utils.config import load_config, register_omegaconf_resolvers from nemo_rl.utils.timer import Timer @@ -117,26 +113,27 @@ def _mock_policy_generation() -> MagicMock: return policy_generation -def test_fused_linear_logprobs_allows_greedy_sampling() -> None: - generation_config = cast( - GenerationConfig, - {"top_k": 5, "top_p": 0.8, "temperature": 0.0}, - ) - - _validate_fused_linear_logprobs_sampling(generation_config) - - -@pytest.mark.parametrize("temperature", [0.5, 1e-2]) -def test_fused_linear_logprobs_reject_non_unit_temperature( - temperature: float, +@pytest.mark.parametrize( + ("sampling_config", "error"), + [ + ({"temperature": 0.5}, "non-unit training-time temperature"), + ({"top_k": 5}, "top-k/top-p training-time filtering"), + ({"top_p": 0.9}, "top-k/top-p training-time filtering"), + ], +) +def test_setup_rejects_fused_linear_logprobs_with_unsupported_sampling( + mock_grpo_components, sampling_config: dict[str, float | int], error: str ) -> None: - generation_config = cast( - GenerationConfig, - {"top_k": None, "top_p": 1.0, "temperature": temperature}, - ) + master_config = mock_grpo_components["master_config"] + master_config.policy["megatron_cfg"] = { + "enabled": True, + "use_fused_linear_logprobs": True, + } + master_config.policy["sequence_packing"] = {"enabled": False} + master_config.policy["generation"].update(sampling_config) - with pytest.raises(AssertionError, match="non-unit training-time temperature"): - _validate_fused_linear_logprobs_sampling(generation_config) + with pytest.raises(AssertionError, match=error): + setup(master_config, MagicMock(), MagicMock(), None) def test_save_async_replay_buffer_checkpoint(tmp_path): From ac7f9ea7d6a0fda3942dec5cbe79084ed42e1c54 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 8 Sep 2026 14:08:45 -0700 Subject: [PATCH 14/17] refactor: keep fused validation in loss setup Signed-off-by: yaoyu-33 --- nemo_rl/algorithms/grpo.py | 76 +++++++++++++++++------------- tests/unit/algorithms/test_grpo.py | 12 ++++- 2 files changed, 55 insertions(+), 33 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 80dec4d0252..d1e65b48545 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -591,38 +591,6 @@ def setup( policy_config["generation"] = generation_config _validate_multimodal_dedup_capability(master_config) - # Fused linear logprobs compute next-token logprobs directly from hidden states - # and therefore cannot apply sampling transforms to materialized logits. - megatron_cfg = policy_config.get("megatron_cfg", {}) - use_fused_linear_logprobs = bool( - megatron_cfg.get("enabled") and megatron_cfg.get("use_fused_linear_logprobs") - ) - if use_fused_linear_logprobs: - assert not policy_config["sequence_packing"]["enabled"], ( - "Linear CE fusion loss is not supported with sequence packing for GRPO. " - "The fused path has not been validated with cu_seqlens-based logprob " - "aggregation. Set policy.megatron_cfg.use_fused_linear_logprobs=false " - "or policy.sequence_packing.enabled=false." - ) - sampling_params = TrainingSamplingParams( - top_k=generation_config["top_k"], - top_p=generation_config["top_p"], - temperature=generation_config["temperature"], - ) - assert sampling_params.temperature == 1.0, ( - "Linear CE fusion loss is not supported with non-unit training-time " - "temperature for GRPO. The fused path computes logprobs before " - "temperature scaling. Set policy.megatron_cfg.use_fused_linear_logprobs=false, " - "or set policy.generation.temperature to 1.0 (or 0.0 for greedy generation)." - ) - assert not need_top_k_or_top_p_filtering(sampling_params), ( - "Linear CE fusion loss is not supported with top-k/top-p training-time " - "filtering for GRPO. The fused path computes logprobs from unfiltered " - "logits. Set policy.megatron_cfg.use_fused_linear_logprobs=false, or " - "disable filtering (policy.generation.top_k=null, " - "policy.generation.top_p=1.0)." - ) - # Validation-only sampling is honored only on the NeMo-Gym vLLM rollout # path; everywhere else validation must sample exactly like training. val_sampling_overridden = ( @@ -776,6 +744,50 @@ def init_train_dataloader(dataset, suffix: str = ""): # ========================== # Loss Function # ========================== + # Fused linear logprobs compute next-token logprobs directly from hidden states + # (chunked over the sequence) and never materialize the full + # [batch, seq_len, vocab_size] logit tensor, which significantly reduces peak + # memory. It is only available on the Megatron backend. + # Both megatron_cfg and use_fused_linear_logprobs are NotRequired, and many + # configs (e.g. nemo_gym, modelopt, non-megatron) omit them -- use .get() with + # a {} fallback to avoid a KeyError. + megatron_cfg = policy_config.get("megatron_cfg", {}) + use_fused_linear_logprobs = bool( + megatron_cfg.get("enabled") and megatron_cfg.get("use_fused_linear_logprobs") + ) + if use_fused_linear_logprobs: + # Sequence packing is not yet validated with the fused path: the fused + # forward rolls labels over the whole (packed) sequence and would mix + # tokens across packed-sequence boundaries. + assert not policy_config["sequence_packing"]["enabled"], ( + "Linear CE fusion loss is not supported with sequence packing for GRPO. " + "The fused path has not been validated with cu_seqlens-based logprob " + "aggregation. Set policy.megatron_cfg.use_fused_linear_logprobs=false " + "or policy.sequence_packing.enabled=false." + ) + sampling_params = TrainingSamplingParams( + top_k=generation_config["top_k"], + top_p=generation_config["top_p"], + temperature=generation_config["temperature"], + ) + assert sampling_params.temperature == 1.0, ( + "Linear CE fusion loss is not supported with non-unit training-time " + "temperature for GRPO. The fused path computes logprobs before " + "temperature scaling. Set policy.megatron_cfg.use_fused_linear_logprobs=false, " + "or set policy.generation.temperature to 1.0 (or 0.0 for greedy generation)." + ) + # The fused forward gathers the logprob of the realized token from the raw + # (unfiltered) logits, so top-k/top-p training-time filtering cannot be + # applied. This also keeps prev/reference logprobs (computed via the fused + # get_logprobs path) consistent with the actor logprobs. + assert not need_top_k_or_top_p_filtering(sampling_params), ( + "Linear CE fusion loss is not supported with top-k/top-p training-time " + "filtering for GRPO. The fused path computes logprobs from unfiltered " + "logits. Set policy.megatron_cfg.use_fused_linear_logprobs=false, or " + "disable filtering (policy.generation.top_k=null, " + "policy.generation.top_p=1.0)." + ) + loss_fn = ClippedPGLossFn( loss_config, use_fused_linear_logprobs=use_fused_linear_logprobs ) diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 8a54df87795..5cfa2f27b42 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -131,8 +131,18 @@ def test_setup_rejects_fused_linear_logprobs_with_unsupported_sampling( } master_config.policy["sequence_packing"] = {"enabled": False} master_config.policy["generation"].update(sampling_config) + master_config.grpo.val_period = 0 + master_config.grpo.batch_multiplier = 1 + master_config.data.update({"shuffle": False, "num_workers": 0}) - with pytest.raises(AssertionError, match=error): + with ( + patch("nemo_rl.algorithms.grpo.Logger"), + patch("nemo_rl.algorithms.grpo.CheckpointManager") as checkpointer_cls, + patch("nemo_rl.algorithms.grpo.StatefulDataLoader"), + pytest.raises(AssertionError, match=error), + ): + checkpointer_cls.return_value.get_latest_checkpoint_path.return_value = None + checkpointer_cls.return_value.load_training_info.return_value = {} setup(master_config, MagicMock(), MagicMock(), None) From 1a6d7ade76aa50298e817570616eea0b945bf916 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 9 Sep 2026 10:40:54 -0700 Subject: [PATCH 15/17] test: align fused sampling validation params Signed-off-by: yaoyu-33 --- tests/unit/algorithms/test_grpo.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 5cfa2f27b42..8ef79b14fde 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -130,7 +130,13 @@ def test_setup_rejects_fused_linear_logprobs_with_unsupported_sampling( "use_fused_linear_logprobs": True, } master_config.policy["sequence_packing"] = {"enabled": False} - master_config.policy["generation"].update(sampling_config) + generation_config = master_config.policy["generation"] + generation_config.update(sampling_config) + # Keep validation sampling aligned so setup reaches the fused-path guard + # exercised by this test instead of the earlier train/validation check. + generation_config.update( + {f"val_{name}": value for name, value in sampling_config.items()} + ) master_config.grpo.val_period = 0 master_config.grpo.batch_multiplier = 1 master_config.data.update({"shuffle": False, "num_workers": 0}) From 435888276285969c8f48a34c6e6adf14c4617188 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 9 Sep 2026 15:19:56 -0700 Subject: [PATCH 16/17] test: run Gym training E2E on GB200 Signed-off-by: yaoyu-33 --- ...ym_Training.sh => L1_Functional_Tests_GB200_Gym_Training.sh} | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename tests/functional/{L1_Functional_Tests_Gym_Training.sh => L1_Functional_Tests_GB200_Gym_Training.sh} (93%) diff --git a/tests/functional/L1_Functional_Tests_Gym_Training.sh b/tests/functional/L1_Functional_Tests_GB200_Gym_Training.sh similarity index 93% rename from tests/functional/L1_Functional_Tests_Gym_Training.sh rename to tests/functional/L1_Functional_Tests_GB200_Gym_Training.sh index c984ee80d0b..873f95e983b 100755 --- a/tests/functional/L1_Functional_Tests_Gym_Training.sh +++ b/tests/functional/L1_Functional_Tests_GB200_Gym_Training.sh @@ -22,7 +22,7 @@ PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") cd "${PROJECT_ROOT}" if [[ "${FAST:-0}" == "1" ]]; then - echo "FAST: Skipping Nano 4B Gym training E2E" + echo "FAST: Skipping Nano 4B Gym training E2E on GB200" else time uv run --no-sync bash ./tests/functional/grpo_nano4b_gym_training_e2e.sh fi From 6ea3c06c30124ef59eacd1b761f3d8bfff74b2b6 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Thu, 10 Sep 2026 21:45:28 -0700 Subject: [PATCH 17/17] fix: address Gym training review feedback Signed-off-by: yaoyu-33 --- nemo_rl/algorithms/grpo.py | 6 ++++-- nemo_rl/models/generation/__init__.py | 21 ++++++++++++------- .../L1_Functional_Tests_GB200_Gym_Training.sh | 5 +++++ .../models/generation/test_vllm_generation.py | 14 ++++++++++--- 4 files changed, 34 insertions(+), 12 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index d1e65b48545..3c62e68e769 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -2538,8 +2538,10 @@ def refit_policy_generation( kv_scales: Optional dictionary of KV cache scales for FP8 quantization. Returns: - Scalar metrics reported by the selected weight synchronizer, or the - number of generation worker groups that acknowledged a direct refit. + Scalar metrics reported by the selected weight synchronizer, or + ``{"generation_workers_updated": n}`` where ``n`` is the number of + generation engine replicas (one TP/PP rank-0 worker per data-parallel + replica) that acknowledged a direct IPC/NCCL refit. """ # Every SGLang deployment reaches its refit through this hook: `setup` # attaches an SGLang synchronizer that owns the whole lifecycle (phase diff --git a/nemo_rl/models/generation/__init__.py b/nemo_rl/models/generation/__init__.py index a7d3e50fc66..9acf4c217e4 100644 --- a/nemo_rl/models/generation/__init__.py +++ b/nemo_rl/models/generation/__init__.py @@ -68,13 +68,20 @@ def configure_generation_config( # vLLM clamps tiny positive temperatures before sampling. Normalize the # shared config at this backend boundary so policy logprob recomputation # receives the same effective temperature without changing other backends. - temperature = config.get("temperature") - if ( - config["backend"] in ("vllm", "dynamo") - and temperature is not None - and 0.0 < temperature < _VLLM_MIN_NON_ZERO_TEMPERATURE - ): - config["temperature"] = _VLLM_MIN_NON_ZERO_TEMPERATURE + if config["backend"] in ("vllm", "dynamo"): + temperature = config.get("temperature") + if ( + temperature is not None + and 0.0 < temperature < _VLLM_MIN_NON_ZERO_TEMPERATURE + ): + config["temperature"] = _VLLM_MIN_NON_ZERO_TEMPERATURE + + val_temperature = config.get("val_temperature") + if ( + val_temperature is not None + and 0.0 < val_temperature < _VLLM_MIN_NON_ZERO_TEMPERATURE + ): + config["val_temperature"] = _VLLM_MIN_NON_ZERO_TEMPERATURE if ( config["backend"] != "vllm" and config.get("worker_extension_cls_fqn") is not None diff --git a/tests/functional/L1_Functional_Tests_GB200_Gym_Training.sh b/tests/functional/L1_Functional_Tests_GB200_Gym_Training.sh index 873f95e983b..833c2149298 100755 --- a/tests/functional/L1_Functional_Tests_GB200_Gym_Training.sh +++ b/tests/functional/L1_Functional_Tests_GB200_Gym_Training.sh @@ -26,3 +26,8 @@ if [[ "${FAST:-0}" == "1" ]]; then else time uv run --no-sync bash ./tests/functional/grpo_nano4b_gym_training_e2e.sh fi + +cd "${PROJECT_ROOT}/tests" +if compgen -G ".coverage*" > /dev/null; then + coverage combine .coverage* +fi diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index cf240c79e1d..b1569db00e4 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -147,18 +147,24 @@ @pytest.mark.parametrize("backend", ["vllm", "dynamo"]) +@pytest.mark.parametrize( + ("temperature", "expected"), + [(1e-4, 1e-2), (0.0, 0.0), (1e-2, 1e-2), (1.0, 1.0)], +) def test_configure_generation_config_matches_vllm_tiny_temperature_clamp( - backend: str, + backend: str, temperature: float, expected: float ): generation_config = deepcopy(basic_vllm_test_config) generation_config["backend"] = backend - generation_config["temperature"] = 1e-4 + generation_config["temperature"] = temperature + generation_config["val_temperature"] = temperature configured = configure_generation_config( generation_config, MagicMock(pad_token_id=0, eos_token_id=1) ) - assert configured["temperature"] == 1e-2 + assert configured["temperature"] == expected + assert configured["val_temperature"] == expected @pytest.mark.parametrize("backend", ["sglang", "trtllm", "megatron"]) @@ -168,12 +174,14 @@ def test_configure_generation_config_preserves_other_backend_tiny_temperature( generation_config = deepcopy(basic_vllm_test_config) generation_config["backend"] = backend generation_config["temperature"] = 1e-4 + generation_config["val_temperature"] = 1e-4 configured = configure_generation_config( generation_config, MagicMock(pad_token_id=0, eos_token_id=1) ) assert configured["temperature"] == 1e-4 + assert configured["val_temperature"] == 1e-4 @pytest.mark.parametrize("async_engine", [False, True])