From cc99a3ba8facaaf9120dfd642212098fa83f4c47 Mon Sep 17 00:00:00 2001 From: ruit Date: Tue, 30 Jun 2026 20:43:21 -0700 Subject: [PATCH 01/69] feat(swe_bench): add GB200 (hsg) SWE2 generation-scaling launcher and config Copy the GB200/aarch64 async-GRPO SWE launcher and its megatron config into examples/swe_bench for the oci-hsg-cs-001 cluster: - run_grpo_swe2_scale_gen_hsg.sh (single-knob NUM_VLLM_REPLICAS, ALIGN_BASELINE switch, LUSTRE_UV_CACHE_SEED default, 4 GPU/node geometry) - grpo_qwen3_30b_async_swe.yaml (incl. megatron_cfg.use_fused_weighted_squared_relu) Signed-off-by: ruit --- .../swe_bench/grpo_qwen3_30b_async_swe.yaml | 429 +++++++++++++++ .../swe_bench/run_grpo_swe2_scale_gen_hsg.sh | 504 ++++++++++++++++++ 2 files changed, 933 insertions(+) create mode 100644 examples/swe_bench/grpo_qwen3_30b_async_swe.yaml create mode 100755 examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh diff --git a/examples/swe_bench/grpo_qwen3_30b_async_swe.yaml b/examples/swe_bench/grpo_qwen3_30b_async_swe.yaml new file mode 100644 index 00000000000..ea485abb375 --- /dev/null +++ b/examples/swe_bench/grpo_qwen3_30b_async_swe.yaml @@ -0,0 +1,429 @@ +# ============================================================================ +# Async GRPO SWE RL Training: Qwen3-30B-A3B-Thinking-2507 (GB200 / aarch64) +# +# Model: Qwen3-30B-A3B-Thinking-2507 (Qwen3MoeForCausalLM, 30B / 3B active, 128 experts top-8, thinking) +# Train data: SWE blend (balanced_language.jsonl) +# Eval data: SWE-bench Verified (swe_public_datasets_val_swebench.jsonl) +# Mode: Async GRPO with non-colocated generation +# Entry: examples/nemo_gym/run_grpo_nemo_gym.py +# Env: swe_agents (OpenHands agent, apptainer sandbox), arm64 SWE sif set +# Cluster: oci-hsg-cs-001 / nemotron_sw_post (GB200, aarch64, 4 GPU/node) +# +# Forked from: examples/swe_bench/grpo_qwen3_30b_async_swe.yaml, adapted for GB200: +# - gpus_per_node 8 -> 4 (ray.sub asserts GPUS_PER_NODE == gres gpu:4) +# - arm64 container_formatter + accessible data paths (x86_64/coreai assets removed) +# - baseline geometry fits 4-GPU nodes; actual geometry/nodes are overridden by +# run_grpo_qwen3_30b_swe_scale_gen.sh. +# ============================================================================ + +checkpointing: + enabled: true + checkpoint_dir: "results/grpo-qwen3-30b-thinking-swe-rl" + metric_name: "train:total_reward/mean" + higher_is_better: true + keep_top_k: 100 + save_period: 5 + checkpoint_must_save_by: "00:03:35:00" + model_save_format: "safetensors" + save_consolidated: false + save_optimizer: true + +grpo: + num_prompts_per_step: 16 + num_generations_per_prompt: 16 + num_val_generations_per_prompt: 1 + max_rollout_turns: 1 + max_num_epochs: 100 + max_num_steps: 1000000 + normalize_rewards: true + use_leave_one_out_baseline: true + advantage_clip_low: -100 + advantage_clip_high: 100 + val_period: 10 + val_at_start: false + val_at_end: false + overlong_filtering: true + max_val_samples: null + val_batch_size: 256 + seed: 42 + invalid_tool_call_strategy: "" + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + penalize_invalid_tool_call: true + invalid_tool_call_advantage: -5.0 + penalize_malformed_thinking: true + malformed_thinking_advantage: -5.0 + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + + seq_logprob_error_threshold: 2 + +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: null + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5.0 + truncated_importance_sampling_ratio_min: null + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + force_on_policy_ratio: true + use_kl_in_reward: false + +policy: + model_name: "/lustre/fsw/portfolios/llmservice/users/igitman/hf_models/Qwen3-30B-A3B-Thinking-2507" + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: + enable_thinking: true + hf_config_overrides: {} + train_global_batch_size: 256 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 131072 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: False + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + gradient_accumulation_fusion: false + empty_unused_memory_level: 1 + activation_checkpointing: true + tensor_model_parallel_size: 4 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 4 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 1 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_aux_loss_coeff: 0.0 + moe_router_enable_expert_bias: true + moe_shared_expert_overlap: false + apply_rope_fusion: True + bias_activation_fusion: False + use_fused_weighted_squared_relu: false + defer_fp32_logits: True + moe_per_layer_logging: True + + optimizer: + optimizer: "adam" + lr: 1.0e-6 + min_lr: 1.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + sgd_momentum: 0.9 + use_distributed_optimizer: true + use_precision_aware_optimizer: true + clip_grad: ${policy.max_grad_norm} + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: 1000000 + lr_warmup_iters: 0 + lr_warmup_init: 0 + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + mtp_loss_scaling_factor: 0.0 + mtp_use_repeated_layer: false + mtp_num_layers: 0 + mtp_detach_heads: false + + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "blockwise" + fp8_param: false + + env_vars: null + + dynamic_batching: + enabled: False + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: True + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + + make_sequence_length_divisible_by: 4 + max_grad_norm: 1.0 + + optimizer: null + scheduler: null + + generation: + port_range_low: 11001 + port_range_high: 15000 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + vllm_cfg: + enable_prefix_caching: true + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 2 + pipeline_parallel_size: 1 + expert_parallel_size: 1 + gpu_memory_utilization: 0.8 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: False + enforce_monotonicity: false + use_deep_gemm: False + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + enable_thinking: true + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: hermes + reasoning_parser: deepseek_r1 + chat_template: | + {%- if tools %} + {{- '<|im_start|>system\n' }} + {%- if messages[0].role == 'system' %} + {{- messages[0].content + '\n\n' }} + {%- endif %} + {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }} + {%- else %} + {%- if messages[0].role == 'system' %} + {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }} + {%- endif %} + {%- endif %} + {%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} + {%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('') and message.content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} + {%- endfor %} + {%- for message in messages %} + {%- if message.content is string %} + {%- set content = message.content %} + {%- else %} + {%- set content = '' %} + {%- endif %} + {%- if (message.role == "user") or (message.role == "system" and not loop.first) %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- if reasoning_content %} + {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- if message.tool_calls %} + {%- for tool_call in message.tool_calls %} + {%- if (loop.first and content) or (not loop.first) %} + {{- '\n' }} + {%- endif %} + {%- if tool_call.function %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {{- '\n{"name": "' }} + {{- tool_call.name }} + {{- '", "arguments": ' }} + {%- if tool_call.arguments is string %} + {{- tool_call.arguments }} + {%- else %} + {{- tool_call.arguments | tojson }} + {%- endif %} + {{- '}\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- endif %} + {%- endfor %} + {%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n\n' }} + {%- endif %} + default_chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + + vllm_kwargs: + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 2 + +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + use_multiple_dataloader: false + train: + data_path: "/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/blends/balanced_language.jsonl" + validation: + data_path: "/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/swe_public_datasets_val_swebench.jsonl" + 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: false + nemo_gym: + skip_venv_if_present: true + port_range_low: 15001 + port_range_high: 20000 + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml + swe_agents_train: + responses_api_agents: + swe_agents: + agent_max_turns: 100 + concurrency: 768 + swebench_agent_timeout: 3600 + run_with_mixed_prompts: true + dataset_path: ${data.train.data_path} + container_formatter: + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/nv_internal/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/r2e_gym/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swegym/sweb.eval.arm64.{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swebench/swe-bench.eval.arm64.{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swerebench/{instance_id}.sif" + swe_agents_val: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: 768 + swebench_agent_timeout: 3600 + dataset_path: ${data.validation.data_path} + container_formatter: + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/nv_internal/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/r2e_gym/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swegym/sweb.eval.arm64.{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swebench/swe-bench.eval.arm64.{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swerebench/{instance_id}.sif" + use_absolute_ip: true + +logger: + log_dir: "logs" + num_val_samples_to_print: 0 + wandb_enabled: true + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "ruit-nemo-rl" + name: "qwen3-30b-thinking-swe-rl-gb200" + tensorboard: {} + mlflow: + experiment_name: "qwen3-30b-thinking-swe-rl-gb200" + run_name: "qwen3-30b-thinking-swe-rl-gb200" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +cluster: + gpus_per_node: 4 + num_nodes: 3 diff --git a/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh b/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh new file mode 100755 index 00000000000..ea161d2160a --- /dev/null +++ b/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh @@ -0,0 +1,504 @@ +#!/bin/bash +# ============================================================================ +# SCALABLE async SWE GRPO launcher for Qwen3-30B-A3B-Thinking-2507 on +# oci-hsg-cs-001 / nemotron_sw_post (GB200, aarch64, 4 GPU/node). +# +# Purpose: confirm the SWE async-GRPO stack runs end-to-end on GB200 with a small +# model before scaling to 235B. Same single-knob design as the 235B launcher. +# +# Single knob: NUM_VLLM_REPLICAS (R) -> everything else is auto-derived. +# INVARIANT held across all R: GPP*PPS/R == GBS/R == SAMPLES_PER_REPLICA (== 2 for 30B). +# GPP is held fixed (GRPO group size); PPS is derived from R to keep the invariant. +# Training geometry is FIXED & small (30B fits 1 node); R only scales generation: +# GEN_NODES = R * VLLM_TP / NUM_GPU +# TRAIN_NODES (default) = (TP*CP*PP) / NUM_GPU = 4/4 = 1 (override w/ TRAIN_NODES=) +# TOTAL_NODES = TRAIN_NODES + GEN_NODES -> sbatch --nodes & cluster.num_nodes +# PPS = SAMPLES_PER_REPLICA * R / GPP (= R/4 with the defaults below) +# GBS = PPS * GPP = SAMPLES_PER_REPLICA * R +# CONCURRENCY = max(768, GBS * max_trajectory_age_steps) +# Non-colocated carve-out (grpo.py:527): cluster.num_nodes - gen_nodes = train_nodes. +# +# Examples (defaults GPP=8, SAMPLES_PER_REPLICA=2 -> GPP*PPS/R = 2): +# NUM_VLLM_REPLICAS=4 bash test_assets/qwen-30B/run_grpo_qwen3_30b_swe_scale_gen.sh +# -> GEN_NODES=2 + TRAIN_NODES=1 = 3 nodes, PPS=1, GPP=8, GBS=8 +# NUM_VLLM_REPLICAS=8 -> GEN_NODES=4 + TRAIN_NODES=1 = 5 nodes, PPS=2, GBS=16 +# NUM_VLLM_REPLICAS=4 SEQLEN=16384 MAX_NUM_STEPS=2 DRY_RUN=1 bash .../run_grpo_qwen3_30b_swe_scale_gen.sh +# SKIP_TRAINING=1 NUM_VLLM_REPLICAS=4 bash .../run_grpo_qwen3_30b_swe_scale_gen.sh # gen-only +# +# ALIGN_BASELINE=1: reproduce baseline dc3m70us GPU/batch geometry. baseline ran on +# x86 8-GPU nodes: 64 gen GPU (32 vLLM replicas) + 64 train GPU, GBS=64. On GB200 +# (4 GPU/node) the same GPU count is 16 gen + 16 train = 32 nodes. The switch just +# pins NUM_VLLM_REPLICAS=32 (-> 16 gen nodes, GBS=64) and TRAIN_NODES=16 (train_DP=16). +# Both remain overridable. Example: +# ALIGN_BASELINE=1 bash .../run_grpo_qwen3_30b_swe_scale_gen.sh # 32 nodes, GBS=64 +# +# Optional env: ALIGN_BASELINE, SKIP_TRAINING, TRAIN_NODES, GPP, WANDB_GROUP, EXP_SUFFIX, +# MODEL_PATH, CONTAINER, MAX_NUM_STEPS, SBATCH_TIME, PERSISTENT_CACHE, +# BASE_LOG_DIR, TP, CP, EP, PP, VLLM_TP, SEQLEN (advanced overrides). +# ============================================================================ + +set -e + +# ============================ Paths ============================ +REPO_ROOT="${REPO_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" +CONFIG_FILE="${CONFIG_FILE:-${REPO_ROOT}/test_assets/qwen-30B/grpo_qwen3_30b_async_swe.yaml}" +CHECKPOINT_ROOT="${CHECKPOINT_ROOT:-${REPO_ROOT}/results}" +TRAIN_DATA_PATH="${TRAIN_DATA_PATH:-/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/blends/balanced_language.jsonl}" +VAL_DATA_PATH="${VAL_DATA_PATH:-/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/swe_public_datasets_val_swebench.jsonl}" +DEFAULT_MODEL_PATH="/lustre/fsw/portfolios/llmservice/users/igitman/hf_models/Qwen3-30B-A3B-Thinking-2507" +MODEL_PATH="${1:-${MODEL_PATH:-${DEFAULT_MODEL_PATH}}}" + +# ================ Container and mount config ================ +# GB200 (aarch64) baked image: apptainer + /opt/nemo_rl_venv with --extra mcore +# (sm_100), built by test_assets/SWE/build_swe_bench_combined.sh. +# export CONTAINER=${CONTAINER:-/lustre/fsw/portfolios/nemotron/users/ruit/enroot-images/ruit-swe_bench-6dc8fabea-aarch64-060426-mcore-apptainer.squashfs} +export CONTAINER=${CONTAINER:-/lustre/fsw/portfolios/nemotron/users/ruit/enroot-images/nemo-rl:nightly-062326.squashfs} +GYM_CODE="${REPO_ROOT}/3rdparty/Gym-workspace/Gym" +export MOUNTS="/lustre:/lustre,$PWD:$PWD,${GYM_CODE}:/opt/nemo-rl/3rdparty/Gym-workspace/Gym" + +# ======================= Cluster / resources ======================= +NUM_GPU=4 # GB200: ray.sub asserts == gres gpu:4 +export GPUS_PER_NODE=${NUM_GPU} +export CPUS_PER_WORKER=${CPUS_PER_WORKER:-96} # GB200: 144 CPUs but ~720GB FreeMem; ray.sub worker srun is + # --exact + CR_CORE_MEMORY so step mem = RealMem*cpus/144. + # 140 -> ~916GB (> free) = intermittent step-creation OOM; 96 -> ~628GB safe. + +# ============================ Parallelism (FIXED, fits 1 node) ============================ +SKIP_TRAINING="${SKIP_TRAINING:-0}" +if [ "${SKIP_TRAINING}" = "1" ]; then + TP="${TP:-4}"; EP="${EP:-4}"; CP="${CP:-1}"; PP="${PP:-1}"; ETP="${ETP:-1}" # model_parallel=4 (1 node) +else + TP="${TP:-4}"; EP="${EP:-4}"; CP="${CP:-1}"; PP="${PP:-1}"; ETP="${ETP:-1}" # model_parallel=4 (1 node) +fi +VLLM_TP="${VLLM_TP:-2}" +MIN_PAD=1 +if [ ${CP} -gt 1 ]; then MIN_PAD=$((MIN_PAD * CP * 2)); fi +if [ ${TP} -gt 1 ]; then MIN_PAD=$((MIN_PAD * TP)); fi +MAKE_SEQ_DIVISIBLE_BY=${MIN_PAD} + +# ================= Generation-scaling: derive all sizes from R ================= +# Invariant held across R: GPP*PPS/R == GBS/R == SAMPLES_PER_REPLICA (constant). +# GPP (GRPO group size) is held fixed; PPS is derived from R so the invariant holds. +# 30B baseline anchors the invariant at 2: GPP=8, SAMPLES_PER_REPLICA=2 +# -> PPS = 2*R/8 = R/4, GBS = 2*R (R=4 => PPS1/GPP8/GBS8). +GPP="${GPP:-8}" # generations per prompt (GRPO group size, fixed) +SAMPLES_PER_REPLICA="${SAMPLES_PER_REPLICA:-2}" # invariant GBS/R = GPP*PPS/R (const = 2) +BASE_CONCURRENCY=768 +MODEL_PARALLEL=$(( TP * CP * PP )) +EXPERT_TMP=$(( ETP * EP * PP )) + +# ALIGN_BASELINE: reproduce baseline dc3m70us GPU/batch geometry (GBS=64, 32 vLLM +# replicas, 64 train GPU). On GB200 (4 GPU/node) that is 16 gen + 16 train = 32 nodes. +# Only seeds defaults -> NUM_VLLM_REPLICAS / TRAIN_NODES still win if set explicitly. +ALIGN_BASELINE="${ALIGN_BASELINE:-0}" +if [ "${ALIGN_BASELINE}" = "1" ]; then + NUM_VLLM_REPLICAS="${NUM_VLLM_REPLICAS:-32}" # 32 replicas -> GEN_NODES=16, GBS=2*32=64 + TRAIN_NODES="${TRAIN_NODES:-16}" # 64 train GPU -> train_DP=64/MODEL_PARALLEL=16 +fi + +NUM_VLLM_REPLICAS="${NUM_VLLM_REPLICAS:-}" +if [ -z "${NUM_VLLM_REPLICAS}" ]; then + echo "ERROR: NUM_VLLM_REPLICAS is required (number of vLLM replicas). e.g. NUM_VLLM_REPLICAS=4" >&2 + exit 1 +fi + +gcd() { local a=$1 b=$2 t; while [ ${b} -ne 0 ]; do t=${b}; b=$(( a % b )); a=${t}; done; echo ${a}; } +lcm() { echo $(( $1 / $(gcd $1 $2) * $2 )); } + +R_STEP_GEN=$(( NUM_GPU / $(gcd ${VLLM_TP} ${NUM_GPU}) )) +R_STEP_PPS=$(( GPP / $(gcd ${SAMPLES_PER_REPLICA} ${GPP}) )) +R_STEP=$(lcm ${R_STEP_GEN} ${R_STEP_PPS}) +if [ $(( NUM_VLLM_REPLICAS % R_STEP )) -ne 0 ] || [ ${NUM_VLLM_REPLICAS} -lt ${R_STEP} ]; then + echo "ERROR: NUM_VLLM_REPLICAS must be a positive multiple of ${R_STEP} (got ${NUM_VLLM_REPLICAS})." >&2 + exit 1 +fi + +GEN_GPUS=$(( NUM_VLLM_REPLICAS * VLLM_TP )) +GEN_NODES=$(( GEN_GPUS / NUM_GPU )) +if [ "${SKIP_TRAINING}" = "1" ]; then + TRAIN_NODES="${TRAIN_NODES:-1}" +else + TRAIN_NODES="${TRAIN_NODES:-$(( MODEL_PARALLEL / NUM_GPU ))}" + if [ ${TRAIN_NODES} -lt 1 ]; then TRAIN_NODES=1; fi +fi +TOTAL_NODES=$(( TRAIN_NODES + GEN_NODES )) +PPS=$(( SAMPLES_PER_REPLICA * NUM_VLLM_REPLICAS / GPP )) +if [ ${PPS} -lt 1 ]; then PPS=1; fi +GBS=$(( PPS * GPP )) +CONCURRENCY=$(( GBS * 1 )) +if [ ${CONCURRENCY} -lt ${BASE_CONCURRENCY} ]; then CONCURRENCY=${BASE_CONCURRENCY}; fi + +TRAIN_WORLD=$(( TRAIN_NODES * NUM_GPU )) +if [ $(( TRAIN_WORLD % MODEL_PARALLEL )) -ne 0 ] || [ $(( TRAIN_WORLD % EXPERT_TMP )) -ne 0 ]; then + echo "ERROR: train world ${TRAIN_WORLD} (TRAIN_NODES=${TRAIN_NODES}) not divisible by model-parallel ${MODEL_PARALLEL} / expert ${EXPERT_TMP}." >&2 + exit 1 +fi +TRAIN_DP=$(( TRAIN_WORLD / MODEL_PARALLEL )) +if [ $(( GBS % TRAIN_DP )) -ne 0 ]; then + echo "ERROR: GBS ${GBS} not divisible by train DP ${TRAIN_DP}." >&2 + exit 1 +fi +PER_GPU_BATCH=$(( GBS / TRAIN_DP )) +PER_REPLICA_SAMPLES=$(( GBS / NUM_VLLM_REPLICAS )) + +# ===================== Sequence length & packing ===================== +SEQLEN="${SEQLEN:-131072}" +SEQUENCE_PACKING=True + +# ================= Sync/Async mode & async GRPO settings ================= +ASYNC_GRPO_ENABLED=True +MAX_TRAJECTORY_AGE_STEPS=1 +FORCE_ON_POLICY_RATIO=True +INFLIGHT_WEIGHT_UPDATE=True +RECOMPUTE_KV_CACHE_AFTER_WEIGHT_UPDATES=False +SEQ_LOGPROB_ERROR_THRESHOLD=null +if [ "${ASYNC_GRPO_ENABLED}" = "True" ]; then + COLOCATED_ENABLED=False + VLLM_GPU_UTIL=0.8 + OVERLAP_GRAD_REDUCE=False + ADVANTAGE_CLIP_LOW=-100 + ADVANTAGE_CLIP_HIGH=100 + TIS_THRESHOLD=5 +else + COLOCATED_ENABLED=True + VLLM_GPU_UTIL=0.5 + OVERLAP_GRAD_REDUCE=True +fi + +# ========================= GRPO / sampling ========================= +NORMALIZE_REWARDS=True +OVERLONG_FILTERING=True + +# ========================== Loss function ========================== +KL=0 +CLIP_MIN=0.2 +CLIP_MAX=0.28 +USE_ON_POLICY_KL_APPROXIMATION=True +IMPORTANCE_SAMPLING_CORRECTION=True +SEQ_LEVEL_IS=False +TOKEN_LEVEL_LOSS=True + +# ============================ Optimizer ============================ +LR="${LR:-1e-06}" + +# =============================== MoE =============================== +MOE_FREEZE_ROUTER=True +MOE_PERMUTE_FUSION=True +MOE_ENABLE_DEEPEP=False +MOE_TOKEN_DISPATCHER_TYPE="alltoall" +MOE_AUX_LOSS_COEFF=0 +MOE_ROUTER_LOAD_BALANCING_TYPE="none" +MOE_ROUTER_BIAS_UPDATE_RATE="1e-3" + +# ======================= Generation / vLLM ======================= +TEMPERATURE=1.0 + +# =================== Checkpointing & validation =================== +SAVE_PERIOD=5 +VAL_PERIOD=1000 +KEEP_TOP_K=2 + +# ============================ SWE agent ============================ +AGENT_MAX_TURNS="${AGENT_MAX_TURNS:-200}" +AGENT_TIMEOUT="${AGENT_TIMEOUT:-1800}" + +# ============================== Logging ============================== +WANDB_PROJ="${WANDB_PROJ:-swe-benchmark}" +WANDB_GROUP="${WANDB_GROUP:-qwen3-30b-gb200-swe-gen-scale}" +LOG_GYM_RESPONSES=true + +# ========================= SLURM submission ========================= +SBATCH_ACCOUNT="nemotron_sw_post" +SBATCH_PARTITION="batch" +SBATCH_TIME="${SBATCH_TIME:-4:0:0}" +MAX_NUM_STEPS="${MAX_NUM_STEPS:-3}" +# Free-form passthrough appended verbatim to COMMAND (extra ++policy.megatron_cfg.* knobs, etc.). Empty by default. +EXTRA_ARGS="${EXTRA_ARGS:-}" + +# ========================= Experiment naming ========================= +if [ "${ASYNC_GRPO_ENABLED}" = "True" ]; then + SYNC_MODE="async-age${MAX_TRAJECTORY_AGE_STEPS}" +else + SYNC_MODE="sync" +fi +EXP_SUFFIX="${EXP_SUFFIX:-qwen3-30b-gb200-swe-genscale-${SYNC_MODE}-genrep${NUM_VLLM_REPLICAS}-nodes${TOTAL_NODES}-tp${TP}cp${CP}ep${EP}pp${PP}-pps${PPS}-gpp${GPP}-gbs${GBS}-seq${SEQLEN}-lr${LR}}" +WANDB_NAME="${EXP_SUFFIX}" +CHECKPOINT_DIR="${CHECKPOINT_ROOT}/${EXP_SUFFIX}" +SNAPSHOT_DIR="${REPO_ROOT}" + +mkdir -p "${CHECKPOINT_DIR}" + +# ============= Unified SLURM/Ray log location ============= +export BASE_LOG_DIR="${BASE_LOG_DIR:-${SNAPSHOT_DIR}/logs/qwen3_30b_swe_scale}" +mkdir -p "${BASE_LOG_DIR}" + +# ========================= Environment variables ========================= +if [ -f "${HOME}/script/export_env_vars.sh" ]; then + # shellcheck disable=SC1090 + source "${HOME}/script/export_env_vars.sh" +fi +export HUGGINGFACE_TOKEN="${HUGGINGFACE_TOKEN:-${HF_TOKEN}}" +export GITLAB_TOKEN="${GITLAB_TOKEN:-}" +export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}" +export UV_CACHE_DIR=/tmp/uv_cache +export LUSTRE_UV_CACHE_SEED="${LUSTRE_UV_CACHE_SEED:-}" +export UV_LOCK_TIMEOUT=3600 +export RAY_DEDUP_LOGS=1 +export SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt +export REQUESTS_CA_BUNDLE=/etc/ssl/certs/ca-certificates.crt +export CURL_CA_BUNDLE=/etc/ssl/certs/ca-certificates.crt +export OMP_NUM_THREADS=16 + +# ========================= Node-local cache config ========================= +# HOME has a 10G quota on this cluster -> persistent caches live on Lustre. +PERSISTENT_CACHE="${PERSISTENT_CACHE:-/lustre/fsw/portfolios/nemotron/users/ruit/.cache/qwen3_30b_thinking_swe_scale}" +export LUSTRE_VLLM_CACHE="${PERSISTENT_CACHE}/vllm_compile_cache" +export LUSTRE_INDUCTOR_CACHE="${PERSISTENT_CACHE}/inductor_cache" +export LUSTRE_TRITON_CACHE="${PERSISTENT_CACHE}/triton_cache" +# Seed the driver's node-local /tmp/uv_cache from a warm lustre copy so `uv run` +# hits the prebuilt transformer-engine wheel instead of recompiling (~20-30min). +# Populate this dir once from a completed head-node /tmp/uv_cache (see repo notes). +export LUSTRE_UV_CACHE_SEED="${LUSTRE_UV_CACHE_SEED:-${PERSISTENT_CACHE}/uv_cache_seed}" +export NRL_VLLM_LOCAL_CACHE_DIR="/tmp/nemo_rl_vllm_cache" +export NRL_VLLM_CACHE_SEED_DIR="/tmp/nemo_rl_vllm_cache_warm" +export INDUCTOR_CACHE_DIR="/tmp/nemo_rl_inductor_cache" +export TRITON_CACHE_DIR="/tmp/nemo_rl_triton_cache" +export CACHE_SYNC_FREQUENCY=120 +mkdir -p "${LUSTRE_VLLM_CACHE}" "${LUSTRE_INDUCTOR_CACHE}" "${LUSTRE_TRITON_CACHE}" + +# ============================== Summary ============================== +echo "==========================================" +echo "Qwen3-30B GB200 SWE generation-scaling | Experiment: ${EXP_SUFFIX}" +echo "Mode: ${SYNC_MODE}, Colocated: ${COLOCATED_ENABLED}, SkipTraining: ${SKIP_TRAINING}, AlignBaseline: ${ALIGN_BASELINE}" +echo "wandb: project=${WANDB_PROJ}, group=${WANDB_GROUP}, name=${WANDB_NAME}" +echo "------------------------------------------" +echo "Scaling input: NUM_VLLM_REPLICAS = ${NUM_VLLM_REPLICAS} (R-step=${R_STEP})" +echo " vllm_tp = ${VLLM_TP} (nodes/replica = ${VLLM_TP}/${NUM_GPU})" +echo " GEN_NODES = ${GEN_NODES}" +echo " TRAIN_NODES = ${TRAIN_NODES} (train_world=${TRAIN_WORLD}, model_parallel=${MODEL_PARALLEL}, train_DP=${TRAIN_DP})" +echo " TOTAL_NODES = ${TOTAL_NODES}" +echo " PPS = ${PPS}" +echo " GPP = ${GPP}" +echo " GBS = ${GBS}" +echo " CONCURRENCY = ${CONCURRENCY}" +echo " invariants : samples/replica=${PER_REPLICA_SAMPLES}, batch/train-GPU=${PER_GPU_BATCH}" +echo "Parallelism: TP=${TP}, EP=${EP}, CP=${CP}, PP=${PP}, ETP=${ETP}, vLLM_TP=${VLLM_TP}, pad=${MAKE_SEQ_DIVISIBLE_BY}" +echo "SeqLen: ${SEQLEN}" +echo "Model: ${MODEL_PATH}" +echo "Container: ${CONTAINER}" +echo "Checkpoint: ${CHECKPOINT_DIR}" +echo "==========================================" + +cd "${SNAPSHOT_DIR}" + +# ================ SETUP_COMMAND (self-skips apptainer install if baked; seed caches) ================ +read -r -d '' SETUP_COMMAND </dev/null 2>&1 || command -v singularity >/dev/null 2>&1; then + echo "[SETUP] singularity/apptainer already available" + RET=0 + break + fi + apt-get update && apt-get install -y git build-essential gcc wget 2>/dev/null || true + cd /tmp && \ + wget --no-check-certificate -q https://github.com/apptainer/apptainer/releases/download/v1.3.1/apptainer_1.3.1_arm64.deb && \ + apt install -y ./apptainer_1.3.1_arm64.deb && \ + ln -sf /usr/bin/apptainer /usr/bin/singularity + if command -v apptainer >/dev/null 2>&1; then + echo "[SETUP] apptainer installed successfully" + RET=0 + break + fi + echo "[SETUP] apptainer install attempt \$attempt failed, retrying..." + sleep 10 +done +if [ \$RET -ne 0 ]; then + echo "[SETUP] WARNING: apptainer not available after \$RETRIES attempts" +fi + +echo "[CACHE SEED] Clearing stale /tmp caches and seeding from Lustre..." +rm -rf /tmp/nemo_rl_vllm_cache /tmp/nemo_rl_vllm_cache_* +rm -rf "${INDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}" +mkdir -p "${INDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}" + +find "${LUSTRE_INDUCTOR_CACHE}" -maxdepth 1 -name '.tmp_*' -mmin +30 -exec rm -rf {} + 2>/dev/null || true +find "${LUSTRE_TRITON_CACHE}" -maxdepth 1 -name '.tmp_*' -mmin +30 -exec rm -rf {} + 2>/dev/null || true + +_seed_cache() { + local lustre="\$1" local_dir="\$2" name="\$3" + if [ -d "\$lustre" ] && [ "\$(ls -A "\$lustre" 2>/dev/null)" ]; then + rsync -a --exclude '.tmp_*' "\$lustre/" "\$local_dir/" 2>/dev/null \ + && echo "[CACHE SEED] \$name: seeded from Lustre" \ + || echo "[CACHE SEED] \$name: seed failed (non-fatal)" + else + echo "[CACHE SEED] \$name: no warm cache on Lustre yet" + fi +} + +_seed_cache "${LUSTRE_INDUCTOR_CACHE}" "${INDUCTOR_CACHE_DIR}" "Inductor" +_seed_cache "${LUSTRE_TRITON_CACHE}" "${TRITON_CACHE_DIR}" "Triton" +mkdir -p /tmp/uv_cache +_seed_cache "${LUSTRE_UV_CACHE_SEED}" "/tmp/uv_cache" "uv (prebuilt transformer-engine)" +echo "[CACHE SEED] Done." + +UV_HTTP_TIMEOUT=3600 \ + uv sync --frozen --extra mcore +SETUPEOF +export SETUP_COMMAND + +# ================ Training command ================ +export COMMAND="NRL_VLLM_USE_V1=1 \ + NRL_WG_USE_RAY_REF=1 \ + WANDB_API_KEY=${WANDB_API_KEY} \ + HUGGINGFACE_TOKEN=${HUGGINGFACE_TOKEN} \ + GITHUB_TOKEN=${GITHUB_TOKEN} \ + GITLAB_TOKEN=${GITLAB_TOKEN} \ + HF_HOME=${HF_HOME} \ + HF_DATASETS_CACHE=${HF_DATASETS_CACHE} \ + UV_CACHE_DIR=${UV_CACHE_DIR} \ + VLLM_ATTENTION_BACKEND=FLASH_ATTN \ + VLLM_CACHE_ROOT=${LUSTRE_VLLM_CACHE} \ + DG_JIT_CACHE_DIR=${LUSTRE_VLLM_CACHE}/deep_gemm \ + VLLM_DEEP_GEMM_WARMUP=skip \ + NRL_FORCE_REBUILD_VENVS=false \ + NRL_IGNORE_VERSION_MISMATCH=1 \ + RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 \ + UV_HTTP_TIMEOUT=3600 \ + UV_LOCK_TIMEOUT=900 \ + TORCH_CUDA_ARCH_LIST='10.0' \ + NEMO_GYM_SKIP_VENV_IF_PRESENT=1 \ + uv run --frozen --extra mcore ./examples/nemo_gym/run_grpo_nemo_gym.py \ + --config=${CONFIG_FILE} \ + cluster.num_nodes=${TOTAL_NODES} \ + cluster.gpus_per_node=${NUM_GPU} \ + ++data.train.data_path=${TRAIN_DATA_PATH} \ + ++data.validation.data_path=${VAL_DATA_PATH} \ + grpo.num_prompts_per_step=${PPS} \ + grpo.num_generations_per_prompt=${GPP} \ + grpo.val_at_start=False \ + grpo.normalize_rewards=${NORMALIZE_REWARDS} \ + grpo.overlong_filtering=${OVERLONG_FILTERING} \ + grpo.val_period=${VAL_PERIOD} \ + grpo.seq_logprob_error_threshold=${SEQ_LOGPROB_ERROR_THRESHOLD} \ + grpo.async_grpo.enabled=${ASYNC_GRPO_ENABLED} \ + grpo.async_grpo.in_flight_weight_updates=${INFLIGHT_WEIGHT_UPDATE} \ + grpo.async_grpo.recompute_kv_cache_after_weight_updates=${RECOMPUTE_KV_CACHE_AFTER_WEIGHT_UPDATES} \ + grpo.async_grpo.max_trajectory_age_steps=${MAX_TRAJECTORY_AGE_STEPS} \ + env.should_log_nemo_gym_responses=${LOG_GYM_RESPONSES} \ + policy.generation.colocated.enabled=${COLOCATED_ENABLED} \ + policy.model_name=${MODEL_PATH} \ + policy.max_total_sequence_length=${SEQLEN} \ + policy.dynamic_batching.enabled=False \ + policy.train_global_batch_size=${GBS} \ + policy.make_sequence_length_divisible_by=${MAKE_SEQ_DIVISIBLE_BY} \ + policy.offload_optimizer_for_logprob=true \ + policy.sequence_packing.enabled=${SEQUENCE_PACKING} \ + policy.megatron_cfg.tensor_model_parallel_size=${TP} \ + policy.megatron_cfg.expert_model_parallel_size=${EP} \ + policy.megatron_cfg.expert_tensor_parallel_size=${ETP} \ + policy.megatron_cfg.context_parallel_size=${CP} \ + policy.megatron_cfg.pipeline_model_parallel_size=${PP} \ + policy.megatron_cfg.sequence_parallel=True \ + policy.megatron_cfg.bias_activation_fusion=False \ + policy.megatron_cfg.distributed_data_parallel_config.overlap_grad_reduce=${OVERLAP_GRAD_REDUCE} \ + policy.megatron_cfg.moe_permute_fusion=${MOE_PERMUTE_FUSION} \ + policy.megatron_cfg.moe_enable_deepep=${MOE_ENABLE_DEEPEP} \ + policy.megatron_cfg.moe_token_dispatcher_type=${MOE_TOKEN_DISPATCHER_TYPE} \ + policy.megatron_cfg.moe_aux_loss_coeff=${MOE_AUX_LOSS_COEFF} \ + policy.megatron_cfg.moe_router_load_balancing_type=${MOE_ROUTER_LOAD_BALANCING_TYPE} \ + policy.megatron_cfg.moe_router_bias_update_rate=${MOE_ROUTER_BIAS_UPDATE_RATE} \ + policy.megatron_cfg.freeze_moe_router=${MOE_FREEZE_ROUTER} \ + policy.megatron_cfg.optimizer.lr=${LR} \ + policy.megatron_cfg.optimizer.min_lr=${LR} \ + policy.megatron_cfg.optimizer.weight_decay=0 \ + policy.megatron_cfg.empty_unused_memory_level=2 \ + policy.megatron_cfg.activation_checkpointing=True \ + policy.generation.temperature=${TEMPERATURE} \ + policy.generation.vllm_cfg.tensor_parallel_size=${VLLM_TP} \ + policy.generation.vllm_cfg.gpu_memory_utilization=${VLLM_GPU_UTIL} \ + policy.generation.vllm_cfg.skip_tokenizer_init=False \ + loss_fn.reference_policy_kl_penalty=${KL} \ + loss_fn.ratio_clip_min=${CLIP_MIN} \ + loss_fn.ratio_clip_max=${CLIP_MAX} \ + loss_fn.use_on_policy_kl_approximation=${USE_ON_POLICY_KL_APPROXIMATION} \ + loss_fn.use_importance_sampling_correction=${IMPORTANCE_SAMPLING_CORRECTION} \ + loss_fn.sequence_level_importance_ratios=${SEQ_LEVEL_IS} \ + loss_fn.token_level_loss=${TOKEN_LEVEL_LOSS} \ + loss_fn.force_on_policy_ratio=${FORCE_ON_POLICY_RATIO} \ + checkpointing.checkpoint_dir=${CHECKPOINT_DIR} \ + checkpointing.save_period=${SAVE_PERIOD} \ + checkpointing.keep_top_k=${KEEP_TOP_K} \ + ++checkpointing.metric_name=train:total_reward/mean \ + ++checkpointing.checkpoint_must_save_by=00:03:35:00 \ + logger.wandb_enabled=True \ + logger.wandb.name=${WANDB_NAME} \ + logger.wandb.project=${WANDB_PROJ} \ + ++logger.wandb.group=${WANDB_GROUP}" + +if [ "${ASYNC_GRPO_ENABLED}" = "True" ]; then + export COMMAND="${COMMAND} \ + policy.generation.colocated.resources.num_nodes=${GEN_NODES} \ + policy.generation.colocated.resources.gpus_per_node=${NUM_GPU} \ + grpo.advantage_clip_low=${ADVANTAGE_CLIP_LOW} \ + grpo.advantage_clip_high=${ADVANTAGE_CLIP_HIGH} \ + loss_fn.truncated_importance_sampling_ratio=${TIS_THRESHOLD} \ + env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.agent_max_turns=${AGENT_MAX_TURNS} \ + env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.swebench_agent_timeout=${AGENT_TIMEOUT} \ + env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.concurrency=${CONCURRENCY} \ + env.nemo_gym.swe_agents_val.responses_api_agents.swe_agents.agent_max_turns=${AGENT_MAX_TURNS} \ + env.nemo_gym.swe_agents_val.responses_api_agents.swe_agents.swebench_agent_timeout=${AGENT_TIMEOUT} \ + env.nemo_gym.swe_agents_val.responses_api_agents.swe_agents.concurrency=${CONCURRENCY}" +fi + +if [ -n "${MAX_NUM_STEPS}" ]; then + export COMMAND="${COMMAND} grpo.max_num_steps=${MAX_NUM_STEPS}" +fi + +if [ "${SKIP_TRAINING}" = "1" ]; then + export COMMAND="${COMMAND} ++grpo.gen_benchmark_skip_training=true checkpointing.enabled=false" +fi + +# Free-form extra Hydra overrides (appended last so they can override anything above). +if [ -n "${EXTRA_ARGS}" ]; then + export COMMAND="${COMMAND} ${EXTRA_ARGS}" +fi + +# ================ Submit job (skipped under DRY_RUN=1) ================ +if [ "${DRY_RUN:-0}" = "1" ]; then + echo "" + echo "[DRY_RUN] Not submitting. Would run:" + echo "[DRY_RUN] sbatch --nodes=${TOTAL_NODES} --account=${SBATCH_ACCOUNT} --partition=${SBATCH_PARTITION} --time=${SBATCH_TIME} --gres=gpu:${NUM_GPU} ... ray.sub" + echo "" + echo "[DRY_RUN] COMMAND:" + echo "${COMMAND}" + cd - > /dev/null + exit 0 +fi + +sbatch \ + --nodes="${TOTAL_NODES}" \ + --account="${SBATCH_ACCOUNT}" \ + --job-name="${WANDB_NAME}" \ + --partition="${SBATCH_PARTITION}" \ + --time="${SBATCH_TIME}" \ + --gres=gpu:${NUM_GPU} \ + --output="${BASE_LOG_DIR}/slurm-%j.out" \ + --exclusive \ + --dependency=singleton \ + --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"180","reason":"data_loading","description":"Async GRPO Qwen3-30B GB200 SWE generation-scaling"}}' \ + ray.sub | tee /dev/stderr | grep -o '[0-9]\+' > latest_30b_scale_gen_job_id.txt + +JOB_ID="$(cat latest_30b_scale_gen_job_id.txt)" +echo "==========================================" +echo "Job submitted: ${EXP_SUFFIX}" +echo "Job ID: ${JOB_ID}" +echo "wandb group: ${WANDB_GROUP}" +echo "Monitor with: squeue -j ${JOB_ID}" +echo "Ray/SLURM logs: ${BASE_LOG_DIR}/${JOB_ID}-logs/" +echo "Checkpoints: ${CHECKPOINT_DIR}/" +echo "==========================================" + +cd - > /dev/null From 9519e04967112ce79067072d7284285f702cf6bb Mon Sep 17 00:00:00 2001 From: ruit Date: Wed, 1 Jul 2026 02:18:42 -0700 Subject: [PATCH 02/69] feat(swe_bench): self-contained hsg SWE launcher (baked-gym image + triton MoE) - rename config to grpo_qwen3_30b_async_swe_hsg.yaml with the GB200 fixes: megatron_cfg.use_fused_weighted_squared_relu + vllm_kwargs.moe_backend=triton (triton MoE backend avoids the weight-refit broadcast_weights_for_collective hang) - run_grpo_swe2_scale_gen_hsg.sh: point CONFIG_FILE at the hsg yaml and default CONTAINER to the baked nightly-063026-gymvenvs image (has /opt/gym_venvs prebuilt so gym spinup skips the concurrent venv build that deadlocks on the uv cache lock) Signed-off-by: ruit --- ..._async_swe.yaml => grpo_qwen3_30b_async_swe_hsg.yaml} | 7 +++++++ examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh | 9 ++++++--- 2 files changed, 13 insertions(+), 3 deletions(-) rename examples/swe_bench/{grpo_qwen3_30b_async_swe.yaml => grpo_qwen3_30b_async_swe_hsg.yaml} (97%) diff --git a/examples/swe_bench/grpo_qwen3_30b_async_swe.yaml b/examples/swe_bench/grpo_qwen3_30b_async_swe_hsg.yaml similarity index 97% rename from examples/swe_bench/grpo_qwen3_30b_async_swe.yaml rename to examples/swe_bench/grpo_qwen3_30b_async_swe_hsg.yaml index ea485abb375..f0d46ad092e 100644 --- a/examples/swe_bench/grpo_qwen3_30b_async_swe.yaml +++ b/examples/swe_bench/grpo_qwen3_30b_async_swe_hsg.yaml @@ -341,6 +341,13 @@ policy: vllm_kwargs: mamba_ssm_cache_dtype: "float32" + # Use triton MoE backend (matches the reference Qwen3-30B SWE configs + # examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe{1,2}.yaml). The default + # flashinfer MoE backend uses a fused/quantized expert weight layout + # (trtllm_bf16_moe) that does not match the megatron weight-refit broadcast, + # which hangs broadcast_weights_for_collective. triton uses the standard + # expert layout that the refit broadcast expects. + moe_backend: triton compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32,64] diff --git a/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh b/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh index ea161d2160a..8657a7c23d8 100755 --- a/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh +++ b/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh @@ -41,7 +41,7 @@ set -e # ============================ Paths ============================ REPO_ROOT="${REPO_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" -CONFIG_FILE="${CONFIG_FILE:-${REPO_ROOT}/test_assets/qwen-30B/grpo_qwen3_30b_async_swe.yaml}" +CONFIG_FILE="${CONFIG_FILE:-${REPO_ROOT}/examples/swe_bench/grpo_qwen3_30b_async_swe_hsg.yaml}" CHECKPOINT_ROOT="${CHECKPOINT_ROOT:-${REPO_ROOT}/results}" TRAIN_DATA_PATH="${TRAIN_DATA_PATH:-/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/blends/balanced_language.jsonl}" VAL_DATA_PATH="${VAL_DATA_PATH:-/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/swe_public_datasets_val_swebench.jsonl}" @@ -51,8 +51,11 @@ MODEL_PATH="${1:-${MODEL_PATH:-${DEFAULT_MODEL_PATH}}}" # ================ Container and mount config ================ # GB200 (aarch64) baked image: apptainer + /opt/nemo_rl_venv with --extra mcore # (sm_100), built by test_assets/SWE/build_swe_bench_combined.sh. -# export CONTAINER=${CONTAINER:-/lustre/fsw/portfolios/nemotron/users/ruit/enroot-images/ruit-swe_bench-6dc8fabea-aarch64-060426-mcore-apptainer.squashfs} -export CONTAINER=${CONTAINER:-/lustre/fsw/portfolios/nemotron/users/ruit/enroot-images/nemo-rl:nightly-062326.squashfs} +# Baked image: nightly-063026 + /opt/gym_venvs (gym server venvs prebuilt) so gym +# spinup hits skip_venv_if_present and does NOT concurrently build venvs (which +# deadlocks on the uv cache lock). Built by test_assets/SWE/build_swe_bench_combined.sh. +# Older non-baked images (e.g. nightly-062326) WILL hang at gym spinup — do not use. +export CONTAINER=${CONTAINER:-/lustre/fsw/portfolios/nemotron/users/ruit/enroot-images/nemo-rl:nightly-063026-gymvenvs.squashfs} GYM_CODE="${REPO_ROOT}/3rdparty/Gym-workspace/Gym" export MOUNTS="/lustre:/lustre,$PWD:$PWD,${GYM_CODE}:/opt/nemo-rl/3rdparty/Gym-workspace/Gym" From 1392b1375d77550fb8f1fbc9ae0028a89ed0adb6 Mon Sep 17 00:00:00 2001 From: ruit Date: Wed, 1 Jul 2026 02:19:39 -0700 Subject: [PATCH 03/69] chore(swe_bench): don't source personal env script; document required creds Remove the hardcoded source of ${HOME}/script/export_env_vars.sh (a personal file). Users export their own HF_HOME/HF_TOKEN/WANDB_API_KEY/etc. before running. Signed-off-by: ruit --- examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh b/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh index 8657a7c23d8..fb7469f6df0 100755 --- a/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh +++ b/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh @@ -236,10 +236,12 @@ export BASE_LOG_DIR="${BASE_LOG_DIR:-${SNAPSHOT_DIR}/logs/qwen3_30b_swe_scale}" mkdir -p "${BASE_LOG_DIR}" # ========================= Environment variables ========================= -if [ -f "${HOME}/script/export_env_vars.sh" ]; then - # shellcheck disable=SC1090 - source "${HOME}/script/export_env_vars.sh" -fi +# Credentials are NOT sourced here. Before running, export your own (e.g. in your +# shell or a personal env script you source yourself): +# HF_HOME=... # HuggingFace cache dir +# HF_TOKEN=... # HuggingFace token (used for HUGGINGFACE_TOKEN below) +# WANDB_API_KEY=... # Weights & Biases API key +# GITHUB_TOKEN=... GITLAB_TOKEN=... # optional, for git-dep rate limits export HUGGINGFACE_TOKEN="${HUGGINGFACE_TOKEN:-${HF_TOKEN}}" export GITLAB_TOKEN="${GITLAB_TOKEN:-}" export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}" From d5ee3df5d2249fa6d9b07b2baaf5e2322f444ed8 Mon Sep 17 00:00:00 2001 From: ruit Date: Wed, 1 Jul 2026 02:26:47 -0700 Subject: [PATCH 04/69] chore(swe_bench): default model/data to SWE1 step_230 + r2e_gym_subset Set DEFAULT_MODEL_PATH to the SWE1 step_230_hf checkpoint and TRAIN/VAL_DATA_PATH to the r2e_gym_subset dataset (train=val) so the hsg launcher reproduces the baseline SWE2 run by default. All three remain overridable via env. Signed-off-by: ruit --- examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh b/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh index fb7469f6df0..ef461cc2940 100755 --- a/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh +++ b/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh @@ -43,9 +43,9 @@ set -e REPO_ROOT="${REPO_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" CONFIG_FILE="${CONFIG_FILE:-${REPO_ROOT}/examples/swe_bench/grpo_qwen3_30b_async_swe_hsg.yaml}" CHECKPOINT_ROOT="${CHECKPOINT_ROOT:-${REPO_ROOT}/results}" -TRAIN_DATA_PATH="${TRAIN_DATA_PATH:-/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/blends/balanced_language.jsonl}" -VAL_DATA_PATH="${VAL_DATA_PATH:-/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/swe_public_datasets_val_swebench.jsonl}" -DEFAULT_MODEL_PATH="/lustre/fsw/portfolios/llmservice/users/igitman/hf_models/Qwen3-30B-A3B-Thinking-2507" +TRAIN_DATA_PATH="${TRAIN_DATA_PATH:-/lustre/fsw/portfolios/nemotron/users/ruit/evolution_rl/test_assets/qwen-30B/data/swe_all_datasets_train_w_agent_ref_r2e_gym_subset.jsonl}" +VAL_DATA_PATH="${VAL_DATA_PATH:-/lustre/fsw/portfolios/nemotron/users/ruit/evolution_rl/test_assets/qwen-30B/data/swe_all_datasets_train_w_agent_ref_r2e_gym_subset.jsonl}" +DEFAULT_MODEL_PATH="/lustre/fsw/portfolios/nemotron/users/ruit/evolution_rl/test_assets/qwen-30B/bihu/qwen3-30b-thinking-swe1-async-age1-pps64-gpp8-gbs512-lr1e-06/step_230_hf" MODEL_PATH="${1:-${MODEL_PATH:-${DEFAULT_MODEL_PATH}}}" # ================ Container and mount config ================ From a6039a06d3e5d30bde3324e2a34e7b1984d9c401 Mon Sep 17 00:00:00 2001 From: ruit Date: Wed, 1 Jul 2026 07:49:57 -0700 Subject: [PATCH 05/69] chore(swe_bench): revert default data to balanced_language + swe_public_val The r2e_gym_subset set has instances with no matching sif on this cluster (missing container -> 500, batch never fills -> step stalls). Revert the launcher defaults to the runnable balanced_language (train) + swe_public_val. Signed-off-by: ruit --- examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh b/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh index ef461cc2940..9e716ecbc00 100755 --- a/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh +++ b/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh @@ -43,8 +43,8 @@ set -e REPO_ROOT="${REPO_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" CONFIG_FILE="${CONFIG_FILE:-${REPO_ROOT}/examples/swe_bench/grpo_qwen3_30b_async_swe_hsg.yaml}" CHECKPOINT_ROOT="${CHECKPOINT_ROOT:-${REPO_ROOT}/results}" -TRAIN_DATA_PATH="${TRAIN_DATA_PATH:-/lustre/fsw/portfolios/nemotron/users/ruit/evolution_rl/test_assets/qwen-30B/data/swe_all_datasets_train_w_agent_ref_r2e_gym_subset.jsonl}" -VAL_DATA_PATH="${VAL_DATA_PATH:-/lustre/fsw/portfolios/nemotron/users/ruit/evolution_rl/test_assets/qwen-30B/data/swe_all_datasets_train_w_agent_ref_r2e_gym_subset.jsonl}" +TRAIN_DATA_PATH="${TRAIN_DATA_PATH:-/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/blends/balanced_language.jsonl}" +VAL_DATA_PATH="${VAL_DATA_PATH:-/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/swe_public_datasets_val_swebench.jsonl}" DEFAULT_MODEL_PATH="/lustre/fsw/portfolios/nemotron/users/ruit/evolution_rl/test_assets/qwen-30B/bihu/qwen3-30b-thinking-swe1-async-age1-pps64-gpp8-gbs512-lr1e-06/step_230_hf" MODEL_PATH="${1:-${MODEL_PATH:-${DEFAULT_MODEL_PATH}}}" From 719590020a45b3440be8015bec83751c6044f416 Mon Sep 17 00:00:00 2001 From: ruit Date: Thu, 2 Jul 2026 02:27:13 -0700 Subject: [PATCH 06/69] feat(swe_bench): nano V3.5 SWE scale-gen launcher with gen-benchmark skip-training MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - add examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh: single-knob (NUM_VLLM_REPLICAS) launcher, 3 modes — ALIGN_BASELINE (32-node baseline, CP16/EP32), default (scale train+gen from the validated CP2/EP8 2-node/DP base), and SKIP_TRAINING (gen benchmark: pin train to 1 node, scale only gen). Segment auto-picks the largest of {16,8,4,2,1} dividing each worker group; wandb staging redirected off the coreai quota; 60min idle-reaper exemption. - add examples/swe_bench/grpo_nano_v3_5_swe_hsg.yaml. - gen-benchmark skip-training support (NRL_GEN_BENCHMARK_SKIP_TRAINING) in grpo.py / lm_policy.py / base_policy_worker.py (per NVIDIA-NeMo/RL#2930): skip optimizer build + policy.train(), NCCL-free keep-alive matmul to avoid the idle-GPU reaper while training is skipped. - MegatronConfig.mtp_num_layers accepts None (int | None) so MTP can be disabled on hybrid (NemotronH) models where 0 still trips the mtp_num_layers>0 assertion. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ruit --- .../swe_bench/grpo_nano_v3_5_swe_hsg.yaml | 476 +++++++++++ .../run_grpo_nano_v3_5_swe_scale_gen_hsg.sh | 748 ++++++++++++++++++ nemo_rl/algorithms/grpo.py | 76 +- nemo_rl/models/policy/__init__.py | 7 +- nemo_rl/models/policy/lm_policy.py | 13 + .../policy/workers/base_policy_worker.py | 29 + 6 files changed, 1336 insertions(+), 13 deletions(-) create mode 100644 examples/swe_bench/grpo_nano_v3_5_swe_hsg.yaml create mode 100755 examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh diff --git a/examples/swe_bench/grpo_nano_v3_5_swe_hsg.yaml b/examples/swe_bench/grpo_nano_v3_5_swe_hsg.yaml new file mode 100644 index 00000000000..f3bc6d547ba --- /dev/null +++ b/examples/swe_bench/grpo_nano_v3_5_swe_hsg.yaml @@ -0,0 +1,476 @@ +# ============================================================================= +# GRPO Ultra E2E — Production SWE Config for GB200 NVL72 +# ============================================================================= +# Production config for Ultra V3 SWE end-to-end GRPO training on +# 128 nodes × 4 GPUs/node (64 gen + 64 train by default). +# +# Static algorithm, loss, and environment settings live here. +# Per-run overrides (model path, data paths, parallelism, precision) +# are applied by the launch script (repro_ultra_e2e.sh) or Hydra CLI. +# +# - gpus_per_node: 4 +# - TP: 8, CP: 8, EP: 32, PP: 1 +# - vLLM TP: 8 +# - Non-colocated async inference with 64 generation nodes +# - Max sequence length: 131072 +# ============================================================================= + +# ============================================================================= +# Cluster — overridden by launch script for parallelism changes +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 128 + segment_size: 16 + +# ============================================================================= +# Checkpointing +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/grpo_ultra_e2e" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 3 + ft_keep_latest_k: 1 + ft_save_period: 1 + checkpoint_must_save_by: "00:03:30:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# GRPO Algorithm +# ============================================================================= +grpo: + num_prompts_per_step: 16 + num_generations_per_prompt: 32 + num_val_generations_per_prompt: 2 + max_rollout_turns: 1 + max_num_epochs: 4 + max_num_steps: 1000000 + normalize_rewards: true + use_leave_one_out_baseline: true + advantage_clip_low: -100 + advantage_clip_high: 100 + val_period: 10000 + val_at_start: false + val_at_end: false + overlong_filtering: false + max_val_samples: null + val_batch_size: 256 + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + penalize_invalid_tool_call: true + invalid_tool_call_advantage: -5.0 + penalize_malformed_thinking: true + malformed_thinking_advantage: -5.0 + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + + use_best_at_k: false + best_at_k_k: 8 + best_at_k_m: 1000 + + use_combined_training: false + combined_training_weight_mode: "auto" + combined_training_best_at_k_weight: 0.2 + combined_training_pass_at_1_weight: 1.0 + + dynamic_sampling_oversample_ratio: 1.0 + seq_logprob_error_threshold: 2 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: null + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + force_on_policy_ratio: true + use_kl_in_reward: false + +# ============================================================================= +# Policy +# ============================================================================= +policy: + model_name: "/lustre/fsw/portfolios/llmservice/users/wdai/megatron-lm-ultra/checkpoints/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-288k-nano-loss-032026/iter_0007000/hf" + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 512 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 131072 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 32 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + use_gloo_process_groups: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + # community_import.py reads megatron_cfg["gradient_accumulation_fusion"] directly + # (no default) on the HF->megatron conversion path; absent here it raises KeyError. + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled + mtp_loss_scaling_factor: 0.0 + mtp_use_repeated_layer: true + # null disables MTP entirely (matches the upstream mtp=0 run's behavior). This + # checkout's hybrid model gates MTP on `mtp_num_layers is not None` + # (hybrid_model.py:544): 0 still enters the MTP block and asserts > 0, and 1 runs + # MTP forward on only the last-stage ranks (suspected cause of the EXPERT_MODEL_ + # PARALLEL NCCL collective timeout). null makes that check False -> MTP is skipped + # on all ranks, no partial-rank EP collective. Requires the schema change making + # MegatronConfig.mtp_num_layers accept None (nemo_rl/models/policy/__init__.py). + mtp_num_layers: null + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 3.0e-6 + min_lr: 3.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 3e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + # FP8 — overridden by precision recipe in launch script + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo + # were dropped from this checkout's megatron-bridge CheckpointConfig (submodule + # 554c7b9); leaving them in raises InstantiationException "Unexpected config keys". + # Removed for the current-repo reproduction. + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + # Must be a multiple of minimum_pad_factor = cp_size*2*tp_size (see + # nemo_rl/models/megatron/data.py). With CP=16 + TP=4 + sequence_parallel that is + # 128; the old value (=TP=4) made get_logprobs assert and abort every training + # step. Formula matches examples/nemo_gym/grpo_nanov3.yaml (TP*CP*2). + make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 11001 + port_range_high: 15000 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 1 + gpu_memory_utilization: 0.8 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + # reasoning_parser_plugin lives at vllm_cfg top level in this checkout + # (vllm_worker_async reads vllm_cfg["reasoning_parser_plugin"]); passing it + # inside http_server_serving_chat_kwargs makes OpenAIServingChat raise + # "unexpected keyword argument 'reasoning_parser_plugin'". Path updated to this + # repo's location (the parser moved out of nemo_rl/utils/). + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + + vllm_kwargs: + attention_backend: FLASH_ATTN + moe_backend: triton + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 64 + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + train: + data_path: null # Set by launch script + validation: + data_path: null # Set by launch script + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + SWE Agents +# ============================================================================= +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + # Reuse pre-built gym venvs. With false, all 3 gym servers (policy_model, + # swe_agents_{train,val}) rebuild the editable nemo-gym package in parallel and + # contend on the same uv distribution-cache lock on Lustre; policy_model then + # times out (default 300s) and NemoGym spinup fails. Mirrors the qwen-30b SWE run. + skip_venv_if_present: true + num_gpu_nodes: 0 + port_range_low: 15001 + port_range_high: 20000 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml + swe_agents_train: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}} + swebench_agent_timeout: 3600 + run_with_mixed_prompts: true + dataset_path: ${data.train.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swerebench/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/nv_internal/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/r2e_gym/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swegym/sweb.eval.arm64.{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swebench/swe-bench.eval.arm64.{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/mercor/swebenchpro_ots/{instance_id}.sif" + swe_agents_val: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}} + swebench_agent_timeout: 3600 + dataset_path: ${data.validation.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter} + use_absolute_ip: true + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "logs" + num_val_samples_to_print: 0 + wandb_enabled: true + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "ultra-v3-swe-e2e" + name: "grpo-ultra-e2e" + tensorboard: {} + mlflow: + experiment_name: "grpo-ultra-e2e" + run_name: "grpo-ultra-e2e" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + +# ============================================================================= +# Token IDs (model-specific, used by token-based penalties) +# ============================================================================= +token_ids: + eos: 2 # + think_open: 12 # + think_close: 13 # + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +penalize_duplicated_reasoning: true # reasoning content == final answer +penalize_empty_final_answer: true # last message output has empty content +penalize_eos_token: true # eos token appears in generation +penalize_malformed_think_tag: true # /<\/think> count != 1 per turn diff --git a/examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh b/examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh new file mode 100755 index 00000000000..178d833b131 --- /dev/null +++ b/examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh @@ -0,0 +1,748 @@ +#!/bin/bash +# ============================================================================ +# nano V3.5 SWE-e2e GRPO reproduction launcher (my conventions). +# +# Adapted from sdevare's repro_nano.sh +# (.../code_snapshots/sdd-swe-e2e-nano-...-gen-nodes-32/scripts/repro_nano.sh) +# but rewritten to match this checkout's launch pattern (see +# test_assets/ultra_SWE/repro_ultra_launch.sh): +# - SLURM account nemotron_sw_post, my container / results / logs paths. +# - Secrets sourced from ~/script/export_env_vars.sh (no hardcoded keys). +# - DRY_RUN=1 default with secret redaction; require_path preflight checks. +# - SIF container_formatter + apptainer setup + compile-cache sidecar. +# +# Run shape (matches the nano gen-nodes-32 snapshot): +# - Model: mopd_ultrav3_to_nanov3_5_repro_v5 step_18 (KD opt full). +# - Geometry: train TP4/CP16/EP32/PP1/ETP1, vLLM TP4/PP1, max_len 196608. +# - Nodes: 32 train + 32 generation (non-colocated, yaml default), gym=0. +# - Batch: PPS=32, GPP=16, GBS=512. +# - Precision: bf16 by default; mxfp8 recipes via PRECISION_RECIPE (see below). +# +# Examples: +# DRY_RUN=1 bash test_assets/nanoV3_5/repro_nano.sh +# MAX_NUM_STEPS=4 DRY_RUN=1 bash test_assets/nanoV3_5/repro_nano.sh +# PRECISION_RECIPE=mxfp8-rollout DRY_RUN=1 bash test_assets/nanoV3_5/repro_nano.sh +# bash test_assets/nanoV3_5/repro_nano.sh # real submit +# SLURM_QOS= SLURM_PARTITION=batch bash test_assets/nanoV3_5/repro_nano.sh +# ============================================================================ + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# nanoV3_5 lives two levels under the repo root (test_assets/nanoV3_5/). +REPO_ROOT="${REPO_ROOT:-$(cd "${SCRIPT_DIR}/../.." && pwd)}" + +# ====================== Source-commit validation ====================== +# This launcher reproduces the W&B run logged in nvidia/ultra-v3-swe-e2e +# (run name == EXP_SUFFIX below: sdd-swe-e2e-nano-ultra-prod-nano-swe-e2e-... +# -gen-nodes-32), which ran from this commit. The submodule revisions are pinned +# deterministically by this superproject commit, so verifying HEAD is sufficient: +# Automodel 1d42deb98169fd94b54c714c0fe4bf308fe7115a +# Gym bd1050cf07b0a0aceb4a63614781f973e74dcdfe +# Megatron-Bridge 4e0209d36597f91026ed5ba42967ad3d9e8ea705 +# Megatron-LM ceb31e61faed0a2e844c1495a877163ddb253694 +# Set ALLOW_COMMIT_MISMATCH=1 to run intentionally from different code. Set +# EXPECTED_COMMIT= (empty) to skip the check entirely. +# +# NOTE (current-repo reproduction): this copy lives in a different checkout/branch +# than the upstream run above (HEAD here is not 9d08c1b2...), so the check defaults +# to empty (skipped). The upstream commit is kept above for provenance. Export +# EXPECTED_COMMIT= to re-enable pinning against a specific commit. +EXPECTED_COMMIT="${EXPECTED_COMMIT-}" +if [ -n "${EXPECTED_COMMIT}" ]; then + CURRENT_COMMIT="$(git -C "${REPO_ROOT}" -c safe.directory="${REPO_ROOT}" rev-parse HEAD 2>/dev/null || echo unknown)" + if [ "${CURRENT_COMMIT}" != "${EXPECTED_COMMIT}" ]; then + echo "ERROR: this run reproduces commit ${EXPECTED_COMMIT}," >&2 + echo " but ${REPO_ROOT} HEAD is ${CURRENT_COMMIT}." >&2 + echo " Fix: git -C ${REPO_ROOT} checkout ${EXPECTED_COMMIT} \\" >&2 + echo " && git -C ${REPO_ROOT} submodule update --init --recursive" >&2 + echo " Or: ALLOW_COMMIT_MISMATCH=1 (proceed) / EXPECTED_COMMIT= (skip check)." >&2 + if [ "${ALLOW_COMMIT_MISMATCH:-0}" != "1" ]; then + exit 1 + fi + echo "WARNING: proceeding despite commit mismatch (ALLOW_COMMIT_MISMATCH=1)." >&2 + fi +fi + +# ============================ Precision recipe ============================ +# bf16 (default) or one of: mxfp8-rollout, mxfp8-train, mxfp8-e2e. +# These emit extra Hydra overrides appended last so they win over the bf16 +# defaults set in the command below. +get_precision_config() { + local PRECISION_RECIPE="$1" + local DISABLE_FP8_LINEAR="$2" + local DISABLE_FP8_MOE="$3" + local ENABLE_FP8_PARAM_IN_TRAIN="$4" + local PRECISION_EXTRA_ARGS="" + + local MXFP8_GEN_EXTRA_ARGS="policy.generation.vllm_cfg.precision=fp8 \ +++policy.generation.vllm_cfg.fp8_cfg.is_mx=true \ +policy.generation.vllm_cfg.gpu_memory_utilization=0.8 \ +policy.generation.vllm_cfg.tensor_parallel_size=4 \ +policy.generation.vllm_cfg.expert_parallel_size=4" + + local IGNORED_LAYER_KWS="\"conv1d\",\"mtp\"" + if [ "$DISABLE_FP8_MOE" == "1" ]; then + IGNORED_LAYER_KWS="$IGNORED_LAYER_KWS,\".experts.\"" + fi + if [ "$DISABLE_FP8_LINEAR" == "1" ]; then + IGNORED_LAYER_KWS="$IGNORED_LAYER_KWS,\"in_proj\",\"out_proj\",\"q_proj\",\"k_proj\",\"v_proj\",\"o_proj\",\"fc1_latent_proj\",\"fc2_latent_proj\",\"shared_experts\"" + fi + MXFP8_GEN_EXTRA_ARGS="$MXFP8_GEN_EXTRA_ARGS +policy.generation.vllm_cfg.quantization_ignored_layer_kws=[$IGNORED_LAYER_KWS]" + + local MXFP8_TRAIN_EXTRA_ARGS="policy.megatron_cfg.fp8_cfg.enabled=true \ +policy.megatron_cfg.fp8_cfg.fp8=\"e4m3\" \ +policy.megatron_cfg.fp8_cfg.fp8_recipe=\"mxfp8\" \ +++policy.megatron_cfg.fp8_cfg.fp8_param=false \ +policy.megatron_cfg.moe_router_dtype=fp32 \ +policy.megatron_cfg.expert_model_parallel_size=64 \ +" + + local MXFP8_PARAM_EXTRA_ARGS="++policy.megatron_cfg.fp8_cfg.fp8_param=true \ ++policy.megatron_cfg.optimizer.reuse_grad_buf_for_mxfp8_param_ag=true \ ++policy.megatron_cfg.optimizer.fp8_recipe=mxfp8 \ ++policy.megatron_cfg.optimizer.overlap_param_gather=true \ +++policy.megatron_cfg.distributed_data_parallel_config.overlap_param_gather=true \ +++policy.megatron_cfg.distributed_data_parallel_config.overlap_grad_reduce=true \ +" + + if [ "$ENABLE_FP8_PARAM_IN_TRAIN" == "1" ]; then + MXFP8_TRAIN_EXTRA_ARGS="$MXFP8_TRAIN_EXTRA_ARGS $MXFP8_PARAM_EXTRA_ARGS" + fi + + if [ "$PRECISION_RECIPE" == "mxfp8-rollout" ]; then + PRECISION_EXTRA_ARGS="$MXFP8_GEN_EXTRA_ARGS" + elif [ "$PRECISION_RECIPE" == "mxfp8-train" ]; then + PRECISION_EXTRA_ARGS="$MXFP8_TRAIN_EXTRA_ARGS" + elif [ "$PRECISION_RECIPE" == "mxfp8-e2e" ]; then + PRECISION_EXTRA_ARGS="$MXFP8_GEN_EXTRA_ARGS $MXFP8_TRAIN_EXTRA_ARGS" + else + PRECISION_EXTRA_ARGS="" + fi + + echo "${PRECISION_EXTRA_ARGS}" +} + +PRECISION_RECIPE="${PRECISION_RECIPE:-bf16}" +DISABLE_FP8_LINEAR="${DISABLE_FP8_LINEAR:-0}" +DISABLE_FP8_MOE="${DISABLE_FP8_MOE:-0}" +ENABLE_FP8_PARAM_IN_TRAIN="${ENABLE_FP8_PARAM_IN_TRAIN:-0}" +PRECISION_EXTRA_ARGS="$(get_precision_config "${PRECISION_RECIPE}" "${DISABLE_FP8_LINEAR}" "${DISABLE_FP8_MOE}" "${ENABLE_FP8_PARAM_IN_TRAIN}")" +# bf16 and mxfp8 compile different subgraphs under the same torch.compile hash, +# so they MUST use separate vLLM compile-cache trees on Lustre. +case "${PRECISION_RECIPE}" in + mxfp8-rollout|mxfp8-e2e) VLLM_CACHE_PRECISION="mxfp8" ;; + *) VLLM_CACHE_PRECISION="bf16" ;; +esac + +# =========== Scaling modes: ALIGN_BASELINE | default (R) | SKIP_TRAINING =========== +# One knob NUM_VLLM_REPLICAS (R) = vLLM gen replicas = gen nodes. The mode picks the +# TRAINING geometry (CP/EP -> nodes-per-DP); nodes/batch/segment are then derived. +# +# ALIGN_BASELINE=1 -> reproduce the validated 32-node baseline: CP=16/EP=32, R +# defaults to 16 -> 16 train + 16 gen = 32 nodes, GBS=256. +# (default, R set) -> real training; scale BOTH train and gen from the validated +# 2-node/DP base CP=2/EP=8 (NODES_PER_DP=2). train grows with R +# (DP=ceil(R/NODES_PER_DP)), gen=R. NOTE: the CP=1/EP=4 1-node +# base OOMs on both GPU and host, so 2 nodes/DP is the stable base. +# SKIP_TRAINING=1 -> generation benchmark (NVIDIA-NeMo/RL#2930). Training pinned to +# a minimal 1-node policy cluster (CP=1/EP=4, no optimizer built -> +# no OOM); train does NOT scale, only gen (=R). Sets +# NRL_GEN_BENCHMARK_SKIP_TRAINING=1. +# +# FIXED across modes: TP=4, PP=1, ETP=1, GPP=16, VLLM_TP=4, MAX_LENGTH=196608. Any of +# CP/EP/NUM_TRAIN_NODES/GBS/PPS/SBATCH_SEGMENT can still be overridden explicitly. +# +# Examples: +# ALIGN_BASELINE=1 bash examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh +# NUM_VLLM_REPLICAS=8 bash examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh +# SKIP_TRAINING=1 NUM_VLLM_REPLICAS=8 bash examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh +NUM_GPU=4 +TP="${TP:-4}" +PP="${PP:-1}" +ETP="${ETP:-1}" +GPP="${GPP:-16}" +VLLM_TP="${VLLM_TP:-4}" +VLLM_PP="${VLLM_PP:-1}" +MAX_LENGTH="${MAX_LENGTH:-196608}" +PER_GEN_REPLICA_BATCH="${PER_GEN_REPLICA_BATCH:-16}" + +ALIGN_BASELINE="${ALIGN_BASELINE:-0}" +SKIP_TRAINING="${SKIP_TRAINING:-0}" +if [ "${ALIGN_BASELINE}" = "1" ] && [ "${SKIP_TRAINING}" = "1" ]; then + echo "ERROR: set only one of ALIGN_BASELINE / SKIP_TRAINING." >&2; exit 1 +fi + +if [ "${ALIGN_BASELINE}" = "1" ]; then + CP="${CP:-16}"; EP="${EP:-32}" # Mode 1: 32-node baseline geometry + NUM_VLLM_REPLICAS="${NUM_VLLM_REPLICAS:-16}" # -> 16 gen + 16 train = 32 nodes +elif [ "${SKIP_TRAINING}" = "1" ]; then + CP="${CP:-1}"; EP="${EP:-4}" # Mode 3: gen benchmark, minimal train + export NRL_GEN_BENCHMARK_SKIP_TRAINING=1 +else + CP="${CP:-2}"; EP="${EP:-8}" # Mode 2: validated 2-node/DP stable base +fi + +NUM_VLLM_REPLICAS="${NUM_VLLM_REPLICAS:-}" +if [ -z "${NUM_VLLM_REPLICAS}" ]; then + echo "ERROR: NUM_VLLM_REPLICAS is required (= vLLM gen replicas = gen nodes)." >&2 + echo " e.g. NUM_VLLM_REPLICAS=8, or ALIGN_BASELINE=1 for the 32-node baseline." >&2 + exit 1 +fi +if ! printf '%s' "${NUM_VLLM_REPLICAS}" | grep -qE '^[0-9]+$' || [ "${NUM_VLLM_REPLICAS}" -lt 1 ]; then + echo "ERROR: NUM_VLLM_REPLICAS must be a positive integer, got ${NUM_VLLM_REPLICAS}" >&2 + exit 1 +fi +R="${NUM_VLLM_REPLICAS}" + +NUM_GEN_NODES="${NUM_GEN_NODES:-${R}}" +NODES_PER_DP=$(( (TP * CP * PP) / NUM_GPU )) # nodes for one train DP replica +[ "${NODES_PER_DP}" -lt 1 ] && NODES_PER_DP=1 +if [ "${SKIP_TRAINING}" = "1" ]; then + # gen benchmark: train pinned to a single minimal DP replica; does NOT scale with R. + NUM_TRAIN_NODES="${NUM_TRAIN_NODES:-${NODES_PER_DP}}" +else + # scale train data-parallelism with R: DP = ceil(R / NODES_PER_DP). + TRAIN_DP_DERIVED=$(( (R + NODES_PER_DP - 1) / NODES_PER_DP )) + [ "${TRAIN_DP_DERIVED}" -lt 1 ] && TRAIN_DP_DERIVED=1 + NUM_TRAIN_NODES="${NUM_TRAIN_NODES:-$(( NODES_PER_DP * TRAIN_DP_DERIVED ))}" +fi +NUM_GYM_NODES="${NUM_GYM_NODES:-0}" +TOTAL_NODES=$(( NUM_TRAIN_NODES + NUM_GEN_NODES + NUM_GYM_NODES )) + +TRAIN_GPUS=$(( NUM_TRAIN_NODES * NUM_GPU )) +ATTN_BASE=$(( TP * CP * PP )) +EXPERT_BASE=$(( ETP * EP * PP )) +if [ $(( TRAIN_GPUS % ATTN_BASE )) -ne 0 ]; then + echo "ERROR: train GPUs ${TRAIN_GPUS} not divisible by TP*CP*PP=${ATTN_BASE}." >&2; exit 1 +fi +if [ $(( TRAIN_GPUS % EXPERT_BASE )) -ne 0 ]; then + echo "ERROR: train GPUs ${TRAIN_GPUS} not divisible by ETP*EP*PP=${EXPERT_BASE}." >&2; exit 1 +fi +TRAIN_DP=$(( TRAIN_GPUS / ATTN_BASE )) + +GBS="${GBS:-$(( PER_GEN_REPLICA_BATCH * R ))}" +if [ $(( GBS % GPP )) -ne 0 ]; then + echo "ERROR: GBS=${GBS} not divisible by GPP=${GPP}." >&2; exit 1 +fi +PPS="${PPS:-$(( GBS / GPP ))}" +if [ $(( PPS % TRAIN_DP )) -ne 0 ]; then + echo "ERROR: PPS=${PPS} must be divisible by train DP=${TRAIN_DP} (prompts split across DP)." >&2 + echo " Pick R so that R is a multiple of ${TRAIN_DP}." >&2; exit 1 +fi +CONCURRENCY="${CONCURRENCY:-$(( 2 * PPS * GPP ))}" + +# Topology segment: largest of {16,8,4,2,1} that DIVIDES *each* worker group's node +# count (train, gen, gym) -- not just the total. nemo-rl's RayVirtualCluster segments +# each group separately and asserts num_nodes % segment_size == 0 per group, so a +# segment that divides the total but not a group (e.g. total=2 but train=1/gen=1) +# fails with "num_nodes (1) must be divisible by segment_size (2)". Override with +# SBATCH_SEGMENT=. +if [ -z "${SBATCH_SEGMENT:-}" ]; then + for s in 16 8 4 2 1; do + if [ $(( NUM_TRAIN_NODES % s )) -eq 0 ] && [ $(( NUM_GEN_NODES % s )) -eq 0 ] \ + && { [ "${NUM_GYM_NODES}" -eq 0 ] || [ $(( NUM_GYM_NODES % s )) -eq 0 ]; }; then + SBATCH_SEGMENT="$s"; break + fi + done +fi +CLUSTER_SEGMENT_SIZE="${CLUSTER_SEGMENT_SIZE:-${SBATCH_SEGMENT}}" + +_MODE="scale train+gen (real training)" +[ "${ALIGN_BASELINE}" = "1" ] && _MODE="ALIGN_BASELINE (32-node baseline)" +[ "${SKIP_TRAINING}" = "1" ] && _MODE="SKIP_TRAINING (gen benchmark; NRL_GEN_BENCHMARK_SKIP_TRAINING=1)" +echo "==========================================" +echo "nano V3.5 SWE scale-gen | mode: ${_MODE} | R=${R}" +echo " nodes: train=${NUM_TRAIN_NODES} (DP=${TRAIN_DP}), gen=${NUM_GEN_NODES}, gym=${NUM_GYM_NODES}, total=${TOTAL_NODES}" +echo " batch: PPS=${PPS}, GPP=${GPP}, GBS=${GBS}, per-replica=$(( GBS / R )), concurrency=${CONCURRENCY}" +echo " parall: TP=${TP}, CP=${CP}, EP=${EP}, PP=${PP}, ETP=${ETP} (nodes/DP=${NODES_PER_DP}); vLLM_TP=${VLLM_TP}, vLLM_PP=${VLLM_PP}; max_length=${MAX_LENGTH}" +echo " segment: slurm=${SBATCH_SEGMENT}, cluster=${CLUSTER_SEGMENT_SIZE} (auto: largest {16..1} dividing each worker group)" +echo "==========================================" + +# ============================ Paths ============================ +EXP_NAME="${EXP_NAME:-nano-v3-5-swe-${USER}-r${R}}" +# In-repo authoritative config (nano_v3 reasoning parser + 6-family container_formatter baked in). +CONFIG_PATH="${CONFIG_PATH:-${REPO_ROOT}/examples/swe_bench/grpo_nano_v3_5_swe_hsg.yaml}" +# nano V3.5 student: mopd ultrav3 -> nanov3_5 repro v5 KD opt full, step_18. +MODEL_PATH="${MODEL_PATH:-/lustre/fsw/portfolios/llmservice/users/pjin/devel/nemo-rl-ultra-v3-nano-opd-dev-20260513/results/mopd_ultrav3_to_nanov3_5_repro_v5_kd_opt_full-hsg-20260524-r1/step_18/hf}" +TRAIN_PATH="${TRAIN_PATH:-/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/blends/large_root_cause_curriculum_with_mercor_ots_plus_singlefile_swerebench_overlap_fix.jsonl}" +VAL_PATH="${VAL_PATH:-/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/swe_public_datasets_val_swebench.jsonl}" +SIF_DIR="${SIF_DIR:-/lustre/fsw/portfolios/llmservice/users/sdevare/images}" + +# Results/logs live under the repo this job is submitted from (REPO_ROOT), +# so a run's artifacts stay with the checkout that produced them. +RESULTS_ROOT="${RESULTS_ROOT:-${REPO_ROOT}/results}" +LOG_ROOT="${LOG_ROOT:-${REPO_ROOT}/logs}" +RESULTS_DIR="${RESULTS_DIR:-${RESULTS_ROOT}/${EXP_NAME}}" +RUN_LOG_DIR="${RUN_LOG_DIR:-${LOG_ROOT}/${EXP_NAME}}" +NEMO_LOG_DIR="${NEMO_LOG_DIR:-${RUN_LOG_DIR}/nemo}" + +RAY_SUB="${RAY_SUB:-${REPO_ROOT}/ray.sub}" +ENTRYPOINT="${ENTRYPOINT:-${REPO_ROOT}/examples/nemo_gym/run_grpo_nemo_gym.py}" +GYM_CODE="${GYM_CODE:-${REPO_ROOT}/3rdparty/Gym-workspace/Gym}" +LATEST_JOB_ID_FILE="${LATEST_JOB_ID_FILE:-${SCRIPT_DIR}/latest_nano_v3_5_job_id.txt}" + +# ========================= Container / mounts ========================= +# Container with gym venvs baked in at /opt/gym_venvs and Ray/Python matching this +# checkout (built from nightly-063026). Override CONTAINER=... for a different image. +CONTAINER="${CONTAINER:-/lustre/fsw/portfolios/nemotron/users/ruit/enroot-images/nemo-rl:nightly-063026-gymvenvs.squashfs}" +SANDBOX_CONTAINER="${SANDBOX_CONTAINER:-/lustre/fsw/portfolios/llmservice/users/igitman/images/nemo-skills-sandbox-b620e79.sqsh}" + +EXTRA_MOUNTS="${EXTRA_MOUNTS:-/lustre:/lustre}" +MOUNTS="${EXTRA_MOUNTS}" +append_mount() { + local src="$1" + local dst="$2" + if [ -d "${src}" ] || [ -f "${src}" ]; then + MOUNTS="${MOUNTS},${src}:${dst}" + else + echo "WARNING: mount source missing, using container built-in if available: ${src}" >&2 + fi +} +append_mount "${REPO_ROOT}" "${REPO_ROOT}" +append_mount "${GYM_CODE}" "/opt/nemo-rl/3rdparty/Gym-workspace/Gym" +export MOUNTS +export CONTAINER +export SANDBOX_CONTAINER + +# ======================= Cluster / resources ======================= +NUM_GPU=4 +NUM_TRAIN_NODES="${NUM_TRAIN_NODES:-32}" +NUM_GEN_NODES="${NUM_GEN_NODES:-32}" +NUM_GYM_NODES="${NUM_GYM_NODES:-0}" +TOTAL_NODES=$((NUM_TRAIN_NODES + NUM_GEN_NODES + NUM_GYM_NODES)) + +export GPUS_PER_NODE="${NUM_GPU}" +export CPUS_PER_WORKER="${CPUS_PER_WORKER:-144}" + +# ============================ nano geometry ============================ +TP="${TP:-4}" +CP="${CP:-16}" +EP="${EP:-32}" +PP="${PP:-1}" +ETP="${ETP:-1}" +GPP="${GPP:-16}" +PPS="${PPS:-32}" +GBS="${GBS:-512}" +VLLM_TP="${VLLM_TP:-4}" +VLLM_PP="${VLLM_PP:-1}" +VLLM_GPU_UTIL="${VLLM_GPU_UTIL:-0.85}" +MAX_LENGTH="${MAX_LENGTH:-196608}" +MAX_NUM_STEPS="${MAX_NUM_STEPS:-1000000}" +CONCURRENCY="${CONCURRENCY:-$((2 * PPS * GPP))}" +USE_MULTIPLE_DATALOADER="${USE_MULTIPLE_DATALOADER:-False}" +# logprob_chunk_size + fuse_loss: aligned to the source nano script (yaml defaults +# to 2048; the prod nano run overrode to 1024 with fused loss). +LOGPROB_CHUNK_SIZE="${LOGPROB_CHUNK_SIZE:-1024}" +FUSE_LOSS="${FUSE_LOSS:-true}" +# Empty = use the yaml value. Set only when a newer code path requires overriding. +GRADIENT_ACCUMULATION_FUSION="${GRADIENT_ACCUMULATION_FUSION:-}" +REASONING_PARSER_PLUGIN="${REASONING_PARSER_PLUGIN:-}" +MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY="${MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY:-}" +if [ -n "${MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY}" ] && ! printf '%s' "${MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY}" | grep -qE '^[0-9]+$'; then + echo "ERROR: MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY must be an integer, got ${MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY}" >&2 + exit 1 +fi + +CHECKPOINTING_SAVE_BY="${CHECKPOINTING_SAVE_BY:-00:03:30:00}" +SAVE_PERIOD="${SAVE_PERIOD:-5}" +SAVE_OPTIMIZER="${SAVE_OPTIMIZER:-true}" +WANDB_PROJ="${WANDB_PROJ:-ruit-nano-v3-5-swe}" +WANDB_GROUP="${WANDB_GROUP:-nano-v3-5-swe}" +# Tag the wandb name as the internal-repo reproduction. +WANDB_NAME="${WANDB_NAME:-${EXP_NAME}-${PRECISION_RECIPE}-internal-repo}" + +# ========================= SLURM submission ========================= +SLURM_ACCOUNT="${SLURM_ACCOUNT-nemotron_sw_post}" +SLURM_PARTITION="${SLURM_PARTITION-batch_long}" +# Default: no explicit QOS -> partition's default QOS (hero-res needs a matching +# reservation/permission, which raises "Invalid qos specification" otherwise). +# Opt in with SLURM_QOS=hero-res SLURM_RESERVATION= when you have one. +SLURM_QOS="${SLURM_QOS-}" +SLURM_RESERVATION="${SLURM_RESERVATION-}" +# Idle-GPU reaper exemption. Async GRPO leaves the training GPUs idle while rollouts +# are collected (~30min for SWE) and during per-step validation, which trips the +# default OccupiedIdleGPUsJobReaper (cancels at 30min idle). This --comment JSON +# raises the exemption to 60min. Override with SLURM_COMMENT=..., disable with +# SLURM_COMMENT= (empty). +DEFAULT_SLURM_COMMENT='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"60","reason":"data_loading","description":"Async GRPO RL training: training GPUs idle during rollout collection (~30min) and validation each step"}}' +SLURM_COMMENT="${SLURM_COMMENT-$DEFAULT_SLURM_COMMENT}" +WALLTIME="${WALLTIME:-4:00:00}" +# SBATCH_SEGMENT / CLUSTER_SEGMENT_SIZE are derived above from the total node count +# (largest divisor of {16,8,4,2,1}); not re-defaulted here to avoid pinning 16. + +# ========================= Environment variables ========================= +# Source the user's env file by default (provides WANDB_API_KEY, HF_TOKEN and +# GITLAB_PAT, the latter needed for the flashinfer internal PyPI index). Set +# SOURCE_USER_ENV=0 to skip -- note this file may also mutate global git config. +if [ "${SOURCE_USER_ENV:-1}" = "1" ] && [ -f "${HOME}/script/export_env_vars.sh" ]; then + # shellcheck disable=SC1090 + source "${HOME}/script/export_env_vars.sh" +fi + +HF_HOME="${HF_HOME:-/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemorl/users/${USER}/hf_home}" +HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}" +PERSISTENT_CACHE="${PERSISTENT_CACHE:-${REPO_ROOT}/.cache/nemotron_nano_v3_5}" +# Default points at the venvs baked into the CONTAINER above (/opt/gym_venvs); with +# skip_venv_if_present=true the gym reuses them and skips the build entirely. +# Do NOT default this to Lustre: building the editable nemo-gym venv there hangs on +# uv's flock. If you use a container WITHOUT baked venvs, set this to a node-local +# path (e.g. /tmp/nemo_gym_venvs) so the runtime build has working flock + fast IO. +NEMO_GYM_VENV_DIR="${NEMO_GYM_VENV_DIR:-/opt/gym_venvs}" +# The uv distribution-cache lock (editable nemo-gym build) is what hangs on Lustre, +# and it lives in the *uv cache*, not the venv dir. So the gym's uv cache must also +# be node-local, else moving only the venv doesn't help. Kept separate from the main +# NeMo-RL UV_CACHE_DIR (which stays on Lustre for package persistence). +NEMO_GYM_UV_CACHE="${NEMO_GYM_UV_CACHE:-/tmp/nemo_gym_uv_cache}" + +# wandb run files + artifact *staging* default to XDG_DATA_HOME (~/.local/share/wandb), +# which on this setup points at the coreai Lustre quota -- that filled up and aborted a +# run mid-artifact-write with "Disk quota exceeded". Redirect wandb to the roomy +# nemotron project space instead. +WANDB_STAGE_ROOT="${WANDB_STAGE_ROOT:-/lustre/fs1/portfolios/nemotron/projects/nemotron_sw_post/users/ruit/wandb_stage}" +WANDB_DIR="${WANDB_DIR:-${WANDB_STAGE_ROOT}/dir}" +WANDB_CACHE_DIR="${WANDB_CACHE_DIR:-${WANDB_STAGE_ROOT}/cache}" +WANDB_DATA_DIR="${WANDB_DATA_DIR:-${WANDB_STAGE_ROOT}/data}" +export WANDB_DIR WANDB_CACHE_DIR WANDB_DATA_DIR + +export HF_HOME +export HF_DATASETS_CACHE +export PERSISTENT_CACHE +export NEMO_GYM_VENV_DIR +export BASE_LOG_DIR="${RUN_LOG_DIR}" +export RAY_LOG_SYNC_FREQUENCY="${RAY_LOG_SYNC_FREQUENCY:-60}" +export CACHE_SYNC_FREQUENCY="${CACHE_SYNC_FREQUENCY:-1800}" + +# vLLM compile cache is precision-scoped (bf16 vs mxfp8 must not share a tree). +LUSTRE_VLLM_CACHE="${PERSISTENT_CACHE}/cache_write/vllm_compile_cache_${VLLM_CACHE_PRECISION}" +LUSTRE_INDUCTOR_CACHE="${PERSISTENT_CACHE}/cache_write/inductor_cache" +LUSTRE_TRITON_CACHE="${PERSISTENT_CACHE}/cache_write/triton_cache" +INDUCTOR_CACHE_DIR="/tmp/nemo_rl_inductor_cache" +TRITON_CACHE_DIR="/tmp/nemo_rl_triton_cache" +VLLM_PRECOMPILED_WHEEL_LOCATION="${VLLM_PRECOMPILED_WHEEL_LOCATION:-https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_aarch64.whl}" + +# NOTE: do NOT mkdir NEMO_GYM_VENV_DIR here. It is a *container-internal* path +# (e.g. /opt/gym_venvs when baked into the sqsh, or node-local /tmp created at +# runtime inside the container) -- creating it from the login node either fails +# (/opt not writable) or is pointless (node-local /tmp differs per node). +mkdir -p "${RESULTS_DIR}" "${RUN_LOG_DIR}" "${NEMO_LOG_DIR}" \ + "${LUSTRE_VLLM_CACHE}" "${LUSTRE_INDUCTOR_CACHE}" "${LUSTRE_TRITON_CACHE}" \ + "${HF_HOME}" "${HF_DATASETS_CACHE}" \ + "${WANDB_DIR}" "${WANDB_CACHE_DIR}" "${WANDB_DATA_DIR}" + +# Record the exact superproject + submodule revisions actually used, for repro. +{ + echo "expected_commit: ${EXPECTED_COMMIT:-}" + echo "superproject: $(git -C "${REPO_ROOT}" -c safe.directory="${REPO_ROOT}" rev-parse HEAD 2>/dev/null || echo unknown)" + echo "submodules (status --recursive):" + git -C "${REPO_ROOT}" -c safe.directory="${REPO_ROOT}" submodule status --recursive 2>/dev/null || echo " " +} > "${RUN_LOG_DIR}/git-revision.txt" 2>/dev/null || true + +require_path() { + local path="$1" + local label="$2" + if [ ! -e "${path}" ]; then + echo "ERROR: missing ${label}: ${path}" >&2 + exit 1 + fi +} + +# vLLM editable source tree: pyproject declares `vllm = {path="3rdparty/vllm", editable}`, +# but that dir is created by the docker build and only exists inside the container at +# /opt/nemo-rl/3rdparty/vllm. In overlay mode we cd into the (mounted) REPO_ROOT and run +# uv there, so REPO_ROOT/3rdparty/vllm must resolve to the container's copy. Create a +# symlink (dangling on the host, resolves inside the container). +if [ ! -e "${REPO_ROOT}/3rdparty/vllm" ] && [ ! -L "${REPO_ROOT}/3rdparty/vllm" ]; then + mkdir -p "${REPO_ROOT}/3rdparty" + ln -s /opt/nemo-rl/3rdparty/vllm "${REPO_ROOT}/3rdparty/vllm" \ + && echo "[INFO] linked ${REPO_ROOT}/3rdparty/vllm -> /opt/nemo-rl/3rdparty/vllm (resolves inside container)" +fi + +require_path "${CONFIG_PATH}" "config" +require_path "${MODEL_PATH}" "model" +require_path "${TRAIN_PATH}" "train data" +require_path "${VAL_PATH}" "validation data" +require_path "${SIF_DIR}" "SIF root" +require_path "${CONTAINER}" "container" +require_path "${RAY_SUB}" "ray.sub" +require_path "${ENTRYPOINT}" "training entrypoint" +require_path "${GYM_CODE}/nemo_gym/__init__.py" "Gym checkout" + +for sif_subdir in swerebench nv_internal r2e_gym swegym swebench mercor/swebenchpro_ots; do + require_path "${SIF_DIR}/${sif_subdir}" "SIF subdir ${sif_subdir}" +done + +if [ "${DRY_RUN:-0}" != "1" ] && [ -z "${WANDB_API_KEY:-}" ]; then + echo "ERROR: WANDB_API_KEY must be set for a real submission." >&2 + exit 1 +fi + +# Per-instance .sif resolution: swe_agents tries each pattern in order and uses the +# first existing file. Mirrors test_assets/ultra_SWE/repro_ultra_launch.sh. +SIF_FORMATTERS="[\"${SIF_DIR}/swerebench/{instance_id}.sif\",\"${SIF_DIR}/nv_internal/{instance_id}.sif\",\"${SIF_DIR}/r2e_gym/{instance_id}.sif\",\"${SIF_DIR}/swegym/sweb.eval.arm64.{instance_id}.sif\",\"${SIF_DIR}/swebench/swe-bench.eval.arm64.{instance_id}.sif\",\"${SIF_DIR}/mercor/swebenchpro_ots/{instance_id}.sif\"]" + +echo "==========================================" +echo "nano V3.5 SWE-e2e reproduction: ${EXP_NAME}" +echo "Repo: ${REPO_ROOT}" +echo "Config: ${CONFIG_PATH}" +echo "Model: ${MODEL_PATH}" +echo "Data train: ${TRAIN_PATH}" +echo "Data val: ${VAL_PATH}" +echo "SIF root: ${SIF_DIR}" +echo "Precision: ${PRECISION_RECIPE} (vLLM cache tree: ${VLLM_CACHE_PRECISION})" +echo "Nodes: train=${NUM_TRAIN_NODES}, gen=${NUM_GEN_NODES}, gym=${NUM_GYM_NODES}, total=${TOTAL_NODES}" +echo "Parallelism: TP=${TP}, CP=${CP}, EP=${EP}, PP=${PP}, ETP=${ETP}, vLLM_TP=${VLLM_TP}, vLLM_PP=${VLLM_PP}" +echo "Batch: PPS=${PPS}, GPP=${GPP}, GBS=${GBS}, concurrency=${CONCURRENCY}" +echo "SeqLen: ${MAX_LENGTH}" +echo "Padding: make_sequence_length_divisible_by=${MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY:-}" +echo "LogprobChunk:${LOGPROB_CHUNK_SIZE}, fuse_loss=${FUSE_LOSS}" +echo "Checkpoints: ${RESULTS_DIR}" +echo "Logs: ${RUN_LOG_DIR}" +echo "Cache: vLLM=${LUSTRE_VLLM_CACHE}; Inductor/Triton=/tmp with ${CACHE_SYNC_FREQUENCY}s rsync to Lustre" +echo "Container: ${CONTAINER}" +echo "Slurm: account=${SLURM_ACCOUNT}, partition=${SLURM_PARTITION}, qos=${SLURM_QOS}, reservation=${SLURM_RESERVATION:-}" +if [ -n "${PRECISION_EXTRA_ARGS}" ]; then + echo "PrecisionArgs: ${PRECISION_EXTRA_ARGS}" +fi +echo "==========================================" + +cd "${REPO_ROOT}" + +read -r -d '' SETUP_COMMAND </dev/null 2>&1 && command -v rsync >/dev/null 2>&1; } || { apt-get update -qq && apt-get install -y -qq zstd rsync; }) 2>/dev/null || true +if ! command -v apptainer >/dev/null 2>&1 && ! command -v singularity >/dev/null 2>&1; then + apt-get update -qq && apt-get install -y -qq git build-essential gcc wget 2>/dev/null || true + cd /tmp + wget --no-check-certificate -q -nc https://github.com/apptainer/apptainer/releases/download/v1.3.1/apptainer_1.3.1_arm64.deb || true + apt install -y ./apptainer_1.3.1_arm64.deb 2>/dev/null || true + ln -sf /usr/bin/apptainer /usr/bin/singularity 2>/dev/null || true +fi + +echo "[CACHE SEED] Clearing stale node-local compile caches..." +rm -rf "${INDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}" +mkdir -p "${INDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}" \ + "${LUSTRE_VLLM_CACHE}" "${LUSTRE_INDUCTOR_CACHE}" "${LUSTRE_TRITON_CACHE}" + +_seed_cache() { + local src="\$1" + local dst="\$2" + local name="\$3" + if [ -d "\$src" ] && [ "\$(ls -A "\$src" 2>/dev/null)" ]; then + rsync -a --exclude '.tmp_*' "\$src/" "\$dst/" 2>/dev/null \ + && echo "[CACHE SEED] \$name: seeded from Lustre" \ + || echo "[CACHE SEED] \$name: seed failed (non-fatal)" + else + echo "[CACHE SEED] \$name: no warm cache yet" + fi +} + +_seed_cache "${LUSTRE_INDUCTOR_CACHE}" "${INDUCTOR_CACHE_DIR}" "Inductor" +_seed_cache "${LUSTRE_TRITON_CACHE}" "${TRITON_CACHE_DIR}" "Triton" +echo "[CACHE SEED] Done." + +_sync_cache_one() { + local src="\$1" + local dst="\$2" + local name="\$3" + mkdir -p "\$dst" + if [ -d "\$src" ] && [ "\$(ls -A "\$src" 2>/dev/null)" ]; then + rsync -a --ignore-existing --exclude '.tmp_*' --exclude 'tmp*' "\$src/" "\$dst/" 2>/dev/null \ + && echo "[CACHE SYNC] \$name: synced node-local cache to Lustre" \ + || echo "[CACHE SYNC] \$name: sync failed (non-fatal)" + else + echo "[CACHE SYNC] \$name: no node-local entries yet" + fi +} + +_sync_compile_caches_to_lustre() { + _sync_cache_one "${INDUCTOR_CACHE_DIR}" "${LUSTRE_INDUCTOR_CACHE}" "Inductor" + _sync_cache_one "${TRITON_CACHE_DIR}" "${LUSTRE_TRITON_CACHE}" "Triton" +} + +_start_cache_sync_sidecar() { + local pidfile="/tmp/nemo_rl_compile_cache_sync_sidecar.pid" + if [ -f "\$pidfile" ]; then + local old_pid + old_pid="\$(cat "\$pidfile" 2>/dev/null || true)" + if [ -n "\$old_pid" ] && kill -0 "\$old_pid" 2>/dev/null; then + echo "[CACHE SYNC] Sidecar already running on this node. pid=\${old_pid}" + return + fi + fi + + local setup_log_dir + setup_log_dir="\$(cd "\$(dirname "\$0")" && pwd)" + local sidecar_log="/tmp/nemo_rl_compile_cache_sync_sidecar.log" + local frequency="${CACHE_SYNC_FREQUENCY}" + + ( + set +e + trap '_sync_compile_caches_to_lustre; exit 0' TERM INT + echo "[CACHE SYNC] Sidecar started. frequency=\${frequency}s log_dir=\${setup_log_dir}" + while true; do + if [ -f "\${setup_log_dir}/ENDED" ]; then + echo "[CACHE SYNC] ENDED detected; final sync." + _sync_compile_caches_to_lustre + exit 0 + fi + sleep "\${frequency}" + _sync_compile_caches_to_lustre + done + ) > "\$sidecar_log" 2>&1 & + echo "\$!" > "\$pidfile" + echo "[CACHE SYNC] Sidecar pid=\$! log=\${sidecar_log}" +} + +if [ "${CACHE_SYNC_FREQUENCY}" -gt 0 ] 2>/dev/null; then + _start_cache_sync_sidecar +else + echo "[CACHE SYNC] Disabled because CACHE_SYNC_FREQUENCY=${CACHE_SYNC_FREQUENCY}" +fi +SETUPEOF +export SETUP_COMMAND + +GITLAB_PASSWORD="${GITLAB_PAT:-${GITLAB_TOKEN:-}}" +export COMMAND="cd ${REPO_ROOT} && \ +trap 'touch ${BASE_LOG_DIR}/\${SLURM_JOB_ID}-logs/ENDED 2>/dev/null || true' EXIT && \ +date && \ +OMP_NUM_THREADS=16 \ +RAY_DEDUP_LOGS=1 \ +NRL_VLLM_USE_V1=1 \ +VLLM_CACHE_ROOT=${LUSTRE_VLLM_CACHE} \ +DG_JIT_CACHE_DIR=${LUSTRE_VLLM_CACHE}/deep_gemm \ +TORCHINDUCTOR_CACHE_DIR=${INDUCTOR_CACHE_DIR} \ +TRITON_CACHE_DIR=${TRITON_CACHE_DIR} \ +UV_CACHE_DIR=${PERSISTENT_CACHE}/uv \ +RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 \ +UV_HTTP_TIMEOUT=10 \ +UV_LOCK_TIMEOUT=${UV_LOCK_TIMEOUT:-1200} \ +NEMO_GYM_VENV_DIR=${NEMO_GYM_VENV_DIR} \ +VLLM_USE_PRECOMPILED=1 \ +VLLM_PRECOMPILED_WHEEL_LOCATION=${VLLM_PRECOMPILED_WHEEL_LOCATION} \ +VLLM_USE_FLASHINFER_MOE_FP8=1 \ +VLLM_FLASHINFER_MOE_BACKEND=latency \ +NRL_VLLM_ASYNC_TIMEOUT_SECONDS=1800 \ +NRL_WG_USE_RAY_REF=1 \ +NRL_USE_FASTOKENS=1 \ +NRL_GEN_BENCHMARK_SKIP_TRAINING=${NRL_GEN_BENCHMARK_SKIP_TRAINING:-0} \ +HF_HOME=${HF_HOME} \ +HF_TOKEN=${HF_TOKEN:-} \ +WANDB_API_KEY=${WANDB_API_KEY:-} \ +WANDB_DIR=${WANDB_DIR} \ +WANDB_CACHE_DIR=${WANDB_CACHE_DIR} \ +WANDB_DATA_DIR=${WANDB_DATA_DIR} \ +UV_INDEX_FLASHINFER_INTERNAL_PYPI_USERNAME=__token__ \ +UV_INDEX_FLASHINFER_INTERNAL_PYPI_PASSWORD=${GITLAB_PASSWORD} \ +uv run ${ENTRYPOINT} \ + --config ${CONFIG_PATH} \ + policy.model_name=${MODEL_PATH} \ + cluster.gpus_per_node=${NUM_GPU} \ + cluster.num_nodes=${TOTAL_NODES} \ + cluster.segment_size=${CLUSTER_SEGMENT_SIZE} \ + grpo.num_prompts_per_step=${PPS} \ + grpo.num_generations_per_prompt=${GPP} \ + policy.train_global_batch_size=${GBS} \ + policy.max_total_sequence_length=${MAX_LENGTH} \ + policy.logprob_chunk_size=${LOGPROB_CHUNK_SIZE} \ + ++policy.sequence_packing.fuse_loss=${FUSE_LOSS} \ + ${MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY:+policy.make_sequence_length_divisible_by=${MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY}} \ + policy.megatron_cfg.tensor_model_parallel_size=${TP} \ + policy.megatron_cfg.context_parallel_size=${CP} \ + policy.megatron_cfg.expert_model_parallel_size=${EP} \ + policy.megatron_cfg.pipeline_model_parallel_size=${PP} \ + policy.megatron_cfg.expert_tensor_parallel_size=${ETP} \ + ${GRADIENT_ACCUMULATION_FUSION:+++policy.megatron_cfg.gradient_accumulation_fusion=${GRADIENT_ACCUMULATION_FUSION}} \ + policy.generation.vllm_cfg.tensor_parallel_size=${VLLM_TP} \ + policy.generation.vllm_cfg.pipeline_parallel_size=${VLLM_PP} \ + policy.generation.vllm_cfg.gpu_memory_utilization=${VLLM_GPU_UTIL} \ + policy.generation.vllm_cfg.max_model_len=${MAX_LENGTH} \ + ${REASONING_PARSER_PLUGIN:+++policy.generation.vllm_cfg.reasoning_parser_plugin=${REASONING_PARSER_PLUGIN}} \ + ${REASONING_PARSER_PLUGIN:+'~policy.generation.vllm_cfg.http_server_serving_chat_kwargs.reasoning_parser_plugin'} \ + policy.generation.colocated.enabled=False \ + policy.generation.colocated.resources.num_nodes=${NUM_GEN_NODES} \ + policy.generation.colocated.resources.gpus_per_node=${NUM_GPU} \ + env.nemo_gym.num_gpu_nodes=${NUM_GYM_NODES} \ + ++env.nemo_gym.uv_venv_dir=${NEMO_GYM_VENV_DIR} \ + ++env.nemo_gym.uv_cache_dir=${NEMO_GYM_UV_CACHE} \ + env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.dataset_path=${TRAIN_PATH} \ + env.nemo_gym.swe_agents_val.responses_api_agents.swe_agents.dataset_path=${VAL_PATH} \ + env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.concurrency=${CONCURRENCY} \ + env.nemo_gym.swe_agents_val.responses_api_agents.swe_agents.concurrency=${CONCURRENCY} \ + 'env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter=${SIF_FORMATTERS}' \ + 'env.nemo_gym.swe_agents_val.responses_api_agents.swe_agents.container_formatter=${SIF_FORMATTERS}' \ + data.train.data_path=${TRAIN_PATH} \ + data.validation.data_path=${VAL_PATH} \ + ++data.use_multiple_dataloader=${USE_MULTIPLE_DATALOADER} \ + checkpointing.checkpoint_dir=${RESULTS_DIR} \ + checkpointing.checkpoint_must_save_by=${CHECKPOINTING_SAVE_BY} \ + checkpointing.save_period=${SAVE_PERIOD} \ + ++checkpointing.save_optimizer=${SAVE_OPTIMIZER} \ + logger.log_dir=${NEMO_LOG_DIR} \ + logger.wandb_enabled=True \ + logger.wandb.name=${WANDB_NAME} \ + logger.wandb.project=${WANDB_PROJ} \ + ++logger.wandb.group=${WANDB_GROUP} \ + grpo.max_num_steps=${MAX_NUM_STEPS} \ + ${PRECISION_EXTRA_ARGS} \ + ${EXTRA_ARGS:-}" + +SBATCH_ARGS=( + --nodes="${TOTAL_NODES}" + --account="${SLURM_ACCOUNT}" + --job-name="${WANDB_NAME}" + --partition="${SLURM_PARTITION}" + --time="${WALLTIME}" + --gres="gpu:${NUM_GPU}" + --exclusive + --mem=0 + --dependency=singleton + --segment="${SBATCH_SEGMENT}" + --output="${RUN_LOG_DIR}/slurm-%j.out" +) + +if [ -n "${SLURM_QOS}" ]; then + SBATCH_ARGS+=(--qos="${SLURM_QOS}") +fi +if [ -n "${SLURM_RESERVATION}" ]; then + SBATCH_ARGS+=(--reservation="${SLURM_RESERVATION}") +fi +if [ -n "${SLURM_COMMENT}" ]; then + SBATCH_ARGS+=(--comment="${SLURM_COMMENT}") +fi + +if [ "${DRY_RUN:-0}" = "1" ]; then + echo "" + echo "[DRY_RUN] Not submitting. Would run:" + printf '[DRY_RUN] sbatch' + printf ' %q' "${SBATCH_ARGS[@]}" + printf ' %q\n' "${RAY_SUB}" + echo "" + echo "[DRY_RUN] COMMAND:" + echo "${COMMAND}" | sed -E \ + -e 's/(WANDB_API_KEY=)[^ ]*/\1/g' \ + -e 's/(HF_TOKEN=)[^ ]*/\1/g' \ + -e 's/(UV_INDEX_FLASHINFER_INTERNAL_PYPI_PASSWORD=)[^ ]*/\1/g' + exit 0 +fi + +SBATCH_OUTPUT="$(sbatch "${SBATCH_ARGS[@]}" "${RAY_SUB}")" +echo "${SBATCH_OUTPUT}" >&2 + +JOB_ID="$(printf '%s\n' "${SBATCH_OUTPUT}" | grep -o '[0-9]\+' | tail -n 1)" +if [ -z "${JOB_ID}" ]; then + echo "ERROR: failed to parse Slurm job id from sbatch output: ${SBATCH_OUTPUT}" >&2 + exit 1 +fi +printf '%s\n' "${JOB_ID}" > "${LATEST_JOB_ID_FILE}" + +echo "==========================================" +echo "Job submitted: ${EXP_NAME}" +echo "Job ID: ${JOB_ID}" +echo "Monitor with: squeue -j ${JOB_ID}" +echo "Ray/Slurm logs: ${RUN_LOG_DIR}/${JOB_ID}-logs/" +echo "Slurm output: ${RUN_LOG_DIR}/slurm-${JOB_ID}.out" +echo "Checkpoints: ${RESULTS_DIR}/" +echo "==========================================" diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index f332c0a6de5..8eb3bdae922 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -838,6 +838,12 @@ def _spinup_nemo_gym(base_urls, model_name): # Define initialization functions that will be used in all paths init_reference_model = loss_config.reference_policy_kl_penalty > 0 + # Benchmark mode performs no training, so there is no KL term and no + # reference model to build. The auto-enable block below then turns on + # skip_reference_policy_logprobs_calculation. + if _gen_benchmark_skip_training(): + init_reference_model = False + # Auto-enable skip_reference_policy_logprobs_calculation when the reference model is not loaded. if not init_reference_model and not grpo_config.get( "skip_reference_policy_logprobs_calculation" @@ -864,7 +870,7 @@ def init_policy(): processor=processor, weights_path=weights_path, optimizer_path=optimizer_path, - init_optimizer=True, + init_optimizer=not _gen_benchmark_skip_training(), init_reference_model=init_reference_model, ) return p, time.perf_counter() - t0 @@ -3321,6 +3327,39 @@ def aggregate_rollout_metrics( return aggregated +def _gen_benchmark_skip_training() -> bool: + """Whether to run async GRPO in generation-benchmark mode. + + Controlled by the ``NRL_GEN_BENCHMARK_SKIP_TRAINING`` environment variable + (truthy values: 1/true/yes/on). In this mode the optimizer is not built and + policy.train() is replaced with zero-valued metrics, while generation, + refit/weight-sync, and the trajectory-collector cadence are preserved so + generation throughput can be benchmarked with training pinned to a minimal + (e.g. single-GPU) policy cluster. Async GRPO only. + """ + return os.environ.get("NRL_GEN_BENCHMARK_SKIP_TRAINING", "").lower() in ( + "1", + "true", + "yes", + "on", + ) + + +def _make_benchmark_dummy_train_results() -> dict[str, Any]: + """Zero-valued stand-in for policy.train() output in gen-benchmark mode. + + Matches the keys the async metrics path consumes: ``loss`` and ``grad_norm`` + are CPU tensors (the consumer calls ``.numpy()`` on them) and + ``all_mb_metrics`` is an empty dict. ``moe_metrics`` / ``mtp_metrics`` are + optional (guarded by ``in`` checks) and intentionally omitted. + """ + return { + "loss": torch.zeros(1), + "grad_norm": torch.zeros(1), + "all_mb_metrics": {}, + } + + def async_grpo_train( policy: ColocatablePolicyInterface, policy_generation: Optional[GenerationInterface], @@ -3541,6 +3580,12 @@ def async_grpo_train( print("📦 Started continuous background trajectory collection") + if _gen_benchmark_skip_training(): + # Training is skipped, so the policy GPUs would look idle between refits. + # Start a tiny NCCL-free keep-alive matmul to avoid idle-GPU reapers. + print("🫀 Starting gen-benchmark keep-alive on policy workers") + policy.start_gen_benchmark_keepalive() + print( f"🚀 Starting async GRPO training with buffer_size={optimal_buffer_size}, max_age={max_trajectory_age_steps} steps" ) @@ -3955,18 +4000,27 @@ def async_grpo_train( train_data["advantages"], master_config.grpo ) - print("▶ Preparing for training...") - with timer.time("training_prep"): - policy.prepare_for_training() + if _gen_benchmark_skip_training(): + # Benchmark mode: skip optimizer/forward/backward entirely and + # return zero-valued metrics. Still mark generation stale so the + # refit below runs every step, preserving the real weight-sync + # cadence we want to measure. + print("▶ Skipping training (gen_benchmark_skip_training)...") POLICY_GENERATION_STALE = True + train_results = _make_benchmark_dummy_train_results() + else: + print("▶ Preparing for training...") + with timer.time("training_prep"): + policy.prepare_for_training() + POLICY_GENERATION_STALE = True - print("▶ Training policy...") - with timer.time("policy_training"): - train_results = policy.train( - train_data, - loss_fn, - timer=timer, - ) + print("▶ Training policy...") + with timer.time("policy_training"): + train_results = policy.train( + train_data, + loss_fn, + timer=timer, + ) print("🔄 Synchronizing policy weights to trajectory collector…") generation_logger_metrics = None diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index e6c9d8e92af..f321ab04e41 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -334,8 +334,11 @@ class MegatronConfig(TypedDict): # Number of tokens per chunk when computing the fused linear CE loss. # Smaller values reduce peak memory further but may decrease throughput. linear_ce_fusion_chunk_size: NotRequired[int] - # When mtp_num_layers=0, Multi-Token Prediction is disabled. - mtp_num_layers: NotRequired[int] + # When mtp_num_layers=0, Multi-Token Prediction is disabled. Set to null/None to + # disable MTP on hybrid (e.g. NemotronH) models: their forward gates MTP on + # `mtp_num_layers is not None`, so 0 still enters the MTP block and asserts + # `mtp_num_layers > 0`; None makes that check False and skips MTP entirely. + mtp_num_layers: NotRequired[int | None] # MTP loss weight added to the main next-token loss (0.0 disables the MTP loss contribution). mtp_loss_scaling_factor: NotRequired[float] # When True, repeat a single MTP layer mtp_num_layers times instead of using distinct layers. diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index f1a5d705c06..30736d4306e 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -413,6 +413,19 @@ def swap_weights_via_reshard(self, *, is_source: bool) -> list[ray.ObjectRef]: is_source=is_source, ) + def start_gen_benchmark_keepalive(self) -> None: + """Start the generation-benchmark keep-alive matmul on all policy workers. + + Used only when ``NRL_GEN_BENCHMARK_SKIP_TRAINING`` is set: training is + skipped, so this issues a tiny NCCL-free local matmul on each worker to + keep otherwise-idle policy GPUs from being reaped. See + ``AbstractPolicyWorker.start_gen_benchmark_keepalive``. + """ + futures = self.worker_group.run_all_workers_single_data( + "start_gen_benchmark_keepalive" + ) + ray.get(futures) + # ── DP-shard helpers ──────────────────────────────────────────────── # DRY for Policy's logprob/train methods only. The data-plane sibling # TQPolicy shards KVBatchMeta via ``shard_meta_for_dp``; the diff --git a/nemo_rl/models/policy/workers/base_policy_worker.py b/nemo_rl/models/policy/workers/base_policy_worker.py index da792d700f4..e6440975691 100644 --- a/nemo_rl/models/policy/workers/base_policy_worker.py +++ b/nemo_rl/models/policy/workers/base_policy_worker.py @@ -48,6 +48,35 @@ def is_alive(self) -> bool: """Check if the worker is alive.""" return True + def start_gen_benchmark_keepalive(self, interval_s: float = 60.0) -> None: + """Start a daemon thread that issues a tiny local matmul periodically. + + Used only in generation-benchmark mode (``NRL_GEN_BENCHMARK_SKIP_TRAINING``), + where ``policy.train()`` is skipped and the policy GPUs would otherwise + look idle between refits and risk being reaped. The matmul is purely + local (no NCCL collectives), so it cannot desync weight-sync. Idempotent: + a second call is a no-op. + """ + import threading + import time + + if getattr(self, "_gen_benchmark_keepalive_started", False): + return + self._gen_benchmark_keepalive_started = True + + device = torch.cuda.current_device() + + def _keepalive_loop() -> None: + while True: + x = torch.randn(256, 256, device=device) + torch.matmul(x, x) + torch.cuda.synchronize(device) + time.sleep(interval_s) + + threading.Thread( + target=_keepalive_loop, name="gen_benchmark_keepalive", daemon=True + ).start() + def reset_peak_memory_stats(self) -> None: """Reset peak memory statistics.""" torch.cuda.reset_peak_memory_stats() From 81f5242954c96065b7b50558eb6d57e69bc12a4a Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Mon, 6 Jul 2026 23:39:11 -0700 Subject: [PATCH 07/69] try to make nano3.5 swe work on main Signed-off-by: Jiaqi Zeng --- 3rdparty/Gym-workspace/Gym | 2 +- examples/configs/grpo_nano_v3_5_swe_hsg.yaml | 479 +++++++++++++++++++ 2 files changed, 480 insertions(+), 1 deletion(-) create mode 100644 examples/configs/grpo_nano_v3_5_swe_hsg.yaml diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index 610a08ab5fe..d3f1bd44235 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit 610a08ab5fe9f8f5fb5fff36b170429ea67f0f92 +Subproject commit d3f1bd44235c88f8c9036d315091e353db518ef9 diff --git a/examples/configs/grpo_nano_v3_5_swe_hsg.yaml b/examples/configs/grpo_nano_v3_5_swe_hsg.yaml new file mode 100644 index 00000000000..9ef0b444a90 --- /dev/null +++ b/examples/configs/grpo_nano_v3_5_swe_hsg.yaml @@ -0,0 +1,479 @@ +# ============================================================================= +# GRPO Ultra E2E — Production SWE Config for GB200 NVL72 +# ============================================================================= +# Production config for Ultra V3 SWE end-to-end GRPO training on +# 128 nodes × 4 GPUs/node (64 gen + 64 train by default). +# +# Static algorithm, loss, and environment settings live here. +# Per-run overrides (model path, data paths, parallelism, precision) +# are applied by the launch script (repro_ultra_e2e.sh) or Hydra CLI. +# +# - gpus_per_node: 4 +# - TP: 8, CP: 8, EP: 32, PP: 1 +# - vLLM TP: 8 +# - Non-colocated async inference with 64 generation nodes +# - Max sequence length: 131072 +# ============================================================================= + +# ============================================================================= +# Cluster — overridden by launch script for parallelism changes +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 128 + segment_size: 16 + +# ============================================================================= +# Checkpointing +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/grpo_ultra_e2e" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 3 + ft_keep_latest_k: 1 + ft_save_period: 1 + save_optimizer: true + checkpoint_must_save_by: "00:03:30:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# GRPO Algorithm +# ============================================================================= +grpo: + num_prompts_per_step: 16 + num_generations_per_prompt: 32 + num_val_generations_per_prompt: 2 + max_rollout_turns: 1 + max_num_epochs: 4 + max_num_steps: 1000000 + normalize_rewards: true + use_leave_one_out_baseline: true + advantage_clip_low: -100 + advantage_clip_high: 100 + val_period: 10000 + val_at_start: false + val_at_end: false + overlong_filtering: false + max_val_samples: null + val_batch_size: 256 + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + penalize_invalid_tool_call: true + invalid_tool_call_advantage: -5.0 + penalize_malformed_thinking: true + malformed_thinking_advantage: -5.0 + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + + use_best_at_k: false + best_at_k_k: 8 + best_at_k_m: 1000 + + use_combined_training: false + combined_training_weight_mode: "auto" + combined_training_best_at_k_weight: 0.2 + combined_training_pass_at_1_weight: 1.0 + + dynamic_sampling_oversample_ratio: 1.0 + seq_logprob_error_threshold: 2 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: null + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + force_on_policy_ratio: true + use_kl_in_reward: false + +# ============================================================================= +# Policy +# ============================================================================= +policy: + model_name: "/lustre/fsw/portfolios/llmservice/users/wdai/megatron-lm-ultra/checkpoints/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-288k-nano-loss-032026/iter_0007000/hf" + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 512 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 131072 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 32 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + use_gloo_process_groups: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + # community_import.py reads megatron_cfg["gradient_accumulation_fusion"] directly + # (no default) on the HF->megatron conversion path; absent here it raises KeyError. + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled + mtp_loss_scaling_factor: 0.0 + mtp_use_repeated_layer: true + # null disables MTP entirely (matches the upstream mtp=0 run's behavior). This + # checkout's hybrid model gates MTP on `mtp_num_layers is not None` + # (hybrid_model.py:544): 0 still enters the MTP block and asserts > 0, and 1 runs + # MTP forward on only the last-stage ranks (suspected cause of the EXPERT_MODEL_ + # PARALLEL NCCL collective timeout). null makes that check False -> MTP is skipped + # on all ranks, no partial-rank EP collective. Requires the schema change making + # MegatronConfig.mtp_num_layers accept None (nemo_rl/models/policy/__init__.py). + mtp_num_layers: null + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 3.0e-6 + min_lr: 3.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 3e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + # FP8 — overridden by precision recipe in launch script + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo + # were dropped from this checkout's megatron-bridge CheckpointConfig (submodule + # 554c7b9); leaving them in raises InstantiationException "Unexpected config keys". + # Removed for the current-repo reproduction. + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + # Must be a multiple of minimum_pad_factor = cp_size*2*tp_size (see + # nemo_rl/models/megatron/data.py). With CP=16 + TP=4 + sequence_parallel that is + # 128; the old value (=TP=4) made get_logprobs assert and abort every training + # step. Formula matches examples/nemo_gym/grpo_nanov3.yaml (TP*CP*2). + make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 11001 + port_range_high: 15000 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 1 + gpu_memory_utilization: 0.8 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + # reasoning_parser_plugin lives at vllm_cfg top level in this checkout + # (vllm_worker_async reads vllm_cfg["reasoning_parser_plugin"]); passing it + # inside http_server_serving_chat_kwargs makes OpenAIServingChat raise + # "unexpected keyword argument 'reasoning_parser_plugin'". Path updated to this + # repo's location (the parser moved out of nemo_rl/utils/). + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + + vllm_kwargs: + attention_backend: FLASH_ATTN + moe_backend: triton + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 64 + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + use_multiple_dataloader: false + train: + data_path: null # Set by launch script + validation: + data_path: null # Set by launch script + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + SWE Agents +# ============================================================================= +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + # Reuse pre-built gym venvs. With false, all 3 gym servers (policy_model, + # swe_agents_{train,val}) rebuild the editable nemo-gym package in parallel and + # contend on the same uv distribution-cache lock on Lustre; policy_model then + # times out (default 300s) and NemoGym spinup fails. Mirrors the qwen-30b SWE run. + skip_venv_if_present: true + num_gpu_nodes: 0 + port_range_low: 15001 + port_range_high: 20000 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml + swe_agents_train: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}} + swebench_agent_timeout: 3600 + run_with_mixed_prompts: true + dataset_path: ${data.train.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swerebench/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/nv_internal/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/r2e_gym/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/sweb.eval.arm64.{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swebench/swe-bench.eval.arm64.{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/swebenchpro_ots/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swegym/sweb.eval.arm64.{instance_id}.sif" + swe_agents_val: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}} + swebench_agent_timeout: 3600 + dataset_path: ${data.validation.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter} + use_absolute_ip: true + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "logs" + num_val_samples_to_print: 0 + wandb_enabled: true + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "ultra-v3-swe-e2e" + name: "grpo-ultra-e2e" + tensorboard: {} + mlflow: + experiment_name: "grpo-ultra-e2e" + run_name: "grpo-ultra-e2e" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + +# ============================================================================= +# Token IDs (model-specific, used by token-based penalties) +# ============================================================================= +token_ids: + eos: 2 # + think_open: 12 # + think_close: 13 # + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +penalize_duplicated_reasoning: true # reasoning content == final answer +penalize_empty_final_answer: true # last message output has empty content +penalize_eos_token: true # eos token appears in generation +penalize_malformed_think_tag: true # /<\/think> count != 1 per turn From b66cb735db7f121b66e9e09fc7cc21bd8cfcf88e Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Tue, 7 Jul 2026 11:07:10 -0700 Subject: [PATCH 08/69] fix: declare missing build deps for custom vLLM container builds Signed-off-by: Jiaqi Zeng --- pyproject.toml | 3 +++ tools/build-custom-vllm.sh | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 25b83808249..66210e837ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -453,6 +453,9 @@ mamba-ssm = [{ requirement = "torch", match-runtime = true }] causal-conv1d = [{ requirement = "torch", match-runtime = true }] fast-hadamard-transform = [{ requirement = "torch", match-runtime = true }] nv-grouped-gemm = [{ requirement = "torch", match-runtime = true }] +# xgrammar (pulled in as a vllm dependency) doesn't declare its own build backend +# (scikit_build_core) in build-system.requires, so uv's build isolation excludes it. +xgrammar = ["scikit_build_core"] # Needed when building from source [[tool.uv.dependency-metadata]] diff --git a/tools/build-custom-vllm.sh b/tools/build-custom-vllm.sh index 71a04479357..4f030227d95 100644 --- a/tools/build-custom-vllm.sh +++ b/tools/build-custom-vllm.sh @@ -74,7 +74,7 @@ uv run --no-project use_existing_torch.py # Install dependencies echo "Installing dependencies..." uv pip install --upgrade pip -uv pip install numpy setuptools setuptools_scm +uv pip install numpy setuptools setuptools_scm setuptools_rust uv pip install torch==2.10.0 --torch-backend=cu129 # Install vLLM using precompiled wheel From c88955d9ed677356b7055f54c40a9ebb4caf1c3c Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Tue, 7 Jul 2026 11:33:57 -0700 Subject: [PATCH 09/69] fix tools/build-custom-vllm.sh Signed-off-by: Jiaqi Zeng --- tools/build-custom-vllm.sh | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tools/build-custom-vllm.sh b/tools/build-custom-vllm.sh index 4f030227d95..ed7085918aa 100644 --- a/tools/build-custom-vllm.sh +++ b/tools/build-custom-vllm.sh @@ -74,7 +74,11 @@ uv run --no-project use_existing_torch.py # Install dependencies echo "Installing dependencies..." uv pip install --upgrade pip -uv pip install numpy setuptools setuptools_scm setuptools_rust +# scikit_build_core: build backend for xgrammar (a vllm runtime dep vllm pins to an +# exact version that may lack a prebuilt wheel here); needed since -e . below runs +# with --no-build-isolation, which builds vllm's deps against this venv's packages +# instead of isolated per-package build envs. +uv pip install numpy setuptools setuptools_scm setuptools_rust scikit_build_core uv pip install torch==2.10.0 --torch-backend=cu129 # Install vLLM using precompiled wheel From afce932a28ad00481029d3de51d3428f1ec4fd02 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Tue, 7 Jul 2026 11:49:52 -0700 Subject: [PATCH 10/69] fix tools/build-custom-vllm.sh, nanobind Signed-off-by: Jiaqi Zeng --- tools/build-custom-vllm.sh | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tools/build-custom-vllm.sh b/tools/build-custom-vllm.sh index ed7085918aa..aef515e4f1b 100644 --- a/tools/build-custom-vllm.sh +++ b/tools/build-custom-vllm.sh @@ -74,11 +74,11 @@ uv run --no-project use_existing_torch.py # Install dependencies echo "Installing dependencies..." uv pip install --upgrade pip -# scikit_build_core: build backend for xgrammar (a vllm runtime dep vllm pins to an -# exact version that may lack a prebuilt wheel here); needed since -e . below runs -# with --no-build-isolation, which builds vllm's deps against this venv's packages -# instead of isolated per-package build envs. -uv pip install numpy setuptools setuptools_scm setuptools_rust scikit_build_core +# scikit_build_core + nanobind: build backend/deps for xgrammar (a vllm runtime dep +# vllm pins to an exact version that may lack a prebuilt wheel here); needed since +# -e . below runs with --no-build-isolation, which builds vllm's deps against this +# venv's packages instead of isolated per-package build envs. +uv pip install numpy setuptools setuptools_scm setuptools_rust scikit_build_core nanobind uv pip install torch==2.10.0 --torch-backend=cu129 # Install vLLM using precompiled wheel From c68f226ad5b2766f950ea2f128fcdd7bacd84f01 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Tue, 7 Jul 2026 11:54:58 -0700 Subject: [PATCH 11/69] fix: restrict tool.uv.environments to aarch64 only The multi-platform lock (aarch64 + x86_64) requires both to jointly resolve. Pinning a custom vLLM whose own transitive deps (e.g. flashinfer-python) differ per version can make the unused x86_64 split unsatisfiable and block the whole lock/sync, even though we only ever build/run on aarch64 (GB200/GB300). Per uv's own hint on this failure. --- pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 66210e837ea..08f95815fbd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -310,7 +310,6 @@ explicit = true [tool.uv] managed = true environments = [ - "python_version == '3.13' and sys_platform == 'linux' and platform_machine == 'x86_64'", "python_version == '3.13' and sys_platform == 'linux' and platform_machine == 'aarch64'", ] preview = true # Enable preview features like extra-build-dependencies From dec7b08ab312dc33f2635b19bf1aa9b9777ac28a Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Wed, 8 Jul 2026 09:29:49 -0700 Subject: [PATCH 12/69] Revert "fix: restrict tool.uv.environments to aarch64 only" This reverts commit c68f226ad5b2766f950ea2f128fcdd7bacd84f01. --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 08f95815fbd..66210e837ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -310,6 +310,7 @@ explicit = true [tool.uv] managed = true environments = [ + "python_version == '3.13' and sys_platform == 'linux' and platform_machine == 'x86_64'", "python_version == '3.13' and sys_platform == 'linux' and platform_machine == 'aarch64'", ] preview = true # Enable preview features like extra-build-dependencies From 64ad3bc7e12c847e5775f4408ae3eaf98029ad69 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Wed, 8 Jul 2026 09:29:49 -0700 Subject: [PATCH 13/69] Revert "fix tools/build-custom-vllm.sh, nanobind" This reverts commit afce932a28ad00481029d3de51d3428f1ec4fd02. --- tools/build-custom-vllm.sh | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tools/build-custom-vllm.sh b/tools/build-custom-vllm.sh index aef515e4f1b..ed7085918aa 100644 --- a/tools/build-custom-vllm.sh +++ b/tools/build-custom-vllm.sh @@ -74,11 +74,11 @@ uv run --no-project use_existing_torch.py # Install dependencies echo "Installing dependencies..." uv pip install --upgrade pip -# scikit_build_core + nanobind: build backend/deps for xgrammar (a vllm runtime dep -# vllm pins to an exact version that may lack a prebuilt wheel here); needed since -# -e . below runs with --no-build-isolation, which builds vllm's deps against this -# venv's packages instead of isolated per-package build envs. -uv pip install numpy setuptools setuptools_scm setuptools_rust scikit_build_core nanobind +# scikit_build_core: build backend for xgrammar (a vllm runtime dep vllm pins to an +# exact version that may lack a prebuilt wheel here); needed since -e . below runs +# with --no-build-isolation, which builds vllm's deps against this venv's packages +# instead of isolated per-package build envs. +uv pip install numpy setuptools setuptools_scm setuptools_rust scikit_build_core uv pip install torch==2.10.0 --torch-backend=cu129 # Install vLLM using precompiled wheel From dffe6a1c3a74cd05cb9ffff90b0cebcaa493e844 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Wed, 8 Jul 2026 09:29:50 -0700 Subject: [PATCH 14/69] Revert "fix tools/build-custom-vllm.sh" This reverts commit c88955d9ed677356b7055f54c40a9ebb4caf1c3c. --- tools/build-custom-vllm.sh | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tools/build-custom-vllm.sh b/tools/build-custom-vllm.sh index ed7085918aa..4f030227d95 100644 --- a/tools/build-custom-vllm.sh +++ b/tools/build-custom-vllm.sh @@ -74,11 +74,7 @@ uv run --no-project use_existing_torch.py # Install dependencies echo "Installing dependencies..." uv pip install --upgrade pip -# scikit_build_core: build backend for xgrammar (a vllm runtime dep vllm pins to an -# exact version that may lack a prebuilt wheel here); needed since -e . below runs -# with --no-build-isolation, which builds vllm's deps against this venv's packages -# instead of isolated per-package build envs. -uv pip install numpy setuptools setuptools_scm setuptools_rust scikit_build_core +uv pip install numpy setuptools setuptools_scm setuptools_rust uv pip install torch==2.10.0 --torch-backend=cu129 # Install vLLM using precompiled wheel From 2b199b7ef38531c8439d205d33610ffc06461eec Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Wed, 8 Jul 2026 09:29:50 -0700 Subject: [PATCH 15/69] Revert "fix: declare missing build deps for custom vLLM container builds" This reverts commit b66cb735db7f121b66e9e09fc7cc21bd8cfcf88e. --- pyproject.toml | 3 --- tools/build-custom-vllm.sh | 2 +- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 66210e837ea..25b83808249 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -453,9 +453,6 @@ mamba-ssm = [{ requirement = "torch", match-runtime = true }] causal-conv1d = [{ requirement = "torch", match-runtime = true }] fast-hadamard-transform = [{ requirement = "torch", match-runtime = true }] nv-grouped-gemm = [{ requirement = "torch", match-runtime = true }] -# xgrammar (pulled in as a vllm dependency) doesn't declare its own build backend -# (scikit_build_core) in build-system.requires, so uv's build isolation excludes it. -xgrammar = ["scikit_build_core"] # Needed when building from source [[tool.uv.dependency-metadata]] diff --git a/tools/build-custom-vllm.sh b/tools/build-custom-vllm.sh index 4f030227d95..71a04479357 100644 --- a/tools/build-custom-vllm.sh +++ b/tools/build-custom-vllm.sh @@ -74,7 +74,7 @@ uv run --no-project use_existing_torch.py # Install dependencies echo "Installing dependencies..." uv pip install --upgrade pip -uv pip install numpy setuptools setuptools_scm setuptools_rust +uv pip install numpy setuptools setuptools_scm uv pip install torch==2.10.0 --torch-backend=cu129 # Install vLLM using precompiled wheel From dc7e4f563ae9872dd09fbda46945a4d4d28bbaaa Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Wed, 8 Jul 2026 09:30:56 -0700 Subject: [PATCH 16/69] use cudagraph_mode: PIECEWISE in vllm 0.20.0 Signed-off-by: Jiaqi Zeng --- examples/configs/grpo_nano_v3_5_swe_hsg.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/configs/grpo_nano_v3_5_swe_hsg.yaml b/examples/configs/grpo_nano_v3_5_swe_hsg.yaml index 9ef0b444a90..ea1b88dee92 100644 --- a/examples/configs/grpo_nano_v3_5_swe_hsg.yaml +++ b/examples/configs/grpo_nano_v3_5_swe_hsg.yaml @@ -348,6 +348,7 @@ policy: moe_backend: triton mamba_ssm_cache_dtype: "float32" compilation_config: + cudagraph_mode: PIECEWISE cudagraph_capture_sizes: [1,2,4,8,16,32,64] pass_config: fuse_allreduce_rms: false From 3ae4c5fbdaed5b8687b1e606dfcd2586b0f067a1 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 9 Jul 2026 15:53:17 -0700 Subject: [PATCH 17/69] default mtp training on Signed-off-by: Jiaqi Zeng --- examples/configs/grpo_nano_v3_5_swe_hsg.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/configs/grpo_nano_v3_5_swe_hsg.yaml b/examples/configs/grpo_nano_v3_5_swe_hsg.yaml index ea1b88dee92..eaa29f79847 100644 --- a/examples/configs/grpo_nano_v3_5_swe_hsg.yaml +++ b/examples/configs/grpo_nano_v3_5_swe_hsg.yaml @@ -202,7 +202,7 @@ policy: scale_loss_by_dp_cp_size: false # MTP — disabled - mtp_loss_scaling_factor: 0.0 + mtp_loss_scaling_factor: 0.3 mtp_use_repeated_layer: true # null disables MTP entirely (matches the upstream mtp=0 run's behavior). This # checkout's hybrid model gates MTP on `mtp_num_layers is not None` @@ -211,7 +211,7 @@ policy: # PARALLEL NCCL collective timeout). null makes that check False -> MTP is skipped # on all ranks, no partial-rank EP collective. Requires the schema change making # MegatronConfig.mtp_num_layers accept None (nemo_rl/models/policy/__init__.py). - mtp_num_layers: null + mtp_num_layers: 5 mtp_detach_heads: true optimizer: From 11b1240f8e51ef2a17d3e2752b94ab3762264672 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 9 Jul 2026 16:29:06 -0700 Subject: [PATCH 18/69] bump gym Signed-off-by: Jiaqi Zeng --- 3rdparty/Gym-workspace/Gym | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index d3f1bd44235..293d3f59b97 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit d3f1bd44235c88f8c9036d315091e353db518ef9 +Subproject commit 293d3f59b97c55f239830367a628647caf8dfebb From cbb4265539531d503f890a0401248fdd25e25e75 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 9 Jul 2026 22:55:53 -0700 Subject: [PATCH 19/69] feat(ppo): add non-colocated generation and train/inference mismatch diagnostics PPO previously required colocated generation and had no visibility into train/inference logprob mismatch, unlike GRPO. Ports both capabilities from GRPO's proven implementations: separate train/inference clusters with NCCL collective weight refit (SGLang guarded as unsupported non-colocated), and the sequence-level mismatch masking, W&B mismatch plot, and gen_kl_error logging GRPO already surfaces. Signed-off-by: Jiaqi Zeng --- docs/guides/ppo.md | 6 +- examples/configs/ppo_math_1B.yaml | 1 + ....5b-gsm8k-1n8g-automodel-noncolocated.yaml | 61 +++ nemo_rl/algorithms/ppo.py | 358 +++++++++++++++--- tests/functional/L1_Functional_Tests_PPO.sh | 1 + tests/functional/ppo_non_colocated.sh | 57 +++ ...-1.5b-gsm8k-1n8g-automodel-noncolocated.sh | 42 ++ tests/test_suites/nightly.txt | 3 + tests/unit/algorithms/test_ppo.py | 272 +++++++++++++ .../ppo_math_1B_megatron.yaml | 1 + 10 files changed, 749 insertions(+), 53 deletions(-) create mode 100644 examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.yaml create mode 100755 tests/functional/ppo_non_colocated.sh create mode 100755 tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.sh diff --git a/docs/guides/ppo.md b/docs/guides/ppo.md index c7cfee08b56..4bee61af2f9 100644 --- a/docs/guides/ppo.md +++ b/docs/guides/ppo.md @@ -38,9 +38,11 @@ We define a [ValueInterface](../../nemo_rl/models/value/interfaces.py) that cont The value model supports the **Megatron-Core backend** (`value.megatron_cfg.enabled: true`) and the **DTensor backend** (`value.dtensor_cfg.enabled: true`). It uses the same architecture and tokenizer as the policy (configured via `value.model_name`), but is trained with a separate MSE loss on GAE returns. -### Colocated Architecture +### Colocated and Non-Colocated Architecture -PPO uses a colocated architecture where the **policy**, **value model**, and **vLLM generation engine** share the same set of GPUs. GPU memory is managed by offloading models to CPU between stages: the value model is loaded to GPU only during its inference and training phases, then offloaded to make room for other components. +By default, PPO uses a colocated architecture where the **policy**, **value model**, and **vLLM/SGLang generation engine** share the same set of GPUs. GPU memory is managed by offloading models to CPU between stages: the value model is loaded to GPU only during its inference and training phases, then offloaded to make room for other components. + +PPO also supports **non-colocated generation**, where the generation engine runs on a dedicated set of GPUs separate from training (`policy.generation.colocated.enabled: false`, with `policy.generation.colocated.resources.{gpus_per_node,num_nodes}` sizing the dedicated generation resources). In this mode the policy and value model still share a single `train_cluster` (offloaded to CPU between each other as in the colocated case), while generation runs concurrently on its own `inference_cluster`; weights are synchronized via an NCCL broadcast collective spanning both clusters instead of the CUDA-IPC/ZMQ path used when colocated. Non-colocated generation currently requires the vLLM backend — `setup()` asserts this explicitly, since SGLang does not yet implement the collective rendezvous or weight refit needed for non-colocated mode. ### Value Model Configuration diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index 689043be54c..76015b58cb2 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -47,6 +47,7 @@ ppo: source_max: 1.0 target_min: -1.0 # DAPO: scale rewards to [-1, 1] target_max: 1.0 + seq_logprob_error_threshold: null loss_fn: disable_ppo_ratio: false diff --git a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.yaml b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.yaml new file mode 100644 index 00000000000..96ee332338c --- /dev/null +++ b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.yaml @@ -0,0 +1,61 @@ +defaults: ../../ppo_math_1B.yaml +ppo: + num_prompts_per_step: 1024 + num_generations_per_prompt: 1 + max_num_epochs: 15 + ppo_epochs: 1 + val_period: 1 + overlong_filtering: true + reward_shaping: + enabled: false + adv_estimator: + gae_lambda_value: 1.0 + gae_lambda_policy: 1 + reward_scaling: + enabled: false +loss_fn: + ratio_clip_max: 0.2 + ratio_clip_c: 3 +value_loss_fn: + scale: 1.0 + cliprange: 0.5 +checkpointing: + checkpoint_dir: results/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated +policy: + model_name: Qwen/Qwen2.5-1.5B-Instruct + train_global_batch_size: 256 + max_total_sequence_length: 1024 + generation: + max_new_tokens: 512 + vllm_cfg: + async_engine: true + gpu_memory_utilization: 0.4 + max_model_len: 1024 + colocated: + enabled: false + resources: + gpus_per_node: 2 + num_nodes: null +value: + model_name: Qwen/Qwen2.5-1.5B-Instruct + train_micro_batch_size: 4 +data: + max_input_seq_length: 512 + train: + dataset_name: gsm8k + split: train + validation: + dataset_name: gsm8k + split: test + default: + system_prompt_file: examples/prompts/gsm8k.txt +env: + math: + math_verify_impl: hf_math_verify +logger: + log_dir: logs/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated + wandb: + project: nemo-rl + name: ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated +cluster: + gpus_per_node: 8 diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index bb8004799bb..46c1de2c5c0 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -15,10 +15,12 @@ import os import time import warnings +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Any, NotRequired, Optional, TypedDict, TypeVar, cast import numpy as np +import ray import torch from pydantic import BaseModel from torchdata.stateful_dataloader import StatefulDataLoader @@ -34,6 +36,7 @@ _should_log_nemo_gym_responses, _should_use_async_rollouts, _should_use_nemo_gym, + compute_and_apply_seq_logprob_error_masking, extract_initial_prompt_messages, refit_policy_generation, scale_rewards, @@ -60,7 +63,12 @@ ) from nemo_rl.data.utils import load_dataloader_state from nemo_rl.distributed.batched_data_dict import BatchedDataDict -from nemo_rl.distributed.virtual_cluster import ClusterConfig, RayVirtualCluster +from nemo_rl.distributed.virtual_cluster import ( + ClusterConfig, + RayVirtualCluster, + get_ray_cluster_topology, + prepare_segment_topology, +) from nemo_rl.environments.interfaces import EnvironmentInterface from nemo_rl.experience.rollouts import ( run_async_multi_turn_rollout, @@ -139,6 +147,9 @@ class PPOConfig(TypedDict): # Value model trains from step 0; policy training is skipped for # total_steps < this value. Default 0 (train from start). policy_training_start_step: NotRequired[int] + # Sequence-level logprob error masking for training stability. If set, mask sequences with mult_prob_error exceeding this threshold (same scale as token_mult_prob_error metric, e.g., 1.5) + # Note that this is slightly different than Masked Importance Sampling (MIS) because this uses the absolute value of the difference between the training and generation logprobs, whereas MIS just uses the difference between the training and generation logprobs. + seq_logprob_error_threshold: float | None class PPOSaveState(TypedDict): @@ -349,16 +360,12 @@ def setup( # ========================== print("\n▶ Setting up compute cluster...", flush=True) colocated_inference = generation_config["colocated"]["enabled"] - assert colocated_inference, ( - "PPO currently requires colocated generation (vLLM / SGLang sharing GPUs " - "with the policy worker). Set policy.generation.colocated.enabled=true. " - "Non-colocated PPO is not yet supported." - ) reward_model_enabled = ( "env_name" in data_config and data_config["env_name"] == "reward_model" ) total_nodes = cluster_config["num_nodes"] + segment_size = cluster_config.get("segment_size") if reward_model_enabled: rm_resource = env_configs["reward_model"]["resources"] rm_nodes = rm_resource["num_nodes"] @@ -376,31 +383,158 @@ def setup( f"policy_nodes:{policy_nodes} + rm_nodes:{rm_nodes} = total_nodes:{total_nodes}" ) - if total_nodes == 1: - policy_gpus_per_node = cluster_config["gpus_per_node"] - rm_gpus_per_node - assert policy_gpus_per_node > 0, ( - "policy.generation.colocated.resources.gpus_per_node must be > 0 " - "when cluster.num_nodes = 1, " - f"but got {policy_gpus_per_node}." + if colocated_inference: + if total_nodes == 1: + policy_gpus_per_node = cluster_config["gpus_per_node"] - rm_gpus_per_node + assert policy_gpus_per_node > 0, ( + "policy.generation.colocated.resources.gpus_per_node must be > 0 " + "when cluster.num_nodes = 1, " + f"but got {policy_gpus_per_node}." + ) + else: + policy_gpus_per_node = cluster_config["gpus_per_node"] + + node_resource_constraints, _, _ = prepare_segment_topology( + segment_size, policy_nodes + ) + cluster = RayVirtualCluster( + name="grpo_policy_cluster", + bundle_ct_per_node_list=[policy_gpus_per_node] * policy_nodes, + use_gpus=True, + num_gpus_per_node=policy_gpus_per_node, + max_colocated_worker_groups=1 + if generation_config["backend"] == "megatron" + else 3, + port_range_low=cluster_config.get("master_port_range_low"), + port_range_high=cluster_config.get("master_port_range_high"), + segment_size=segment_size, + node_resource_constraints=node_resource_constraints, + ) + train_cluster = cluster + inference_cluster = cluster + print( + f" ✓ Ray cluster for policy initialized with {policy_nodes} nodes", + flush=True, ) else: - policy_gpus_per_node = cluster_config["gpus_per_node"] - - cluster = RayVirtualCluster( - name="grpo_policy_cluster", - bundle_ct_per_node_list=[policy_gpus_per_node] * policy_nodes, - use_gpus=True, - num_gpus_per_node=policy_gpus_per_node, - max_colocated_worker_groups=1 - if generation_config["backend"] == "megatron" - else 3, - ) - train_cluster = cluster - inference_cluster = cluster - print( - f" ✓ Ray cluster for policy initialized with {policy_nodes} nodes", - flush=True, - ) + # Train resources are reduced below to carve out dedicated generation + # (inference) resources from the same overall cluster. + train_gpus_per_node = cluster_config["gpus_per_node"] + train_nodes = policy_nodes + + inference_resources = generation_config["colocated"]["resources"] + inference_gpus_per_node = inference_resources["gpus_per_node"] + inference_nodes = inference_resources["num_nodes"] + + if policy_nodes == 1: + # Train and inference share the single node. + assert ( + inference_gpus_per_node is not None and inference_gpus_per_node > 0 + ), ( + "policy.generation.colocated.resources.gpus_per_node must be explicitly set to a value > 0 " + "when policy_nodes = 1 and inference is non-colocated, " + f"but got {inference_gpus_per_node}." + ) + assert inference_nodes is None or inference_nodes == 1, ( + "policy.generation.colocated.resources.num_nodes must be 1 or set to null " + "when policy_nodes = 1 and inference is non-colocated, " + f"but got {inference_nodes}." + ) + inference_nodes = 1 + reward_gpus_to_subtract = ( + rm_gpus_per_node if total_nodes == 1 and reward_model_enabled else 0 + ) + train_gpus_per_node -= inference_gpus_per_node + reward_gpus_to_subtract + assert train_gpus_per_node > 0, ( + "Not enough GPUs for training: " + f"train_gpus_per_node:{train_gpus_per_node} = cluster_config['gpus_per_node']:{cluster_config['gpus_per_node']} - inference_gpus_per_node:{inference_gpus_per_node}" + + ( + f" - rm_gpus_per_node:{rm_gpus_per_node}" + if total_nodes == 1 and reward_model_enabled + else "" + ) + ) + else: + # Train, inference, and reward model each get dedicated whole nodes. + assert inference_nodes > 0, ( + "policy.generation.colocated.resources.num_nodes must be > 0 " + "when cluster.num_nodes > 1 and inference is non-colocated, " + f"but got {inference_nodes}." + ) + assert ( + inference_gpus_per_node is not None + and inference_gpus_per_node == cluster_config["gpus_per_node"] + ), ( + "policy.generation.colocated.resources.gpus_per_node must be explicitly set and equal to cluster.gpus_per_node " + "when cluster.num_nodes > 1 and inference is non-colocated, " + f"but got inference_gpus_per_node={inference_gpus_per_node}, cluster.gpus_per_node={cluster_config['gpus_per_node']}." + ) + train_nodes -= inference_nodes + + assert train_nodes > 0 and inference_nodes > 0, ( + "Non-colocated mode requires train_nodes > 0 and inference_nodes > 0, " + f"got train_nodes={train_nodes}, inference_nodes={inference_nodes}" + ) + + # NVLink-domain-aware node constraints for the training cluster only; + # unlike GRPO, PPO does not (yet) split a single generation instance + # across multiple nodes, so the inference cluster below gets no + # topology constraints and falls back to default PACK placement. + node_resource_constraints = None + if segment_size is not None: + topology = get_ray_cluster_topology() + num_alive_nodes = len(topology) + required_nodes = train_nodes + inference_nodes + assert num_alive_nodes >= required_nodes, ( + "Not enough alive Ray nodes for all roles: " + f"need {required_nodes} (train={train_nodes} + inference={inference_nodes}), " + f"but only {num_alive_nodes} alive nodes found" + ) + node_resource_constraints, _, _ = prepare_segment_topology( + segment_size, train_nodes, topology=topology, role="training" + ) + + # Value shares train_cluster with policy (2 worker groups timesharing + # the same GPUs, offloaded/reloaded in turn); generation lives on its + # own inference_cluster. + train_cluster = RayVirtualCluster( + name="ppo_train_cluster", + bundle_ct_per_node_list=[train_gpus_per_node] * train_nodes, + use_gpus=True, + num_gpus_per_node=train_gpus_per_node, + max_colocated_worker_groups=2, + port_range_low=cluster_config.get("master_port_range_low"), + port_range_high=cluster_config.get("master_port_range_high"), + segment_size=segment_size, + node_resource_constraints=node_resource_constraints, + ) + if node_resource_constraints is not None: + train_cluster.get_placement_groups() + print( + f" ✓ Ray train cluster initialized with {train_nodes} nodes with {train_gpus_per_node} GPUs per node", + flush=True, + ) + + if segment_size is not None and inference_nodes > 1: + print( + f" ⚠ segment_size={segment_size} is set but PPO does not (yet) apply " + f"NVLink-domain-aware placement to a {inference_nodes}-node inference " + "cluster; falling back to default PACK placement for generation.", + flush=True, + ) + inference_cluster = RayVirtualCluster( + name="ppo_inference_cluster", + bundle_ct_per_node_list=[inference_gpus_per_node] * inference_nodes, + use_gpus=True, + num_gpus_per_node=inference_gpus_per_node, + max_colocated_worker_groups=1, + port_range_low=cluster_config.get("master_port_range_low"), + port_range_high=cluster_config.get("master_port_range_high"), + ) + print( + f" ✓ Ray inference cluster initialized with {inference_nodes} nodes with {inference_gpus_per_node} GPUs per node", + flush=True, + ) # ========================== # Training and Inference @@ -531,46 +665,82 @@ def init_sglang(): pg.finish_generation() return pg, time.perf_counter() - t0 + def init_policy_and_value(): + """Initialize policy then value, sequentially, freeing GPU between each. + + Used as one unit of work (against ``train_cluster``) both in the + colocated sequential path and as the "train" half of the + non-colocated parallel path below. + """ + policy, policy_time = init_policy() + # Block until the policy worker's __init__ completes and offload to + # CPU, freeing GPU for value model initialization. Policy will be + # reloaded before the vLLM refit step below. + policy.offload_to_cpu() + + print(" ⚙️ Initializing value model for GAE...", flush=True) + value_model, value_time = init_value() + # Block until the value worker's __init__ completes and offload + # model + optimizer to CPU. Without this, __init__ runs asynchronously + # in the Ray actor and may overlap with generation, causing GPU OOM. + value_model.finish_training() + print(f" ✓ Value model initialized in {value_time:.2f}s", flush=True) + + return policy, policy_time, value_model, value_time + def initialize_generation_with_policy( init_generation_fn, generation_name: str, init_time_key: str, + colocated_inference: bool, worker_init_timing_metrics: dict, ): - """Generic function to initialize a generation engine (vLLM or SGLang) along with policy. + """Generic function to initialize a generation engine (vLLM or SGLang) along with policy and value. Args: init_generation_fn: Function that initializes the generation engine (init_vllm or init_sglang) generation_name: Name of the generation engine ("vLLM" or "SGLang") init_time_key: Key name for storing initialization time in metrics ("vllm_init_time_s" or "sglang_init_time_s") + colocated_inference: Whether inference shares GPUs with training worker_init_timing_metrics: Dictionary to store timing metrics Returns: - Tuple of (policy_generation, policy) + Tuple of (policy_generation, policy, value_model) """ - # Initialize generation engine first so it claims its GPU memory - # before policy/value workers are constructed; then policy, then value. - print(" ⚙️ Initializing workers (colocated mode)", flush=True) + if colocated_inference: + # Initialize generation engine first so it claims its GPU memory + # before policy/value workers are constructed; then policy, then value. + print(" ⚙️ Initializing workers (colocated mode)", flush=True) - policy_generation, generation_time = init_generation_fn() - worker_init_timing_metrics[init_time_key] = generation_time + policy_generation, generation_time = init_generation_fn() + worker_init_timing_metrics[init_time_key] = generation_time - policy, policy_time = init_policy() - # Block until the policy worker's __init__ completes and offload to - # CPU, freeing GPU for value model initialization. Policy will be - # reloaded before the vLLM refit step below. - policy.offload_to_cpu() - worker_init_timing_metrics["policy_init_time_s"] = policy_time + policy, policy_time, value_model, value_time = init_policy_and_value() + worker_init_timing_metrics["policy_init_time_s"] = policy_time + worker_init_timing_metrics["value_init_time_s"] = value_time + else: + # Generation lives on a disjoint inference_cluster, so it can + # initialize in parallel with the policy+value chain on + # train_cluster instead of waiting for GPU memory to free up. + print( + " ⚡ Using parallel worker initialization (non-colocated mode)", + flush=True, + ) + parallel_start_time = time.perf_counter() + with ThreadPoolExecutor(max_workers=2) as executor: + generation_future = executor.submit(init_generation_fn) + policy_value_future = executor.submit(init_policy_and_value) + policy_generation, generation_time = generation_future.result() + policy, policy_time, value_model, value_time = ( + policy_value_future.result() + ) + parallel_wall_time = time.perf_counter() - parallel_start_time - print(" ⚙️ Initializing value model for GAE...", flush=True) - value_model, value_time = init_value() - # Block until the value worker's __init__ completes and offload - # model + optimizer to CPU. Without this, __init__ runs asynchronously - # in the Ray actor and may overlap with vLLM generation, causing - # GPU OOM. - value_model.finish_training() - worker_init_timing_metrics["value_init_time_s"] = value_time - print(f" ✓ Value model initialized in {value_time:.2f}s", flush=True) + worker_init_timing_metrics[init_time_key] = generation_time + worker_init_timing_metrics["policy_init_time_s"] = policy_time + worker_init_timing_metrics["value_init_time_s"] = value_time + worker_init_timing_metrics["parallel_wall_time_s"] = parallel_wall_time + worker_init_timing_metrics["parallel_init_enabled"] = True return policy_generation, policy, value_model @@ -578,6 +748,13 @@ def initialize_generation_with_policy( f"PPO requires vllm or sglang generation backend; got {backend!r}. " "The megatron generation backend is not supported." ) + assert colocated_inference or backend == "vllm", ( + "Non-colocated PPO currently requires the vLLM generation backend. " + "SGLangGeneration.init_collective() is a no-op, so the training side's " + "NCCL collective rendezvous would hang waiting for peers that never " + "join. Set policy.generation.backend=vllm or " + "policy.generation.colocated.enabled=true." + ) if backend == "vllm": # vLLM generation: setup config, then initialize with policy @@ -612,6 +789,7 @@ def initialize_generation_with_policy( init_generation_fn=init_vllm, generation_name="vLLM", init_time_key="vllm_init_time_s", + colocated_inference=colocated_inference, worker_init_timing_metrics=worker_init_timing_metrics, ) @@ -631,6 +809,7 @@ def initialize_generation_with_policy( init_generation_fn=init_sglang, generation_name="SGLang", init_time_key="sglang_init_time_s", + colocated_inference=colocated_inference, worker_init_timing_metrics=worker_init_timing_metrics, ) @@ -645,6 +824,28 @@ def initialize_generation_with_policy( # print the node IP and GPU ID of the policy workers for debugging policy.print_node_ip_and_gpu_id() + # Non-colocated inference has no shared GPU memory for weight refit, so + # establish an NCCL collective spanning train_cluster + inference_cluster + # up front; refit_policy_generation() broadcasts weights over it later. + if not colocated_inference: + t0 = time.perf_counter() + ip, port = train_cluster.get_master_address_and_port() + print(f"Using ip: {ip}, port: {port} for collective communication", flush=True) + train_world_size = train_cluster.world_size() + inference_world_size = inference_nodes * inference_gpus_per_node + world_size = train_world_size + inference_world_size + + futures_train = policy.init_collective( + ip, port, world_size, train_world_size=train_world_size + ) + futures_inference = policy_generation.init_collective( + ip, port, world_size, train_world_size=train_world_size + ) + ray.get(futures_train + futures_inference) + worker_init_timing_metrics["collective_init_time_s"] = ( + time.perf_counter() - t0 + ) + # Reload policy weights to GPU before refit (they may have been offloaded # during setup to free GPU for value model initialization). policy.prepare_for_training() @@ -1226,6 +1427,40 @@ def ppo_train( policy.finish_inference() + # Seq-level logprob error metrics/masking (train/inference mismatch + # diagnostics; always computed so a threshold can be tuned from the + # logged metrics even before it's enabled). + # + # NOTE: when the threshold masks a sequence, this only zeroes + # train_data["sample_mask"], which takes effect at the loss level + # (ClippedPGLossFn/MseValueLossFn both read sample_mask directly). + # It does NOT get combined into the `mask` GAE receives below + # (adv_kwargs["mask"] = train_data["token_mask"] alone) — the same + # pre-existing characteristic overlong_filtering's loss_multiplier + # masking already has in PPO, unlike GRPO which explicitly combines + # token_mask * sample_mask before advantage computation. So a masked + # sequence's tokens still contribute to GAE's whitening/bootstrap + # statistics even though they're excluded from the loss. + # + # Unlike GRPO, PPO never skips real prev_logprobs computation (no + # force_on_policy_ratio placeholder-zero path), so it's always safe + # to call this unconditionally. If PPO ever grows such a skip path, + # port GRPO's companion guard (grpo.py's force_on_policy_ratio + + # seq_logprob_error_threshold conflict warning) alongside it. + seq_logprob_error_threshold = master_config.ppo.get( + "seq_logprob_error_threshold", None + ) + seq_error_result = compute_and_apply_seq_logprob_error_masking( + train_data=train_data, + rewards=rewards, + seq_logprob_error_threshold=seq_logprob_error_threshold, + ) + seq_logprob_error_metrics = seq_error_result + if "num_masked_seqs" in seq_logprob_error_metrics: + seq_logprob_error_metrics["num_masked_seqs_by_logprob_error"] = ( + seq_logprob_error_metrics.pop("num_masked_seqs") + ) + # Build prompt IDs for advantage estimation (groups responses from same prompt). # Use the token-length-based extractor so multi-turn prompts containing # assistant messages still resolve to the original prompt only. @@ -1489,6 +1724,9 @@ def ppo_train( if "global_valid_toks" in metrics: total_valid_tokens += metrics["global_valid_toks"] + # Always log sequence-level error metrics (useful for deciding threshold) + metrics.update(seq_logprob_error_metrics) + ## Checkpointing consumed_samples += master_config.ppo["num_prompts_per_step"] timeout.mark_iteration() @@ -1610,11 +1848,29 @@ def ppo_train( reduction_op="sum" ) # type: ignore + # track example with high token mult prob error above 1.05 + # (metrics["token_mult_prob_error"] is only populated when + # train_results is not None, i.e. outside critic-only warmup steps) + if metrics.get("token_mult_prob_error", 0) > 1.05: + logger.log_plot_token_mult_prob_error( + { + "prompt_lengths": repeated_batch["length"], + "full_lengths": input_lengths, + "generation_logprobs": train_data["generation_logprobs"], + "prev_logprobs": train_data["prev_logprobs"], + "token_mask": train_data["token_mask"], + "sample_mask": train_data["sample_mask"], + }, + total_steps + 1, + name="train/token_mult_prob_error_plot_sample", + ) + del train_data print("\n📊 Training Results:") if train_results is not None: print(f" • Policy Loss: {metrics.get('loss', 'N/A')}") + print(f" • Generation KL Error: {metrics.get('gen_kl_error', 'N/A')}") if value_results is not None: print(f" • Critic Loss: {metrics.get('critic/loss', 'N/A')}") print(f" • Critic Grad Norm: {metrics.get('critic/grad_norm', 'N/A')}") diff --git a/tests/functional/L1_Functional_Tests_PPO.sh b/tests/functional/L1_Functional_Tests_PPO.sh index c398722dbc5..8865c2157dc 100755 --- a/tests/functional/L1_Functional_Tests_PPO.sh +++ b/tests/functional/L1_Functional_Tests_PPO.sh @@ -36,6 +36,7 @@ run_test() { run_test fast uv run --no-sync bash ./tests/functional/ppo_automodel.sh run_test fast uv run --no-sync bash ./tests/functional/ppo_megatron.sh +run_test uv run --no-sync bash ./tests/functional/ppo_non_colocated.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/ppo_non_colocated.sh b/tests/functional/ppo_non_colocated.sh new file mode 100755 index 00000000000..010761991e8 --- /dev/null +++ b/tests/functional/ppo_non_colocated.sh @@ -0,0 +1,57 @@ +#!/bin/bash + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) +# Mark the current repo as safe, since wandb fetches metadata about the repo +git config --global --add safe.directory $PROJECT_ROOT + +set -eou pipefail + +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 +export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} + +rm -rf $EXP_DIR $LOG_DIR +mkdir -p $EXP_DIR $LOG_DIR + +cd $PROJECT_ROOT +uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ + $PROJECT_ROOT/examples/run_ppo.py \ + policy.model_name=Qwen/Qwen2.5-0.5B \ + value.model_name=Qwen/Qwen2.5-0.5B \ + ppo.num_prompts_per_step=2 \ + ppo.num_generations_per_prompt=4 \ + ppo.ppo_epochs=2 \ + policy.train_global_batch_size=4 \ + policy.logprob_batch_size=4 \ + policy.train_micro_batch_size=1 \ + policy.generation.colocated.enabled=false \ + policy.generation.colocated.resources.gpus_per_node=1 \ + policy.generation.vllm_cfg.async_engine=true \ + value.train_global_batch_size=4 \ + value.train_micro_batch_size=1 \ + cluster.gpus_per_node=2 \ + ppo.max_num_steps=2 \ + logger.tensorboard_enabled=true \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=false \ + logger.monitor_gpus=true \ + checkpointing.enabled=false \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/token_mult_prob_error"]) < 1.05' \ + 'min(data["train/probs_ratio_clamped_min"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_min"]) < 1.21' \ + 'min(data["train/probs_ratio_clamped_max"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_max"]) < 1.29' \ + 'max(data["train/critic/loss"]) < 6.0' \ + 'min(data["train/critic/loss"]) >= 0' \ + 'max(data["train/critic/explained_var"]) <= 1.0001' \ + 'max(data["train/critic/grad_norm"]) < 350' diff --git a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.sh b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.sh new file mode 100755 index 00000000000..c156ce816be --- /dev/null +++ b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.sh @@ -0,0 +1,42 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=1 +STEPS_PER_RUN=40 +MAX_STEPS=40 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=60 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +# Run the experiment +cd $PROJECT_ROOT +uv run examples/run_ppo.py \ + --config $CONFIG_PATH \ + ppo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + $@ \ + 2>&1 | tee $RUN_LOG + +# Convert tensorboard logs to json +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Only run metrics if the target step is reached +if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then + uv run tests/check_metrics.py $JSON_METRICS \ + 'data["train/reward"]["40"] > 0.75' \ + 'data["validation/accuracy"]["40"] > 0.65' + + # Clean up checkpoint directory after successful run to save space. + rm -rf "$CKPT_DIR" +fi diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index 81e8ebc90c4..c16cdaceac7 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -264,6 +264,9 @@ tests/test_suites/llm/distillation-xtoken-off-policy-multiteacher-qwen3-4b-llama # DTensor PPO with value sequence parallelism (TP2+SP) (Qwen2.5-1.5B, GSM8K) tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-valuetp2sp.sh +# DTensor PPO with non-colocated generation (Qwen2.5-1.5B, GSM8K) +tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.sh + # Megatron PPO with value sequence parallelism (TP2+SP) and dynamic batching (Qwen2.5-1.5B, GSM8K) tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-megatron-valuetp2sp-dynbatch.sh diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 5682f574bfd..4c3a43b20a8 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -693,3 +693,275 @@ def test_create_advantage_estimator_requires_adv_estimator_key(): with pytest.raises(KeyError): _create_advantage_estimator(master_config) + + +# =============================================================================== +# Non-colocated generation (setup()) +# =============================================================================== + + +def _build_ppo_master_config(): + """Minimal PPO MasterConfig sufficient to drive setup() end to end. + + Uses model_construct (like test_grpo.py's mock_grpo_components) so nested + sub-configs can stay plain dicts instead of satisfying full pydantic/TypedDict + validation. Individual tests mutate fields in place before calling setup(). + """ + from nemo_rl.algorithms.ppo import MasterConfig + + return MasterConfig.model_construct( + **{ + "policy": { + "model_name": "fake-model", + "train_global_batch_size": 1, + "train_micro_batch_size": 1, + "max_total_sequence_length": 128, + "generation": { + "backend": "vllm", + "model_name": "fake-model", + "colocated": { + "enabled": True, + "resources": {"gpus_per_node": None, "num_nodes": None}, + }, + "vllm_cfg": { + "precision": "bfloat16", + "kv_cache_dtype": "auto", + }, + "vllm_kwargs": {}, + }, + }, + "value": { + "megatron_cfg": {"enabled": False}, + "dtensor_cfg": {"enabled": True, "context_parallel_size": 1}, + "sequence_packing": {"enabled": False}, + "dynamic_batching": {"enabled": False}, + }, + "loss_fn": ClippedPGLossConfig(), + "value_loss_fn": MseValueLossConfig(), + "env": {}, + "data": {"shuffle": False, "num_workers": 0}, + "ppo": { + "seed": 42, + "batch_multiplier": 1, + "num_prompts_per_step": 1, + "use_dynamic_sampling": False, + "val_period": 0, + "val_at_start": False, + "val_at_end": False, + "ppo_epochs": 1, + "max_num_steps": 1, + "max_num_epochs": 1, + }, + "logger": {}, + "cluster": {"num_nodes": 1, "gpus_per_node": 2}, + "checkpointing": {}, + } + ) + + +@pytest.mark.parametrize( + "num_nodes,inference_num_nodes", + [ + pytest.param(1, None, id="single_node"), + pytest.param(2, 1, id="multi_node"), + ], +) +def test_ppo_noncolocated_requires_explicit_gpus_per_node( + num_nodes, inference_num_nodes +): + """Non-colocated PPO must set an explicit inference GPU count, whether + train and inference share a single node or each get dedicated nodes.""" + from unittest.mock import MagicMock, patch + + from nemo_rl.algorithms.ppo import setup + + master_config = _build_ppo_master_config() + master_config.policy["generation"]["colocated"] = { + "enabled": False, + "resources": {"gpus_per_node": None, "num_nodes": inference_num_nodes}, + } + master_config.cluster["num_nodes"] = num_nodes + master_config.cluster["gpus_per_node"] = 8 + + tokenizer = MagicMock() + dataset = MagicMock() + dataset.__len__ = MagicMock(return_value=10) + + with ( + patch("nemo_rl.algorithms.ppo.Logger"), + patch("nemo_rl.algorithms.ppo.CheckpointManager") as mock_checkpointer, + patch("nemo_rl.algorithms.ppo.StatefulDataLoader"), + pytest.raises( + AssertionError, + match="policy.generation.colocated.resources.gpus_per_node must be explicitly set", + ), + ): + mock_checkpointer.return_value.get_latest_checkpoint_path.return_value = None + mock_checkpointer.return_value.load_training_info.return_value = None + setup(master_config, tokenizer, dataset, None) + + +def test_ppo_noncolocated_rejects_sglang_backend(): + """SGLangGeneration.init_collective() is a no-op, so non-colocated PPO must + fail loudly at setup() rather than hang forever waiting for the training + side's NCCL collective to be joined by peers that never connect.""" + from unittest.mock import MagicMock, patch + + from nemo_rl.algorithms.ppo import setup + + master_config = _build_ppo_master_config() + master_config.policy["generation"]["backend"] = "sglang" + master_config.policy["generation"]["colocated"] = { + "enabled": False, + "resources": {"gpus_per_node": 1, "num_nodes": None}, + } + master_config.cluster["num_nodes"] = 1 + master_config.cluster["gpus_per_node"] = 2 + + tokenizer = MagicMock() + dataset = MagicMock() + dataset.__len__ = MagicMock(return_value=10) + + with ( + patch("nemo_rl.algorithms.ppo.Logger"), + patch("nemo_rl.algorithms.ppo.CheckpointManager") as mock_checkpointer, + patch("nemo_rl.algorithms.ppo.StatefulDataLoader"), + pytest.raises( + AssertionError, + match="Non-colocated PPO currently requires the vLLM generation backend", + ), + ): + mock_checkpointer.return_value.get_latest_checkpoint_path.return_value = None + mock_checkpointer.return_value.load_training_info.return_value = None + setup(master_config, tokenizer, dataset, None) + + +@pytest.mark.parametrize("colocated", [True, False]) +def test_ppo_setup_cluster_split_matches_colocation_mode(monkeypatch, colocated): + """train_cluster/inference_cluster identity and worker-group budget by mode. + + Colocated: train_cluster is inference_cluster (single shared pool of + generation+policy+value). Non-colocated: they are distinct clusters, and + train_cluster must budget for 2 co-timesharing worker groups (policy + + value) since generation now lives on its own inference_cluster. + """ + from unittest.mock import MagicMock + + from nemo_rl.algorithms import ppo as ppo_mod + + class DummyLogger: + def log_hyperparams(self, *_args, **_kwargs): + pass + + def log_metrics(self, *_args, **_kwargs): + pass + + class DummyCheckpointer: + def get_latest_checkpoint_path(self): + return None + + def load_training_info(self, _path): + return None + + class DummyLoader: + def __init__(self, *_args, **_kwargs): + pass + + def __len__(self): + return 1 + + class DummyCluster: + instances = [] + + def __init__(self, *_args, max_colocated_worker_groups=1, **_kwargs): + self.max_colocated_worker_groups = max_colocated_worker_groups + DummyCluster.instances.append(self) + + def world_size(self): + return 1 + + def get_master_address_and_port(self): + return "127.0.0.1", 1234 + + def get_placement_groups(self): + return [] + + class DummyPolicy: + def offload_to_cpu(self): + pass + + def print_node_ip_and_gpu_id(self): + pass + + def init_collective(self, *_args, **_kwargs): + return [] + + def prepare_for_training(self): + pass + + def prepare_refit_info(self): + return {} + + class DummyValue: + def __init__(self, *_args, **_kwargs): + pass + + def finish_training(self): + pass + + class DummyVllmGeneration: + def __init__(self, *_args, **_kwargs): + pass + + def finish_generation(self): + pass + + def prepare_refit_info(self, _state): + pass + + def init_collective(self, *_args, **_kwargs): + return [] + + DummyCluster.instances = [] + monkeypatch.setattr(ppo_mod, "Logger", lambda *_a, **_k: DummyLogger()) + monkeypatch.setattr( + ppo_mod, "CheckpointManager", lambda *_a, **_k: DummyCheckpointer() + ) + monkeypatch.setattr(ppo_mod, "StatefulDataLoader", DummyLoader) + monkeypatch.setattr(ppo_mod, "RayVirtualCluster", DummyCluster) + monkeypatch.setattr(ppo_mod, "Policy", lambda *_a, **_k: DummyPolicy()) + monkeypatch.setattr(ppo_mod, "Value", lambda *_a, **_k: DummyValue()) + monkeypatch.setattr( + ppo_mod, "VllmGeneration", lambda *_a, **_k: DummyVllmGeneration() + ) + monkeypatch.setattr(ppo_mod.ray, "get", lambda x: x) + + master_config = _build_ppo_master_config() + if colocated: + master_config.policy["generation"]["colocated"] = { + "enabled": True, + "resources": {"gpus_per_node": None, "num_nodes": None}, + } + else: + master_config.policy["generation"]["colocated"] = { + "enabled": False, + "resources": {"gpus_per_node": 1, "num_nodes": None}, + } + master_config.cluster["num_nodes"] = 1 + master_config.cluster["gpus_per_node"] = 2 + + tokenizer = MagicMock() + dataset = MagicMock() + dataset.__len__ = MagicMock(return_value=1) + + _, _, _, (train_cluster, inference_cluster), *_ = ppo_mod.setup( + master_config, tokenizer, dataset, None + ) + + if colocated: + assert train_cluster is inference_cluster + assert train_cluster.max_colocated_worker_groups == 3 + else: + assert train_cluster is not inference_cluster + assert train_cluster.max_colocated_worker_groups == 2 + assert inference_cluster.max_colocated_worker_groups == 1 diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 00c1fb02845..d20060f88b7 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -47,6 +47,7 @@ ppo: source_max: 1.0 target_min: -1.0 # DAPO: scale rewards to [-1, 1] target_max: 1.0 + seq_logprob_error_threshold: null loss_fn: disable_ppo_ratio: false From 04011ac18f7c752c098925690979a44b7c3b6fa3 Mon Sep 17 00:00:00 2001 From: Sugam Devare Date: Thu, 9 Jul 2026 15:45:51 -0700 Subject: [PATCH 20/69] feat: add mask_sample --- nemo_rl/algorithms/grpo.py | 38 ++++++++++++++++++++++++++++++++++ nemo_rl/experience/rollouts.py | 13 ++++++++++++ 2 files changed, 51 insertions(+) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 981e0a02bf3..375899a7152 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -2690,6 +2690,25 @@ def grpo_train( loss_multiplier[truncated] = 0 repeated_batch["loss_multiplier"] = loss_multiplier + + # Mask samples flagged by the environment (keep for advantage, mask for gradient) + if "mask_sample" in repeated_batch: + loss_multiplier = repeated_batch["loss_multiplier"].clone() + mask_sample = repeated_batch["mask_sample"] + + if isinstance(mask_sample, list): + mask_sample = torch.tensor(mask_sample, dtype=torch.bool) + mask_sample_bool = mask_sample.bool() + + num_masked = int(mask_sample_bool.sum().item()) + if num_masked > 0: + print( + f" 📊 mask_sample filtering: masking {num_masked}/{len(mask_sample_bool)} env-flagged samples", + flush=True, + ) + loss_multiplier[mask_sample_bool] = 0 + repeated_batch["loss_multiplier"] = loss_multiplier + add_grpo_token_loss_masks_and_generation_logprobs( repeated_batch["message_log"] ) @@ -4000,6 +4019,25 @@ def async_grpo_train( loss_multiplier[truncated] = 0 repeated_batch["loss_multiplier"] = loss_multiplier + # Mask samples flagged by the environment (keep for advantage, mask for gradient) + with timer.time("mask_sample_filter"): + if "mask_sample" in repeated_batch: + loss_multiplier = repeated_batch["loss_multiplier"].clone() + mask_sample = repeated_batch["mask_sample"] + + if isinstance(mask_sample, list): + mask_sample = torch.tensor(mask_sample, dtype=torch.bool) + mask_sample_bool = mask_sample.bool() + + num_masked = int(mask_sample_bool.sum().item()) + if num_masked > 0: + print( + f" 📊 mask_sample filtering: masking {num_masked}/{len(mask_sample_bool)} env-flagged samples", + flush=True, + ) + loss_multiplier[mask_sample_bool] = 0 + repeated_batch["loss_multiplier"] = loss_multiplier + # Add loss mask to each message # Only unmask assistant messages that were actually generated (have generation_logprobs), # not assistant messages that were part of the prompt history diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index e4b7c0350f4..4ac564557e2 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -1974,6 +1974,19 @@ def run_async_nemo_gym_rollout( "truncated": torch.tensor( [m["hit_max_tokens"] for m in all_sample_metrics], dtype=torch.bool ), + # Agent/env-driven mask flag — True means this sample should be masked + # from the GRPO gradient (kept for advantage computation). + "mask_sample": torch.tensor( + [ + bool( + (r["full_result"].get("instance_config") or {}).get( + "mask_sample", False + ) + ) + for r in results + ], + dtype=torch.bool, + ), } ) From 0a5adde916f45e87ba4ed2c4add9b9c63e54c924 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 10 Jul 2026 20:26:43 -0700 Subject: [PATCH 21/69] fix hsg config Signed-off-by: Jiaqi Zeng --- examples/configs/grpo_nano_v3_5_swe_hsg.yaml | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/examples/configs/grpo_nano_v3_5_swe_hsg.yaml b/examples/configs/grpo_nano_v3_5_swe_hsg.yaml index eaa29f79847..84efb1795ae 100644 --- a/examples/configs/grpo_nano_v3_5_swe_hsg.yaml +++ b/examples/configs/grpo_nano_v3_5_swe_hsg.yaml @@ -414,13 +414,13 @@ env: dataset_path: ${data.train.data_path} apptainer_memory_limit_mb: 65536 container_formatter: - - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swerebench/{instance_id}.sif" - - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/nv_internal/{instance_id}.sif" - - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/r2e_gym/{instance_id}.sif" - - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/sweb.eval.arm64.{instance_id}.sif" - - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swebench/swe-bench.eval.arm64.{instance_id}.sif" - - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/swebenchpro_ots/{instance_id}.sif" - - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swegym/sweb.eval.arm64.{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swerebench/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/nv_internal/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/r2e_gym/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swegym/sweb.eval.arm64.{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swebench/swe-bench.eval.arm64.{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/mercor/swebenchpro_ots/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/mercor/{instance_id}.sif" swe_agents_val: responses_api_agents: swe_agents: From 165a41e067a3708464dbf1eaf8063b985b9f469f Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 10 Jul 2026 20:37:41 -0700 Subject: [PATCH 22/69] add cmh swe config Signed-off-by: Jiaqi Zeng --- examples/configs/grpo_nano_v3_5_swe_cmh.yaml | 480 +++++++++++++++++++ 1 file changed, 480 insertions(+) create mode 100644 examples/configs/grpo_nano_v3_5_swe_cmh.yaml diff --git a/examples/configs/grpo_nano_v3_5_swe_cmh.yaml b/examples/configs/grpo_nano_v3_5_swe_cmh.yaml new file mode 100644 index 00000000000..eaa29f79847 --- /dev/null +++ b/examples/configs/grpo_nano_v3_5_swe_cmh.yaml @@ -0,0 +1,480 @@ +# ============================================================================= +# GRPO Ultra E2E — Production SWE Config for GB200 NVL72 +# ============================================================================= +# Production config for Ultra V3 SWE end-to-end GRPO training on +# 128 nodes × 4 GPUs/node (64 gen + 64 train by default). +# +# Static algorithm, loss, and environment settings live here. +# Per-run overrides (model path, data paths, parallelism, precision) +# are applied by the launch script (repro_ultra_e2e.sh) or Hydra CLI. +# +# - gpus_per_node: 4 +# - TP: 8, CP: 8, EP: 32, PP: 1 +# - vLLM TP: 8 +# - Non-colocated async inference with 64 generation nodes +# - Max sequence length: 131072 +# ============================================================================= + +# ============================================================================= +# Cluster — overridden by launch script for parallelism changes +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 128 + segment_size: 16 + +# ============================================================================= +# Checkpointing +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/grpo_ultra_e2e" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 3 + ft_keep_latest_k: 1 + ft_save_period: 1 + save_optimizer: true + checkpoint_must_save_by: "00:03:30:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# GRPO Algorithm +# ============================================================================= +grpo: + num_prompts_per_step: 16 + num_generations_per_prompt: 32 + num_val_generations_per_prompt: 2 + max_rollout_turns: 1 + max_num_epochs: 4 + max_num_steps: 1000000 + normalize_rewards: true + use_leave_one_out_baseline: true + advantage_clip_low: -100 + advantage_clip_high: 100 + val_period: 10000 + val_at_start: false + val_at_end: false + overlong_filtering: false + max_val_samples: null + val_batch_size: 256 + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + penalize_invalid_tool_call: true + invalid_tool_call_advantage: -5.0 + penalize_malformed_thinking: true + malformed_thinking_advantage: -5.0 + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + + use_best_at_k: false + best_at_k_k: 8 + best_at_k_m: 1000 + + use_combined_training: false + combined_training_weight_mode: "auto" + combined_training_best_at_k_weight: 0.2 + combined_training_pass_at_1_weight: 1.0 + + dynamic_sampling_oversample_ratio: 1.0 + seq_logprob_error_threshold: 2 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: null + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + force_on_policy_ratio: true + use_kl_in_reward: false + +# ============================================================================= +# Policy +# ============================================================================= +policy: + model_name: "/lustre/fsw/portfolios/llmservice/users/wdai/megatron-lm-ultra/checkpoints/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-288k-nano-loss-032026/iter_0007000/hf" + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 512 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 131072 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 32 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + use_gloo_process_groups: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + # community_import.py reads megatron_cfg["gradient_accumulation_fusion"] directly + # (no default) on the HF->megatron conversion path; absent here it raises KeyError. + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled + mtp_loss_scaling_factor: 0.3 + mtp_use_repeated_layer: true + # null disables MTP entirely (matches the upstream mtp=0 run's behavior). This + # checkout's hybrid model gates MTP on `mtp_num_layers is not None` + # (hybrid_model.py:544): 0 still enters the MTP block and asserts > 0, and 1 runs + # MTP forward on only the last-stage ranks (suspected cause of the EXPERT_MODEL_ + # PARALLEL NCCL collective timeout). null makes that check False -> MTP is skipped + # on all ranks, no partial-rank EP collective. Requires the schema change making + # MegatronConfig.mtp_num_layers accept None (nemo_rl/models/policy/__init__.py). + mtp_num_layers: 5 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 3.0e-6 + min_lr: 3.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 3e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + # FP8 — overridden by precision recipe in launch script + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo + # were dropped from this checkout's megatron-bridge CheckpointConfig (submodule + # 554c7b9); leaving them in raises InstantiationException "Unexpected config keys". + # Removed for the current-repo reproduction. + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + # Must be a multiple of minimum_pad_factor = cp_size*2*tp_size (see + # nemo_rl/models/megatron/data.py). With CP=16 + TP=4 + sequence_parallel that is + # 128; the old value (=TP=4) made get_logprobs assert and abort every training + # step. Formula matches examples/nemo_gym/grpo_nanov3.yaml (TP*CP*2). + make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 11001 + port_range_high: 15000 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 1 + gpu_memory_utilization: 0.8 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + # reasoning_parser_plugin lives at vllm_cfg top level in this checkout + # (vllm_worker_async reads vllm_cfg["reasoning_parser_plugin"]); passing it + # inside http_server_serving_chat_kwargs makes OpenAIServingChat raise + # "unexpected keyword argument 'reasoning_parser_plugin'". Path updated to this + # repo's location (the parser moved out of nemo_rl/utils/). + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + + vllm_kwargs: + attention_backend: FLASH_ATTN + moe_backend: triton + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_mode: PIECEWISE + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 64 + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + use_multiple_dataloader: false + train: + data_path: null # Set by launch script + validation: + data_path: null # Set by launch script + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + SWE Agents +# ============================================================================= +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + # Reuse pre-built gym venvs. With false, all 3 gym servers (policy_model, + # swe_agents_{train,val}) rebuild the editable nemo-gym package in parallel and + # contend on the same uv distribution-cache lock on Lustre; policy_model then + # times out (default 300s) and NemoGym spinup fails. Mirrors the qwen-30b SWE run. + skip_venv_if_present: true + num_gpu_nodes: 0 + port_range_low: 15001 + port_range_high: 20000 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml + swe_agents_train: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}} + swebench_agent_timeout: 3600 + run_with_mixed_prompts: true + dataset_path: ${data.train.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swerebench/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/nv_internal/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/r2e_gym/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/sweb.eval.arm64.{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swebench/swe-bench.eval.arm64.{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/swebenchpro_ots/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swegym/sweb.eval.arm64.{instance_id}.sif" + swe_agents_val: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}} + swebench_agent_timeout: 3600 + dataset_path: ${data.validation.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter} + use_absolute_ip: true + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "logs" + num_val_samples_to_print: 0 + wandb_enabled: true + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "ultra-v3-swe-e2e" + name: "grpo-ultra-e2e" + tensorboard: {} + mlflow: + experiment_name: "grpo-ultra-e2e" + run_name: "grpo-ultra-e2e" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + +# ============================================================================= +# Token IDs (model-specific, used by token-based penalties) +# ============================================================================= +token_ids: + eos: 2 # + think_open: 12 # + think_close: 13 # + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +penalize_duplicated_reasoning: true # reasoning content == final answer +penalize_empty_final_answer: true # last message output has empty content +penalize_eos_token: true # eos token appears in generation +penalize_malformed_think_tag: true # /<\/think> count != 1 per turn From a3d2a49beb3e2d1c5ee2d1e433176c89e9dc3a20 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 10 Jul 2026 20:38:47 -0700 Subject: [PATCH 23/69] fix gym merge Signed-off-by: Jiaqi Zeng --- 3rdparty/Gym-workspace/Gym | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index 293d3f59b97..8b145ad5693 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit 293d3f59b97c55f239830367a628647caf8dfebb +Subproject commit 8b145ad5693a50740fcbb12fe89bd56eebcdb6d4 From ea57ce0d7f79b72b7bda0009887ba7dbc1cf55b2 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Sat, 11 Jul 2026 13:40:14 -0700 Subject: [PATCH 24/69] feat(ppo): asynchronous PPO with non-colocated generation Add async PPO training (async_ppo_train) mirroring async GRPO: an AsyncTrajectoryCollector + ReplayBuffer Ray-actor pipeline with a weight-version staleness window, critic-warmup trajectory banking (warmup_max_trajectory_age_steps), in-flight weight updates, and freeze-aware trajectory-age diagnostics. Also route async rollouts through validate(), fix replay-buffer and value-optimizer checkpoint/resume, and add unit + functional (ppo_async.sh) coverage. Signed-off-by: Jiaqi Zeng --- docs/guides/ppo.md | 23 + examples/configs/ppo_math_1B.yaml | 18 + examples/run_ppo.py | 93 +- .../algorithms/async_utils/replay_buffer.py | 69 +- .../async_utils/trajectory_collector.py | 143 +- nemo_rl/algorithms/ppo.py | 1402 ++++++++++++++++- ray.sub | 23 +- tests/functional/L1_Functional_Tests_PPO.sh | 1 + tests/functional/ppo_async.sh | 64 + tests/unit/algorithms/test_async_utils.py | 32 + tests/unit/algorithms/test_ppo.py | 325 ++++ .../ppo_math_1B_megatron.yaml | 15 + 12 files changed, 2064 insertions(+), 144 deletions(-) create mode 100755 tests/functional/ppo_async.sh diff --git a/docs/guides/ppo.md b/docs/guides/ppo.md index 4bee61af2f9..6d4ff7ae745 100644 --- a/docs/guides/ppo.md +++ b/docs/guides/ppo.md @@ -44,6 +44,29 @@ By default, PPO uses a colocated architecture where the **policy**, **value mode PPO also supports **non-colocated generation**, where the generation engine runs on a dedicated set of GPUs separate from training (`policy.generation.colocated.enabled: false`, with `policy.generation.colocated.resources.{gpus_per_node,num_nodes}` sizing the dedicated generation resources). In this mode the policy and value model still share a single `train_cluster` (offloaded to CPU between each other as in the colocated case), while generation runs concurrently on its own `inference_cluster`; weights are synchronized via an NCCL broadcast collective spanning both clusters instead of the CUDA-IPC/ZMQ path used when colocated. Non-colocated generation currently requires the vLLM backend — `setup()` asserts this explicitly, since SGLang does not yet implement the collective rendezvous or weight refit needed for non-colocated mode. +### Asynchronous PPO + +PPO supports an **asynchronous** training mode (`ppo.async_ppo.enabled: true`) that mirrors [async GRPO](grpo.md): a background `AsyncTrajectoryCollector` continuously generates trajectories into a replay buffer while the driver trains, so generation and training overlap instead of running in lockstep. Async PPO requires non-colocated vLLM generation (`policy.generation.colocated.enabled: false`, `vllm_cfg.async_engine: true`) and importance-sampling correction (`loss_fn.use_importance_sampling_correction: true`) — all enforced by asserts in `async_ppo_train`. Enable it via `examples/run_ppo.py`, which dispatches to `async_ppo_train` when `ppo.async_ppo.enabled` is set. + +Each async step: sample a fixed batch from the replay buffer → compute fresh critic **values** at train time → compute fresh policy/reference logprobs → compute GAE advantages → run the `ppo_epochs` inner loop (critic then actor) → perform a **single** weight refit to the generation engine and bump the replay-buffer weight version. Computing values at train time (rather than stashing them when the trajectory was generated) keeps the critic side as fresh as possible. + +**Critic-staleness caveat.** Off-policy trajectories are corrected on the *actor* side by importance sampling (`exp(prev_logprobs - generation_logprobs)`), exactly as in async GRPO. PPO's GAE, however, recursively bootstraps value estimates across each trajectory, so bias from a stale trajectory's actions/rewards compounds along the sequence in a way GRPO's memoryless reward-only advantage does not, and there is **no** corresponding critic-side correction. This bias is *bounded* (not eliminated) by keeping `ppo.async_ppo.max_trajectory_age_steps: 1` (at most one policy version stale), which is the recommended and validated setting. Values above 1 are permitted but only warned about; a rigorous fix (folding importance ratios into the GAE recursion, V-trace style) is future work. + +**Critic warmup** (`ppo.policy_training_start_step > 0`) is supported: during warmup the policy is frozen, so each step's weight refit skips the actual weight transfer (generation already holds the correct initial weights) but still advances the replay-buffer weight version so the async pipeline keeps making progress. + +Because the actor is frozen at its initial policy `π₀` all the way **through** step `W = ppo.policy_training_start_step` (it first trains *during* step `W`, and the refit that publishes `π₁` to generation happens at the *end* of step `W`), **every rollout — whatever its generation-version tag — was produced by the same model `π₀`**. The generation-version counter still increments each step (so the collector's lead can advance), but it does *not* track policy staleness during warmup: gen-version `g ≤ W` ⇒ policy `π₀`, and `g > W` ⇒ policy `π_{g−W}`. This lets the collector bank cheap frozen rollouts far ahead for free: set `ppo.async_ppo.warmup_max_trajectory_age_steps` above `max_trajectory_age_steps` (`A_t`) and it generates that far ahead while the critic pretrains. + +The snap-back is governed by **two distinct boundaries** (this is what makes the knob correct and hang-free): + +- The **collector's generation-lead** drops to `A_t` at step `W`, so from step `W+1` it stops over-banking `π₀` and regenerates its lead targets against the freshly-trained policy (`π₁`, `π₂`, …). +- The **buffer's eviction age** stays elevated through step `W + A_t`. A frozen (`π₀`) rollout is within `A_t` *policy*-steps of the actor for every step `s ≤ W + A_t` (its policy-age is `s − W`), so those banked rollouts are admitted as legitimate lag-≤`A_t` data — IS-corrected at train time exactly like normal async — and are only evicted once the actor has genuinely moved more than `A_t` steps past `π₀`. (Snapping the eviction age at `W` instead would discard the still-on-policy boundary batch and **deadlock**, since the collector's lead has already advanced past that target and never regenerates it.) + +Note this only improves *throughput* when warmup is **generation-bound** (the collector can actually bank ahead). When critic training is the bottleneck it is a throughput no-op — but it is correct and hang-free either way. + +**Reading the age metrics.** `avg_trajectory_age` reports the *generation-version* age (`current_weight_version − gen_version`), which **overcounts** off-policyness across the warmup boundary: every gen-version `≤ W` is the same frozen `π₀`, so a rollout banked at gen-version 0 and consumed at step `W+1` shows `avg_trajectory_age ≈ A_w` even though it is only **1 policy step** off-policy. The metric that actually bounds the importance-sampling correction is `avg_trajectory_policy_age` / `max_trajectory_policy_age` = `max(0, s−W) − max(0, g−W)`, which stays **≤ `max_trajectory_age_steps`** at every step (equal to the gen-version age when there is no warmup). A spike in `avg_trajectory_age` at the boundary with `max_trajectory_policy_age` still at the bound is expected and safe; a `max_trajectory_policy_age` *above* the bound would indicate a real regression. Absent/`null` ⇒ same as `max_trajectory_age_steps`. + +**v1 limitations.** Async PPO currently requires the vLLM backend (no SGLang/Megatron generation, no colocated inference). DAPO-style dynamic sampling, reward scaling, and reward shaping are also unsupported in async mode (rejected by `examples/run_ppo.py`). + ### Value Model Configuration ```yaml diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index 76015b58cb2..ff48399b2c2 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -49,6 +49,24 @@ ppo: target_max: 1.0 seq_logprob_error_threshold: null + async_ppo: + enabled: false # Set to true for async PPO (requires non-colocated vLLM generation) + # Max age (in training steps) for trajectories drawn from the replay buffer. + # Recommended/validated value is 1; higher values increase critic-staleness + # bias (warned, not forbidden). See AsyncPPOConfig in nemo_rl/algorithms/ppo.py. + max_trajectory_age_steps: 1 + # Max age used ONLY during critic warmup (step < policy_training_start_step), + # where the frozen actor makes any-age trajectories on-policy for free. null + # => same as max_trajectory_age_steps. Only helps throughput when warmup is + # generation-bound. See AsyncPPOConfig. + warmup_max_trajectory_age_steps: null + in_flight_weight_updates: true # Set to true to enable in-flight weight updates + recompute_kv_cache_after_weight_updates: false + # Heartbeat frequency for the collector/replay-buffer per-rollout progress + # prints (they flood the log at large num_prompts_per_step). Log every Nth + # event (plus the final one per target); 0 silences them; unset => every event. + log_every: 500 + loss_fn: disable_ppo_ratio: false reference_policy_kl_penalty: 0.0 # DAPO: no KL penalty diff --git a/examples/run_ppo.py b/examples/run_ppo.py index dc8fab3f2a5..5977df7b837 100644 --- a/examples/run_ppo.py +++ b/examples/run_ppo.py @@ -18,7 +18,7 @@ from omegaconf import OmegaConf -from nemo_rl.algorithms.ppo import MasterConfig, ppo_train, setup +from nemo_rl.algorithms.ppo import MasterConfig, async_ppo_train, ppo_train, setup from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.data.utils import setup_response_data from nemo_rl.distributed.virtual_cluster import init_ray @@ -44,6 +44,33 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]: return args, overrides +def _async_ppo_enabled(config: MasterConfig) -> bool: + """Whether async PPO is requested in the config.""" + return "async_ppo" in config.ppo and config.ppo["async_ppo"]["enabled"] + + +def _validate_async_ppo_config(config: MasterConfig) -> None: + """Reject async-incompatible config up front (before setup()/vLLM init). + + Async PPO does not support DAPO-style dynamic sampling / reward scaling / + reward shaping, nor multiple dataloaders. + """ + if not _async_ppo_enabled(config): + return + for feature in ("use_dynamic_sampling", "reward_scaling", "reward_shaping"): + if feature not in config.ppo: + continue + if feature == "use_dynamic_sampling": + if config.ppo[feature]: + raise NotImplementedError(f"{feature} is not supported with async PPO") + elif config.ppo[feature]["enabled"]: + raise NotImplementedError(f"{feature} is not supported with async PPO") + if config.data.get("use_multiple_dataloader"): + raise NotImplementedError( + "use_multiple_dataloader is not supported with async PPO" + ) + + def main() -> None: """Main entry point.""" # Parse arguments @@ -66,6 +93,10 @@ def main() -> None: config = MasterConfig(**config) print("Applied CLI overrides") + # Fail fast (before the expensive setup()/vLLM init) on async-incompatible + # config, so users don't wait through a full worker init to be rejected. + _validate_async_ppo_config(config) + # Print config print("Final config:") pprint.pprint(config) @@ -112,25 +143,47 @@ def main() -> None: master_config, ) = setup(config, tokenizer, dataset, val_dataset) - print("🚀 Running synchronous PPO training") - - # Run standard PPO training - ppo_train( - policy, - policy_generation, - value_model, - dataloader, - val_dataloader, - tokenizer, - loss_fn, - value_loss_fn, - task_to_env, - val_task_to_env, - logger, - checkpointer, - ppo_state, - master_config, - ) + # Dispatch to async PPO when enabled, otherwise standard synchronous PPO. + if _async_ppo_enabled(config): + print("🚀 Running asynchronous PPO training") + async_config = config.ppo["async_ppo"] + async_ppo_train( + policy, + policy_generation, + value_model, + dataloader, + val_dataloader, + tokenizer, + loss_fn, + value_loss_fn, + task_to_env, + val_task_to_env, + logger, + checkpointer, + ppo_state, + master_config, + max_trajectory_age_steps=async_config["max_trajectory_age_steps"], + ) + else: + print("🚀 Running synchronous PPO training") + + # Run standard PPO training + ppo_train( + policy, + policy_generation, + value_model, + dataloader, + val_dataloader, + tokenizer, + loss_fn, + value_loss_fn, + task_to_env, + val_task_to_env, + logger, + checkpointer, + ppo_state, + master_config, + ) if __name__ == "__main__": diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 22939bf72b3..3a72273d899 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import threading as _threading from collections import Counter from typing import Any, Iterable, Optional @@ -20,6 +21,17 @@ from nemo_rl.algorithms.async_utils.interfaces import ReplayBufferProtocol +# Per-add / per-sample debug prints fire once per prompt group, so they flood the +# log at large ppo.num_prompts_per_step (e.g. 8192 => tens of thousands of lines +# per step). Throttle to a periodic heartbeat: print every Nth event. N is +# configured in YAML (ppo.async_ppo.log_every / grpo.async_grpo.log_every) and +# passed to the buffer constructor: +# N > 0 -> print every Nth event +# N == 0 -> silence the throttled lines +# unset -> print every event (no throttling) +# NRL_ASYNC_VERBOSE=1 additionally enables the O(n) full-list dumps. +_ASYNC_VERBOSE = os.environ.get("NRL_ASYNC_VERBOSE", "0").lower() in ("1", "true") + # Classes with @ray.remote can't be inherited from, so we split the implementation out. class ReplayBufferImpl(ReplayBufferProtocol): @@ -29,10 +41,14 @@ class ReplayBufferImpl(ReplayBufferProtocol): grpo.num_generations_per_prompt (required to compute per-prompt advantages). """ - def __init__(self, max_size: int): + def __init__(self, max_size: int, log_every: Optional[int] = None): if max_size <= 0: raise ValueError(f"max_size must be positive, got {max_size}") self.max_size = max_size + # Heartbeat frequency for the throttled progress prints, from YAML + # (ppo.async_ppo.log_every). N>0 => every Nth event; 0 => silence; unset + # (None) => every event (log_every_n == 1, no throttling). + self._log_every_n = 1 if log_every is None else int(log_every) self.trajectories = [] # List[dict[str, Any]] # If trajectory_version is 1 and target_weight_version is 4 it means that weight version 1 was used for generating a trajectory and this trajectory will be used for training when weight version is 4. self.trajectory_versions = [] # it is the weight-version used for generation of a trajectory @@ -41,6 +57,10 @@ def __init__(self, max_size: int): self.last_target_weight_already_generated = -1 self._lock = _threading.Lock() + def _should_log(self, n: int) -> bool: + """True on every ``log_every``-th event (always if ``NRL_ASYNC_VERBOSE``).""" + return _ASYNC_VERBOSE or (self._log_every_n > 0 and n % self._log_every_n == 0) + def add( self, trajectory: dict[str, Any], @@ -58,15 +78,24 @@ def add( if len(self.trajectories) >= self.max_size: return "full" - print("🔍 ReplayBuffer.add: Adding trajectory") self.trajectories.append(trajectory) self.trajectory_versions.append(weight_version) self.target_weight_versions.append(target_weight_version) # Do not advance last_target_weight_already_generated here. A target # is only safe to skip once training consumes a complete batch for it. - print( - f"ReplayBuffer state: {len(self.trajectories)} groups, versions={self.trajectory_versions}, targets={self.target_weight_versions}, last_target_weight_already_generated={self.last_target_weight_already_generated}" - ) + n = len(self.trajectories) + if _ASYNC_VERBOSE: + print( + f"🔍 ReplayBuffer.add: {n} groups, " + f"versions={self.trajectory_versions}, " + f"targets={self.target_weight_versions}, " + f"last_target_weight_already_generated={self.last_target_weight_already_generated}" + ) + elif self._should_log(n): + print( + f"🔍 ReplayBuffer: {n} groups buffered " + f"(last_target_already_generated={self.last_target_weight_already_generated})" + ) return "success" def get_debug_info(self) -> dict: @@ -115,13 +144,15 @@ def sample( return None total_trajectories = len(self.trajectories) - print("🔍 ReplayBuffer sampling debug:") - print(f" {current_weight_version=}, {max_age_steps=}") - print(f" {self.trajectory_versions=}") - - # For debugging: check for unexpected old trajectories + # Counter is a compact summary; the raw per-group version list is O(n) + # and floods at large buffers, so only dump it under NRL_ASYNC_VERBOSE. version_counts = Counter(self.trajectory_versions) - print(f" {version_counts=}") + print( + f"🔍 ReplayBuffer sampling: {current_weight_version=}, " + f"{max_age_steps=}, {total_trajectories=}, {version_counts=}" + ) + if _ASYNC_VERBOSE: + print(f" {self.trajectory_versions=}") # Compute minimum valid version based on age window # max_age_steps=1 means trajectories from the last 1 step are valid @@ -220,13 +251,24 @@ def sample( f"(consumed batch for step {current_weight_version})" ) + _new_targets = ( + self.target_weight_versions + if _ASYNC_VERBOSE + else f"<{len(self.target_weight_versions)} groups>" + ) print( - f"🗑️ Consumed and removed {len(selected)} groups from buffer, old buffer size: {total_trajectories}, new buffer size: {len(self.trajectories)}, new target weight versions {self.target_weight_versions}" + f"🗑️ Consumed and removed {len(selected)} groups from buffer, old buffer size: {total_trajectories}, new buffer size: {len(self.trajectories)}, new target weight versions {_new_targets}" ) return { "trajectories": sampled_items, "avg_trajectory_age": avg_trajectory_age, + # Per-group generation weight-versions of the sampled trajectories. + # avg_trajectory_age above is the GEN-version age; a consumer that + # knows the freeze boundary (e.g. async PPO's critic warmup) can use + # these to compute the true POLICY-age instead. Extra key — ignored + # by callers that don't need it (e.g. async GRPO). + "generation_weight_versions": sampled_weights, } def size(self) -> int: @@ -640,4 +682,7 @@ def sample( return { "trajectories": sampled_items, "avg_trajectory_age": avg_trajectory_age, + # Keep the return schema identical to ReplayBufferImpl.sample so + # consumers can compute freeze-aware policy-age uniformly. + "generation_weight_versions": sampled_weights, } diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index 14c60f9dcec..85088115793 100644 --- a/nemo_rl/algorithms/async_utils/trajectory_collector.py +++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py @@ -15,6 +15,7 @@ from __future__ import annotations import concurrent.futures +import os import threading as _threading import time from collections import defaultdict @@ -38,6 +39,17 @@ TokenizerType = PreTrainedTokenizerBase +# Per-worker / per-group progress prints fire once per prompt group, so they flood +# the log at large ppo.num_prompts_per_step (e.g. 8192 => tens of thousands of lines +# per step). Throttle to a periodic heartbeat: print every Nth event, plus the final +# one per target. N is configured in YAML (ppo.async_ppo.log_every / +# grpo.async_grpo.log_every): +# N > 0 -> print every Nth event +# N == 0 -> silence the throttled lines +# unset -> print every event (no throttling) +# NRL_ASYNC_VERBOSE=1 additionally enables the O(n) full-list dumps. +_ASYNC_VERBOSE = os.environ.get("NRL_ASYNC_VERBOSE", "0").lower() in ("1", "true") + @ray.remote # pragma: no cover class AsyncTrajectoryCollector: @@ -59,6 +71,43 @@ def __init__( self.tokenizer = tokenizer self.task_to_env = task_to_env self.master_config = master_config + # Resolve the algorithm config block generically so this collector works + # for both GRPO (master_config.grpo / "async_grpo") and PPO + # (master_config.ppo / "async_ppo"). The shared trajectory-collection + # logic (num_prompts_per_step, num_generations_per_prompt, + # max_rollout_turns, and the async sub-block) is identical between them; + # only the config key differs. `_algo_cfg`/`_async_cfg` are read below in + # place of the former `master_config.grpo[...]` / `["async_grpo"]` reads. + self._algo_cfg = getattr(master_config, "grpo", None) + if self._algo_cfg is None: + self._algo_cfg = master_config.ppo + self._async_cfg = ( + self._algo_cfg.get("async_grpo") + or self._algo_cfg.get("async_ppo") + or {} + ) + # Effective max trajectory age used for the generation-lead decisions + # below. The driver may override it per phase via set_max_trajectory_age() + # — e.g. a larger value during critic warmup, where the frozen actor makes + # arbitrarily-old trajectories on-policy (importance ratio == 1). Defaults + # to the configured value, so callers that never call the setter (e.g. + # async GRPO) behave exactly as before. + self._current_max_age = int(self._async_cfg["max_trajectory_age_steps"]) + # Highest age the collector may be switched to (warmup age if larger), + # used to size the fixed in-flight semaphore up front. Absent for GRPO + # (async_grpo has no warmup key) → falls back to the train age. + self._max_age_ceiling = max( + self._current_max_age, + int( + self._async_cfg.get("warmup_max_trajectory_age_steps") + or self._current_max_age + ), + ) + # Heartbeat frequency for throttled progress prints, from YAML + # (ppo.async_ppo.log_every / grpo.async_grpo.log_every). N>0 => every Nth + # event; 0 => silence; unset (None) => every event (no throttling). + _log_every_cfg = self._async_cfg.get("log_every") + self._log_every_n = 1 if _log_every_cfg is None else int(_log_every_cfg) self.replay_buffer = replay_buffer self.teacher_worker_groups = teacher_worker_groups or {} self.alias_to_group_alias = alias_to_group_alias or {} @@ -100,11 +149,19 @@ def __init__( self._inflight_threads: set[_threading.Thread] = set() self._threads_lock: _threading.Lock = _threading.Lock() - # Limit in-flight generator requests to num_prompts_per_step * max_trajectory_age_steps - # This value limits the parallelism of the generation requests. + # Limit in-flight generator requests to num_prompts_per_step * max age, + # sized for the ceiling (warmup age if larger) so deep warmup generation + # isn't throttled. The semaphore size is fixed for the collector's life. + # + # CAUTION: this bounds CONCURRENT in-flight rollouts, and each rollout runs + # its own asyncio/uvloop event loop (several file descriptors). So the peak + # concurrency is num_prompts_per_step * _max_age_ceiling — keep the product + # under the process FD limit (`ulimit -n`). E.g. num_prompts_per_step=8192 + # with warmup age 10 => 81920 concurrent event loops => ~300k FDs, which + # overruns a default limit (OSError: [Errno 24] Too many open files). Either + # keep warmup_max_trajectory_age_steps modest or raise `ulimit -n`. max_inflight = ( - int(self.master_config.grpo["num_prompts_per_step"]) - * int(self.master_config.grpo["async_grpo"]["max_trajectory_age_steps"]) + int(self._algo_cfg["num_prompts_per_step"]) * self._max_age_ceiling ) or 1 self._inflight_sema = _threading.Semaphore(max_inflight) @@ -122,6 +179,10 @@ def __init__( # Timer for efficiency metrics self._efficiency_timer = ThreadSafeTimer(context={"worker": "collector"}) + def _should_log(self, n: int) -> bool: + """True on every ``log_every``-th event (always if ``NRL_ASYNC_VERBOSE``).""" + return _ASYNC_VERBOSE or (self._log_every_n > 0 and n % self._log_every_n == 0) + def _calculate_target_weights(self, generation_weight_version: int) -> list[int]: """Calculate target weight versions for given generation weight version. @@ -137,9 +198,8 @@ def _calculate_target_weights(self, generation_weight_version: int) -> list[int] Returns: [11, 12, 13, 14] # Meaning this generation server can create trajectories for training step 11, 12, 13, 14 """ - # Read async config strictly from grpo.async_grpo - async_cfg = self.master_config.grpo.get("async_grpo", {}) - max_trajectory_age = async_cfg["max_trajectory_age_steps"] + # Effective (possibly warmup-elevated) max age, see set_max_trajectory_age. + max_trajectory_age = self._current_max_age if generation_weight_version == self.initial_weight_version: return [ i @@ -156,10 +216,8 @@ def _get_next_target_for_generation( ) -> Optional[int]: """Get the next target weight that needs generation (if any).""" target_weights = self._calculate_target_weights(generation_weight_version) - num_prompts = int(self.master_config.grpo["num_prompts_per_step"]) - max_age_steps = int( - self.master_config.grpo["async_grpo"]["max_trajectory_age_steps"] - ) + num_prompts = int(self._algo_cfg["num_prompts_per_step"]) + max_age_steps = int(self._current_max_age) last_consumed_target = ray.get( self.replay_buffer.get_last_target_weight_already_generated.remote() ) @@ -202,14 +260,26 @@ def set_weight_version(self, version: int) -> None: else: print(f"🔄 Updated weight version to {version}") + def set_max_trajectory_age(self, age: int) -> None: + """Override the effective max trajectory age for generation-lead decisions. + + Used by async PPO to run a larger age during critic warmup (where the + frozen actor makes old trajectories on-policy) and snap back to the + configured training age once the actor starts training. The replay + buffer's own eviction (driven by the age the *driver* passes to + ``sample``) is what actually clamps staleness, so lowering this at the + warmup→training transition simply stops the collector from generating + further ahead than the training age allows. + """ + self._current_max_age = int(age) + print(f"🔧 Collector max trajectory age set to {self._current_max_age}") + def _should_pause_for_generation_limits(self) -> bool: """Check if collection should be paused due to generation limits.""" try: target_weights = self._calculate_target_weights(self.current_weight_version) - num_prompts = int(self.master_config.grpo["num_prompts_per_step"]) - max_age_steps = int( - self.master_config.grpo["async_grpo"]["max_trajectory_age_steps"] - ) + num_prompts = int(self._algo_cfg["num_prompts_per_step"]) + max_age_steps = int(self._current_max_age) last_consumed_target = ray.get( self.replay_buffer.get_last_target_weight_already_generated.remote() ) @@ -271,8 +341,7 @@ def _collection_loop(self): if self._should_pause_for_generation_limits() and self.running: # Only log warning once per weight version if self._last_limit_warning_version != self.current_weight_version: - async_cfg = self.master_config.grpo.get("async_grpo", {}) - max_trajectory_age = async_cfg["max_trajectory_age_steps"] + max_trajectory_age = self._current_max_age target_weights = [ self.current_weight_version + i for i in range(max_trajectory_age) @@ -312,12 +381,10 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None: target_weight: Optional[int] = None try: generation_weight_version = self.current_weight_version - num_generations = self.master_config.grpo["num_generations_per_prompt"] + num_generations = self._algo_cfg["num_generations_per_prompt"] num_prompts_in_batch = batch.size - num_prompts_per_step = int(self.master_config.grpo["num_prompts_per_step"]) - max_age_steps = int( - self.master_config.grpo["async_grpo"]["max_trajectory_age_steps"] - ) + num_prompts_per_step = int(self._algo_cfg["num_prompts_per_step"]) + max_age_steps = int(self._current_max_age) # Get the next target weight that needs generation target_weight = self._get_next_target_for_generation( @@ -425,10 +492,11 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None: self._inflight_sema.release() raise started += 1 - print( - f"📊 Started worker {started}/{num_prompts_to_generate} for " - f"target_weight={target_weight} ({spawned_count} total)" - ) + if self._should_log(started) or started == num_prompts_to_generate: + print( + f"📊 Started worker {started}/{num_prompts_to_generate} " + f"for target_weight={target_weight} ({spawned_count} total)" + ) finally: if started < num_prompts_to_generate: print( @@ -454,6 +522,9 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None: def get_weight_version(self) -> int: return self.current_weight_version + def get_max_trajectory_age(self) -> int: + return self._current_max_age + def pause(self) -> None: """Pause trajectory collection.""" self._manual_pause_cleared.clear() # Signal collection to pause @@ -493,7 +564,7 @@ def prepare_for_refit(self) -> None: ) else: is_async_engine = False - in_flight_weight_updates = self.master_config.grpo.get("async_grpo", {}).get( + in_flight_weight_updates = self._async_cfg.get( "in_flight_weight_updates", False ) @@ -524,7 +595,7 @@ def resume_after_refit(self) -> None: # Invalidate&recompute vLLM caches after the in-flight weight updates if # recompute_kv_cache_after_weight_updates is True (AREAL-style implementation). # Otherwise, keep using the stale KV caches (Magistral-style implementation). - async_cfg = self.master_config.grpo.get("async_grpo", {}) + async_cfg = self._async_cfg if async_cfg.get("in_flight_weight_updates", False) and async_cfg.get( "recompute_kv_cache_after_weight_updates", False ): @@ -769,7 +840,7 @@ def _run_prompt_group_worker( tokenizer=self.tokenizer, task_to_env=self.task_to_env, max_seq_len=self.master_config.policy["max_total_sequence_length"], - max_rollout_turns=self.master_config.grpo["max_rollout_turns"], + max_rollout_turns=self._algo_cfg["max_rollout_turns"], greedy=False, ) @@ -840,11 +911,15 @@ def _run_prompt_group_worker( "idle/buffer_full_backoff", time.perf_counter() - backoff_start, ) - print( - f"📦 Buffered per-prompt group (prompt_idx {prompt_idx}, " - f"target_weight {target_weight_version}) " - f"[{buffered_count}/{spawned_count} buffered]" - ) + if ( + self._should_log(buffered_count) + or buffered_count == spawned_count + ): + print( + f"📦 Buffered per-prompt group (prompt_idx {prompt_idx}, " + f"target_weight {target_weight_version}) " + f"[{buffered_count}/{spawned_count} buffered]" + ) break elif status == "full": if backoff_start is None: diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 46c1de2c5c0..4bf17860122 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -36,6 +36,8 @@ _should_log_nemo_gym_responses, _should_use_async_rollouts, _should_use_nemo_gym, + _write_latest_checkpoint_status, + aggregate_rollout_metrics, compute_and_apply_seq_logprob_error_masking, extract_initial_prompt_messages, refit_policy_generation, @@ -52,7 +54,11 @@ RewardShapingConfig, apply_reward_shaping, ) -from nemo_rl.algorithms.utils import print_performance_metrics, set_seed +from nemo_rl.algorithms.utils import ( + print_efficiency_summary, + print_performance_metrics, + set_seed, +) from nemo_rl.data import DataConfig from nemo_rl.data.collate_fn import rl_collate_fn from nemo_rl.data.datasets import AllTaskProcessedDataset @@ -63,6 +69,7 @@ ) from nemo_rl.data.utils import load_dataloader_state from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env from nemo_rl.distributed.virtual_cluster import ( ClusterConfig, RayVirtualCluster, @@ -89,6 +96,7 @@ from nemo_rl.utils.memory_tracker import MemoryTracker from nemo_rl.utils.nsys import maybe_gpu_profile_step from nemo_rl.utils.timer import TimeoutChecker, Timer +from nemo_rl.utils.venvs import create_local_venv_on_each_node # =============================================================================== # Configuration @@ -96,6 +104,55 @@ TokenizerType = TypeVar("TokenizerType", bound=PreTrainedTokenizerBase) +class AsyncPPOConfig(TypedDict): + """Configuration for asynchronous PPO training (mirrors AsyncGRPOConfig).""" + + enabled: bool + # Maximum trajectory age in training steps for samples drawn from the async + # replay buffer. Trajectories older than this are excluded during sampling; + # buffer sizing also scales with this value. + # + # NOTE: values > 1 are ALLOWED but only warned about, not forbidden. PPO's + # GAE recursively bootstraps value estimates across each trajectory, so + # stale-trajectory bias compounds along the sequence in a way GRPO's + # memoryless reward-only advantage does not. The recommended/validated + # value is 1 (at most one policy version stale); higher values trade more + # critic-staleness bias for throughput and are the user's call. + max_trajectory_age_steps: int + # Maximum trajectory age used during critic warmup. Through step + # W = ppo.policy_training_start_step the actor is FROZEN at its initial policy + # pi_0 (it first trains DURING step W), so every rollout — however old its + # generation-version tag — was produced by the same fixed policy: the + # importance ratio is exactly 1 and there is NO off-policy staleness. The + # collector may therefore bank rollouts far ahead for free while the critic + # pretrains. Two boundaries govern the snap-back (see async_ppo_train): + # * the collector's generation-lead drops to max_trajectory_age_steps at + # step W (so it regenerates fresh rollouts against pi_1, pi_2, ...); + # * the buffer's eviction age stays elevated through step W + A_t, because a + # frozen (pi_0) rollout is only within A_t POLICY-steps of the actor until + # then — so those banked rollouts are admitted as valid lag-<=A_t data and + # only evicted once pi_0 is genuinely too stale. + # Absent OR null => same as max_trajectory_age_steps (no special warmup + # behavior); read as `.get(...) or max_trajectory_age_steps`, so None is a valid + # "unset" value and the type must allow it. + # + # NOTE: this only helps THROUGHPUT when warmup is generation-bound (the + # collector can actually bank ahead). When critic training is the bottleneck it + # is a no-op on throughput, but it is still correct/hang-free either way. + warmup_max_trajectory_age_steps: NotRequired[int | None] + # Broadcast weights to the generation engine without waiting for in-flight + # generations to drain first (only supported by async-capable engines). + in_flight_weight_updates: NotRequired[bool] + # Recompute the KV cache after in-flight weight updates. + recompute_kv_cache_after_weight_updates: NotRequired[bool] + # Heartbeat frequency for the collector/replay-buffer per-rollout progress + # prints: log every Nth event (plus the final one per target). These fire once + # per prompt group and flood the log at large num_prompts_per_step, so throttle + # them. N>0 => every Nth; 0 => silence; unset/null => every event. Default 500 + # (YAML). Read with a None fallback, so the type must allow null. + log_every: NotRequired[int | None] + + class AdvEstimatorConfig(TypedDict): """Configuration for PPO advantage estimator (GAE or raw_reward).""" @@ -149,7 +206,13 @@ class PPOConfig(TypedDict): policy_training_start_step: NotRequired[int] # Sequence-level logprob error masking for training stability. If set, mask sequences with mult_prob_error exceeding this threshold (same scale as token_mult_prob_error metric, e.g., 1.5) # Note that this is slightly different than Masked Importance Sampling (MIS) because this uses the absolute value of the difference between the training and generation logprobs, whereas MIS just uses the difference between the training and generation logprobs. - seq_logprob_error_threshold: float | None + # NotRequired (read via .get(..., None)): configs that omit it disable the + # masking rather than fail MasterConfig validation — e.g. a config whose + # `defaults:` base predates this field. + seq_logprob_error_threshold: NotRequired[float | None] + # Asynchronous PPO (replay-buffer, non-colocated generation). Absent/disabled + # runs synchronous PPO. + async_ppo: NotRequired[AsyncPPOConfig] class PPOSaveState(TypedDict): @@ -196,6 +259,36 @@ class MasterConfig(BaseModel, extra="allow"): # =============================================================================== +def _resolve_resume_optimizer_path( + optim_dir: Path, + weights_path: Optional[Path], + model_config: dict[str, Any], +) -> Optional[Path]: + """Resolve ``optimizer_path`` for resuming the optimizer + LR scheduler. + + The two training backends persist optimizer state differently, and + ``optimizer_path`` means different things to each: + + * **DTensor** writes/reads a SEPARATE ``optimizer/`` directory — resume it + only if that dir exists. + * **Megatron** bundles the optimizer AND the LR scheduler INSIDE the weights + dist-checkpoint (``weights/iter_*/*.distcp``) and never creates a separate + ``optimizer/`` dir; there ``optimizer_path`` is merely the ``load_optim`` + flag (the load reads from ``weights_path``). So point it at ``weights_path`` + whenever weights exist — otherwise ``load_optim`` stays False and the + optimizer + LR-warmup scheduler are NOT restored, so the scheduler restarts + from step 0 on every resume (the tell-tale V-shaped ``critic/lr`` / + ``policy/lr``, plus silently discarded Adam moments). + """ + if optim_dir.exists(): + return optim_dir + if weights_path is not None and model_config.get("megatron_cfg", {}).get( + "enabled", False + ): + return weights_path + return None + + def setup( master_config: MasterConfig, tokenizer: TokenizerType, @@ -554,7 +647,9 @@ def setup( _policy_weights = Path(last_checkpoint_path) / "policy" / "weights" _policy_optim = Path(last_checkpoint_path) / "policy" / "optimizer" weights_path = _policy_weights if _policy_weights.exists() else None - optimizer_path = _policy_optim if _policy_optim.exists() else None + optimizer_path = _resolve_resume_optimizer_path( + _policy_optim, weights_path, policy_config + ) if weights_path is None: print( f" ⚠ Policy weights not found in checkpoint {last_checkpoint_path} " @@ -620,7 +715,9 @@ def init_value(): _value_weights = Path(last_checkpoint_path) / "value" / "weights" _value_optim = Path(last_checkpoint_path) / "value" / "optimizer" value_weights_path = _value_weights if _value_weights.exists() else None - value_optimizer_path = _value_optim if _value_optim.exists() else None + value_optimizer_path = _resolve_resume_optimizer_path( + _value_optim, value_weights_path, value_config + ) if value_weights_path is None: print( f" ⚠ Value weights not found in checkpoint {last_checkpoint_path} " @@ -1073,6 +1170,53 @@ def _create_advantage_estimator(master_config: MasterConfig): return adv_estimator +def _compute_critic_metrics(value_results: dict[str, Any]) -> dict[str, Any]: + """Build the ``critic/*`` metrics dict from a value model's train results. + + Shared by the synchronous (:func:`ppo_train`) and asynchronous + (:func:`async_ppo_train`) loops so the critic metric aggregation and + explained-variance computation stay in exactly one place. + """ + value_mb_metrics = value_results.get("all_mb_metrics", {}) + critic_metrics: dict[str, Any] = { + "critic/grad_norm": value_results["grad_norm"].numpy(), + "critic/loss": value_results["loss"].numpy(), + } + + for k, v in value_mb_metrics.items(): + if k in { + "lr", + "wd", + "global_valid_seqs", + "global_valid_toks", + "grad_norm", + }: + critic_metrics["critic/" + k] = np.mean(v).item() + elif k in {"values_min"}: + critic_metrics["critic/" + k] = np.min(v).item() + elif k in {"values_max"}: + critic_metrics["critic/" + k] = np.max(v).item() + elif isinstance(v, (np.ndarray, list)): + critic_metrics["critic/" + k] = np.sum(v).item() + else: + raise ValueError( + f"Unknown metric for value don't know how to handle: {k}" + ) + + # Compute explained variance from sufficient statistics: + # EV = 1 - Var(returns - values) / Var(returns) + r_mean = critic_metrics.get("critic/returns_mean", 0) + v_mean = critic_metrics.get("critic/values_mean", 0) + r_sq = critic_metrics.get("critic/returns_sq_mean", 0) + res_sq = critic_metrics.get("critic/residual_sq_mean", 0) + var_returns = r_sq - r_mean**2 + var_residual = res_sq - (r_mean - v_mean) ** 2 + critic_metrics["critic/explained_var"] = 1.0 - var_residual / max( + var_returns, 1e-8 + ) + return critic_metrics + + # =============================================================================== # Training & Validation # =============================================================================== @@ -1633,45 +1777,7 @@ def ppo_train( # Extract critic metrics from value training results if value_results is not None: - value_mb_metrics = value_results.get("all_mb_metrics", {}) - critic_metrics = { - "critic/grad_norm": value_results["grad_norm"].numpy(), - "critic/loss": value_results["loss"].numpy(), - } - - for k, v in value_mb_metrics.items(): - if k in { - "lr", - "wd", - "global_valid_seqs", - "global_valid_toks", - "grad_norm", - }: - critic_metrics["critic/" + k] = np.mean(v).item() - elif k in {"values_min"}: - critic_metrics["critic/" + k] = np.min(v).item() - elif k in {"values_max"}: - critic_metrics["critic/" + k] = np.max(v).item() - elif isinstance(v, (np.ndarray, list)): - critic_metrics["critic/" + k] = np.sum(v).item() - else: - raise ValueError( - f"Unknown metric for value don't know how to handle: {k}" - ) - - # Compute explained variance from sufficient statistics: - # EV = 1 - Var(returns - values) / Var(returns) - r_mean = critic_metrics.get("critic/returns_mean", 0) - v_mean = critic_metrics.get("critic/values_mean", 0) - r_sq = critic_metrics.get("critic/returns_sq_mean", 0) - res_sq = critic_metrics.get("critic/residual_sq_mean", 0) - var_returns = r_sq - r_mean**2 - var_residual = res_sq - (r_mean - v_mean) ** 2 - critic_metrics["critic/explained_var"] = 1.0 - var_residual / max( - var_returns, 1e-8 - ) - - metrics.update(critic_metrics) + metrics.update(_compute_critic_metrics(value_results)) metrics.update( { "reward": rewards.numpy(), @@ -1954,49 +2060,1201 @@ def ppo_train( current_step = 0 -def validate( - policy_generation: GenerationInterface, +def _async_warmup_collector_lead_age( + step: int, + policy_training_start_step: int, + train_age: int, + warmup_age: int, +) -> int: + """Collector generation-lead age at ``step`` (boundary at W). + + During critic warmup the actor is FROZEN at its initial policy ``pi_0`` through + step ``W = policy_training_start_step`` (it first trains DURING step W). The + collector banks ``warmup_age``-deep while frozen (``step <= W``), then drops to + ``train_age`` from step ``W+1`` so it stops over-banking ``pi_0`` and instead + regenerates its lead targets against the freshly-trained policy (pi_1, pi_2, …). + """ + return warmup_age if step <= policy_training_start_step else train_age + + +def _async_warmup_sample_max_age( + step: int, + policy_training_start_step: int, + train_age: int, + warmup_age: int, +) -> int: + """Driver buffer-eviction age at ``step`` (boundary at W + train_age). + + A frozen (``pi_0``) rollout banked during warmup carries a large *generation*-age + but a small *policy*-age: gen-version ``g <= W`` is always ``pi_0``, and at step + ``s`` the actor is ``pi_{max(0, s-W)}``, so a ``pi_0`` rollout consumed at step + ``s`` is only ``s - W`` policy-steps stale. It therefore stays within the allowed + ``train_age`` policy-steps for every ``s <= W + train_age`` and is legitimate + lag-<=``train_age`` data there (IS-corrected at train time, exactly like normal + async). Keeping the elevated age through ``W + train_age`` admits those banked + rollouts; dropping only afterwards correctly evicts ``pi_0`` once the actor has + genuinely moved more than ``train_age`` steps past it. + + Snapping this boundary at ``W`` instead would evict the still-on-policy boundary + batch and DEADLOCK: the collector's lead has already advanced past that target + (its window is ``[current+1, …]``, never ``current`` itself) so it never + regenerates it. + """ + return warmup_age if step <= policy_training_start_step + train_age else train_age + + +def _async_trajectory_policy_age( + gen_version: int, + current_weight_version: int, + policy_training_start_step: int, +) -> int: + """True off-policy staleness (in POLICY updates) of a sampled rollout. + + The replay buffer's ``avg_trajectory_age`` is a *generation*-version age + (``current_weight_version - gen_version``), which OVERCOUNTS staleness during + critic warmup: the actor is frozen at ``pi_0`` through step + ``W = policy_training_start_step``, so every gen-version ``<= W`` is the same + model ``pi_0``. The policy that actually produced a gen-version-``g`` rollout is + ``pi_{max(0, g-W)}`` and the actor at weight-version ``s`` is ``pi_{max(0, s-W)}``, + so the number of policy updates between them — i.e. how far the importance ratio + must reach — is:: + + max(0, s - W) - max(0, g - W) + + This is ``<= max_trajectory_age_steps`` by construction (enforced by the + two-boundary eviction in ``_async_warmup_sample_max_age``). With no warmup + (``W == 0``) it reduces to the plain gen-version age ``s - g``. + """ + return max(0, current_weight_version - policy_training_start_step) - max( + 0, gen_version - policy_training_start_step + ) + + +def async_ppo_train( + policy: ColocatablePolicyInterface, + policy_generation: Optional[GenerationInterface], + value_model: ValueInterface, + dataloader: StatefulDataLoader, val_dataloader: Optional[StatefulDataLoader], - tokenizer, + tokenizer: TokenizerType, + loss_fn: LossFunction, + value_loss_fn: LossFunction, + task_to_env: dict[str, EnvironmentInterface], val_task_to_env: Optional[dict[str, EnvironmentInterface]], - step: int, + logger: Logger, + checkpointer: CheckpointManager, + ppo_save_state: PPOSaveState, master_config: MasterConfig, - logger: Optional[Logger] = None, -) -> tuple[dict[str, Any], dict[str, Any]]: - """Run validation on the validation dataset.""" - if val_dataloader is None: - assert val_dataloader is not None or master_config.ppo["val_period"] == 0, ( - "val_dataloader is None, so ppo.val_period must be 0" + max_trajectory_age_steps: int = 1, +) -> None: + """Run asynchronous PPO training with a replay buffer. + + Ported from :func:`nemo_rl.algorithms.grpo.async_grpo_train`, with PPO's + value/critic model spliced in. Per outer step, the driver: + + 1. samples a fixed batch of trajectories from the async replay buffer + (continuously filled by a background ``AsyncTrajectoryCollector``), + 2. computes fresh values for the batch (critic forward, train-time), + 3. computes fresh policy/reference logprobs, + 4. computes GAE advantages/returns using those values, + 5. runs the ``ppo_epochs`` inner loop (critic train, then actor train), + 6. performs a single weight refit to the generation engine and bumps the + replay-buffer weight version. + + The value is computed once per outer step at train time (not stashed at + generation time) so it is as fresh as possible; see the AsyncPPOConfig + docstring for the residual critic-staleness caveat. + + Critic warmup (``ppo.policy_training_start_step > 0``) is supported: during + warmup the policy is frozen (never trained), so the per-step weight refit + skips the actual weight transfer — generation already holds the correct + (initial) weights — but still runs the collector coordination and advances + the replay-buffer weight version so the pipeline keeps making progress. + + v1 limitations (documented in docs/guides/ppo.md): + - vLLM non-colocated generation only (SGLang/Megatron generation and + colocated inference are unsupported for async PPO). + """ + # ------------------------------------------------------------------ + # Entry guards (fail loud at startup, not deep in the loop) + # ------------------------------------------------------------------ + generation_config = master_config.policy["generation"] + backend = generation_config.get("backend", "") if generation_config else "" + assert backend == "vllm" and _should_use_async_rollouts(master_config), ( + "Async PPO requires an async vLLM generation engine. " + "Set policy.generation.backend=vllm and " + "policy.generation.vllm_cfg.async_engine=true." + ) + assert master_config.loss_fn.use_importance_sampling_correction, ( + "Importance sampling correction must be enabled for async PPO for good " + "convergence due to off-policy samples " + "(loss_fn.use_importance_sampling_correction=true)." + ) + colocated_inference = master_config.policy["generation"]["colocated"]["enabled"] + assert not colocated_inference, ( + "Colocated inference is not supported for async PPO. Use non-colocated " + "generation (policy.generation.colocated.enabled=false)." + ) + + async_cfg = master_config.ppo["async_ppo"] + policy_training_start_step = master_config.ppo["policy_training_start_step"] + assert master_config.ppo["ppo_epochs"] >= 1, ( + f"ppo.ppo_epochs must be >= 1 (got {master_config.ppo['ppo_epochs']})." + ) + # NeMo-Gym rollout is not yet wired for async PPO. Without this guard the + # collector's gym branch would AttributeError on the (absent) reward_penalties + # field inside a worker thread, silently stalling the buffer instead of + # failing loudly here. + assert not _should_use_nemo_gym(master_config), ( + "NeMo-Gym rollout is not yet supported for async PPO " + "(env.should_use_nemo_gym must be false)." + ) + if max_trajectory_age_steps > 1: + if not async_cfg.get("in_flight_weight_updates", False): + print( + "⚠️ WARNING: in_flight_weight_updates is recommended for async PPO " + "with max_trajectory_age_steps > 1; without it, a larger age gives " + "no throughput benefit." + ) + print( + "⚠️ WARNING: max_trajectory_age_steps > 1 increases critic-staleness " + "bias (GAE bootstraps stale value estimates across each trajectory). " + "The validated/recommended value is 1." ) - print(" ⚠️ No validation dataloader provided, skipping validation", flush=True) - return {}, {} - timer = Timer() - with timer.time("total_validation_time"): - print(f"▶ Starting validation at step {step}...", flush=True) + # Import async utilities only when needed (heavy Ray actors). + from nemo_rl.algorithms.async_utils import AsyncTrajectoryCollector, ReplayBuffer - total_rewards = [] - total_lengths = [] - all_message_logs = [] # Collect all message logs + timer = Timer(context={"worker": "driver"}) + training_wall_start = time.perf_counter() + timeout = TimeoutChecker( + timeout=master_config.checkpointing["checkpoint_must_save_by"], + fit_last_save_time=True, + ) + timeout.start_iterations() - max_batches = ( - master_config.ppo["max_val_samples"] // master_config.ppo["val_batch_size"] + # PPO async always uses non-colocated vLLM generation, so a refit is always + # required and the generation engine is a real (non-None) actor. + NEED_REFIT = True + POLICY_GENERATION_STALE = True + assert policy_generation is not None + + # ------------------------------------------------------------------ + # Training state. `step` is the global monotonic training step; it is what + # max_num_steps bounds and what the replay-buffer weight versioning tracks. + # ------------------------------------------------------------------ + step = ppo_save_state["total_steps"] + weight_version = step + consumed_samples = ppo_save_state["consumed_samples"] + total_valid_tokens = ppo_save_state.get("total_valid_tokens", 0) + max_num_steps = master_config.ppo["max_num_steps"] + ppo_epochs = master_config.ppo["ppo_epochs"] + val_period = master_config.ppo["val_period"] + val_at_start = master_config.ppo["val_at_start"] + val_at_end = master_config.ppo["val_at_end"] + num_prompts_per_step = master_config.ppo["num_prompts_per_step"] + + # During critic warmup the actor is FROZEN at its initial policy pi_0 all the + # way through step W = policy_training_start_step (it first trains DURING step + # W, and the refit that publishes pi_1 to generation happens at the END of step + # W). So although the collector bumps a fresh generation-version every step, + # every gen-version 0..W is produced by the SAME model pi_0. Concretely: + # gen-version g <= W -> policy pi_0 + # gen-version g > W -> policy pi_{g - W} + # at training step s -> actor is pi_{max(0, s - W)} + # This means a frozen (gen-version <= W) trajectory consumed at step W+k has a + # true POLICY-age of k, not a gen-version age of (W+k - g). Raising the age + # during warmup lets the collector bank cheap frozen rollouts ahead so training + # never stalls on generation variance. Defaults to the training age (no special + # warmup behavior). + warmup_max_trajectory_age_steps = ( + async_cfg.get("warmup_max_trajectory_age_steps") or max_trajectory_age_steps + ) + + # Two DISTINCT boundaries (module-level pure fns below, unit-tested) — decoupling + # these is what makes the warmup-age knob correct and hang-free. + def _collector_lead_age(s: int) -> int: + return _async_warmup_collector_lead_age( + s, + policy_training_start_step, + max_trajectory_age_steps, + warmup_max_trajectory_age_steps, ) - for batch_idx, val_batch in enumerate(val_dataloader): - if batch_idx >= max_batches: - break - additional_metrics_to_report = dict() + def _sample_max_age(s: int) -> int: + return _async_warmup_sample_max_age( + s, + policy_training_start_step, + max_trajectory_age_steps, + warmup_max_trajectory_age_steps, + ) + + adv_estimator = _create_advantage_estimator(master_config) + + # ------------------------------------------------------------------ + # Spin up the replay buffer + trajectory collector Ray actors. + # ------------------------------------------------------------------ + _replay_py_exec = get_actor_python_env( + "nemo_rl.algorithms.async_utils.ReplayBuffer" + ) + if _replay_py_exec.startswith("uv"): + _replay_py_exec = create_local_venv_on_each_node( + _replay_py_exec, + "nemo_rl.algorithms.async_utils.ReplayBuffer", + ) + _replay_py_venv = os.path.dirname(os.path.dirname(_replay_py_exec)) + _replay_runtime_env = { + "py_executable": _replay_py_exec, + "env_vars": { + **os.environ, + "VIRTUAL_ENV": _replay_py_venv, + "UV_PROJECT_ENVIRONMENT": _replay_py_venv, + }, + } + + late_arrival_slack = 2 + # Size for the largest age the run will use (warmup age if larger) so the + # deeper warmup buffer fits. + buffer_age = max(max_trajectory_age_steps, warmup_max_trajectory_age_steps) + optimal_buffer_size = num_prompts_per_step * buffer_age * late_arrival_slack + min_trajectories_needed = num_prompts_per_step + + print("📊 Async PPO buffer requirements:") + print(f" - num_prompts_per_step: {num_prompts_per_step}") + print(f" - max_trajectory_age_steps: {max_trajectory_age_steps}") + print(f" - warmup_max_trajectory_age_steps: {warmup_max_trajectory_age_steps}") + print(f" - optimal_buffer_size: {optimal_buffer_size}") + + replay_buffer = ReplayBuffer.options(runtime_env=_replay_runtime_env).remote( + max_size=optimal_buffer_size, + log_every=async_cfg.get("log_every"), + ) + + last_checkpoint_path = checkpointer.get_latest_checkpoint_path() + if last_checkpoint_path is not None: + replay_buffer_path = os.path.join(last_checkpoint_path, "replay_buffer.pt") + if os.path.exists(replay_buffer_path): + print(f"📦 Restoring replay buffer from checkpoint: {replay_buffer_path}") + # weights_only=False: trajectories are pickled BatchedDataDict/dicts, + # not plain tensors. The checkpoint is a trusted same-job artifact. + replay_buffer_state = torch.load(replay_buffer_path, weights_only=False) + ray.get( + replay_buffer.load_state_dict.remote( + replay_buffer_state, + num_prompts_per_step=num_prompts_per_step, + current_training_step=step, + max_age_steps=_sample_max_age(step), + ) + ) + print("✅ Replay buffer restored from checkpoint") + else: + print( + f"⚠️ No replay buffer checkpoint found at {replay_buffer_path}. " + "Starting with an empty replay buffer." + ) + + _tc_py_exec = get_actor_python_env( + "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector" + ) + if _tc_py_exec.startswith("uv"): + _tc_py_exec = create_local_venv_on_each_node( + _tc_py_exec, + "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector", + ) + _tc_py_venv = os.path.dirname(os.path.dirname(_tc_py_exec)) + _tc_runtime_env = { + "py_executable": _tc_py_exec, + "env_vars": { + **os.environ, + "VIRTUAL_ENV": _tc_py_venv, + "UV_PROJECT_ENVIRONMENT": _tc_py_venv, + }, + } + + # The collector resolves the algorithm config block generically + # (master_config.ppo / "async_ppo"), so the PPO master_config is passed + # through directly. + trajectory_collector = AsyncTrajectoryCollector.options( + runtime_env=_tc_runtime_env + ).remote( + policy_generation=policy_generation, + tokenizer=tokenizer, + task_to_env=task_to_env, + master_config=master_config, + replay_buffer=replay_buffer, + start_step=step, + ) - val_batch, gen_metrics = run_multi_turn_rollout( + collection_task = trajectory_collector.start_collection.remote(dataloader) # noqa: F841 + trajectory_collector.set_weight_version.remote(weight_version) + # Match the collector's generation-lead to the current phase (a larger age + # during warmup, if configured). set_max_trajectory_age is a no-op-equivalent + # for GRPO (never called) — it only affects async PPO. + trajectory_collector.set_max_trajectory_age.remote(_collector_lead_age(step)) + print("📦 Started continuous background trajectory collection") + + # ------------------------------------------------------------------ + # Initial refit so the generation engine holds real trained weights. + # After setup(), the policy params are on GPU and the value model is + # offloaded, so the (non-colocated NCCL) broadcast can read policy weights + # directly. We then offload the policy so the first step's value forward has + # room on the shared train_cluster GPUs. + # ------------------------------------------------------------------ + print("⏳ Preparing policy generation for training (initial refit)...") + refit_policy_generation(policy, policy_generation, colocated_inference) + POLICY_GENERATION_STALE = False + policy.offload_to_cpu() + + if val_at_start and step == 0: + print("\n🔍 Running initial validation...") + trajectory_collector.pause.remote() + try: + val_metrics, validation_timings = validate( policy_generation, - val_batch, + val_dataloader, tokenizer, val_task_to_env, - max_seq_len=master_config.policy["max_total_sequence_length"], - max_rollout_turns=master_config.ppo["max_rollout_turns"], - greedy=False, + step=0, + master_config=master_config, + logger=logger, + ) + policy_generation.finish_generation() + logger.log_metrics(val_metrics, step, prefix="validation") + logger.log_metrics(validation_timings, step, prefix="timing/validation") + finally: + trajectory_collector.resume.remote() + + if policy_generation is not None: + policy_generation.clear_logger_metrics() + + # Wait for the buffer to hold a full step's worth of trajectories. + print(f"⏳ Waiting for replay buffer to be ready for step {step}...") + timer.start("init/total") + wait_iterations = 0 + while True: + current_step_ready = ray.get( + replay_buffer.has_complete_batch.remote( + step, num_prompts_per_step, _sample_max_age(step) ) + ) + if current_step_ready: + break + if wait_iterations % 10 == 0: + buffer_size_current = ray.get(replay_buffer.size.remote()) + print( + f" Wait iteration {wait_iterations}: buffer_size={buffer_size_current}, " + f"step {step} ready={current_step_ready}" + ) + wait_iterations += 1 + time.sleep(1.0) + timer.stop("init/total") + print(f"✅ Buffer ready for step {step}! Starting async PPO training loop...") + + # ------------------------------------------------------------------ + # Main loop + # ------------------------------------------------------------------ + try: + while step < max_num_steps: + print(f"\n{'=' * 25} Step {step + 1}/{max_num_steps} {'=' * 25}") + maybe_gpu_profile_step(policy, step + 1) + if policy != policy_generation: + maybe_gpu_profile_step(policy_generation, step + 1) + + metrics: dict[str, Any] = {} + val_metrics, validation_timings = None, None + + with timer.time("total_step_time"): + # ---- 1. Sample a fixed batch of trajectories from the buffer ---- + print("📦 Sampling from replay buffer...") + with timer.time("exposed_generation"): + # _sample_max_age keeps the elevated age through step W + A_t so + # frozen (pi_0) rollouts banked during warmup are admitted while + # they are still within A_t POLICY-steps of the actor, then drops + # to the training age (evicting genuinely-stale pi_0). Passing it + # to sample() is what actually clamps staleness. + sample_result = ray.get( + replay_buffer.sample.remote( + num_prompt_groups=num_prompts_per_step, + current_weight_version=weight_version, + max_age_steps=_sample_max_age(step), + ) + ) + if ( + sample_result is None + or len(sample_result["trajectories"]) != num_prompts_per_step + ): + print( + "⏳ Buffer empty or not enough groups for a full step, " + "waiting..." + ) + with timer.time("idle/buffer_starvation"): + time.sleep(0.5) + continue + + trajectories = sample_result["trajectories"] + avg_trajectory_age = sample_result["avg_trajectory_age"] + # Freeze-aware POLICY-age: the number that actually bounds the + # importance-sampling correction (see _async_trajectory_policy_age). + # During critic warmup the actor is frozen at pi_0, so a rollout's + # *generation*-version age overcounts its off-policyness; the true + # off-policy distance is <= max_trajectory_age_steps by construction + # (what the two-boundary eviction enforces), so avg_trajectory_age + # can legitimately read higher across the warmup boundary without + # any extra off-policyness. For GRPO / no-warmup the two are equal. + gen_versions = sample_result.get("generation_weight_versions", []) + _policy_ages = [ + _async_trajectory_policy_age( + g, weight_version, policy_training_start_step + ) + for g in gen_versions + ] + avg_trajectory_policy_age = ( + sum(_policy_ages) / len(_policy_ages) + if _policy_ages + else avg_trajectory_age + ) + max_trajectory_policy_age = max(_policy_ages, default=0) + print( + f"✅ Sampled {len(trajectories)} trajectory groups " + f"(gen-version age: {avg_trajectory_age:.2f} steps | " + f"policy-age: {avg_trajectory_policy_age:.2f} avg / " + f"{max_trajectory_policy_age} max, " + f"bound {max_trajectory_age_steps})" + ) + if max_trajectory_policy_age > max_trajectory_age_steps: + print( + "⚠️ WARNING: sampled policy-age " + f"{max_trajectory_policy_age} exceeds " + f"max_trajectory_age_steps={max_trajectory_age_steps} — " + "this should not happen; the buffer admitted a rollout " + "more off-policy than configured." + ) + + per_prompt_batches = [t["batch"] for t in trajectories] + repeated_batch = BatchedDataDict.from_batches(per_prompt_batches) + + per_group_metrics: dict[str, list] = {} + for t in trajectories: + for k, v in t["rollout_metrics"].items(): + per_group_metrics.setdefault(k, []).append(v) + rollout_metrics = aggregate_rollout_metrics(per_group_metrics) + + expected_batch_size = ( + master_config.ppo["num_prompts_per_step"] + * master_config.ppo["num_generations_per_prompt"] + ) + if repeated_batch.size != expected_batch_size: + print( + f"❌ Unexpected training batch size: got {repeated_batch.size}, " + f"expected {expected_batch_size}. Waiting for correct buffer " + "content." + ) + time.sleep(0.5) + continue + + # ---- 2. Build PPO training data (rewards + inline loss mask) ---- + print("▶ Processing rewards...") + with timer.time("data_processing"): + rewards = repeated_batch["total_reward"] + + use_overlong_filtering = master_config.ppo["overlong_filtering"] + if use_overlong_filtering: + loss_multiplier = repeated_batch["loss_multiplier"].clone() + truncated = repeated_batch["truncated"] + if isinstance(truncated, list): + truncated = torch.tensor(truncated, dtype=torch.bool) + loss_multiplier[truncated] = 0 + repeated_batch["loss_multiplier"] = loss_multiplier + + # PPO's inline loss-mask setup (unmask all assistant messages), + # matching sync ppo_train — deliberately NOT GRPO's helper, + # which only unmasks generated assistant messages. + for message_log in repeated_batch["message_log"]: + for message in message_log: + if message["role"] == "assistant": + message["token_loss_mask"] = torch.ones_like( + message["token_ids"] + ) + else: + message["token_loss_mask"] = torch.zeros_like( + message["token_ids"] + ) + if "generation_logprobs" not in message: + message["generation_logprobs"] = torch.zeros_like( + message["token_ids"], dtype=torch.float32 + ) + + flat_messages, input_lengths = batched_message_log_to_flat_message( + repeated_batch["message_log"], + pad_value_dict={"token_ids": tokenizer.pad_token_id}, + make_sequence_length_divisible_by=master_config.policy[ + "make_sequence_length_divisible_by" + ], + ) + + train_data = BatchedDataDict[ClippedPGLossDataDict]( + { + "input_ids": flat_messages["token_ids"], + "input_lengths": input_lengths, + "generation_logprobs": flat_messages["generation_logprobs"], + "rewards": repeated_batch["total_reward"], + "token_mask": flat_messages["token_loss_mask"], + "sample_mask": repeated_batch["loss_multiplier"], + } + ) + extra_multimodal_data = flat_messages.get_multimodal_dict( + as_tensors=False + ) + train_data.update(extra_multimodal_data) + train_data.to("cpu") + + # ---- 3. Value forward (critic on GPU, then offloaded) ---- + # GPU state entering here: policy OFF, value OFF (see refit/step + # end below). Load value only. + print("▶ Computing values...") + with timer.time("value_inference"): + value_model.prepare_for_inference() + train_data["values"] = value_model.get_values(train_data)[ + "values" + ].squeeze(-1) + value_model.finish_inference() + + # ---- 4. Policy / reference logprobs (policy on GPU, then off) ---- + print("▶ Computing logprobs...") + with timer.time("logprob_inference_prep"): + policy.prepare_for_lp_inference() + with timer.time("policy_and_reference_logprobs"): + logprob_data = BatchedDataDict[ClippedPGLossDataDict]( + { + "input_ids": train_data["input_ids"], + "input_lengths": train_data["input_lengths"], + **extra_multimodal_data, + } + ) + train_data["prev_logprobs"] = policy.get_logprobs( + logprob_data, timer=timer + )["logprobs"] + if not master_config.ppo.get( + "skip_reference_policy_logprobs_calculation" + ): + train_data["reference_policy_logprobs"] = ( + policy.get_reference_policy_logprobs( + logprob_data, + timer=timer, + )["reference_logprobs"] + ) + del logprob_data + del extra_multimodal_data + policy.finish_inference() + + # ---- 5. Sequence-level train/inference mismatch diagnostics ---- + seq_logprob_error_threshold = master_config.ppo.get( + "seq_logprob_error_threshold", None + ) + seq_error_result = compute_and_apply_seq_logprob_error_masking( + train_data=train_data, + rewards=rewards, + seq_logprob_error_threshold=seq_logprob_error_threshold, + ) + seq_logprob_error_metrics = seq_error_result + if "num_masked_seqs" in seq_logprob_error_metrics: + seq_logprob_error_metrics["num_masked_seqs_by_logprob_error"] = ( + seq_logprob_error_metrics.pop("num_masked_seqs") + ) + + # ---- 6. GAE advantages/returns (uses fresh values) ---- + with timer.time("advantage_calculation"): + print("▶ Computing advantages...") + initial_prompt_message_logs = extract_initial_prompt_messages( + repeated_batch["message_log"], + repeated_batch["length"], + ) + prompt_batched_flat, _ = batched_message_log_to_flat_message( + initial_prompt_message_logs, + pad_value_dict={"token_ids": tokenizer.pad_token_id}, + ) + prompt_ids_for_adv = prompt_batched_flat["token_ids"] + del initial_prompt_message_logs + del prompt_batched_flat + + adv_kwargs = dict( + prompt_ids=prompt_ids_for_adv, + rewards=train_data["rewards"], + mask=train_data["token_mask"], + reference_logprobs=train_data.get("reference_policy_logprobs"), + logprobs=train_data["prev_logprobs"], + ) + if "values" in train_data: + adv_kwargs["values"] = train_data["values"] + result = adv_estimator.compute_advantage(**adv_kwargs) + if isinstance(result, tuple): + advantages, returns = result + else: + advantages, returns = result, None + del prompt_ids_for_adv + train_data["advantages"] = advantages + if returns is not None: + train_data["returns"] = returns + + # ---- 7. ppo_epochs inner loop (critic, then actor) ---- + # Each epoch: value on GPU -> train -> off. Then, once past critic + # warmup, policy on GPU -> train -> off (except the last epoch, + # which leaves the policy on GPU for the refit broadcast below). + # During warmup (step < policy_training_start_step) the policy is + # frozen: it is never loaded/trained here, exactly as in sync + # ppo_train, so train_results stays None for the step. + is_policy_training_step = step >= policy_training_start_step + train_results = None + value_results = None + for epoch in range(ppo_epochs): + print(f"▶ PPO epoch {epoch + 1}/{ppo_epochs}...") + with timer.time("value_training_prep"): + value_model.prepare_for_training() + with timer.time("value_training"): + value_results = value_model.train( + train_data, + value_loss_fn, + timer=timer, + ) + value_model.finish_training() + + if is_policy_training_step: + if ( + step == policy_training_start_step + and policy_training_start_step > 0 + and epoch == 0 + ): + print( + f" ✓ Critic warmup complete ({policy_training_start_step} " + "steps). Starting policy training.", + flush=True, + ) + with timer.time("training_prep"): + policy.prepare_for_training() + POLICY_GENERATION_STALE = True + with timer.time("policy_training"): + train_results = policy.train( + train_data, loss_fn, timer=timer + ) + if epoch < ppo_epochs - 1: + policy.offload_to_cpu() + + # ---- 8. Single weight refit after the whole ppo_epochs loop ---- + # When the policy trained this step: GPU state is policy ON (last + # epoch did not offload), value OFF; the non-colocated broadcast + # reads policy weights directly, then we offload the policy so the + # next step's value forward has room on the shared train_cluster. + # + # During critic warmup the policy is frozen and stayed OFF this + # step, and generation already holds the correct (initial) weights + # from the last real refit, so we SKIP the weight transfer. We + # still run the same collector coordination (pause -> bump weight + # version -> resume) so the replay buffer/collector keep advancing. + generation_logger_metrics = None + if NEED_REFIT: + print("🔄 Coordinating with trajectory collector before refit...") + with timer.time("idle/refit_bubble"): + with timer.time("exposed_generation"): + ray.get(trajectory_collector.prepare_for_refit.remote()) + if policy_generation is not None: + generation_logger_metrics = ( + policy_generation.get_logger_metrics() + ) + with timer.time("weight_sync"): + if is_policy_training_step: + refit_policy_generation( + policy, policy_generation, colocated_inference + ) + else: + print( + "▶ Critic warmup: skipping policy weight transfer " + "(policy frozen; generation already up to date)" + ) + POLICY_GENERATION_STALE = False + weight_version += 1 + trajectory_collector.set_weight_version.remote( + weight_version + ) + # At the warmup->training boundary (end of step W) the + # collector's generation-lead drops A_w -> A_t. This both + # stops over-banking frozen rollouts AND is required for + # correctness: from step W+1 the collector must regenerate + # its lead targets against the freshly-trained policy + # (pi_1, pi_2, ...) rather than reuse banked pi_0. The + # driver's own eviction age (_sample_max_age) drops later, + # at W + A_t, so the banked boundary batch is still + # admitted as valid lag-<=A_t data. + next_age = _collector_lead_age(step + 1) + if next_age != _collector_lead_age(step): + trajectory_collector.set_max_trajectory_age.remote( + next_age + ) + trajectory_collector.resume_after_refit.remote() + # Only the policy-training path leaves the policy resident on GPU; + # during warmup it is already offloaded, so skip the redundant call. + if is_policy_training_step: + policy.offload_to_cpu() + + if policy_generation is not None: + policy_generation.clear_logger_metrics() + + # ---- Validation ---- + is_last_step = step + 1 == max_num_steps + if (val_period > 0 and (step + 1) % val_period == 0) or ( + val_at_end and is_last_step + ): + with timer.time("idle/validation"): + trajectory_collector.pause.remote() + # Policy weights are already synced to generation from the + # refit above (POLICY_GENERATION_STALE is False), so just + # enter generation mode. + policy_generation.prepare_for_generation() + val_metrics, validation_timings = validate( + policy_generation, + val_dataloader, + tokenizer, + val_task_to_env, + step=step + 1, + master_config=master_config, + logger=logger, + ) + policy_generation.finish_generation() + logger.log_metrics( + validation_timings, step + 1, prefix="timing/validation" + ) + logger.log_metrics(val_metrics, step + 1, prefix="validation") + gc.collect() + torch.cuda.empty_cache() + trajectory_collector.resume.remote() + + # ---- Metrics ---- + flat_advantages = train_data["advantages"] + flat_token_mask = flat_messages["token_loss_mask"] + flat_messages_content = flat_messages.get("content", []) + del flat_messages + response_advantages = torch.masked_select( + flat_advantages, flat_token_mask.bool() + ) + + metrics.update( + { + "reward": rewards.numpy(), + "mean_prompt_length": repeated_batch["length"].numpy(), + "total_num_tokens": input_lengths.numpy(), + "advantages/mean": torch.mean(response_advantages) + .detach() + .item() + if response_advantages.numel() > 0 + else 0.0, + "advantages/max": torch.max(response_advantages).detach().item() + if response_advantages.numel() > 0 + else 0.0, + "advantages/min": torch.min(response_advantages).detach().item() + if response_advantages.numel() > 0 + else 0.0, + } + ) + # Policy metrics are absent during critic warmup (train_results is + # None because the policy was not trained this step). + if train_results is not None: + metrics["loss"] = train_results["loss"].numpy() + metrics["grad_norm"] = train_results["grad_norm"].numpy() + if "moe_metrics" in train_results: + metrics.update( + { + f"moe/{k}": v + for k, v in train_results["moe_metrics"].items() + } + ) + metrics.update(train_results["all_mb_metrics"]) + if value_results is not None: + metrics.update(_compute_critic_metrics(value_results)) + + for k, v in metrics.items(): + if k in {"probs_ratio_min", "probs_ratio_clamped_min"}: + valid_values = [x for x in v if not np.isinf(x)] + metrics[k] = ( + np.min(valid_values).item() if valid_values else -1.0 + ) + elif k in {"probs_ratio_max", "probs_ratio_clamped_max"}: + valid_values = [x for x in v if not np.isinf(x)] + metrics[k] = ( + np.max(valid_values).item() if valid_values else -1.0 + ) + elif k in { + "lr", + "wd", + "reward", + "global_valid_seqs", + "global_valid_toks", + "mean_prompt_length", + }: + metrics[k] = np.mean(v).item() + elif isinstance(v, (np.ndarray, list)): + metrics[k] = np.sum(v).item() + + metrics.update(rollout_metrics) + if generation_logger_metrics is not None: + metrics["generation_logger_metrics"] = generation_logger_metrics + if "global_valid_toks" in metrics: + total_valid_tokens += metrics["global_valid_toks"] + # Always log seq-level error metrics (useful for tuning threshold). + metrics.update(seq_logprob_error_metrics) + + # ---- Checkpointing ---- + consumed_samples += master_config.ppo["num_prompts_per_step"] + timeout.mark_iteration() + should_save_by_step = ( + is_last_step + or (step + 1) % master_config.checkpointing["save_period"] == 0 + ) + should_save_by_timeout = timeout.check_save() + if master_config.checkpointing["enabled"] and ( + should_save_by_step or should_save_by_timeout + ): + ppo_save_state["current_step"] = step + 1 + ppo_save_state["total_steps"] = step + 1 + ppo_save_state["total_valid_tokens"] = total_valid_tokens + if val_metrics is not None: + ppo_save_state["val_reward"] = val_metrics["accuracy"] + elif "val_reward" in ppo_save_state: + del ppo_save_state["val_reward"] + ppo_save_state["consumed_samples"] = consumed_samples + + # Record the top-k ranking metric into the save state so + # get_best_checkpoint_path / top-k pruning work (parity with + # sync ppo_train and async_grpo_train). + full_metric_name = master_config.checkpointing["metric_name"] + if full_metric_name is not None: + assert full_metric_name.startswith( + "train:" + ) or full_metric_name.startswith("val:"), ( + f"metric_name={full_metric_name} must start with 'val:' or 'train:',\n" + f'followed by the corresponding name in the "val" or "train" metrics dictionary.' + ) + prefix, metric_name = full_metric_name.split(":", 1) + metrics_source = metrics if prefix == "train" else val_metrics + if not metrics_source: + warnings.warn( + f"You asked to save checkpoints based on {metric_name} but no {prefix} metrics were collected. " + "This checkpoint will not be saved as top-k.", + stacklevel=2, + ) + if full_metric_name in ppo_save_state: + del ppo_save_state[full_metric_name] + elif metric_name not in metrics_source: + raise ValueError( + f"Metric {metric_name} not found in {prefix} metrics" + ) + else: + ppo_save_state[full_metric_name] = metrics_source[ + metric_name + ] + + with timer.time("checkpointing"): + print(f"Saving checkpoint for step {step + 1}...") + checkpoint_path = checkpointer.init_tmp_checkpoint( + step + 1, ppo_save_state, master_config + ) + # Policy first (its presence marks a trained policy), + # then value. Both are offloaded at this point, so load + # each for saving and offload again after. During critic + # warmup the policy optimizer has no state yet, so skip + # the policy checkpoint entirely (matching sync ppo_train); + # the resume path falls back to the base model weights. + if is_policy_training_step: + policy.prepare_for_training() + policy.save_checkpoint( + weights_path=os.path.join( + checkpoint_path, "policy", "weights" + ), + optimizer_path=os.path.join( + checkpoint_path, "policy", "optimizer" + ), + tokenizer_path=os.path.join( + checkpoint_path, "policy", "tokenizer" + ), + checkpointing_cfg=master_config.checkpointing, + ) + policy.offload_to_cpu() + else: + print( + f"Skipping policy checkpoint (critic warmup: " + f"step {step} < {policy_training_start_step})", + flush=True, + ) + + value_model.prepare_for_training() + value_model.save_checkpoint( + weights_path=os.path.join( + checkpoint_path, "value", "weights" + ), + optimizer_path=os.path.join( + checkpoint_path, "value", "optimizer" + ), + tokenizer_path=os.path.join( + checkpoint_path, "value", "tokenizer" + ), + checkpointing_cfg=master_config.checkpointing, + ) + value_model.finish_training() + + actual_dataloader_state = ray.get( + trajectory_collector.get_dataloader_state.remote() + ) + torch.save( + actual_dataloader_state, + os.path.join(checkpoint_path, "train_dataloader.pt"), + ) + print("📦 Saving replay buffer state...") + replay_buffer_state = ray.get(replay_buffer.state_dict.remote()) + torch.save( + replay_buffer_state, + os.path.join(checkpoint_path, "replay_buffer.pt"), + ) + checkpointer.finalize_checkpoint(checkpoint_path) + _write_latest_checkpoint_status( + checkpointer, last_checkpoint_step=step + 1 + ) + + # ---- Logging ---- + if not _should_log_nemo_gym_responses(master_config): + log_data = {} + log_data["content"] = flat_messages_content + log_data["rewards"] = rewards.tolist() + log_data["input_lengths"] = input_lengths.tolist() + log_data["token_ids"] = train_data["input_ids"].tolist() + log_data["token_loss_mask"] = train_data["token_mask"].tolist() + log_data["sample_loss_mask"] = train_data["sample_mask"].tolist() + log_data["advantages"] = train_data["advantages"].tolist() + log_data["generation_logprobs"] = train_data[ + "generation_logprobs" + ].tolist() + log_data["prev_logprobs"] = train_data["prev_logprobs"].tolist() + logger.log_batched_dict_as_jsonl( + log_data, f"train_data_step{step + 1}.jsonl" + ) + del log_data + del flat_messages_content + + timing_metrics: dict[str, float] = timer.get_timing_metrics( + reduction_op="sum" + ) # type: ignore + + buffer_size_current = ray.get(replay_buffer.size.remote()) + metrics["buffer_size"] = buffer_size_current + metrics["avg_trajectory_age"] = avg_trajectory_age + # Freeze-aware off-policy staleness (bounded by max_trajectory_age_steps); + # equals avg_trajectory_age when there is no critic warmup. + metrics["avg_trajectory_policy_age"] = avg_trajectory_policy_age + metrics["max_trajectory_policy_age"] = max_trajectory_policy_age + + # Track the worst-mismatch example plot (parity with sync PPO). + if metrics.get("token_mult_prob_error", 0) > 1.05: + logger.log_plot_token_mult_prob_error( + { + "prompt_lengths": repeated_batch["length"], + "full_lengths": input_lengths, + "generation_logprobs": train_data["generation_logprobs"], + "prev_logprobs": train_data["prev_logprobs"], + "token_mask": train_data["token_mask"], + "sample_mask": train_data["sample_mask"], + }, + step + 1, + name="train/token_mult_prob_error_plot_sample", + ) + del train_data + + print("\n📊 Training Results:") + if "loss" in metrics: + print(f" • Loss: {metrics['loss']:.4f}") + print( + f" • Generation KL Error: {metrics.get('gen_kl_error', 'N/A')}" + ) + else: + print(" • (critic warmup: policy not trained this step)") + if "critic/loss" in metrics: + print(f" • Critic Loss: {metrics['critic/loss']:.4f}") + print(f" • Avg Reward: {np.mean(rewards.numpy()):.4f}") + print(f" • Buffer Size: {buffer_size_current}") + print( + f" • Avg Trajectory Age (gen-version): {avg_trajectory_age:.2f} steps" + ) + print( + f" • Avg Trajectory Policy-Age (freeze-aware, off-policy staleness): " + f"{avg_trajectory_policy_age:.2f} avg / {max_trajectory_policy_age} max " + f"(bound: max_trajectory_age_steps={max_trajectory_age_steps})" + ) + + total_time = timing_metrics.get("total_step_time", 0) + total_num_gpus = ( + master_config.cluster["num_nodes"] + * master_config.cluster["gpus_per_node"] + ) + if total_time > 0 and "global_valid_toks" in metrics: + timing_metrics["valid_tokens_per_sec_per_gpu"] = ( + metrics["global_valid_toks"] / total_time / total_num_gpus + ) + performance_metrics = print_performance_metrics( + train_results if train_results is not None else (value_results or {}), + metrics, + timing_metrics, + master_config, + ) + + collector_efficiency = ray.get( + trajectory_collector.get_efficiency_metrics.remote() + ) + driver_efficiency = { + cat: timer.reduce(cat, "sum") + for cat in [ + "init/total", + "idle/buffer_starvation", + "idle/refit_bubble", + "idle/validation", + ] + if cat in timer._timers + } + merged_efficiency = {**driver_efficiency} + for cat, dur in collector_efficiency.items(): + merged_efficiency[cat] = merged_efficiency.get(cat, 0.0) + dur + total_wall_time = time.perf_counter() - training_wall_start + efficiency_loggable = print_efficiency_summary( + merged_efficiency, total_wall_time, step + 1 + ) + + logger.log_metrics(performance_metrics, step + 1, prefix="performance") + logger.log_metrics(metrics, step + 1, prefix="train") + logger.log_metrics(efficiency_loggable, step + 1, prefix="") + logger.log_metrics( + timing_metrics, + step + 1, + prefix="timing/train", + step_finished=True, + ) + + timer.reset() + step += 1 + if should_save_by_timeout: + print("Timeout has been reached, stopping training early", flush=True) + return + if step >= max_num_steps: + print( + "Max number of steps has been reached, stopping training early", + flush=True, + ) + return + + except Exception as e: + print(f"❌ Error in async PPO loop: {e}") + import traceback + + traceback.print_exc() + + finally: + print("🛑 Stopping trajectory collection...") + try: + ray.kill(trajectory_collector) + except Exception as e: + print(f"Error stopping trajectory collector: {e}") + try: + ray.kill(replay_buffer) + except Exception as e: + print(f"Error stopping replay buffer: {e}") + + # Shut down environments before generation workers: they may hold + # in-flight HTTP requests to the vLLM endpoints. + for env_dict in (task_to_env, val_task_to_env): + if env_dict is None: + continue + for task_name, env in env_dict.items(): + print(f"🛑 Shutting down environment {task_name}...") + try: + ray.get(env.shutdown.remote(), timeout=10) + except Exception: + try: + ray.kill(env) + except Exception as e: + print(f"Error shutting down environment {task_name}: {e}") + + print("🛑 Shutting down generation workers...") + try: + policy_generation.shutdown() + except Exception as e: + print(f"Error shutting down generation workers: {e}") + if policy is not policy_generation: + print("🛑 Shutting down policy workers...") + try: + policy.shutdown() + except Exception as e: + print(f"Error shutting down policy workers: {e}") + print("🛑 Shutting down value workers...") + try: + value_model.shutdown() + except Exception as e: + print(f"Error shutting down value workers: {e}") + print("Async PPO training complete!") + + +def validate( + policy_generation: GenerationInterface, + val_dataloader: Optional[StatefulDataLoader], + tokenizer, + val_task_to_env: Optional[dict[str, EnvironmentInterface]], + step: int, + master_config: MasterConfig, + logger: Optional[Logger] = None, +) -> tuple[dict[str, Any], dict[str, Any]]: + """Run validation on the validation dataset.""" + if val_dataloader is None: + assert val_dataloader is not None or master_config.ppo["val_period"] == 0, ( + "val_dataloader is None, so ppo.val_period must be 0" + ) + print(" ⚠️ No validation dataloader provided, skipping validation", flush=True) + return {}, {} + + timer = Timer() + with timer.time("total_validation_time"): + print(f"▶ Starting validation at step {step}...", flush=True) + + total_rewards = [] + total_lengths = [] + all_message_logs = [] # Collect all message logs + + max_batches = ( + master_config.ppo["max_val_samples"] // master_config.ppo["val_batch_size"] + ) + for batch_idx, val_batch in enumerate(val_dataloader): + if batch_idx >= max_batches: + break + + additional_metrics_to_report = dict() + + # Async PPO uses the vLLM async engine (async_engine=true), whose + # generation worker exposes only the async rollout path — the sync + # `run_multi_turn_rollout` -> policy_generation.generate() would raise + # AttributeError ('generate') on the async worker. Dispatch the same way + # the training loop does (see async_ppo_train / grpo.validate). + if _should_use_async_rollouts(master_config): + val_batch, gen_metrics = run_async_multi_turn_rollout( + policy_generation, + val_batch, + tokenizer, + val_task_to_env, + max_seq_len=master_config.policy["max_total_sequence_length"], + max_rollout_turns=master_config.ppo["max_rollout_turns"], + greedy=False, + ) + else: + val_batch, gen_metrics = run_multi_turn_rollout( + policy_generation, + val_batch, + tokenizer, + val_task_to_env, + max_seq_len=master_config.policy["max_total_sequence_length"], + max_rollout_turns=master_config.ppo["max_rollout_turns"], + greedy=False, + ) total_rewards.extend(val_batch["total_reward"].tolist()) total_lengths.append(gen_metrics["mean_gen_tokens_per_sample"]) diff --git a/ray.sub b/ray.sub index e71a90ebbb5..b575ba3276b 100644 --- a/ray.sub +++ b/ray.sub @@ -117,13 +117,24 @@ export RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 # Setting ulimit is recommended by ray best practices page # @ https://docs.ray.io/en/latest/cluster/vms/user-guides/large-cluster-best-practices.html -# It's session based and won't affect the system outside the script -# Ensure that the soft limit isn't above the hard limit -if [[ $(ulimit -Hn) == "unlimited" ]] || [[ 65535 -lt $(ulimit -Hn) ]]; then - ulimit -Sn 65535 -elif [[ $(ulimit -Hn) != "unlimited" ]] && [[ $(ulimit -Hn) -lt 65535 ]]; then - echo "[WARNING]: Cannot increase ulimit on file descriptors to 65535 according ray recommendation: https://docs.ray.io/en/latest/cluster/vms/user-guides/large-cluster-best-practices.html. Speak to cluster admins to increase, otherwise ray may crash unexpectedly." +# It's session based and won't affect the system outside the script. +# 65535 is Ray's *minimum* recommendation. Async RL opens one asyncio/uvloop event +# loop (several file descriptors) per in-flight generation rollout, so with a large +# ppo.num_prompts_per_step the collector can need far more than 65535 and hit +# "OSError: [Errno 24] Too many open files". Raise the soft limit as high as the +# hard limit permits (capped at 1048576, which covers e.g. 8192 prompts * age 10). +_NRL_TARGET_NOFILE=1048576 +_NRL_HARD_NOFILE=$(ulimit -Hn) +if [[ "$_NRL_HARD_NOFILE" == "unlimited" ]] || [[ "$_NRL_HARD_NOFILE" -ge "$_NRL_TARGET_NOFILE" ]]; then + ulimit -Sn "$_NRL_TARGET_NOFILE" +elif [[ "$_NRL_HARD_NOFILE" -ge 65535 ]]; then + # Hard limit below our target but still >= Ray's minimum: raise soft to the hard cap. + ulimit -Sn "$_NRL_HARD_NOFILE" +else + echo "[WARNING]: hard FD limit ($_NRL_HARD_NOFILE) is below Ray's recommended 65535. Async RL with a large ppo.num_prompts_per_step may crash with 'Too many open files' (Errno 24). Ask cluster admins to raise the nofile hard limit." + ulimit -Sn "$_NRL_HARD_NOFILE" fi +echo "[INFO] FD limit (ulimit -n): soft=$(ulimit -Sn) hard=$(ulimit -Hn)" # Worker port range must NOT overlap with the OS ephemeral range (32768-60999) # to prevent TOCTOU collisions. Ray's Raylet uses a probe-and-release pattern diff --git a/tests/functional/L1_Functional_Tests_PPO.sh b/tests/functional/L1_Functional_Tests_PPO.sh index 8865c2157dc..e6b9b23276e 100755 --- a/tests/functional/L1_Functional_Tests_PPO.sh +++ b/tests/functional/L1_Functional_Tests_PPO.sh @@ -37,6 +37,7 @@ run_test() { run_test fast uv run --no-sync bash ./tests/functional/ppo_automodel.sh run_test fast uv run --no-sync bash ./tests/functional/ppo_megatron.sh run_test uv run --no-sync bash ./tests/functional/ppo_non_colocated.sh +run_test uv run --no-sync bash ./tests/functional/ppo_async.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/ppo_async.sh b/tests/functional/ppo_async.sh new file mode 100755 index 00000000000..0d67feecb9f --- /dev/null +++ b/tests/functional/ppo_async.sh @@ -0,0 +1,64 @@ +#!/bin/bash +# Async PPO smoke test on a single 2-GPU node (1 train / 1 inference, +# non-colocated). policy_training_start_step=1 + max_num_steps=3 deliberately +# exercises the critic-warmup path (step 0), the warmup->training transition +# (step 1), and a normal async step (step 2) in one run. + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) +# Mark the current repo as safe, since wandb fetches metadata about the repo +git config --global --add safe.directory $PROJECT_ROOT + +set -eou pipefail + +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 +export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} + +rm -rf $EXP_DIR $LOG_DIR +mkdir -p $EXP_DIR $LOG_DIR + +cd $PROJECT_ROOT +uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ + $PROJECT_ROOT/examples/run_ppo.py \ + policy.model_name=Qwen/Qwen2.5-0.5B \ + value.model_name=Qwen/Qwen2.5-0.5B \ + ppo.num_prompts_per_step=2 \ + ppo.num_generations_per_prompt=4 \ + ppo.ppo_epochs=2 \ + ppo.policy_training_start_step=1 \ + policy.train_global_batch_size=4 \ + policy.logprob_batch_size=4 \ + policy.train_micro_batch_size=1 \ + policy.generation.colocated.enabled=false \ + policy.generation.colocated.resources.gpus_per_node=1 \ + policy.generation.colocated.resources.num_nodes=1 \ + policy.generation.vllm_cfg.async_engine=true \ + ppo.async_ppo.enabled=true \ + ppo.async_ppo.max_trajectory_age_steps=1 \ + ppo.async_ppo.in_flight_weight_updates=false \ + loss_fn.use_importance_sampling_correction=true \ + value.train_global_batch_size=4 \ + value.train_micro_batch_size=1 \ + cluster.gpus_per_node=2 \ + ppo.max_num_steps=3 \ + logger.tensorboard_enabled=true \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=false \ + logger.monitor_gpus=true \ + checkpointing.enabled=false \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/token_mult_prob_error"]) < 1.05' \ + 'max(data["train/critic/loss"]) < 6.0' \ + 'min(data["train/critic/loss"]) >= 0' \ + 'max(data["train/critic/explained_var"]) <= 1.0001' \ + 'min(data["train/buffer_size"]) >= 0' \ + 'max(data["train/avg_trajectory_age"]) <= 1.0' diff --git a/tests/unit/algorithms/test_async_utils.py b/tests/unit/algorithms/test_async_utils.py index b7cff7eecd7..d7bdc617159 100644 --- a/tests/unit/algorithms/test_async_utils.py +++ b/tests/unit/algorithms/test_async_utils.py @@ -1184,6 +1184,38 @@ def test_async_trajectory_collector_weight_version_updates(self): ray.kill(buffer) ray.kill(mock_env) + def test_async_trajectory_collector_max_trajectory_age_override(self): + """Collector max trajectory age defaults to config and is settable. + + Async PPO raises this during critic warmup (frozen actor => any-age + trajectories are on-policy) and lowers it at the warmup->training + boundary. The mock config has async_grpo.max_trajectory_age_steps=2 and + no warmup key, so the default is 2. + """ + buffer = ReplayBuffer.remote(max_size=10) + mock_generation = MockGenerationInterface() + mock_tokenizer = mock.MagicMock() + mock_env = MockEnvironment.remote(rewards=[1.0, 2.0]) + task_to_env = {"test": mock_env} + master_config = self.create_mock_config() + + collector = AsyncTrajectoryCollector.remote( + policy_generation=mock_generation, + tokenizer=mock_tokenizer, + task_to_env=task_to_env, + master_config=master_config, + replay_buffer=buffer, + start_step=0, + ) + + assert ray.get(collector.get_max_trajectory_age.remote()) == 2 + ray.get(collector.set_max_trajectory_age.remote(8)) + assert ray.get(collector.get_max_trajectory_age.remote()) == 8 + + ray.kill(collector) + ray.kill(buffer) + ray.kill(mock_env) + def test_async_trajectory_collector_pause_resume(self): """Test pause and resume functionality.""" buffer = ReplayBuffer.remote(max_size=10) diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 4c3a43b20a8..0a62c687aa1 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -965,3 +965,328 @@ def init_collective(self, *_args, **_kwargs): assert train_cluster is not inference_cluster assert train_cluster.max_colocated_worker_groups == 2 assert inference_cluster.max_colocated_worker_groups == 1 + + +# =============================================================================== +# Async PPO entry guards (async_ppo_train) +# =============================================================================== + + +def _build_async_ppo_master_config(): + """PPO MasterConfig pre-configured for a valid async run. + + Individual guard tests mutate one field to trip a specific assertion. All + guards fire before any Ray actor is created, so the other async_ppo_train + arguments can be plain mocks. + """ + master_config = _build_ppo_master_config() + master_config.policy["generation"]["backend"] = "vllm" + master_config.policy["generation"]["colocated"] = { + "enabled": False, + "resources": {"gpus_per_node": 1, "num_nodes": 1}, + } + master_config.policy["generation"]["vllm_cfg"]["async_engine"] = True + master_config.loss_fn = ClippedPGLossConfig(use_importance_sampling_correction=True) + master_config.ppo["policy_training_start_step"] = 0 + master_config.ppo["num_generations_per_prompt"] = 1 + master_config.ppo["overlong_filtering"] = False + master_config.ppo["async_ppo"] = { + "enabled": True, + "max_trajectory_age_steps": 1, + "in_flight_weight_updates": False, + } + return master_config + + +def _call_async_ppo_train(master_config): + from unittest.mock import MagicMock + + from nemo_rl.algorithms.ppo import async_ppo_train + + async_ppo_train( + policy=MagicMock(), + policy_generation=MagicMock(), + value_model=MagicMock(), + dataloader=MagicMock(), + val_dataloader=None, + tokenizer=MagicMock(), + loss_fn=MagicMock(), + value_loss_fn=MagicMock(), + task_to_env={}, + val_task_to_env=None, + logger=MagicMock(), + checkpointer=MagicMock(), + ppo_save_state=MagicMock(), + master_config=master_config, + max_trajectory_age_steps=master_config.ppo["async_ppo"][ + "max_trajectory_age_steps" + ], + ) + + +def test_async_ppo_rejects_non_vllm_backend(): + master_config = _build_async_ppo_master_config() + master_config.policy["generation"]["backend"] = "sglang" + with pytest.raises(AssertionError, match="async vLLM generation engine"): + _call_async_ppo_train(master_config) + + +def test_async_ppo_requires_async_engine(): + master_config = _build_async_ppo_master_config() + master_config.policy["generation"]["vllm_cfg"]["async_engine"] = False + with pytest.raises(AssertionError, match="async vLLM generation engine"): + _call_async_ppo_train(master_config) + + +def test_async_ppo_requires_importance_sampling_correction(): + master_config = _build_async_ppo_master_config() + master_config.loss_fn = ClippedPGLossConfig( + use_importance_sampling_correction=False + ) + with pytest.raises(AssertionError, match="Importance sampling correction"): + _call_async_ppo_train(master_config) + + +def test_async_ppo_rejects_colocated_inference(): + master_config = _build_async_ppo_master_config() + master_config.policy["generation"]["colocated"]["enabled"] = True + with pytest.raises(AssertionError, match="Colocated inference is not supported"): + _call_async_ppo_train(master_config) + + +def test_async_ppo_requires_positive_ppo_epochs(): + """ppo_epochs == 0 would leave train_results unset; guard rejects it.""" + master_config = _build_async_ppo_master_config() + master_config.ppo["ppo_epochs"] = 0 + with pytest.raises(AssertionError, match="ppo_epochs must be >= 1"): + _call_async_ppo_train(master_config) + + +def test_async_ppo_rejects_nemo_gym(): + """NeMo-Gym rollout is not wired for async PPO; guard fails loud at startup.""" + master_config = _build_async_ppo_master_config() + master_config.env["should_use_nemo_gym"] = True + # _should_use_nemo_gym also requires the http server to be exposed. + master_config.policy["generation"]["vllm_cfg"]["expose_http_server"] = True + with pytest.raises(AssertionError, match="NeMo-Gym rollout is not yet supported"): + _call_async_ppo_train(master_config) + + +# --------------------------------------------------------------------------- +# Warmup trajectory-age boundaries (the two-boundary fix for the frozen-actor +# critic-warmup phase). See async_ppo_train's _collector_lead_age / +# _sample_max_age and the pi_0 policy-version analysis. +# --------------------------------------------------------------------------- +def test_async_warmup_age_boundaries_no_elevation(): + """warmup_age == train_age (the default) => constant age at every step, so + behaviour is identical to plain async PPO regardless of the boundary math.""" + from nemo_rl.algorithms.ppo import ( + _async_warmup_collector_lead_age, + _async_warmup_sample_max_age, + ) + + W, train_age = 20, 1 + for s in range(0, W + 10): + assert _async_warmup_collector_lead_age(s, W, train_age, train_age) == train_age + assert _async_warmup_sample_max_age(s, W, train_age, train_age) == train_age + + +def test_async_warmup_collector_lead_age_boundary_at_W(): + """Collector generation-lead is elevated through step W (frozen actor), then + drops to the training age from W+1 so it regenerates against the trained policy.""" + from nemo_rl.algorithms.ppo import _async_warmup_collector_lead_age + + W, train_age, warmup_age = 20, 1, 8 + assert _async_warmup_collector_lead_age(W - 1, W, train_age, warmup_age) == warmup_age + assert _async_warmup_collector_lead_age(W, W, train_age, warmup_age) == warmup_age + assert _async_warmup_collector_lead_age(W + 1, W, train_age, warmup_age) == train_age + + +def test_async_warmup_sample_max_age_boundary_at_W_plus_train_age(): + """Driver eviction age stays elevated through W + train_age: a frozen (pi_0) + rollout is within train_age POLICY-steps of the actor there, so it is admitted + as valid lag-<=train_age data instead of being wrongly evicted (which would + deadlock, since the collector never regenerates that target).""" + from nemo_rl.algorithms.ppo import _async_warmup_sample_max_age + + W, warmup_age = 20, 8 + # train_age = 1: elevated through W+1 (policy-age 1), drops at W+2. + assert _async_warmup_sample_max_age(W, W, 1, warmup_age) == warmup_age + assert _async_warmup_sample_max_age(W + 1, W, 1, warmup_age) == warmup_age + assert _async_warmup_sample_max_age(W + 2, W, 1, warmup_age) == 1 + # train_age = 2: pi_0 stays within 2 policy-steps through W+2, drops at W+3. + assert _async_warmup_sample_max_age(W + 2, W, 2, warmup_age) == warmup_age + assert _async_warmup_sample_max_age(W + 3, W, 2, warmup_age) == 2 + + +def test_async_trajectory_policy_age_is_freeze_aware(): + """The gen-version age overcounts staleness during warmup; policy-age is the + true off-policy distance and must stay <= max_trajectory_age_steps. + + Reproduces the smoke-test boundary (W=3, A_w=5, A_t=1): the frozen pi_0 banked + at gen-version 0 is consumed at steps 0..4, then a fresh pi_1 (gen-version 4) at + step 5. Its gen-version age spikes to 4 at step 4, but its POLICY-age is only 1. + """ + from nemo_rl.algorithms.ppo import _async_trajectory_policy_age + + W, train_age = 3, 1 + # (weight_version s, gen_version g) actually consumed each step of the smoke run. + consumed = [(0, 0), (1, 0), (2, 0), (3, 0), (4, 0), (5, 4)] + expected_gen_age = [0, 1, 2, 3, 4, 1] + expected_policy_age = [0, 0, 0, 0, 1, 1] + for (s, g), gen_age, pol_age in zip( + consumed, expected_gen_age, expected_policy_age + ): + assert (s - g) == gen_age # what avg_trajectory_age reports + got = _async_trajectory_policy_age(g, s, W) + assert got == pol_age, (s, g, got, pol_age) + assert got <= train_age # the invariant that actually matters + + # No warmup (W == 0): policy-age reduces to the plain gen-version age. + for s in range(6): + for g in range(s + 1): + assert _async_trajectory_policy_age(g, s, 0) == s - g + + +def test_replay_buffer_admits_banked_frozen_rollout_at_boundary(): + """End-to-end buffer check of the boundary: a frozen rollout banked deep during + warmup (low gen-version) targeting step W+1 must be ADMITTED when sampled with + the elevated (warmup) age at step W+1, and EVICTED once the age snaps back. + + This is the exact scenario that deadlocked with a single W boundary: the + collector had already advanced its lead past target W+1, so an eviction there + is unrecoverable. + """ + from nemo_rl.algorithms.async_utils.replay_buffer import ReplayBufferImpl + + W, train_age, warmup_age = 20, 1, 8 + num_prompts = 2 + + def _fresh_buffer(): + buf = ReplayBufferImpl(max_size=64) + # Two groups for target W+1, generated during warmup at gen-version W+1-A_w + # (the frozen actor pi_0). gen-version age is warmup_age, policy-age is 1. + for _ in range(num_prompts): + buf.add({"batch": None}, W + 1 - warmup_age, W + 1) + return buf + + # At step W+1 with the elevated (warmup) age -> admitted (valid lag-1 pi_0). + buf = _fresh_buffer() + result = buf.sample( + num_prompt_groups=num_prompts, + current_weight_version=W + 1, + max_age_steps=warmup_age, + ) + assert result is not None + assert len(result["trajectories"]) == num_prompts + + # Same rollout at step W+2 with the training age -> gen-version W+1-A_w is now + # older than (W+2 - train_age), so it is correctly evicted (pi_0 genuinely stale). + buf = _fresh_buffer() + # relabel the (identical) banked groups to target W+2 to model the surplus that + # the driver would try to consume one step later. + buf.target_weight_versions = [W + 2, W + 2] + result = buf.sample( + num_prompt_groups=num_prompts, + current_weight_version=W + 2, + max_age_steps=train_age, + ) + assert result is None # evicted as stale; collector must regenerate fresh + + +# --------------------------------------------------------------------------- +# validate() rollout dispatch. Async PPO runs the vLLM async engine, whose +# generation worker has no sync `generate` method — validate() must take the +# async rollout path or it raises AttributeError ('ActorHandle' has no +# attribute 'generate') at the first validation step. +# --------------------------------------------------------------------------- +def _run_validate_with_mocked_rollouts(master_config): + """Drive validate() through one batch with both rollout fns mocked. + + Returns (async_rollout_mock, sync_rollout_mock) for call assertions. + """ + from unittest.mock import MagicMock, patch + + from nemo_rl.algorithms.ppo import validate + + # Fields validate() reads to run one batch and summarize. + master_config.ppo["max_val_samples"] = 1 + master_config.ppo["val_batch_size"] = 1 + master_config.ppo["max_rollout_turns"] = 1 + master_config.logger = {"num_val_samples_to_print": 0} + + rollout_return = ( + {"total_reward": torch.tensor([1.0]), "message_log": []}, + {"mean_gen_tokens_per_sample": 5.0}, + ) + with ( + patch( + "nemo_rl.algorithms.ppo.run_async_multi_turn_rollout", + return_value=rollout_return, + ) as async_mock, + patch( + "nemo_rl.algorithms.ppo.run_multi_turn_rollout", + return_value=rollout_return, + ) as sync_mock, + ): + validate( + policy_generation=MagicMock(), + val_dataloader=[MagicMock()], # one validation batch + tokenizer=MagicMock(), + val_task_to_env={}, + step=5, + master_config=master_config, + logger=MagicMock(), + ) + return async_mock, sync_mock + + +def test_validate_dispatches_to_async_rollout_for_async_engine(): + """Regression: async PPO (vLLM async_engine) validation must use the ASYNC + rollout path. The sync path calls policy_generation.generate(), which the + async worker lacks -> AttributeError at the first validation step.""" + master_config = _build_async_ppo_master_config() # vllm async_engine=True + async_mock, sync_mock = _run_validate_with_mocked_rollouts(master_config) + async_mock.assert_called_once() + sync_mock.assert_not_called() + + +def test_validate_dispatches_to_sync_rollout_when_not_async(): + """Non-async (colocated/sync) PPO validation uses the synchronous rollout.""" + master_config = _build_async_ppo_master_config() + master_config.policy["generation"]["vllm_cfg"]["async_engine"] = False + async_mock, sync_mock = _run_validate_with_mocked_rollouts(master_config) + sync_mock.assert_called_once() + async_mock.assert_not_called() + + +# --------------------------------------------------------------------------- +# Resume optimizer-path resolution. Megatron bundles the optimizer + LR +# scheduler inside weights/iter_*/*.distcp (no separate optimizer/ dir), so a +# resume must point optimizer_path at the weights or load_optim stays False and +# the LR-warmup scheduler restarts on every resume (V-shaped critic/lr). +# --------------------------------------------------------------------------- +def test_resolve_resume_optimizer_path(tmp_path): + from nemo_rl.algorithms.ppo import _resolve_resume_optimizer_path + + weights = tmp_path / "weights" + weights.mkdir() + optim = tmp_path / "optimizer" + missing_optim = tmp_path / "optimizer_absent" + megatron = {"megatron_cfg": {"enabled": True}} + dtensor = {"megatron_cfg": {"enabled": False}} + + # DTensor writes a separate optimizer/ dir -> use it when present. + optim.mkdir() + assert _resolve_resume_optimizer_path(optim, weights, dtensor) == optim + # A present separate dir wins even on megatron. + assert _resolve_resume_optimizer_path(optim, weights, megatron) == optim + + # Megatron: no separate optimizer/ dir (bundled in weights) -> weights path, + # so load_optim=True and the optimizer + scheduler actually resume. + assert _resolve_resume_optimizer_path(missing_optim, weights, megatron) == weights + # Megatron during warmup: policy weights not saved -> nothing to resume. + assert _resolve_resume_optimizer_path(missing_optim, None, megatron) is None + # DTensor without a separate optimizer/ dir -> None (don't misread weights as + # an optimizer path for a backend that stores it separately). + assert _resolve_resume_optimizer_path(missing_optim, weights, dtensor) is None diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index d20060f88b7..7f055de3ad2 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -49,6 +49,21 @@ ppo: target_max: 1.0 seq_logprob_error_threshold: null + async_ppo: + enabled: false # Set to true for async PPO (requires non-colocated vLLM generation) + # Max age (in training steps) for trajectories drawn from the replay buffer. + # Recommended/validated value is 1; higher values increase critic-staleness + # bias (warned, not forbidden). See AsyncPPOConfig in nemo_rl/algorithms/ppo.py. + max_trajectory_age_steps: 1 + # Max age used ONLY during critic warmup (step < policy_training_start_step), + # where the frozen actor makes any-age trajectories on-policy for free. null + # => same as max_trajectory_age_steps. Only helps throughput when warmup is + # generation-bound. See AsyncPPOConfig. + warmup_max_trajectory_age_steps: null + in_flight_weight_updates: false # Set to true to enable in-flight weight updates + recompute_kv_cache_after_weight_updates: false + log_every: 500 # heartbeat frequency for collector progress prints (0 silences) + loss_fn: disable_ppo_ratio: false reference_policy_kl_penalty: 0.0 # DAPO: no KL penalty From 4a8acb7b2d1124378442bbdb08a5e8e03c6ca88f Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Sat, 11 Jul 2026 13:45:28 -0700 Subject: [PATCH 25/69] fix(ppo): drop survivorship-biased incomplete replay-buffer targets on resume MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The partially-generated frontier target saved in the async replay buffer holds only the fast-completing (short) rollouts, so keeping and gap-filling it on resume trains the first post-resume step on a systematically shorter, higher-reward batch — a dip in mean_gen_tokens_per_sample and a reward spike after every resume. Add ppo.async_ppo.drop_incomplete_targets_on_restore (default false = historical gap-fill behavior). When true, incomplete restored targets are dropped and regenerated fresh (unbiased batch, at the cost of a one-target generation bubble at resume). Complete banked targets are always kept and replayed exactly. Signed-off-by: Jiaqi Zeng --- examples/configs/ppo_math_1B.yaml | 7 ++ .../algorithms/async_utils/replay_buffer.py | 73 ++++++++++++-- nemo_rl/algorithms/ppo.py | 20 ++++ tests/unit/algorithms/test_async_utils.py | 99 ++++++++++++++++++- .../ppo_math_1B_megatron.yaml | 3 + 5 files changed, 191 insertions(+), 11 deletions(-) diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index ff48399b2c2..e29c582bf47 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -62,6 +62,13 @@ ppo: warmup_max_trajectory_age_steps: null in_flight_weight_updates: true # Set to true to enable in-flight weight updates recompute_kv_cache_after_weight_updates: false + # On resume, how to treat the INCOMPLETE (partially-generated) frontier target + # in the restored replay buffer. false (default) keeps the partial batch and + # gap-fills the missing groups (fast resume, but survivorship-biased toward + # short rollouts => a dip in mean_gen_tokens_per_sample + reward spike after + # each resume). true drops it and regenerates fresh (unbiased, one-target + # bubble at resume). See AsyncPPOConfig in nemo_rl/algorithms/ppo.py. + drop_incomplete_targets_on_restore: false # Heartbeat frequency for the collector/replay-buffer per-rollout progress # prints (they flood the log at large num_prompts_per_step). Log every Nth # event (plus the final one per target); 0 silences them; unset => every event. diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 3a72273d899..ae1c3090d56 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -41,7 +41,12 @@ class ReplayBufferImpl(ReplayBufferProtocol): grpo.num_generations_per_prompt (required to compute per-prompt advantages). """ - def __init__(self, max_size: int, log_every: Optional[int] = None): + def __init__( + self, + max_size: int, + log_every: Optional[int] = None, + drop_incomplete_targets_on_restore: bool = False, + ): if max_size <= 0: raise ValueError(f"max_size must be positive, got {max_size}") self.max_size = max_size @@ -49,6 +54,18 @@ def __init__(self, max_size: int, log_every: Optional[int] = None): # (ppo.async_ppo.log_every). N>0 => every Nth event; 0 => silence; unset # (None) => every event (log_every_n == 1, no throttling). self._log_every_n = 1 if log_every is None else int(log_every) + # How to treat an INCOMPLETE (partially-generated) restored target step on + # resume — see _prepare_for_training_step for the full rationale. + # False (default): keep the partial survivors and let the collector + # gap-fill only the missing groups. This is the historical behavior. + # True: drop the incomplete target(s) and regenerate them fresh. Fixes a + # survivorship bias (the partial batch holds only the fast/short + # rollouts that finished before the checkpoint) at the cost of a + # one-target generation bubble per resume. + # Configured via YAML (ppo.async_ppo.drop_incomplete_targets_on_restore). + self._drop_incomplete_targets_on_restore = bool( + drop_incomplete_targets_on_restore + ) self.trajectories = [] # List[dict[str, Any]] # If trajectory_version is 1 and target_weight_version is 4 it means that weight version 1 was used for generating a trajectory and this trajectory will be used for training when weight version is 4. self.trajectory_versions = [] # it is the weight-version used for generation of a trajectory @@ -432,15 +449,57 @@ def _prepare_for_training_step( " Complete targets: " f"{sorted(complete_targets) if complete_targets else 'none'}" ) - for target in sorted(incomplete_targets): + + # How to handle an INCOMPLETE (partially-generated) restored target. This + # matters because a partial (frontier) target in the checkpoint is + # SURVIVORSHIP-BIASED toward SHORT responses: the buffer was saved while the + # collector was mid-generation, so only the fast-completing (shorter) + # rollouts made it in — the slow/long ones were still generating. The first + # step trained after each resume then sees a systematically shorter + + # higher-reward batch (short math answers score higher), biasing the + # critic/policy toward short generations. Complete banked targets are + # unbiased either way and are always kept/replayed exactly. + # + # Two behaviors, selected by drop_incomplete_targets_on_restore (default + # False, i.e. the historical behavior): + # * False -> KEEP the partial survivors; the collector gap-fills only the + # missing groups (fast resume, but reintroduces the bias above). + # * True -> DROP the incomplete target(s) so the collector regenerates + # them from scratch (unbiased batch, at the cost of a one-target + # generation bubble at resume). + if incomplete_targets and self._drop_incomplete_targets_on_restore: print( - f" Incomplete target {target}: " - f"{target_counts[target]}/{num_prompts_per_step}" + " Dropping survivorship-biased incomplete restored target(s) " + "(will regenerate fresh): " + + ", ".join( + f"{t}={target_counts[t]}/{num_prompts_per_step}" + for t in sorted(incomplete_targets) + ) + ) + keep = [ + i + for i, t in enumerate(self.target_weight_versions) + if t not in incomplete_targets + ] + self.trajectories = [self.trajectories[i] for i in keep] + self.trajectory_versions = [self.trajectory_versions[i] for i in keep] + self.target_weight_versions = [ + self.target_weight_versions[i] for i in keep + ] + elif incomplete_targets: + print( + " Keeping incomplete restored target(s) for gap-fill " + "(set ppo.async_ppo.drop_incomplete_targets_on_restore=true to " + "regenerate fresh instead): " + + ", ".join( + f"{t}={target_counts[t]}/{num_prompts_per_step}" + for t in sorted(incomplete_targets) + ) ) - # Let the collector ask each target from current_step onward how many - # trajectories are still needed, so incomplete restored batches can be - # gap-filled and complete batches can be skipped. + # Collector gap-fills each target from current_step onward: complete + # restored batches (and kept incomplete ones) are topped up only by their + # missing groups; dropped (incomplete) ones regenerate fully. self.last_target_weight_already_generated = current_step - 1 @staticmethod diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 4bf17860122..fe1ddb86200 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -145,6 +145,19 @@ class AsyncPPOConfig(TypedDict): in_flight_weight_updates: NotRequired[bool] # Recompute the KV cache after in-flight weight updates. recompute_kv_cache_after_weight_updates: NotRequired[bool] + # How to treat the INCOMPLETE (partially-generated) frontier target step found + # in a restored replay-buffer checkpoint on resume: + # * false (default): keep the partial survivors and let the collector gap-fill + # only the missing groups. Fast resume, but the partial batch is + # SURVIVORSHIP-BIASED toward SHORT rollouts (only the fast-completing ones + # were saved), so the first step after each resume trains on a shorter + + # higher-reward batch — visible as a dip in mean_gen_tokens_per_sample and a + # reward spike right after every resume. + # * true: drop the incomplete target and regenerate it fresh (unbiased batch), + # at the cost of a one-target generation bubble at resume. Recommended when + # resuming frequently. Complete banked targets are unbiased and always kept. + # Absent/null => false. Read with a None-safe default, so allow null. + drop_incomplete_targets_on_restore: NotRequired[bool | None] # Heartbeat frequency for the collector/replay-buffer per-rollout progress # prints: log every Nth event (plus the final one per target). These fire once # per prompt group and flood the log at large num_prompts_per_step, so throttle @@ -2328,6 +2341,13 @@ def _sample_max_age(s: int) -> int: replay_buffer = ReplayBuffer.options(runtime_env=_replay_runtime_env).remote( max_size=optimal_buffer_size, log_every=async_cfg.get("log_every"), + # Default False => historical gap-fill behavior; True drops the + # survivorship-biased incomplete frontier target on resume (regenerate + # fresh). None (absent) is coerced to False by the buffer constructor. + drop_incomplete_targets_on_restore=async_cfg.get( + "drop_incomplete_targets_on_restore" + ) + or False, ) last_checkpoint_path = checkpointer.get_latest_checkpoint_path() diff --git a/tests/unit/algorithms/test_async_utils.py b/tests/unit/algorithms/test_async_utils.py index d7bdc617159..9a3972e6b22 100644 --- a/tests/unit/algorithms/test_async_utils.py +++ b/tests/unit/algorithms/test_async_utils.py @@ -125,7 +125,13 @@ def _state( "max_size": max_size, } - def test_local_restore_prepares_current_step_for_gap_fill(self): + def test_local_restore_default_keeps_incomplete_targets_for_gap_fill(self): + # DEFAULT behavior (drop_incomplete_targets_on_restore=False): the restored + # buffer holds past target 1 (dropped), COMPLETE target 2 (count 2 == + # num_prompts_per_step), and an INCOMPLETE frontier target 3 (count 1 < 2) — + # the target the collector was mid-generating when the checkpoint was + # written. By default that partial target is KEPT and the collector gap-fills + # only the missing group (historical behavior). buffer = ReplayBufferImpl(max_size=10) state = self._state( trajectory_versions=[0, 1, 1, 2], @@ -139,14 +145,48 @@ def test_local_restore_prepares_current_step_for_gap_fill(self): current_training_step=2, ) - assert buffer.get_debug_info()["trajectory_versions"] == [1, 1, 2] + # target 1 dropped (past); complete target 2 and the partial target 3 kept. assert buffer.get_debug_info()["target_weight_versions"] == [2, 2, 3] + assert buffer.get_debug_info()["trajectory_versions"] == [1, 1, 2] assert buffer.get_last_target_weight_already_generated() == 1 assert buffer.has_complete_batch(2, 2) assert not buffer.has_complete_batch(3, 2) assert buffer.get_trajectories_needed(2, 2) == 0 + # kept -> only the missing group is gap-filled, not the whole batch. assert buffer.get_trajectories_needed(3, 2) == 1 + def test_local_restore_drops_incomplete_targets_for_fresh_regen(self): + # With drop_incomplete_targets_on_restore=True: the INCOMPLETE frontier + # target 3 (count 1 < 2) is survivorship-biased toward SHORT rollouts (short + # responses finish first), so on resume it must be DROPPED and regenerated + # fresh, NOT gap-filled from the biased subset (which would train the first + # post-resume step on a systematically shorter + higher-reward batch). + buffer = ReplayBufferImpl( + max_size=10, drop_incomplete_targets_on_restore=True + ) + state = self._state( + trajectory_versions=[0, 1, 1, 2], + target_weight_versions=[1, 2, 2, 3], + last_target_weight_already_generated=3, + ) + + buffer.load_state_dict( + state, + num_prompts_per_step=2, + current_training_step=2, + ) + + # target 1 dropped (past); incomplete target 3 dropped (regenerate fresh); + # complete target 2 kept and replayed exactly. + assert buffer.get_debug_info()["target_weight_versions"] == [2, 2] + assert buffer.get_debug_info()["trajectory_versions"] == [1, 1] + assert buffer.get_last_target_weight_already_generated() == 1 + assert buffer.has_complete_batch(2, 2) + assert not buffer.has_complete_batch(3, 2) + assert buffer.get_trajectories_needed(2, 2) == 0 + # dropped -> fully regenerated (needs the whole batch), NOT gap-filled (1) + assert buffer.get_trajectories_needed(3, 2) == 2 + def test_local_restore_empty_state_resets_generation_watermark(self): buffer = ReplayBufferImpl(max_size=10) state = self._state( @@ -820,8 +860,12 @@ def test_replay_buffer_load_state_dict_inconsistent_lengths(self): ray.kill(buffer) - def test_replay_buffer_restore_for_training_step_gap_fill_accounting(self): - """Test resume cleanup keeps incomplete future targets for gap filling.""" + def test_replay_buffer_restore_default_keeps_incomplete_targets(self): + """DEFAULT resume cleanup keeps incomplete future targets for gap filling. + + With drop_incomplete_targets_on_restore unset (False), the partial frontier + target is kept and the collector gap-fills only the missing group. + """ buffer = ReplayBuffer.remote(max_size=10) state = { @@ -846,16 +890,63 @@ def test_replay_buffer_restore_for_training_step_gap_fill_accounting(self): ) debug_info = ray.get(buffer.get_debug_info.remote()) + # complete target 2 and partial target 3 kept; past target 1 dropped. assert debug_info["trajectory_versions"] == [1, 1, 2] assert debug_info["target_weight_versions"] == [2, 2, 3] assert ray.get(buffer.has_complete_batch.remote(2, 2)) assert not ray.get(buffer.has_complete_batch.remote(3, 2)) assert ray.get(buffer.get_trajectories_needed.remote(2, 2)) == 0 + # kept -> only the missing group is gap-filled. assert ray.get(buffer.get_trajectories_needed.remote(3, 2)) == 1 assert ray.get(buffer.get_last_target_weight_already_generated.remote()) == 1 ray.kill(buffer) + def test_replay_buffer_restore_drops_incomplete_targets_when_enabled(self): + """With drop_incomplete_targets_on_restore=True, incomplete targets drop. + + The incomplete frontier target the collector was mid-generating at save + time is survivorship-biased toward short rollouts, so it is dropped and + regenerated rather than gap-filled from the biased subset. + """ + buffer = ReplayBuffer.remote( + max_size=10, drop_incomplete_targets_on_restore=True + ) + + state = { + "trajectories": [ + {"batch": {"data": "past"}}, + {"batch": {"data": "step2_a"}}, + {"batch": {"data": "step2_b"}}, + {"batch": {"data": "step3_a"}}, + ], + "trajectory_versions": [0, 1, 1, 2], + "target_weight_versions": [1, 2, 2, 3], + "last_target_weight_already_generated": 3, + "max_size": 10, + } + + ray.get( + buffer.load_state_dict.remote( + state, + num_prompts_per_step=2, + current_training_step=2, + ) + ) + + debug_info = ray.get(buffer.get_debug_info.remote()) + # complete target 2 kept; past target 1 and incomplete target 3 dropped. + assert debug_info["trajectory_versions"] == [1, 1] + assert debug_info["target_weight_versions"] == [2, 2] + assert ray.get(buffer.has_complete_batch.remote(2, 2)) + assert not ray.get(buffer.has_complete_batch.remote(3, 2)) + assert ray.get(buffer.get_trajectories_needed.remote(2, 2)) == 0 + # dropped -> fully regenerated, not gap-filled (would be 1) + assert ray.get(buffer.get_trajectories_needed.remote(3, 2)) == 2 + assert ray.get(buffer.get_last_target_weight_already_generated.remote()) == 1 + + ray.kill(buffer) + def test_replay_buffer_remove_incomplete_resets_watermark_before_first_remaining_target( self, ): diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 7f055de3ad2..d596091bac7 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -62,6 +62,9 @@ ppo: warmup_max_trajectory_age_steps: null in_flight_weight_updates: false # Set to true to enable in-flight weight updates recompute_kv_cache_after_weight_updates: false + # false (default) keeps the partial frontier target on resume and gap-fills it + # (survivorship-biased toward short rollouts); true drops + regenerates it fresh. + drop_incomplete_targets_on_restore: false log_every: 500 # heartbeat frequency for collector progress prints (0 silences) loss_fn: From c41b70227f8d039f427a1c7755b8eb46aa1e6832 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Mon, 13 Jul 2026 14:25:40 -0700 Subject: [PATCH 26/69] feat(ppo): NeMo-Gym rollout support (sync + async) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Finish the partial NeMo-Gym port for PPO so it reaches GRPO parity: - setup(): spin up the NeMo-Gym actor via the vLLM deferred-load path, bundling policy+value+vLLM as one GPU task while the gym overlaps (CPU/HTTP); add nemo_gym_actor to the return tuple. - ppo_train(): complete the gym rollout call (max_seq_len, effort_config, reward_penalty_config, thinking_tags, stop_token_ids/stop_strings=None). - validate(): add the gym branch (3-way dispatch) — closes the NemoGym.step() NotImplementedError crash gap. - async_ppo_train(): drop the gym-off guard; gym rollouts run inside the shared AsyncTrajectoryCollector, which now reads master_config.reward_penalties. - reward_penalties (reward-zeroing, pre-GAE) are wired functional and GAE-compatible; mask_sample is honored in both training loops (mirrors GRPO); the advantage-overwrite penalties (invalid_tool_call_advantage / malformed_thinking_advantage) are explicitly rejected for PPO (they'd break GAE advantage/return consistency). - Harden the collector gym branch to clear stop_token_ids/stop_strings by construction (avoids a swallowed assert -> silent buffer stall). Add examples/nemo_gym/run_ppo_nemo_gym.py + ppo_math_rlvr_nemo_gym.yaml (Math/RLVR), unit coverage for the validate gym dispatch, and functional tests ppo_nemo_gym.sh (sync) + ppo_async_gym.sh (async). Both smoke tests pass on GB200 (DAPO17k, Qwen2.5-0.5B) end to end. Signed-off-by: Jiaqi Zeng --- examples/nemo_gym/ppo_math_rlvr_nemo_gym.yaml | 83 +++++ examples/nemo_gym/run_ppo_nemo_gym.py | 236 +++++++++++++ examples/run_ppo.py | 1 + .../async_utils/trajectory_collector.py | 12 +- nemo_rl/algorithms/ppo.py | 313 ++++++++++++++++-- tests/functional/L1_Functional_Tests_Gym.sh | 2 + tests/functional/ppo_async_gym.sh | 102 ++++++ tests/functional/ppo_nemo_gym.sh | 99 ++++++ tests/unit/algorithms/test_ppo.py | 53 ++- 9 files changed, 862 insertions(+), 39 deletions(-) create mode 100644 examples/nemo_gym/ppo_math_rlvr_nemo_gym.yaml create mode 100644 examples/nemo_gym/run_ppo_nemo_gym.py create mode 100755 tests/functional/ppo_async_gym.sh create mode 100755 tests/functional/ppo_nemo_gym.sh diff --git a/examples/nemo_gym/ppo_math_rlvr_nemo_gym.yaml b/examples/nemo_gym/ppo_math_rlvr_nemo_gym.yaml new file mode 100644 index 00000000000..c1642318fd0 --- /dev/null +++ b/examples/nemo_gym/ppo_math_rlvr_nemo_gym.yaml @@ -0,0 +1,83 @@ +# PPO + NeMo-Gym (Math / RLVR) +# +# Reuses the PPO-DAPO training setup from examples/configs/ppo_math_1B.yaml +# (ppo / value / loss_fn / value_loss_fn / policy blocks) and swaps in a +# NeMo-Gym math_with_judge (RLVR) rollout environment. Run with: +# +# uv run examples/nemo_gym/run_ppo_nemo_gym.py \ +# --config examples/nemo_gym/ppo_math_rlvr_nemo_gym.yaml +# +# NOTE: the `env.nemo_gym.config_paths` and `data.{train,validation}.data_path` +# entries are specific to your NeMo-Gym install / prepared data — edit them to +# match your environment (see examples/nemo_gym/grpo_nanov3.yaml). +defaults: ../configs/ppo_math_1B.yaml + +policy: + generation: + # NeMo-Gym drives an HTTP-exposed async vLLM server (also force-set by + # setup_nemo_gym_config, but kept explicit here for clarity). + vllm_cfg: + async_engine: true + expose_http_server: true + +ppo: + # NeMo-Gym principle: the validation set is used verbatim (no hidden pre/post + # processing). run_ppo_nemo_gym.py rejects a preset max_val_samples and sets + # both to len(val_dataset), so they must start null. + max_val_samples: null + val_batch_size: null + # Use NeMo-Gym RLVR rewards as-is: turn OFF the DAPO reward rescaling and + # length shaping inherited from ppo_math_1B (every NeMo-Gym recipe disables + # both). They are applied unconditionally after the rollout, so leaving them on + # would silently remap / length-penalize the gym reward. reward_penalties + # (reward-zeroing) handle response quality instead. + reward_scaling: + enabled: false + reward_shaping: + enabled: false + # For async PPO + NeMo-Gym: set async_ppo.enabled=true AND + # policy.generation.colocated.enabled=false (async requires non-colocated + # generation). reward_scaling / reward_shaping / use_dynamic_sampling must stay + # off (the async entry path rejects them) — they already are, above. + async_ppo: + enabled: false + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + +# Fully replace the parent's math (native-env) data block with the NeMo-Gym one. +data: + _override_: true + max_input_seq_length: null + shuffle: false + num_workers: 1 + use_multiple_dataloader: false + train: + data_path: "/path/to/train.jsonl" # EDIT: ng_prepare_data output + validation: + data_path: "/path/to/validation.jsonl" # EDIT: ng_prepare_data output + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# Fully replace the parent's native math env with the NeMo-Gym env. +env: + _override_: true + should_use_nemo_gym: true + # true: skip expensive train_data_step*.jsonl; false: write full jsonl. + should_log_nemo_gym_responses: true + nemo_gym: # Passed into NeMo-Gym as the initial_global_config_dict. + # Port range for Gym HTTP servers (kept below the OS ephemeral range and + # non-overlapping with NeMo-RL / vLLM ports — see ray.sub). + port_range_low: 15001 + port_range_high: 20000 + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml # Required; must be *_for_training + # Base math_with_judge RLVR env (math-verify). It defines the + # math_with_judge_simple_agent that ng-prepared math records (incl. DAPO17k) + # reference via agent_ref, and already sets judge_model_server.name=policy_model + # + should_use_judge=false. Swap for another resources_servers/*/configs/*.yaml + # to target a different env. + - resources_servers/math_with_judge/configs/math_with_judge.yaml diff --git a/examples/nemo_gym/run_ppo_nemo_gym.py b/examples/nemo_gym/run_ppo_nemo_gym.py new file mode 100644 index 00000000000..be8e0afc036 --- /dev/null +++ b/examples/nemo_gym/run_ppo_nemo_gym.py @@ -0,0 +1,236 @@ +# Copyright (c) 2025, 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. + +import argparse +import os +import pprint +import time + +# Increase the W&B single object size warning threshold. Initially 100_000 (100 KB) -> 10_000_000 (10 MB) +import wandb.util + +wandb.util.VALUE_BYTES_LIMIT = 10_000_000 + +from omegaconf import OmegaConf + +from nemo_rl.algorithms.grpo import _should_use_nemo_gym +from nemo_rl.algorithms.ppo import ( + MasterConfig, + async_ppo_train, + ppo_train, + setup, +) +from nemo_rl.algorithms.utils import get_tokenizer +from nemo_rl.data.utils import setup_response_data +from nemo_rl.distributed.virtual_cluster import init_ray +from nemo_rl.environments.nemo_gym import setup_nemo_gym_config +from nemo_rl.models.generation import configure_generation_config +from nemo_rl.utils.config import ( + load_config, + parse_hydra_overrides, + register_omegaconf_resolvers, +) +from nemo_rl.utils.logger import get_next_experiment_dir, log_container_init_timing +from nemo_rl.utils.timer import Timer + + +def parse_args() -> tuple[argparse.Namespace, list[str]]: + """Parse command line arguments.""" + parser = argparse.ArgumentParser(description="Run PPO training with NeMo-Gym") + parser.add_argument( + "--config", type=str, default=None, help="Path to YAML config file" + ) + + # Parse known args for the script + args, overrides = parser.parse_known_args() + + return args, overrides + + +def main() -> None: + """Main entry point.""" + main_start = time.perf_counter() + log_container_init_timing() + rl_init_timer = Timer(context={"worker": "driver"}) + + register_omegaconf_resolvers() + args, overrides = parse_args() + + if not args.config: + args.config = os.path.join( + os.path.dirname(__file__), + "ppo_math_rlvr_nemo_gym.yaml", + ) + + with rl_init_timer.time("config"): + config = load_config(args.config) + print(f"Loaded configuration from: {args.config}") + + if overrides: + print(f"Overrides: {overrides}") + config = parse_hydra_overrides(config, overrides) + + config = OmegaConf.to_container(config, resolve=True) + config = MasterConfig(**config) + print("Applied CLI overrides") + + # Get the next experiment directory with incremented ID + config.logger["log_dir"] = get_next_experiment_dir(config.logger["log_dir"]) + print(f"📊 Using log directory: {config.logger['log_dir']}") + if config.checkpointing["enabled"]: + print( + f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}" + ) + + with rl_init_timer.time("tokenizer"): + tokenizer = get_tokenizer(config.policy["tokenizer"]) + assert config.policy["generation"] is not None, ( + "A generation config is required for PPO" + ) + config.policy["generation"] = configure_generation_config( + config.policy["generation"], tokenizer + ) + + # NeMo-Gym specific config setup (forces async_engine / expose_http_server, + # nulls stop_strings / stop_token_ids). + setup_nemo_gym_config(config, tokenizer) + + # Assert here since this is right after the final config has been materialized. + assert _should_use_nemo_gym(config) + + # NeMo-Gym env needs dp_openai_server_base_urls from policy_generation, so the + # gym actor is created inside setup(); we don't build the env here. + with rl_init_timer.time("data"): + print("\n▶ Setting up data...") + train_dataset, val_dataset = setup_response_data( + tokenizer, config.data, env_configs=None + ) + + # Validation dataset config setup. Gym principle: what you pass in is used + # verbatim (no hidden pre/post processing), so a preset max_val_samples is + # rejected and the full val set is used as one batch. + if config.ppo["max_val_samples"] is not None: + raise ValueError( + """A non-null `ppo.max_val_samples` parameter is not supported. + +Gym principle is that there is no hidden data pre or post processing from you. What you see is what you get. + +The validation set you pass in will directly be used for validation with no additional preprocessing. If you want to have some number of repetitions, please include that in your dataset, via ``num_repeats``, in your dataset config and `ng_prepare_data` will prepare it accordingly.""" + ) + + if val_dataset is not None: + print( + f"Setting `ppo.max_val_samples` and `ppo.val_batch_size` to the length of the validation dataset, which is {len(val_dataset)}" + ) + config.ppo["max_val_samples"] = len(val_dataset) + config.ppo["val_batch_size"] = config.ppo["max_val_samples"] + + # Print config + print("Final config:") + pprint.pprint(config) + + with rl_init_timer.time("ray_connect"): + init_ray() + + with rl_init_timer.time("setup"): + ( + policy, + policy_generation, + nemo_gym, + value_model, + cluster, + dataloader, + val_dataloader, + loss_fn, + value_loss_fn, + logger, + checkpointer, + ppo_state, + master_config, + ) = setup(config, tokenizer, train_dataset, val_dataset) + + rl_init_timer.record("total", time.perf_counter() - main_start) + rl_init_metrics = rl_init_timer.get_timing_metrics(reduction_op="sum") + print("\n" + "=" * 60) + print(" " * 14 + "RL INIT TIMING BREAKDOWN") + for label, value in sorted(rl_init_metrics.items()): + if isinstance(value, (int, float)): + print(f" {label}: {value:.1f}s") + print("=" * 60 + "\n", flush=True) + + # NeMo-Gym is spun up inside setup() (overlapped with vLLM model load). + # Bind task_to_env / val_task_to_env for the nemo_gym env. Hardcode here to + # match `run_async_nemo_gym_rollout`. + task_to_env = {"nemo_gym": nemo_gym} + val_task_to_env = task_to_env + + if "async_ppo" in config.ppo and config.ppo["async_ppo"]["enabled"]: + # Async PPO does not support dynamic sampling / reward scaling / reward shaping. + unsupported_features = [ + "use_dynamic_sampling", + "reward_scaling", + "reward_shaping", + ] + for feature in unsupported_features: + if feature not in config.ppo: + continue + if feature == "use_dynamic_sampling": + if config.ppo[feature]: + raise NotImplementedError( + f"{feature} is not supported with async PPO" + ) + elif config.ppo[feature]["enabled"]: + raise NotImplementedError(f"{feature} is not supported with async PPO") + + print("🚀 Running async PPO training with NeMo-Gym") + async_config = config.ppo["async_ppo"] + async_ppo_train( + policy, + policy_generation, + value_model, + dataloader, + val_dataloader, + tokenizer, + loss_fn, + value_loss_fn, + task_to_env, + val_task_to_env, + logger, + checkpointer, + ppo_state, + master_config, + max_trajectory_age_steps=async_config["max_trajectory_age_steps"], + ) + else: + print("🚀 Running synchronous PPO training with NeMo-Gym") + ppo_train( + policy, + policy_generation, + value_model, + dataloader, + val_dataloader, + tokenizer, + loss_fn, + value_loss_fn, + task_to_env, + val_task_to_env, + logger, + checkpointer, + ppo_state, + master_config, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/run_ppo.py b/examples/run_ppo.py index 5977df7b837..b0d9dde507b 100644 --- a/examples/run_ppo.py +++ b/examples/run_ppo.py @@ -131,6 +131,7 @@ def main() -> None: ( policy, policy_generation, + _nemo_gym, # None on this path (NeMo-Gym has its own entry: run_ppo_nemo_gym.py) value_model, cluster, dataloader, diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index 85088115793..4fc1b03187a 100644 --- a/nemo_rl/algorithms/async_utils/trajectory_collector.py +++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py @@ -818,7 +818,17 @@ def _run_prompt_group_worker( # Async engine supports concurrent generation; avoid locking # Check if we should use nemo_gym (similar to synchronous GRPO) if _should_use_nemo_gym(self.master_config): - generation_config = self.master_config.policy["generation"] + # NeMo-Gym manages its own stop criteria; run_async_nemo_gym_rollout + # asserts stop_token_ids/stop_strings are unset. setup_nemo_gym_config + # nulls them globally, but clear them here too (on a copy) so this + # collector path is safe by construction, not just by convention — + # a stray auto-filled stop token would otherwise trip the assert + # inside this worker thread and silently stall the buffer. + generation_config = { + **self.master_config.policy["generation"], + "stop_token_ids": None, + "stop_strings": None, + } nemo_gym_rollout_result = run_async_nemo_gym_rollout( policy_generation=self.policy_generation, input_batch=repeated_batch, diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index fe1ddb86200..a7cf6b39cbe 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -22,7 +22,8 @@ import numpy as np import ray import torch -from pydantic import BaseModel +from pydantic import BaseModel, Field +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy from torchdata.stateful_dataloader import StatefulDataLoader from transformers import AutoProcessor from transformers.tokenization_utils_base import PreTrainedTokenizerBase @@ -32,7 +33,10 @@ RawRewardAdvantageEstimator, ) from nemo_rl.algorithms.grpo import ( + RewardPenaltyConfig, RewardScalingConfig, + _get_effort_config, + _raise_if_reward_penalties_enabled_without_nemo_gym, _should_log_nemo_gym_responses, _should_use_async_rollouts, _should_use_nemo_gym, @@ -77,7 +81,14 @@ prepare_segment_topology, ) from nemo_rl.environments.interfaces import EnvironmentInterface +from nemo_rl.environments.nemo_gym import ( + NemoGym, + NemoGymConfig, + get_nemo_gym_uv_cache_dir, + get_nemo_gym_venv_dir, +) from nemo_rl.experience.rollouts import ( + get_nemo_gym_thinking_tags, run_async_multi_turn_rollout, run_async_nemo_gym_rollout, run_multi_turn_rollout, @@ -188,12 +199,12 @@ class PPOConfig(TypedDict): max_num_steps: int max_rollout_turns: int val_period: int - val_batch_size: int + val_batch_size: int | None # None for NeMo-Gym compatibility val_at_start: bool # Whether to run validation on the last training step. Setting this to True ensures the # final checkpoint has validation metrics, which is required for get_best_checkpoint_path(). val_at_end: bool - max_val_samples: int + max_val_samples: int | None # None for NeMo-Gym compatibility skip_reference_policy_logprobs_calculation: NotRequired[bool] seed: int overlong_filtering: bool @@ -265,6 +276,12 @@ class MasterConfig(BaseModel, extra="allow"): logger: PPOLoggerConfig cluster: ClusterConfig checkpointing: CheckpointingConfig + # Reward-zeroing penalties applied to NeMo-Gym rollout results (see + # RewardPenaltyConfig). They set the reward to 0.0 for pathological responses + # BEFORE advantage/value computation, so they flow through GAE naturally. + # Only usable on the NeMo-Gym path (enforced by + # _raise_if_reward_penalties_enabled_without_nemo_gym). Defaults to all-off. + reward_penalties: RewardPenaltyConfig = Field(default_factory=RewardPenaltyConfig) # =============================================================================== @@ -311,6 +328,7 @@ def setup( ) -> tuple[ ColocatablePolicyInterface, Optional[GenerationInterface], + Optional[NemoGym], ValueInterface, tuple[RayVirtualCluster, RayVirtualCluster], StatefulDataLoader, @@ -325,9 +343,10 @@ def setup( """Main entry point for running PPO algorithm. Returns: - tuple of (policy, policy_generation, value_model, clusters, - dataloader, val_dataloader, loss_fn, value_loss_fn, logger, - checkpointer, ppo_save_state, master_config). + tuple of (policy, policy_generation, nemo_gym_actor, value_model, + clusters, dataloader, val_dataloader, loss_fn, value_loss_fn, logger, + checkpointer, ppo_save_state, master_config). ``nemo_gym_actor`` is None + unless NeMo-Gym is enabled (env.should_use_nemo_gym). """ # Start timing the entire setup process setup_start_time = time.perf_counter() @@ -854,10 +873,99 @@ def initialize_generation_with_policy( return policy_generation, policy, value_model + # ------------------------------------------------------------------------- + # NeMo-Gym gating. Initialized inside setup() (rather than by the caller) so + # its CPU/HTTP spinup can overlap with vLLM model loading via deferred model + # load (see the vLLM branch below). vLLM-only: gym requires an HTTP-exposed + # async generation server, which SGLang/native paths don't provide here. + # ------------------------------------------------------------------------- + enable_nemo_gym = _should_use_nemo_gym(master_config) + _raise_if_reward_penalties_enabled_without_nemo_gym( + master_config, enable_nemo_gym=enable_nemo_gym + ) + # Message-level advantage-OVERWRITE penalties (invalid_tool_call_advantage / + # malformed_thinking_advantage) set advantage spans directly post-hoc. That is + # GRPO-specific: in PPO it would overwrite GAE advantages and break the + # advantage<->return consistency GAE relies on. Deferred; the reward-level + # reward_penalties (reward-zeroing, pre-GAE) are the PPO-compatible path. + for _unsupported in ("invalid_tool_call_advantage", "malformed_thinking_advantage"): + if master_config.ppo.get(_unsupported) is not None: + raise NotImplementedError( + f"ppo.{_unsupported} (message-level advantage-overwrite penalty) is " + "not supported for PPO: it overwrites GAE advantages and breaks the " + "advantage/return consistency GAE relies on. Use reward_penalties " + "(reward-zeroing, applied pre-GAE) instead." + ) + if enable_nemo_gym: + nemo_gym_num_nodes = master_config.env.get("nemo_gym", {}).get( + "num_gpu_nodes", 0 + ) + ray_cur_node_id = ray.get_runtime_context().get_node_id() + else: + nemo_gym_num_nodes = 0 + ray_cur_node_id = None + nemo_gym_actor = None + + def _spinup_nemo_gym(base_urls, model_name): + """Spin up the NeMo-Gym actor against the given vLLM server URLs.""" + t0 = time.perf_counter() + nemo_gym_py_exec = get_actor_python_env( + "nemo_rl.environments.nemo_gym.NemoGym" + ) + if nemo_gym_py_exec.startswith("uv"): + nemo_gym_py_exec = create_local_venv_on_each_node( + nemo_gym_py_exec, "nemo_rl.environments.nemo_gym.NemoGym" + ) + nemo_gym_dict = dict(master_config.env["nemo_gym"]) + # NeMo-RL-side detection knobs are top-level NemoGymConfig fields (where + # the detector reads them), not part of Gym's global config. + invalid_tool_call_patterns = nemo_gym_dict.pop( + "invalid_tool_call_patterns", None + ) + thinking_tags = nemo_gym_dict.pop("thinking_tags", None) + # Reuse image-baked cache + venv dirs so the gym doesn't rebuild them. + uv_cache_dir = get_nemo_gym_uv_cache_dir() + if uv_cache_dir is not None: + nemo_gym_dict.setdefault("uv_cache_dir", uv_cache_dir) + uv_venv_dir = get_nemo_gym_venv_dir() + if uv_venv_dir is not None: + nemo_gym_dict.setdefault("uv_venv_dir", uv_venv_dir) + nemo_gym_cfg = NemoGymConfig( + model_name=model_name, + base_urls=base_urls, + invalid_tool_call_patterns=invalid_tool_call_patterns, + thinking_tags=thinking_tags, + # PPO has no Megatron generation backend, so router-replay / routed- + # experts preservation is N/A (this gym path is vLLM-only). + require_routed_experts=False, + initial_global_config_dict=nemo_gym_dict, + ) + nemo_gym_opts = {} + if nemo_gym_num_nodes: + nemo_gym_opts["scheduling_strategy"] = NodeAffinitySchedulingStrategy( + node_id=ray_cur_node_id, + soft=True, + ) + nemo_gym_opts["runtime_env"] = { + "py_executable": nemo_gym_py_exec, + "env_vars": { + **os.environ, + "VIRTUAL_ENV": nemo_gym_py_exec, + "UV_PROJECT_ENVIRONMENT": nemo_gym_py_exec, + }, + } + actor = NemoGym.options(**nemo_gym_opts).remote(nemo_gym_cfg) + ray.get(actor._spinup.remote()) + return actor, time.perf_counter() - t0 + assert backend in ("vllm", "sglang"), ( f"PPO requires vllm or sglang generation backend; got {backend!r}. " "The megatron generation backend is not supported." ) + assert not (enable_nemo_gym and backend != "vllm"), ( + "NeMo-Gym requires the vLLM generation backend (HTTP-exposed async " + f"server); got backend={backend!r}." + ) assert colocated_inference or backend == "vllm", ( "Non-colocated PPO currently requires the vLLM generation backend. " "SGLangGeneration.init_collective() is a no-op, so the training side's " @@ -895,13 +1003,86 @@ def initialize_generation_with_policy( "hf_config_overrides", {} ) - policy_generation, policy, value_model = initialize_generation_with_policy( - init_generation_fn=init_vllm, - generation_name="vLLM", - init_time_key="vllm_init_time_s", - colocated_inference=colocated_inference, - worker_init_timing_metrics=worker_init_timing_metrics, - ) + if enable_nemo_gym: + # Reserve vLLM ports up-front so we can hand the server URLs to + # NeMo-Gym and spin it up (CPU/HTTP) WHILE vLLM loads weights and the + # policy+value workers initialize. + print( + " ⚡ Deferred vLLM load: reserving ports for overlapped NeMo-Gym init", + flush=True, + ) + deferred_vllm = VllmGeneration( + cluster=inference_cluster, + config=generation_config, + defer_model_load=True, + ) + print( + f" ✓ Reserved {len(deferred_vllm.dp_openai_server_base_urls)} vLLM " + f"server URLs: {deferred_vllm.dp_openai_server_base_urls}", + flush=True, + ) + + def init_vllm_deferred(): + """Complete the deferred vLLM model load started above.""" + t0 = time.perf_counter() + deferred_vllm.load_and_start() + deferred_vllm.finish_generation() + return deferred_vllm, time.perf_counter() - t0 + + def init_nemo_gym(): + return _spinup_nemo_gym( + deferred_vllm.dp_openai_server_base_urls, + generation_config["model_name"], + ) + + # Colocated: vLLM + policy + value share GPUs, so run them as one + # sequential GPU task; non-colocated: vLLM (inference_cluster) and + # policy+value (train_cluster) run in parallel. NeMo-Gym overlaps + # either way (CPU/HTTP, no training-GPU contention). + init_tasks = {} + if colocated_inference: + + def init_vllm_then_policy_value(): + pg, vllm_t = init_vllm_deferred() + p, p_t, v, v_t = init_policy_and_value() + return pg, vllm_t, p, p_t, v, v_t + + init_tasks["vllm_policy_value"] = init_vllm_then_policy_value + else: + init_tasks["vllm"] = init_vllm_deferred + init_tasks["policy_value"] = init_policy_and_value + init_tasks["nemo_gym"] = init_nemo_gym + + print(f" ⚡ Init tasks: {', '.join(init_tasks.keys())}", flush=True) + with ThreadPoolExecutor(max_workers=len(init_tasks)) as executor: + submitted = {k: executor.submit(fn) for k, fn in init_tasks.items()} + results = {k: f.result() for k, f in submitted.items()} + + if colocated_inference: + ( + policy_generation, + vllm_time, + policy, + policy_time, + value_model, + value_time, + ) = results["vllm_policy_value"] + else: + policy_generation, vllm_time = results["vllm"] + policy, policy_time, value_model, value_time = results["policy_value"] + nemo_gym_actor, nemo_gym_time = results["nemo_gym"] + worker_init_timing_metrics["vllm_init_time_s"] = vllm_time + worker_init_timing_metrics["policy_init_time_s"] = policy_time + worker_init_timing_metrics["value_init_time_s"] = value_time + worker_init_timing_metrics["nemo_gym_init_time_s"] = nemo_gym_time + else: + policy_generation, policy, value_model = initialize_generation_with_policy( + init_generation_fn=init_vllm, + generation_name="vLLM", + init_time_key="vllm_init_time_s", + colocated_inference=colocated_inference, + worker_init_timing_metrics=worker_init_timing_metrics, + ) print( f" ✓ Using vLLM backend for generation with {policy_config['model_name']}", @@ -1001,6 +1182,7 @@ def initialize_generation_with_policy( return ( policy, policy_generation, + nemo_gym_actor, value_model, (train_cluster, inference_cluster), dataloader, @@ -1416,18 +1598,32 @@ def ppo_train( policy_generation.clear_logger_metrics() if _should_use_nemo_gym(master_config): - generation_config = master_config.policy["generation"] + # configure_generation_config auto-fills stop_token_ids from + # the EOS token, but run_async_nemo_gym_rollout asserts these + # are unset (NeMo-Gym manages its own stop criteria). Clear + # them on a copy so the assertion reflects user intent. + generation_config = { + **master_config.policy["generation"], + "stop_token_ids": None, + "stop_strings": None, + } nemo_gym_rollout_result = run_async_nemo_gym_rollout( policy_generation=policy_generation, input_batch=repeated_batch, tokenizer=tokenizer, task_to_env=task_to_env, - max_seq_len=None, + max_seq_len=master_config.policy[ + "max_total_sequence_length" + ], generation_config=generation_config, max_rollout_turns=None, greedy=False, + effort_config=_get_effort_config(master_config), + reward_penalty_config=master_config.reward_penalties, + thinking_tags=get_nemo_gym_thinking_tags( + master_config.env + ), ) - input_ids = nemo_gym_rollout_result.input_ids repeated_batch = nemo_gym_rollout_result.final_batch rollout_metrics = nemo_gym_rollout_result.rollout_metrics del nemo_gym_rollout_result @@ -1496,6 +1692,25 @@ def ppo_train( loss_multiplier[truncated] = 0 repeated_batch["loss_multiplier"] = loss_multiplier + # Mask samples the environment flagged (e.g. NeMo-Gym): keep them + # for advantage/value targets but zero their gradient. Only present + # on the gym path (final_batch["mask_sample"]); mirrors GRPO. + if "mask_sample" in repeated_batch: + loss_multiplier = repeated_batch["loss_multiplier"].clone() + mask_sample = repeated_batch["mask_sample"] + if isinstance(mask_sample, list): + mask_sample = torch.tensor(mask_sample, dtype=torch.bool) + mask_sample_bool = mask_sample.bool() + num_masked = int(mask_sample_bool.sum().item()) + if num_masked > 0: + print( + f" 📊 mask_sample filtering: masking {num_masked}/" + f"{len(mask_sample_bool)} env-flagged samples", + flush=True, + ) + loss_multiplier[mask_sample_bool] = 0 + repeated_batch["loss_multiplier"] = loss_multiplier + for i, message_log in enumerate(repeated_batch["message_log"]): for j, message in enumerate(message_log): if message["role"] == "assistant": @@ -2214,14 +2429,11 @@ def async_ppo_train( assert master_config.ppo["ppo_epochs"] >= 1, ( f"ppo.ppo_epochs must be >= 1 (got {master_config.ppo['ppo_epochs']})." ) - # NeMo-Gym rollout is not yet wired for async PPO. Without this guard the - # collector's gym branch would AttributeError on the (absent) reward_penalties - # field inside a worker thread, silently stalling the buffer instead of - # failing loudly here. - assert not _should_use_nemo_gym(master_config), ( - "NeMo-Gym rollout is not yet supported for async PPO " - "(env.should_use_nemo_gym must be false)." - ) + # NeMo-Gym for async PPO is handled inside the AsyncTrajectoryCollector: its + # _run_prompt_group_worker checks _should_use_nemo_gym and calls + # run_async_nemo_gym_rollout with master_config.reward_penalties (now a real + # field). reward_penalties zero rewards pre-GAE, so they flow through the value + # + GAE stage below like any other reward. No extra wiring is needed here. if max_trajectory_age_steps > 1: if not async_cfg.get("in_flight_weight_updates", False): print( @@ -2586,6 +2798,25 @@ def _sample_max_age(s: int) -> int: loss_multiplier[truncated] = 0 repeated_batch["loss_multiplier"] = loss_multiplier + # Mask samples the environment flagged (e.g. NeMo-Gym): keep them + # for advantage/value targets but zero their gradient. Only present + # on the gym path (final_batch["mask_sample"]); mirrors GRPO. + if "mask_sample" in repeated_batch: + loss_multiplier = repeated_batch["loss_multiplier"].clone() + mask_sample = repeated_batch["mask_sample"] + if isinstance(mask_sample, list): + mask_sample = torch.tensor(mask_sample, dtype=torch.bool) + mask_sample_bool = mask_sample.bool() + num_masked = int(mask_sample_bool.sum().item()) + if num_masked > 0: + print( + f" 📊 mask_sample filtering: masking {num_masked}/" + f"{len(mask_sample_bool)} env-flagged samples", + flush=True, + ) + loss_multiplier[mask_sample_bool] = 0 + repeated_batch["loss_multiplier"] = loss_multiplier + # PPO's inline loss-mask setup (unmask all assistant messages), # matching sync ppo_train — deliberately NOT GRPO's helper, # which only unmasks generated assistant messages. @@ -3250,12 +3481,44 @@ def validate( additional_metrics_to_report = dict() + # NeMo-Gym runs its own async rollout loop; cascade it first (it also + # requires async generation). Without this branch a gym config would + # fall into run_multi_turn_rollout -> NemoGym.step() (NotImplementedError). + if _should_use_nemo_gym(master_config): + # NeMo-Gym manages its own stop criteria; clear the auto-filled + # stop_token_ids/stop_strings on a copy (asserted unset by the rollout). + generation_config = { + **master_config.policy["generation"], + "stop_token_ids": None, + "stop_strings": None, + } + nemo_gym_rollout_result = run_async_nemo_gym_rollout( + policy_generation=policy_generation, + input_batch=val_batch, + tokenizer=tokenizer, + task_to_env=val_task_to_env, + max_seq_len=master_config.policy["max_total_sequence_length"], + generation_config=generation_config, + max_rollout_turns=None, + greedy=False, + effort_config=_get_effort_config(master_config), + reward_penalty_config=master_config.reward_penalties, + thinking_tags=get_nemo_gym_thinking_tags(master_config.env), + ) + val_batch = nemo_gym_rollout_result.final_batch + gen_metrics = nemo_gym_rollout_result.rollout_metrics + # NeMo-Gym responses can be huge; strip full_result unless opted in. + if not _should_log_nemo_gym_responses(master_config): + for key in list(gen_metrics): + if "full_result" in key: + gen_metrics.pop(key) + additional_metrics_to_report = gen_metrics # Async PPO uses the vLLM async engine (async_engine=true), whose # generation worker exposes only the async rollout path — the sync # `run_multi_turn_rollout` -> policy_generation.generate() would raise # AttributeError ('generate') on the async worker. Dispatch the same way # the training loop does (see async_ppo_train / grpo.validate). - if _should_use_async_rollouts(master_config): + elif _should_use_async_rollouts(master_config): val_batch, gen_metrics = run_async_multi_turn_rollout( policy_generation, val_batch, diff --git a/tests/functional/L1_Functional_Tests_Gym.sh b/tests/functional/L1_Functional_Tests_Gym.sh index 992eab603bd..05b05ce502c 100644 --- a/tests/functional/L1_Functional_Tests_Gym.sh +++ b/tests/functional/L1_Functional_Tests_Gym.sh @@ -35,6 +35,8 @@ run_test() { } run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym.sh +run_test fast uv run --no-sync bash ./tests/functional/ppo_nemo_gym.sh +run_test fast uv run --no-sync bash ./tests/functional/ppo_async_gym.sh run_test fast uv run --no-sync bash ./tests/functional/distillation_nemo_gym.sh cd ${PROJECT_ROOT}/tests diff --git a/tests/functional/ppo_async_gym.sh b/tests/functional/ppo_async_gym.sh new file mode 100755 index 00000000000..c6dd8e2cdce --- /dev/null +++ b/tests/functional/ppo_async_gym.sh @@ -0,0 +1,102 @@ +#!/bin/bash +# Async PPO + NeMo-Gym (Math / RLVR) smoke test. Non-colocated + async, which +# exercises the AsyncTrajectoryCollector gym rollout path feeding async_ppo_train's +# replay buffer + value/GAE stage. The sync counterpart is ppo_nemo_gym.sh; see +# grpo_async_gym.sh for the GRPO analog and the NeMo-Gym data-prep prerequisites. + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) +# Mark the current repo as safe, since wandb fetches metadata about the repo +git config --global --add safe.directory $PROJECT_ROOT + +set -eou pipefail + +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 +CHECKPOINT_DIR=$EXP_DIR/checkpoints +DATA_DIR=$EXP_DIR/data +export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} + +rm -rf $EXP_DIR $LOG_DIR +mkdir -p $EXP_DIR $LOG_DIR $CHECKPOINT_DIR $DATA_DIR + +# clean up checkpoint directory on exit +trap "rm -rf $CHECKPOINT_DIR" EXIT + +cd $PROJECT_ROOT + +# Follow the NeMo-Gym instructions to obtain the Gym workspace + data: +# https://docs.nvidia.com/nemo/gym/0.1.0/tutorials/nemo-rl-grpo/setup.html +cd 3rdparty/Gym-workspace/Gym + +# We need HF_TOKEN to download the data from huggingface +if [[ ! -f env.yaml ]]; then + if [[ -z "${HF_TOKEN:-}" ]]; then + echo "[ERROR] HF_TOKEN is not set" + exit 1 + fi + echo "hf_token: $HF_TOKEN" >> env.yaml +fi + +uv run ng_prepare_data "+config_paths=[resources_servers/math_with_judge/configs/math_with_judge.yaml]" \ + +output_dirpath=data/math_with_judge \ + +mode=train_preparation \ + +should_download=true \ + +data_source=huggingface +cd - + +TRAIN_PATH=$DATA_DIR/math_with_judge_train.jsonl +VALIDATION_PATH=$DATA_DIR/math_with_judge_validation.jsonl +cp 3rdparty/Gym-workspace/Gym/data/math_with_judge/train.jsonl $TRAIN_PATH +cp 3rdparty/Gym-workspace/Gym/data/math_with_judge/validation.jsonl $VALIDATION_PATH + +uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ + $PROJECT_ROOT/examples/nemo_gym/run_ppo_nemo_gym.py \ + --config $PROJECT_ROOT/examples/nemo_gym/ppo_math_rlvr_nemo_gym.yaml \ + policy.model_name=Qwen/Qwen3-0.6B \ + policy.dtensor_cfg.enabled=true \ + policy.megatron_cfg.enabled=false \ + value.model_name=Qwen/Qwen3-0.6B \ + value.dtensor_cfg.enabled=true \ + value.megatron_cfg.enabled=false \ + policy.generation.vllm_cfg.tensor_parallel_size=1 \ + policy.generation.vllm_cfg.async_engine=true \ + policy.generation.vllm_cfg.expose_http_server=true \ + policy.max_total_sequence_length=512 \ + policy.generation.colocated.enabled=false \ + policy.generation.colocated.resources.num_nodes=1 \ + policy.generation.colocated.resources.gpus_per_node=1 \ + ppo.num_prompts_per_step=4 \ + ppo.num_generations_per_prompt=2 \ + ppo.max_num_steps=10 \ + ppo.val_period=5 \ + ppo.policy_training_start_step=0 \ + ppo.async_ppo.enabled=true \ + ppo.async_ppo.max_trajectory_age_steps=1 \ + ppo.async_ppo.in_flight_weight_updates=true \ + policy.train_global_batch_size=4 \ + policy.train_micro_batch_size=1 \ + cluster.gpus_per_node=2 \ + loss_fn.use_importance_sampling_correction=true \ + logger.tensorboard_enabled=true \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=false \ + logger.monitor_gpus=true \ + checkpointing.enabled=true \ + checkpointing.save_period=5 \ + checkpointing.checkpoint_dir=$CHECKPOINT_DIR \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Smoke thresholds: the async collector fed the buffer (training progressed to +# step 10) and the mismatch diagnostic stayed bounded. +uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/gen_kl_error"]) < 1.3' \ + 'data["validation/accuracy"]["10"] >= 0.0' diff --git a/tests/functional/ppo_nemo_gym.sh b/tests/functional/ppo_nemo_gym.sh new file mode 100755 index 00000000000..d2fd8607975 --- /dev/null +++ b/tests/functional/ppo_nemo_gym.sh @@ -0,0 +1,99 @@ +#!/bin/bash +# PPO + NeMo-Gym (Math / RLVR) smoke test. Colocated + synchronous PPO, which +# exercises the three new gym paths: setup() gym spinup (vLLM deferred load), +# ppo_train() gym rollout, and validate() gym dispatch. See grpo_async_gym.sh +# for the GRPO analog and the NeMo-Gym data-prep prerequisites. + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) +# Mark the current repo as safe, since wandb fetches metadata about the repo +git config --global --add safe.directory $PROJECT_ROOT + +set -eou pipefail + +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 +CHECKPOINT_DIR=$EXP_DIR/checkpoints +DATA_DIR=$EXP_DIR/data +export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} + +rm -rf $EXP_DIR $LOG_DIR +mkdir -p $EXP_DIR $LOG_DIR $CHECKPOINT_DIR $DATA_DIR + +# clean up checkpoint directory on exit +trap "rm -rf $CHECKPOINT_DIR" EXIT + +cd $PROJECT_ROOT + +# Follow the NeMo-Gym instructions to obtain the Gym workspace + data: +# https://docs.nvidia.com/nemo/gym/0.1.0/tutorials/nemo-rl-grpo/setup.html +cd 3rdparty/Gym-workspace/Gym + +# We need HF_TOKEN to download the data from huggingface +if [[ ! -f env.yaml ]]; then + if [[ -z "${HF_TOKEN:-}" ]]; then + echo "[ERROR] HF_TOKEN is not set" + exit 1 + fi + echo "hf_token: $HF_TOKEN" >> env.yaml +fi + +uv run ng_prepare_data "+config_paths=[resources_servers/math_with_judge/configs/math_with_judge.yaml]" \ + +output_dirpath=data/math_with_judge \ + +mode=train_preparation \ + +should_download=true \ + +data_source=huggingface +cd - + +TRAIN_PATH=$DATA_DIR/math_with_judge_train.jsonl +VALIDATION_PATH=$DATA_DIR/math_with_judge_validation.jsonl +cp 3rdparty/Gym-workspace/Gym/data/math_with_judge/train.jsonl $TRAIN_PATH +cp 3rdparty/Gym-workspace/Gym/data/math_with_judge/validation.jsonl $VALIDATION_PATH + +uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ + $PROJECT_ROOT/examples/nemo_gym/run_ppo_nemo_gym.py \ + --config $PROJECT_ROOT/examples/nemo_gym/ppo_math_rlvr_nemo_gym.yaml \ + policy.model_name=Qwen/Qwen3-0.6B \ + policy.dtensor_cfg.enabled=true \ + policy.megatron_cfg.enabled=false \ + value.model_name=Qwen/Qwen3-0.6B \ + value.dtensor_cfg.enabled=true \ + value.megatron_cfg.enabled=false \ + policy.generation.vllm_cfg.tensor_parallel_size=1 \ + policy.generation.vllm_cfg.async_engine=true \ + policy.generation.vllm_cfg.expose_http_server=true \ + policy.max_total_sequence_length=512 \ + policy.generation.colocated.enabled=true \ + ppo.num_prompts_per_step=4 \ + ppo.num_generations_per_prompt=2 \ + ppo.max_num_steps=10 \ + ppo.val_period=5 \ + ppo.policy_training_start_step=0 \ + ppo.reward_scaling.enabled=false \ + ppo.reward_shaping.enabled=false \ + policy.train_global_batch_size=4 \ + policy.train_micro_batch_size=1 \ + cluster.gpus_per_node=2 \ + loss_fn.use_importance_sampling_correction=true \ + logger.tensorboard_enabled=true \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=false \ + logger.monitor_gpus=true \ + checkpointing.enabled=true \ + checkpointing.save_period=5 \ + checkpointing.checkpoint_dir=$CHECKPOINT_DIR \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Smoke thresholds: mismatch diagnostic stays bounded and the critic learns a +# finite value loss (explained_variance is logged by PPO's value stage). +uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/gen_kl_error"]) < 1.3' \ + 'data["validation/accuracy"]["10"] >= 0.0' diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 0a62c687aa1..3ecb22f57a9 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -1062,14 +1062,10 @@ def test_async_ppo_requires_positive_ppo_epochs(): _call_async_ppo_train(master_config) -def test_async_ppo_rejects_nemo_gym(): - """NeMo-Gym rollout is not wired for async PPO; guard fails loud at startup.""" - master_config = _build_async_ppo_master_config() - master_config.env["should_use_nemo_gym"] = True - # _should_use_nemo_gym also requires the http server to be exposed. - master_config.policy["generation"]["vllm_cfg"]["expose_http_server"] = True - with pytest.raises(AssertionError, match="NeMo-Gym rollout is not yet supported"): - _call_async_ppo_train(master_config) +# NeMo-Gym IS supported for async PPO: the AsyncTrajectoryCollector runs the gym +# rollout internally (checks _should_use_nemo_gym, passes master_config.reward_penalties), +# so async_ppo_train no longer guards it out. Validate-side dispatch is covered by +# test_validate_dispatches_to_nemo_gym_rollout below. # --------------------------------------------------------------------------- @@ -1201,9 +1197,9 @@ def _fresh_buffer(): # attribute 'generate') at the first validation step. # --------------------------------------------------------------------------- def _run_validate_with_mocked_rollouts(master_config): - """Drive validate() through one batch with both rollout fns mocked. + """Drive validate() through one batch with all three rollout fns mocked. - Returns (async_rollout_mock, sync_rollout_mock) for call assertions. + Returns (async_rollout_mock, sync_rollout_mock, nemo_gym_rollout_mock). """ from unittest.mock import MagicMock, patch @@ -1219,6 +1215,10 @@ def _run_validate_with_mocked_rollouts(master_config): {"total_reward": torch.tensor([1.0]), "message_log": []}, {"mean_gen_tokens_per_sample": 5.0}, ) + # NeMo-Gym rollout returns a result object (not a tuple). + gym_result = MagicMock() + gym_result.final_batch = {"total_reward": torch.tensor([1.0]), "message_log": []} + gym_result.rollout_metrics = {"mean_gen_tokens_per_sample": 5.0} with ( patch( "nemo_rl.algorithms.ppo.run_async_multi_turn_rollout", @@ -1228,6 +1228,13 @@ def _run_validate_with_mocked_rollouts(master_config): "nemo_rl.algorithms.ppo.run_multi_turn_rollout", return_value=rollout_return, ) as sync_mock, + patch( + "nemo_rl.algorithms.ppo.run_async_nemo_gym_rollout", + return_value=gym_result, + ) as gym_mock, + # Isolate the dispatch decision from gym-config internals. + patch("nemo_rl.algorithms.ppo._get_effort_config", return_value=None), + patch("nemo_rl.algorithms.ppo.get_nemo_gym_thinking_tags", return_value=[]), ): validate( policy_generation=MagicMock(), @@ -1238,7 +1245,7 @@ def _run_validate_with_mocked_rollouts(master_config): master_config=master_config, logger=MagicMock(), ) - return async_mock, sync_mock + return async_mock, sync_mock, gym_mock def test_validate_dispatches_to_async_rollout_for_async_engine(): @@ -1246,18 +1253,38 @@ def test_validate_dispatches_to_async_rollout_for_async_engine(): rollout path. The sync path calls policy_generation.generate(), which the async worker lacks -> AttributeError at the first validation step.""" master_config = _build_async_ppo_master_config() # vllm async_engine=True - async_mock, sync_mock = _run_validate_with_mocked_rollouts(master_config) + async_mock, sync_mock, gym_mock = _run_validate_with_mocked_rollouts(master_config) async_mock.assert_called_once() sync_mock.assert_not_called() + gym_mock.assert_not_called() def test_validate_dispatches_to_sync_rollout_when_not_async(): """Non-async (colocated/sync) PPO validation uses the synchronous rollout.""" master_config = _build_async_ppo_master_config() master_config.policy["generation"]["vllm_cfg"]["async_engine"] = False - async_mock, sync_mock = _run_validate_with_mocked_rollouts(master_config) + async_mock, sync_mock, gym_mock = _run_validate_with_mocked_rollouts(master_config) sync_mock.assert_called_once() async_mock.assert_not_called() + gym_mock.assert_not_called() + + +def test_validate_dispatches_to_nemo_gym_rollout(): + """Regression for the validation crash gap: under a NeMo-Gym config, validate() + must take the gym rollout path. Otherwise it falls into run_multi_turn_rollout + -> NemoGym.step() (NotImplementedError).""" + master_config = _build_async_ppo_master_config() # vllm async_engine=True + master_config.env = { + "should_use_nemo_gym": True, + "should_log_nemo_gym_responses": True, + "nemo_gym": {}, + } + # _should_use_nemo_gym also requires the http server to be exposed. + master_config.policy["generation"]["vllm_cfg"]["expose_http_server"] = True + async_mock, sync_mock, gym_mock = _run_validate_with_mocked_rollouts(master_config) + gym_mock.assert_called_once() + async_mock.assert_not_called() + sync_mock.assert_not_called() # --------------------------------------------------------------------------- From 82af1026fabcbf8539c5a9958dff0dd8693832e4 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Mon, 13 Jul 2026 15:36:26 -0700 Subject: [PATCH 27/69] fix(async ppo): cycle the collector dataloader across epochs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The AsyncTrajectoryCollector iterated the training dataloader with a single, non-cycling `for batch in self.dataloader`. A finite dataset is consumed once and the loop returns, stopping the collector and silently stalling the replay buffer (the driver waits on buffer_size=0 forever). This bites hardest on RESUME: the checkpoint saves the collector's real dataloader position, and a restored StatefulDataLoader whose samples_yielded is already at the epoch end yields ZERO batches on the first epoch — so the run hangs immediately after resume (observed at step 195 of a 218-batch DAPOMath17K run). Wrap the loop in `while self.running:` so it cycles epochs, mirroring sync ppo_train/grpo_train's `while epoch < max_num_epochs: for batch in dataloader`. StatefulDataLoader resume-once semantics (next_iter_state cleared after the first iterator) mean the resumed partial epoch yields its remainder and fresh full epochs follow; a two-empty-epoch guard fails loudly on a genuinely empty dataset instead of busy-spinning. Add unit tests for the resume-exhaustion recovery, multi-epoch cycling, and the empty-dataset guard. Shared with async GRPO. Signed-off-by: Jiaqi Zeng --- .../async_utils/trajectory_collector.py | 115 ++++++++++++------ tests/unit/algorithms/test_async_utils.py | 105 ++++++++++++++++ 2 files changed, 182 insertions(+), 38 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index 4fc1b03187a..3cd127105ca 100644 --- a/nemo_rl/algorithms/async_utils/trajectory_collector.py +++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py @@ -322,51 +322,90 @@ def start_collection(self, dataloader: StatefulDataLoader) -> None: def _collection_loop(self): """Run the collection loop in background thread.""" try: - for batch in self.dataloader: - if not self.running: - break - - # Check if manually paused and wait - if not self._manual_pause_cleared.is_set() and self.running: - self._manual_pause_cleared.wait() - - # Check if refit is in progress and wait - if not self._refit_pause_cleared.is_set() and self.running: - print("⏸️ Pausing collection for refit...") - with self._efficiency_timer.time("idle/refit_event_wait"): - self._refit_pause_cleared.wait() - print("▶️ Refit completed, resuming collection") - - # Check if generation limits require pausing collection - if self._should_pause_for_generation_limits() and self.running: - # Only log warning once per weight version - if self._last_limit_warning_version != self.current_weight_version: - max_trajectory_age = self._current_max_age - target_weights = [ - self.current_weight_version + i - for i in range(max_trajectory_age) - ] + # Cycle the dataloader across epochs. A single `for batch in + # self.dataloader` is ONE epoch, but the collector must feed the + # buffer continuously for the whole (possibly multi-epoch) run — + # mirroring sync ppo_train/grpo_train's + # `while epoch < max_num_epochs: for batch in dataloader`. Without + # this outer loop the finite dataset is consumed once and the loop + # returns, stopping the collector and silently stalling the buffer. + # This bites hardest on RESUME: the restored StatefulDataLoader can + # already be at the epoch end (samples_yielded == dataset size), so + # the very first epoch yields ZERO batches and the run hangs. + consecutive_empty_epochs = 0 + while self.running: + produced_this_epoch = 0 + for batch in self.dataloader: + if not self.running: + break - print( - f"⏸️ Pausing collection: all target weights {target_weights} for weight version {self.current_weight_version} " - f"already exist in buffer. Waiting for weight update..." - ) - self._last_limit_warning_version = self.current_weight_version + # Check if manually paused and wait + if not self._manual_pause_cleared.is_set() and self.running: + self._manual_pause_cleared.wait() - self._generation_limit_cleared.clear() # Clear the event to pause + # Check if refit is in progress and wait + if not self._refit_pause_cleared.is_set() and self.running: + print("⏸️ Pausing collection for refit...") + with self._efficiency_timer.time("idle/refit_event_wait"): + self._refit_pause_cleared.wait() + print("▶️ Refit completed, resuming collection") - # Efficiently wait for generation limits to be cleared (no polling!) - with self._efficiency_timer.time("idle/generation_limit_pause"): - self._generation_limit_cleared.wait() + # Check if generation limits require pausing collection + if self._should_pause_for_generation_limits() and self.running: + # Only log warning once per weight version + if ( + self._last_limit_warning_version + != self.current_weight_version + ): + max_trajectory_age = self._current_max_age + target_weights = [ + self.current_weight_version + i + for i in range(max_trajectory_age) + ] + + print( + f"⏸️ Pausing collection: all target weights {target_weights} for weight version {self.current_weight_version} " + f"already exist in buffer. Waiting for weight update..." + ) + self._last_limit_warning_version = ( + self.current_weight_version + ) + + self._generation_limit_cleared.clear() # Clear the event to pause + + # Efficiently wait for generation limits to be cleared (no polling!) + with self._efficiency_timer.time( + "idle/generation_limit_pause" + ): + self._generation_limit_cleared.wait() + + # Double-check we're still running after being woken up + if not self.running: + break - # Double-check we're still running after being woken up if not self.running: break - if not self.running: - break - - self._process_batch(batch) + self._process_batch(batch) + produced_this_epoch += 1 + + # Epoch exhausted; the next `for` starts a fresh epoch + # (StatefulDataLoader resume-once semantics: the loaded position + # only applies to the first epoch after load_state_dict). A + # resumed first epoch can legitimately yield 0 batches; the fresh + # epoch that follows will not. If TWO consecutive epochs yield + # nothing, the dataset is genuinely empty — fail loudly instead of + # busy-spinning the outer while. + if produced_this_epoch == 0: + consecutive_empty_epochs += 1 + if consecutive_empty_epochs >= 2: + print( + "❌ Dataloader yielded no batches for two consecutive " + "epochs (empty dataset?); stopping collection." + ) + break + else: + consecutive_empty_epochs = 0 except Exception as e: print(f"❌ Error in trajectory collection: {e}") import traceback diff --git a/tests/unit/algorithms/test_async_utils.py b/tests/unit/algorithms/test_async_utils.py index 9a3972e6b22..5e4acfd80a7 100644 --- a/tests/unit/algorithms/test_async_utils.py +++ b/tests/unit/algorithms/test_async_utils.py @@ -1165,6 +1165,25 @@ def test_sample_consumes_selected_rows(self): ray.kill(buf) +class _EpochListLoader: + """Fake dataloader for collection-loop tests: each ``iter()`` is one epoch + yielding ``sizes[i]`` batches (the last size repeats for further epochs). + + Models StatefulDataLoader resume-once semantics: the FIRST epoch may be + short/empty (a resumed loader replays only the remainder from the restored + ``samples_yielded``), and later ``iter()`` calls start fresh full epochs. + """ + + def __init__(self, sizes): + self._sizes = list(sizes) + self.iter_count = 0 + + def __iter__(self): + n = self._sizes[min(self.iter_count, len(self._sizes) - 1)] + self.iter_count += 1 + return iter([{"batch": i} for i in range(n)]) + + class TestAsyncTrajectoryCollector: """Test cases for AsyncTrajectoryCollector.""" @@ -1560,6 +1579,92 @@ def test_dataloader_state_retrieval(self): ray.kill(buffer) ray.kill(mock_env) + # ------------------------------------------------------------------ + # Dataloader epoch-cycling in the async collection loop. A single + # `for batch in self.dataloader` is ONE epoch; the collector must cycle + # epochs to feed the buffer for the whole run. Without it an exhausted + # dataset — most acutely a RESUMED StatefulDataLoader whose restored + # samples_yielded is already at the epoch end — makes the loop return + # immediately, stop the collector, and silently stall the buffer forever. + # ------------------------------------------------------------------ + def _prime_collector_for_loop(self, collector, dataloader): + """Wire a collector so _collection_loop's body reduces to _process_batch. + + __init__ leaves the manual/refit pause events SET (cleared), so with + generation-limit pausing stubbed off the loop body is just + `if not running: break; _process_batch(batch)`. + """ + collector.dataloader = dataloader + collector.running = True + collector._should_pause_for_generation_limits = lambda: False + return collector + + def _run_collection_loop(self, collector, timeout=15.0): + """Run _collection_loop to completion in a watchdog thread.""" + t = threading.Thread(target=collector._collection_loop, daemon=True) + t.start() + t.join(timeout) + assert not t.is_alive(), "collection loop did not terminate (possible hang)" + + def test_collection_loop_recovers_from_exhausted_first_epoch_on_resume(self): + """Regression for the resume-hang: a resumed dataloader whose FIRST epoch + yields 0 batches must NOT stop the collector. The outer epoch loop + re-iterates a fresh full epoch and keeps feeding the buffer (otherwise the + driver waits on an empty buffer forever).""" + collector = self.create_local_collector() + # First iter() -> 0 batches (resumed at epoch end); next iter() -> full. + loader = _EpochListLoader(sizes=[0, 3]) + self._prime_collector_for_loop(collector, loader) + + processed = [] + + def _fake_process(batch): + processed.append(batch) + if len(processed) >= 3: + collector.running = False # stop after the fresh epoch + + collector._process_batch = _fake_process + self._run_collection_loop(collector) + + # The buggy single-pass loop processes 0 and stops; cycling re-iterates + # past the empty first epoch and processes the fresh epoch's 3 batches. + assert processed == [{"batch": 0}, {"batch": 1}, {"batch": 2}] + assert loader.iter_count >= 2 # re-iterated past the empty first epoch + + def test_collection_loop_cycles_multiple_epochs(self): + """A finite dataset must be re-iterated so the collector can feed more + steps than a single dataset pass provides.""" + collector = self.create_local_collector() + loader = _EpochListLoader(sizes=[2]) # 2 batches every epoch + self._prime_collector_for_loop(collector, loader) + + processed = [] + + def _fake_process(batch): + processed.append(batch) + if len(processed) >= 5: # 5 > 2/epoch -> must cross epoch boundaries + collector.running = False + + collector._process_batch = _fake_process + self._run_collection_loop(collector) + + assert len(processed) == 5 + assert loader.iter_count >= 3 # 2 + 2 + 1 spans three epochs + + def test_collection_loop_stops_on_empty_dataset(self): + """A genuinely empty dataloader (0 batches for two consecutive epochs) + must stop the loop, not busy-spin the outer while forever.""" + collector = self.create_local_collector() + loader = _EpochListLoader(sizes=[0]) # always empty + self._prime_collector_for_loop(collector, loader) + collector._process_batch = mock.MagicMock() + + self._run_collection_loop(collector) # must return (watchdog asserts) + + collector._process_batch.assert_not_called() + assert collector.running is False + assert loader.iter_count == 2 # stopped after two empty epochs + class TestAsyncUtilsIntegration: """Integration tests for async utilities working together.""" From f70c08ca0a58aa33aad5c8846a6d0eaa7eea9498 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Tue, 14 Jul 2026 10:09:23 -0700 Subject: [PATCH 28/69] bump gym Signed-off-by: Jiaqi Zeng --- 3rdparty/Gym-workspace/Gym | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index 8b145ad5693..67f7955902a 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit 8b145ad5693a50740fcbb12fe89bd56eebcdb6d4 +Subproject commit 67f7955902aa7a54678e22fa680858dc1fe7e613 From 3ee78be4e710a59e751d36e33df32bb7f4720605 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Tue, 14 Jul 2026 10:10:03 -0700 Subject: [PATCH 29/69] fix survivorship bias for grpo and add dapo grpo baseline Signed-off-by: Jiaqi Zeng --- .../configs/grpo_math_1B_dapo_megatron.yaml | 35 +++++++++++++++++++ nemo_rl/algorithms/grpo.py | 22 +++++++++++- 2 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 examples/configs/grpo_math_1B_dapo_megatron.yaml diff --git a/examples/configs/grpo_math_1B_dapo_megatron.yaml b/examples/configs/grpo_math_1B_dapo_megatron.yaml new file mode 100644 index 00000000000..66082adffa8 --- /dev/null +++ b/examples/configs/grpo_math_1B_dapo_megatron.yaml @@ -0,0 +1,35 @@ +# GRPO (Megatron-Core) on DAPO-Math-17K — async baseline for the PPO async run +# (repro_ppo_mcore_7b_math_async.sh / ppo_math_1B.yaml). +# +# Inherits grpo_math_1B_megatron.yaml and swaps the `data:` and `env:` blocks so +# the data setup AND the reward scorer are byte-for-byte the same as +# ppo_math_1B.yaml (apple-to-apple): +# - full DAPOMath17K train set. DAPOMath17KDataset takes **kwargs and ignores +# the inherited split_validation_size, so NO 5% train-side split happens — +# the whole training set is used, exactly like PPO. +# - explicit DAPOMathAIME2024 validation (overrides the parent's +# `validation: null`, which would otherwise fall back to a train split). +# - prompt_file: null (no CoT template) to match PPO; the parent grpo config +# uses examples/prompts/cot.txt. +# - math_verify_impl: dapo_math_verify to match PPO. The grpo default is the +# more lenient hf_math_verify, which scores the SAME completions higher and +# was inflating GRPO's reward/accuracy vs PPO (~0.45 vs ~0.29 val at step 0). +# Everything else (loss_fn, policy/megatron, cluster, grpo sizes, async_grpo) is +# inherited and overridden on the command line by the launch script. +defaults: "grpo_math_1B_megatron.yaml" + +data: + train: + dataset_name: DAPOMath17K + split_validation_size: null # no train-side split; use the explicit validation set below (matches PPO) + validation: + dataset_name: DAPOMathAIME2024 + default: + prompt_file: null + system_prompt_file: null + processor: "math_hf_data_processor" + env_name: "math" + +env: + math: + math_verify_impl: "dapo_math_verify" # match PPO's stricter verifier (grpo default is hf_math_verify) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 375899a7152..44d4047730a 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -151,6 +151,19 @@ class AsyncGRPOConfig(TypedDict): in_flight_weight_updates: NotRequired[bool] # Recomputes the KV cache after the in-flight weight updates. recompute_kv_cache_after_weight_updates: NotRequired[bool] + # On resume, how to treat an INCOMPLETE (partially-generated) restored target + # step in the replay-buffer checkpoint: + # * false (default): keep the partial survivors and let the collector gap-fill + # only the missing groups. Fast resume, but the partial batch is + # SURVIVORSHIP-BIASED toward SHORT rollouts (only the fast-completing ones + # were saved), so the first step after each resume trains on a shorter + + # higher-reward batch — visible as a dip in mean_gen_tokens_per_sample and a + # reward spike right after every resume. + # * true: drop the incomplete target and regenerate it fresh (unbiased batch), + # at the cost of a one-target generation bubble at resume. Recommended when + # resuming frequently. Complete banked targets are unbiased and always kept. + # Absent/null => false. Read with a None-safe default, so allow null. + drop_incomplete_targets_on_restore: NotRequired[bool | None] class AdvEstimatorConfig(TypedDict): @@ -3691,7 +3704,14 @@ def async_grpo_train( ) replay_buffer = ReplayBuffer.options(runtime_env=_replay_runtime_env).remote( - max_size=optimal_buffer_size + max_size=optimal_buffer_size, + # Default False => historical gap-fill behavior; True drops the + # survivorship-biased incomplete frontier target on resume (regenerate + # fresh). None (absent) is coerced to False by the buffer constructor. + drop_incomplete_targets_on_restore=master_config.grpo["async_grpo"].get( + "drop_incomplete_targets_on_restore" + ) + or False, ) last_checkpoint_path = checkpointer.get_latest_checkpoint_path() From 021b0adbd929d1cfbe664532cd6d93973552b9de Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 16 Jul 2026 11:54:43 -0700 Subject: [PATCH 30/69] fix yaml Signed-off-by: Jiaqi Zeng --- examples/configs/grpo_ultra_256n4g_bf16.yaml | 676 +++++++++++++++++++ 1 file changed, 676 insertions(+) create mode 100644 examples/configs/grpo_ultra_256n4g_bf16.yaml diff --git a/examples/configs/grpo_ultra_256n4g_bf16.yaml b/examples/configs/grpo_ultra_256n4g_bf16.yaml new file mode 100644 index 00000000000..b9c3c69bc67 --- /dev/null +++ b/examples/configs/grpo_ultra_256n4g_bf16.yaml @@ -0,0 +1,676 @@ +# ============================================================================= +# GRPO Ultra V3 — 256-node GB200 NVL72 Config (bf16) +# ============================================================================= +# Config for GRPO training on 256 nodes × 4 GPUs/node (1024 GPUs). +# Full batch sizes and sequence lengths for convergence runs. +# +# Node allocation (256 total, set via launch script env vars): +# - Training: 64 nodes (256 GPUs) — 4 segments of 16 +# - vLLM: 182 nodes (728 GPUs) — 91 instances at TP=8 EP=8 (2 nodes each) +# - Gym/Judge: 10 nodes ( 40 GPUs) — judges scaled for production throughput +# +# Generation-heavy split follows the SuperV3 production ratio (~25/71/4). +# +# Training parallelism (256 GPUs = 64 nodes): +# - TP: 8 +# - EP: 64 +# - CP: 8 +# - PP: 1 +# - SP: true +# +# vLLM parallelism (bf16, TP=8, EP=8): +# NeMo Gym requires async_engine=true, but vLLM DP+EP (EP > TP) requires +# async_engine=false (see https://github.com/NVIDIA-NeMo/RL/issues/1101). +# This forces EP <= TP. With EP=8 (=TP), vllm_dp_size = 8/8 = 1 so +# async_engine=true works. +# ============================================================================= + +# ============================================================================= +# Cluster — overridden by launch script +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 256 + segment_size: 16 + +# ============================================================================= +# Checkpointing +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/grpo_ultra_v3" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_optimizer: true + save_period: 10 + ft_keep_latest_k: 1 + ft_save_period: 1 + checkpoint_must_save_by: "00:03:30:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# GRPO Algorithm +# ============================================================================= +grpo: + num_prompts_per_step: 256 + num_generations_per_prompt: 16 + num_val_generations_per_prompt: 2 + max_rollout_turns: 1 + max_num_epochs: 1 + max_num_steps: 1000000 + normalize_rewards: true + use_leave_one_out_baseline: true + advantage_clip_low: -50 + advantage_clip_high: 50 + val_period: -1 + val_at_start: false + val_at_end: false + overlong_filtering: false + max_val_samples: null + val_batch_size: 256 + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + penalize_invalid_tool_call: true + invalid_tool_call_advantage: -5.0 + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + + use_best_at_k: false + best_at_k_k: 8 + best_at_k_m: 1000 + + use_combined_training: false + combined_training_weight_mode: "auto" + combined_training_best_at_k_weight: 0.2 + combined_training_pass_at_1_weight: 1.0 + + dynamic_sampling_oversample_ratio: 1.0 + seq_logprob_error_threshold: 2 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: null + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + force_on_policy_ratio: true + use_kl_in_reward: false + +# ============================================================================= +# Policy +# ============================================================================= +policy: + model_name: "/lustre/fsw/portfolios/llmservice/users/soumyes/sft-runs/eval_and_sleep/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-mar13-blend-512k-filt-1e-5/iter_0001900/hf" + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 4096 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 65536 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + # EP=64: 512 experts / 64 = 8 experts per EP rank + # All-to-all spans 64 GPUs (16 nodes), fits within one NVLink domain + expert_model_parallel_size: 64 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "flex" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled + mtp_loss_scaling_factor: 0.3 + mtp_use_repeated_layer: true + mtp_num_layers: 5 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 3.0e-6 + min_lr: 3.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 3e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + # FP8 — disabled for bf16 runs. Enable for mxfp8 validation. + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + use_gloo_process_groups: false + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + fully_parallel_save_process_group: "ep_dp" + fully_parallel_load_process_group: "ep_dp" + fully_parallel_load_exchange_algo: "broadcast" + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 3000 + port_range_high: 4999 + backend: "vllm" + max_new_tokens: 49152 + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. + # Same memory as EP=1 but uses all-to-all for expert routing. + # EP > TP blocked by https://github.com/NVIDIA-NeMo/RL/issues/1101. + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 8 + gpu_memory_utilization: 0.85 + max_model_len: 49152 + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + # nano_v3 is a CUSTOM reasoning parser; its plugin file is registered at vllm_cfg + # top level (read by vllm_worker_async.py:481 -> ReasoningParserManager.import_reasoning_parser). + # It must NOT sit inside http_server_serving_chat_kwargs: that dict is splatted into + # OpenAIServingChat(**kwargs), which has no reasoning_parser_plugin arg -> TypeError. + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + + vllm_kwargs: + attention_backend: FLASH_ATTN + moe_backend: triton + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_mode: PIECEWISE + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 182 # Overridden by launch script + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + use_multiple_dataloader: false + train: + data_path: null # Set by launch script + validation: + data_path: null # Set by launch script + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + Judge Models +# ============================================================================= +env: + should_use_nemo_gym: true + # true: skip expensive train_data_step*.jsonl (recommended for large Gym runs); false: write full jsonl. + should_log_nemo_gym_responses: true + nemo_gym: + nemo_gym_log_dir: "logs/nemo_gym" + skip_venv_if_present: true + port_range_low: 5000 + port_range_high: 5999 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/math_with_judge/configs/math_with_judge.yaml + - resources_servers/code_gen/configs/code_gen.yaml + - resources_servers/workplace_assistant/configs/workplace_assistant.yaml + - resources_servers/mcqa/configs/mcqa.yaml + - resources_servers/instruction_following/configs/instruction_following.yaml + - resources_servers/equivalence_llm_judge/configs/lc_judge.yaml + - resources_servers/calendar/configs/calendar.yaml + - resources_servers/genrm_compare/configs/genrm_compare.yaml + - resources_servers/equivalence_llm_judge/configs/nl2bash-equivalency.yaml + - resources_servers/equivalence_llm_judge/configs/equivalence_llm_judge.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml + - resources_servers/reasoning_gym/configs/reasoning_gym.yaml + - resources_servers/terminus_judge/configs/terminus_judge_string_only.yaml + - resources_servers/ns_tools/configs/ns_tools.yaml + - resources_servers/math_formal_lean/configs/math_formal_lean_multi_turn.yaml + # swerl_gen disabled: requires Apptainer/Singularity (not available on aarch64) + # - resources_servers/swerl_gen/configs/swerl_gen.yaml + - resources_servers/multichallenge/configs/multichallenge.yaml + - resources_servers/inverse_if/configs/inverse_if.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/search_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/toolcall_schema_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/swe_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/abstention/configs/abstention.yaml + - resources_servers/nvarc/configs/inductive.yaml + - resources_servers/nvarc/configs/transductive.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/droid_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/equivalence_rule/configs/lc.yaml + - resources_servers/ether0/configs/ether0.yaml + - resources_servers/structured_outputs/configs/structured_outputs_json_yaml_xml_v1.yaml + - resources_servers/structured_outputs/configs/structured_outputs_v3.yaml + - resources_servers/format_verification/configs/freeform_formatting.yaml + - resources_servers/format_verification/configs/citation_format.yaml + - resources_servers/rdkit_chemistry/configs/rdkit_chemistry.yaml + - resources_servers/jailbreak_detection/configs/jailbreak_detection_nemotron_combined_reward_tp8.yaml + - resources_servers/indirect_prompt_injection/configs/indirect_prompt_injection.yaml + + abstention: + resources_servers: + abstention: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + # Safety Model: 4B — TP=4 ensures each PG claims a full node, + # avoiding GPU fragmentation that can block larger-TP models. + jailbreak_detection: + resources_servers: + jailbreak_detection: + judge_model_server: + type: responses_api_models + name: safety_judge_model + + safety_judge_model: + responses_api_models: + local_vllm_model: + entrypoint: app.py + model: null # Set by launch script + return_token_id_information: false + uses_reasoning_parser: false + debug: true + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + # The gym `local_vllm_model` venv has no deep_gemm (nemo_rl's policy venv + # does, but this stock-vLLM judge venv doesn't). On an FP8 model, vLLM's + # deep_gemm warmup then hard-fails. Disable it -> fall back to vLLM's other + # FP8 kernels, consistent with the policy running use_deep_gemm=false. + VLLM_USE_DEEP_GEMM: "0" + + vllm_serve_kwargs: + attention_backend: TRITON_ATTN + tensor_parallel_size: 4 + data_parallel_size: 2 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + gpu_memory_utilization: 0.85 + max_model_len: 96000 + max_num_seqs: 16 + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16] + + + # nl2bash / General Judge: TP=4 on GB200 192GB + nl2bash_judge_model: + responses_api_models: + local_vllm_model: + entrypoint: app.py + model: null # Set by launch script + return_token_id_information: false + uses_reasoning_parser: false + debug: true + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + # The gym `local_vllm_model` venv has no deep_gemm (nemo_rl's policy venv + # does, but this stock-vLLM judge venv doesn't). On an FP8 model, vLLM's + # deep_gemm warmup then hard-fails. Disable it -> fall back to vLLM's other + # FP8 kernels, consistent with the policy running use_deep_gemm=false. + VLLM_USE_DEEP_GEMM: "0" + + vllm_serve_kwargs: + attention_backend: FLASH_ATTN + tensor_parallel_size: 4 + data_parallel_size: 4 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + enable_expert_parallel: true + enable_auto_tool_choice: true + tool_call_parser: hermes + gpu_memory_utilization: 0.85 + max_model_len: 131072 + max_num_seqs: 256 + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + inverse_if: + resources_servers: + inverse_if: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + + multichallenge: + resources_servers: + multichallenge: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + equivalence_llm_judge: + resources_servers: + equivalence_llm_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + # GenRM: TP=4 on GB200 192GB + genrm_compare_resources_server: + resources_servers: + genrm_compare: + num_rollouts_per_prompt: ${grpo.num_generations_per_prompt} + genrm_model_server: + type: responses_api_models + name: genrm_model + genrm_responses_create_params: + max_output_tokens: 24576 + temperature: 1.0 + top_p: 0.95 + comparison_strategy: "circular" + use_golden_anchor: true + num_judges_per_comparison: 1 + use_principle: true + default_principle: "You will be given one or more evaluation criteria (rubrics).\nEvaluate both responses on EACH criterion individually first, then synthesize an overall judgment.\nCriteria:\n\n1. Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user prompt. Begin your evaluation by generating your own answer to the prompt. You must provide your answer before judging any answers. When evaluating the assistants' answers, compare both assistants' answers with your answer. You must identify and correct any mistakes or inaccurate information. Then consider if the assistant's answers are helpful, relevant, and concise. Helpful means the answer correctly responds to the prompt or follows the instructions. Note when user prompt has any ambiguity or more than one interpretation, it is more helpful and appropriate to ask for clarifications or more information from the user than providing an answer based on assumptions. Relevant means all parts of the response closely connect or are appropriate to what is being asked. Concise means the response is clear and not verbose or excessive. Then consider the creativity and novelty of the assistant's answers when needed. Finally, identify any missing important information in the assistants' answers that would be beneficial to include when responding to the user prompt." + aggregator_method: "simple_tiebreaker" + reasoning_bonus: 0.5 + answer_bonus: 0.5 + top_percentile: 0.2 + group_reasoning_length_penalty_coeff: 0.12 + group_answer_length_penalty_coeff: 0.12 + group_style_penalty_coeff: 0.0 + default_score: 3.0 + default_ranking: 3.5 + genrm_parse_retries: 1 + + # GenRM: external server managed by genrm_server_manager.sh + # Setting base_url triggers external mode — no local vLLM launch. + genrm_model: + responses_api_models: + genrm_model: + entrypoint: app.py + base_url: http://10.109.28.142:9213/v1 # Set by launch script: env.nemo_gym.genrm_model.responses_api_models.genrm_model.base_url=http://... + model: "model" # Must match --served-model-name in external vLLM server + uses_reasoning_parser: true + return_token_id_information: false + debug: true + vllm_serve_env_vars: {} + vllm_serve_kwargs: + tensor_parallel_size: 1 + data_parallel_size: 1 + pipeline_parallel_size: 1 + + lc_judge: + resources_servers: + equivalence_llm_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + math_with_judge: + resources_servers: + math_with_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + should_use_judge: true + + code_gen: + resources_servers: + code_gen: + num_processes: 2048 + unit_test_timeout_secs: 10 + debug: false + + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "/lustre/fsw/portfolios/llmservice/users/ansubramania/ultra_v3/runs/logs" + num_val_samples_to_print: 0 + wandb_enabled: false + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "grpo-ultra-v3" + name: "grpo-ultra-v3-256n" + tensorboard: {} + mlflow: + experiment_name: "grpo-ultra-v3" + run_name: "grpo-ultra-v3-256n" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + +# ============================================================================= +# Token IDs (model-specific, used by token-based penalties) +# ============================================================================= +token_ids: + eos: 2 # + think_open: 12 # + think_close: 13 # + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +penalize_duplicated_reasoning: true # reasoning content == final answer +penalize_empty_final_answer: true # last message output has empty content +penalize_eos_token: true # eos token appears in generation +penalize_malformed_think_tag: true # /<\/think> count != 1 per turn + +# Reward-zeroing penalties applied to NeMo-Gym rollout results. +reward_penalties: + penalize_duplicated_reasoning: true + penalize_empty_final_answer: true + penalize_unwanted_tokens: true + penalize_malformed_think_tag: true + # Optional model/tokenizer-specific token IDs. Add token IDs that should + # be penalized when emitted by the model. Unwanted token IDs must be + # specified explicitly; think-tag IDs are inferred when each tag encodes + # to one token. Example: {unwanted: [2], think_open: 12, think_close: 13} + token_ids: {unwanted: [2], think_open: 12, think_close: 13} \ No newline at end of file From af6b8769526a70a6c8a04910361609ab3ca8afba Mon Sep 17 00:00:00 2001 From: Daniele Molinari Date: Fri, 17 Jul 2026 01:09:56 +0200 Subject: [PATCH 31/69] PPO: add token-level logging (#3252) Signed-off-by: Daniele Molinari --- examples/configs/ppo_math_1B.yaml | 6 + nemo_rl/algorithms/ppo.py | 195 ++++++++++++++++-- tests/unit/algorithms/test_ppo.py | 158 ++++++++++++++ .../ppo_math_1B_megatron.yaml | 2 + 4 files changed, 349 insertions(+), 12 deletions(-) diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index e29c582bf47..6a554da5474 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -21,6 +21,12 @@ ppo: dynamic_sampling_max_gen_batches: 10 batch_multiplier: 1 skip_reference_policy_logprobs_calculation: true # No KL, so skip ref logprobs + # Per-token rollout dump for offline analysis: writes packed token ids / + # critic values / GAE advantages / logprobs (plus decoded text and per-sample + # context) to {logger.log_dir}/ppo_rollout_dump_step{N}.pt. Tensors reflect + # the state entering the PPO-epoch loop (before any update this step). + log_rollout_dump: false + rollout_dump_period: 1 # dump every N PPO steps (when log_rollout_dump is true) reward_shaping: enabled: true diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index a7cf6b39cbe..f69e8faea85 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -234,6 +234,15 @@ class PPOConfig(TypedDict): # masking rather than fail MasterConfig validation — e.g. a config whose # `defaults:` base predates this field. seq_logprob_error_threshold: NotRequired[float | None] + # Optional per-token rollout dump for offline analysis: writes packed + # token ids / critic values / GAE advantages / logprobs (plus decoded text + # and per-sample context) to {logger.log_dir}/ppo_rollout_dump_step{N}.pt. + # Tensors reflect the state entering the PPO-epoch loop (before any update + # this step). NotRequired (read via .get()): configs that omit it disable + # the dump rather than fail MasterConfig validation. + log_rollout_dump: NotRequired[bool] + # Dump every N PPO steps. Must be set when log_rollout_dump is true. + rollout_dump_period: NotRequired[int] # Asynchronous PPO (replay-buffer, non-colocated generation). Absent/disabled # runs synchronous PPO. async_ppo: NotRequired[AsyncPPOConfig] @@ -1394,9 +1403,7 @@ def _compute_critic_metrics(value_results: dict[str, Any]) -> dict[str, Any]: elif isinstance(v, (np.ndarray, list)): critic_metrics["critic/" + k] = np.sum(v).item() else: - raise ValueError( - f"Unknown metric for value don't know how to handle: {k}" - ) + raise ValueError(f"Unknown metric for value don't know how to handle: {k}") # Compute explained variance from sufficient statistics: # EV = 1 - Var(returns - values) / Var(returns) @@ -1406,12 +1413,138 @@ def _compute_critic_metrics(value_results: dict[str, Any]) -> dict[str, Any]: res_sq = critic_metrics.get("critic/residual_sq_mean", 0) var_returns = r_sq - r_mean**2 var_residual = res_sq - (r_mean - v_mean) ** 2 - critic_metrics["critic/explained_var"] = 1.0 - var_residual / max( - var_returns, 1e-8 - ) + critic_metrics["critic/explained_var"] = 1.0 - var_residual / max(var_returns, 1e-8) return critic_metrics +def _should_log_ppo_rollout_dump(master_config: MasterConfig, step: int) -> bool: + """Whether to write the packed per-token rollout dump for this (1-based) step.""" + if not master_config.ppo.get("log_rollout_dump"): + return False + + period = master_config.ppo.get("rollout_dump_period") + if period is None: + raise ValueError( + "ppo.log_rollout_dump is enabled but ppo.rollout_dump_period is not set." + ) + return step % max(int(period), 1) == 0 + + +def _build_ppo_rollout_dump_payload( + *, + step: int, + num_generations_per_prompt: int, + tokenizer: PreTrainedTokenizerBase, + train_data: BatchedDataDict[ClippedPGLossDataDict], + prompt_lengths: torch.Tensor, + content: list[str], + repeated_batch: BatchedDataDict[DatumSpec], +) -> dict[str, Any]: + """Build the packed per-token rollout dump payload for torch.save. + + Response tokens (``token_mask == 1``) from the whole batch are packed into + flat tensors; ``token_sample_index`` / ``token_sequence_position`` map each + packed token back to its sample and position. All tensors reflect the state + entering the PPO-epoch loop, i.e. before any policy/value update this step. + + Args: + step: 1-based PPO step number (matches the filename suffix). + num_generations_per_prompt: Generations per prompt; used to derive + prompt-group/generation indices. + tokenizer: Tokenizer used to decode per-token text fragments. + train_data: Fully populated training batch (values, advantages, + logprobs, masks). + prompt_lengths: Per-sample prompt length, shape ``[batch]``. + content: Per-sample decoded conversation text. + repeated_batch: Rollout batch; optional per-sample metadata + (``task_name``, ``idx``, ``truncated``) is copied when present. + + Returns: + Payload dict of packed per-token tensors, per-sample arrays, and + metadata, ready for ``torch.save``. + """ + token_mask = train_data["token_mask"].detach().bool().cpu() + batch_size = token_mask.shape[0] + token_count = token_mask.sum(dim=-1) + + token_coords = token_mask.nonzero(as_tuple=False) + token_sample_index = token_coords[:, 0].to(torch.int32) + token_sequence_position = token_coords[:, 1].to(torch.int32) + if token_sample_index.numel() > 0: + token_response_position = torch.cat( + [ + torch.arange(int(count), dtype=torch.int32) + for count in token_count.tolist() + ] + ) + else: + token_response_position = torch.empty(0, dtype=torch.int32) + + def _pack_float(field: str) -> torch.Tensor: + return train_data[field].detach().float().cpu()[token_mask] + + token_ids = train_data["input_ids"].detach().cpu()[token_mask] + sample_indices = torch.arange(batch_size, dtype=torch.int32) + + payload: dict[str, Any] = { + "format_version": 1, + "description": ( + "Packed response-token PPO rollout dump. Per-token tensors are " + "flattened over token_mask; entry i belongs to sample " + "token_sample_index[i] at sequence position " + "token_sequence_position[i] and describes token token_ids[i]. " + "prev_logprobs are the policy logprobs before any update this " + "step (the PPO-ratio denominator, pi_old); generation_logprobs " + "come from the rollout backend; values are the critic estimates " + "GAE consumed; returns are the critic regression targets." + ), + "step": int(step), + "num_generations_per_prompt": int(num_generations_per_prompt), + # Per-sample arrays, shape [batch]. + "sample_index": sample_indices, + "reward": train_data["rewards"].detach().float().cpu(), + "sample_loss_mask": train_data["sample_mask"].detach().float().cpu(), + "input_length": train_data["input_lengths"].detach().cpu(), + "prompt_length": prompt_lengths.detach().cpu(), + "num_response_tokens": token_count, + "content": list(content), + # Packed per-token tensors, shape [num_response_tokens_total]. + "token_sample_index": token_sample_index, + "token_sequence_position": token_sequence_position, + "token_response_position": token_response_position, + "token_ids": token_ids, + "token_text": tokenizer.convert_ids_to_tokens(token_ids.tolist()), + "values": _pack_float("values"), + "advantages": _pack_float("advantages"), + "generation_logprobs": _pack_float("generation_logprobs"), + "prev_logprobs": _pack_float("prev_logprobs"), + } + if "returns" in train_data: + payload["returns"] = _pack_float("returns") + if "reference_policy_logprobs" in train_data: + payload["reference_policy_logprobs"] = _pack_float("reference_policy_logprobs") + + if num_generations_per_prompt > 0: + payload["prompt_group_index"] = ( + sample_indices // num_generations_per_prompt + ).to(torch.int32) + payload["generation_index"] = (sample_indices % num_generations_per_prompt).to( + torch.int32 + ) + + for key in ("task_name", "idx", "truncated"): + if key in repeated_batch: + value = repeated_batch[key] + if isinstance(value, torch.Tensor): + value = value.detach().cpu() + else: + value = list(value) + if len(value) == batch_size: + payload[key] = value + + return payload + + # =============================================================================== # Training & Validation # =============================================================================== @@ -1620,9 +1753,7 @@ def ppo_train( greedy=False, effort_config=_get_effort_config(master_config), reward_penalty_config=master_config.reward_penalties, - thinking_tags=get_nemo_gym_thinking_tags( - master_config.env - ), + thinking_tags=get_nemo_gym_thinking_tags(master_config.env), ) repeated_batch = nemo_gym_rollout_result.final_batch rollout_metrics = nemo_gym_rollout_result.rollout_metrics @@ -1870,6 +2001,27 @@ def ppo_train( if returns is not None: train_data["returns"] = returns + if _should_log_ppo_rollout_dump(master_config, total_steps + 1): + with timer.time("rollout_dump"): + rollout_dump_path = os.path.join( + logger.base_log_dir, + f"ppo_rollout_dump_step{total_steps + 1}.pt", + ) + rollout_dump_payload = _build_ppo_rollout_dump_payload( + step=total_steps + 1, + num_generations_per_prompt=master_config.ppo[ + "num_generations_per_prompt" + ], + tokenizer=tokenizer, + train_data=train_data, + prompt_lengths=repeated_batch["length"], + content=flat_messages["content"], + repeated_batch=repeated_batch, + ) + torch.save(rollout_dump_payload, rollout_dump_path) + print(f" 📝 Dumped rollout data to {rollout_dump_path}") + del rollout_dump_payload + # PPO: Multiple training steps per rollout memory_tracker.snapshot_start_of_stage("Policy train", dir()) for step in range(ppo_epochs): @@ -3286,6 +3438,27 @@ def _sample_max_age(s: int) -> int: log_data, f"train_data_step{step + 1}.jsonl" ) del log_data + + if _should_log_ppo_rollout_dump(master_config, step + 1): + with timer.time("rollout_dump"): + rollout_dump_path = os.path.join( + logger.base_log_dir, + f"ppo_rollout_dump_step{step + 1}.pt", + ) + rollout_dump_payload = _build_ppo_rollout_dump_payload( + step=step + 1, + num_generations_per_prompt=master_config.ppo[ + "num_generations_per_prompt" + ], + tokenizer=tokenizer, + train_data=train_data, + prompt_lengths=repeated_batch["length"], + content=flat_messages_content, + repeated_batch=repeated_batch, + ) + torch.save(rollout_dump_payload, rollout_dump_path) + print(f" 📝 Dumped rollout data to {rollout_dump_path}") + del rollout_dump_payload del flat_messages_content timing_metrics: dict[str, float] = timer.get_timing_metrics( @@ -3319,9 +3492,7 @@ def _sample_max_age(s: int) -> int: print("\n📊 Training Results:") if "loss" in metrics: print(f" • Loss: {metrics['loss']:.4f}") - print( - f" • Generation KL Error: {metrics.get('gen_kl_error', 'N/A')}" - ) + print(f" • Generation KL Error: {metrics.get('gen_kl_error', 'N/A')}") else: print(" • (critic warmup: policy not trained this step)") if "critic/loss" in metrics: diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 3ecb22f57a9..4e593bbd471 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -1317,3 +1317,161 @@ def test_resolve_resume_optimizer_path(tmp_path): # DTensor without a separate optimizer/ dir -> None (don't misread weights as # an optimizer path for a backend that stores it separately). assert _resolve_resume_optimizer_path(missing_optim, weights, dtensor) is None + + +# --------------------------------------------------------------------------- +# Per-token rollout dump (ppo.log_rollout_dump) +# --------------------------------------------------------------------------- +class _FakeDumpTokenizer: + def convert_ids_to_tokens(self, ids: list[int]) -> list[str]: + return [f"" for i in ids] + + +def test_should_log_ppo_rollout_dump_gate(): + from types import SimpleNamespace + + from nemo_rl.algorithms.ppo import _should_log_ppo_rollout_dump + + # Absent / disabled -> never dump. + assert not _should_log_ppo_rollout_dump(SimpleNamespace(ppo={}), 1) + assert not _should_log_ppo_rollout_dump( + SimpleNamespace(ppo={"log_rollout_dump": False, "rollout_dump_period": 1}), 1 + ) + + # Enabled without a period -> loud failure, not a hidden default. + with pytest.raises(ValueError, match="rollout_dump_period"): + _should_log_ppo_rollout_dump(SimpleNamespace(ppo={"log_rollout_dump": True}), 1) + + # Enabled with period 2 -> fires on even (1-based) steps only. + cfg = SimpleNamespace(ppo={"log_rollout_dump": True, "rollout_dump_period": 2}) + assert not _should_log_ppo_rollout_dump(cfg, 1) + assert _should_log_ppo_rollout_dump(cfg, 2) + assert not _should_log_ppo_rollout_dump(cfg, 3) + assert _should_log_ppo_rollout_dump(cfg, 4) + + +def test_build_ppo_rollout_dump_payload_packing(tmp_path): + from nemo_rl.algorithms.ppo import _build_ppo_rollout_dump_payload + + # 2 samples x 5 positions; sample 0 has 3 response tokens, sample 1 has 2. + token_mask = torch.tensor([[0, 0, 1, 1, 1], [0, 0, 0, 1, 1]], dtype=torch.float32) + input_ids = torch.arange(10).reshape(2, 5) + values = torch.arange(10, dtype=torch.float32).reshape(2, 5) * 0.1 + advantages = torch.arange(10, dtype=torch.float32).reshape(2, 5) * -1.0 + returns = torch.arange(10, dtype=torch.float32).reshape(2, 5) * 2.0 + generation_logprobs = torch.arange(10, dtype=torch.float32).reshape(2, 5) - 20.0 + prev_logprobs = torch.arange(10, dtype=torch.float32).reshape(2, 5) - 30.0 + reference_policy_logprobs = ( + torch.arange(10, dtype=torch.float32).reshape(2, 5) - 40.0 + ) + + train_data = BatchedDataDict( + { + "input_ids": input_ids, + "input_lengths": torch.tensor([5, 5]), + "token_mask": token_mask, + "sample_mask": torch.tensor([1.0, 0.0]), + "rewards": torch.tensor([1.0, -1.0]), + "values": values, + "advantages": advantages, + "returns": returns, + "generation_logprobs": generation_logprobs, + "prev_logprobs": prev_logprobs, + "reference_policy_logprobs": reference_policy_logprobs, + } + ) + repeated_batch = BatchedDataDict( + { + "length": torch.tensor([2, 3]), + "task_name": ["math", "math"], + "idx": torch.tensor([7, 8]), + "truncated": [False, True], + } + ) + + payload = _build_ppo_rollout_dump_payload( + step=3, + num_generations_per_prompt=2, + tokenizer=_FakeDumpTokenizer(), + train_data=train_data, + prompt_lengths=repeated_batch["length"], + content=["convo A", "convo B"], + repeated_batch=repeated_batch, + ) + + assert payload["format_version"] == 1 + assert payload["step"] == 3 + + # Packed token coordinates. + mask = token_mask.bool() + assert payload["token_sample_index"].tolist() == [0, 0, 0, 1, 1] + assert payload["token_sequence_position"].tolist() == [2, 3, 4, 3, 4] + assert payload["token_response_position"].tolist() == [0, 1, 2, 0, 1] + + # Per-token tensors align with the masked positions. + assert torch.equal(payload["token_ids"], input_ids[mask]) + assert torch.equal(payload["values"], values[mask]) + assert torch.equal(payload["advantages"], advantages[mask]) + assert torch.equal(payload["returns"], returns[mask]) + assert torch.equal(payload["generation_logprobs"], generation_logprobs[mask]) + assert torch.equal(payload["prev_logprobs"], prev_logprobs[mask]) + assert torch.equal( + payload["reference_policy_logprobs"], reference_policy_logprobs[mask] + ) + assert payload["token_text"] == [f"" for i in input_ids[mask].tolist()] + + # Per-sample context. + assert payload["reward"].tolist() == [1.0, -1.0] + assert payload["sample_loss_mask"].tolist() == [1.0, 0.0] + assert payload["input_length"].tolist() == [5, 5] + assert payload["prompt_length"].tolist() == [2, 3] + assert payload["num_response_tokens"].tolist() == [3, 2] + assert payload["content"] == ["convo A", "convo B"] + assert payload["prompt_group_index"].tolist() == [0, 0] + assert payload["generation_index"].tolist() == [0, 1] + assert payload["task_name"] == ["math", "math"] + assert payload["idx"].tolist() == [7, 8] + assert payload["truncated"] == [False, True] + + # Round-trips through torch.save/load (what the dump actually does). + dump_path = tmp_path / "ppo_rollout_dump_step3.pt" + torch.save(payload, dump_path) + reloaded = torch.load(dump_path, weights_only=False) + assert torch.equal(reloaded["advantages"], payload["advantages"]) + assert reloaded["token_text"] == payload["token_text"] + + +def test_build_ppo_rollout_dump_payload_optional_fields_absent(): + from nemo_rl.algorithms.ppo import _build_ppo_rollout_dump_payload + + train_data = BatchedDataDict( + { + "input_ids": torch.tensor([[1, 2, 3]]), + "input_lengths": torch.tensor([3]), + "token_mask": torch.tensor([[0.0, 1.0, 1.0]]), + "sample_mask": torch.tensor([1.0]), + "rewards": torch.tensor([0.5]), + "values": torch.zeros(1, 3), + "advantages": torch.zeros(1, 3), + "generation_logprobs": torch.zeros(1, 3), + "prev_logprobs": torch.zeros(1, 3), + } + ) + + payload = _build_ppo_rollout_dump_payload( + step=1, + num_generations_per_prompt=1, + tokenizer=_FakeDumpTokenizer(), + train_data=train_data, + prompt_lengths=torch.tensor([1]), + content=["convo"], + repeated_batch=BatchedDataDict({"length": torch.tensor([1])}), + ) + + # No returns / reference logprobs / batch metadata -> keys omitted, not None. + assert "returns" not in payload + assert "reference_policy_logprobs" not in payload + assert "task_name" not in payload + assert "idx" not in payload + assert "truncated" not in payload + assert payload["num_response_tokens"].tolist() == [2] diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index d596091bac7..66c2284e040 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -21,6 +21,8 @@ ppo: dynamic_sampling_max_gen_batches: 10 batch_multiplier: 1 skip_reference_policy_logprobs_calculation: true # No KL, so skip ref logprobs + log_rollout_dump: false + rollout_dump_period: 1 reward_shaping: enabled: true From cb182244d6fdbf88f4c0d2ac470d79fe17019b15 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 16 Jul 2026 16:12:09 -0700 Subject: [PATCH 32/69] privileged critic Signed-off-by: Jiaqi Zeng --- examples/configs/ppo_math_1B.yaml | 2 +- nemo_rl/algorithms/ppo.py | 155 +++++++++++- nemo_rl/algorithms/privileged_critic.py | 239 ++++++++++++++++++ .../unit/algorithms/test_privileged_critic.py | 189 ++++++++++++++ 4 files changed, 581 insertions(+), 4 deletions(-) create mode 100644 nemo_rl/algorithms/privileged_critic.py create mode 100644 tests/unit/algorithms/test_privileged_critic.py diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index e29c582bf47..f16653841e5 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -104,7 +104,7 @@ checkpointing: checkpoint_dir: "results/ppo_dapo" metric_name: "val:accuracy" higher_is_better: true - keep_top_k: 5 + keep_top_k: 500000 save_period: 10 checkpoint_must_save_by: null model_save_format: "safetensors" diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index a7cf6b39cbe..ca619621192 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -72,6 +72,10 @@ get_keys_from_message_log, ) from nemo_rl.data.utils import load_dataloader_state +from nemo_rl.algorithms.privileged_critic import ( + build_privileged_value_inputs, + remap_by_response_mask, +) from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env from nemo_rl.distributed.virtual_cluster import ( @@ -354,6 +358,49 @@ def setup( # Extract individual configs for easier access policy_config = master_config.policy value_config = master_config.value + + # Privileged (answer-conditioned) critic is currently wired for synchronous PPO + # only; fail loudly rather than silently falling back to a blind critic on the + # async path (which does not build the answer-augmented value batch). + _privileged_critic = value_config.get("privileged_critic") + if _privileged_critic is not None and _privileged_critic.get("enabled"): + assert not ( + "async_ppo" in master_config.ppo + and master_config.ppo["async_ppo"].get("enabled") + ), ( + "value.privileged_critic is supported only for synchronous PPO for now " + "(the async_ppo value stage does not build the answer-augmented batch). " + "Disable async_ppo or value.privileged_critic.enabled." + ) + # The critic scores [prompt + answer + response], which is longer than the + # policy's [prompt + response] by up to max_answer_tokens (+ grader-note + # template overhead). Give the VALUE model that much headroom so every + # answer-augmented sample fits its sequence-packing bins. Only ever RAISES the + # budget; if the user already configured enough, this is a no-op. + _needed = ( + policy_config["max_total_sequence_length"] + + int(_privileged_critic.get("max_answer_tokens", 256) or 0) + + 128 # grader-note template + chat re-render slack + ) + if value_config["max_total_sequence_length"] < _needed: + print( + " ↑ privileged critic: raising value.max_total_sequence_length " + f"{value_config['max_total_sequence_length']} -> {_needed}", + flush=True, + ) + value_config["max_total_sequence_length"] = _needed + for _bcfg_key in ("sequence_packing", "dynamic_batching"): + _bcfg = value_config.get(_bcfg_key) or {} + if not _bcfg.get("enabled"): + continue + for _tok_key in ("train_mb_tokens", "logprob_mb_tokens"): + if _bcfg.get(_tok_key) is not None and _bcfg[_tok_key] < _needed: + print( + f" ↑ privileged critic: raising value.{_bcfg_key}.{_tok_key} " + f"{_bcfg[_tok_key]} -> {_needed}", + flush=True, + ) + _bcfg[_tok_key] = _needed generation_config = master_config.policy["generation"] env_configs = master_config.env loss_config: ClippedPGLossConfig = master_config.loss_fn @@ -1412,6 +1459,48 @@ def _compute_critic_metrics(value_results: dict[str, Any]) -> dict[str, Any]: return critic_metrics +def _positional_value_metrics( + values: torch.Tensor, + returns: torch.Tensor, + token_mask: torch.Tensor, + n_bins: int = 3, +) -> dict[str, float]: + """Explained variance of V(s_t) vs the GAE return, bucketed by RELATIVE position + within each response (early / mid / late thirds). + + Answers "where does the critic do its job well?" — expected to rise early->late, + since near the end the outcome is nearly determined. Comparing a privileged + (answer-conditioned) run against a blind one, the golden answer should lift the + EARLY/MID buckets most (a large early-bucket gap = the privileged-critic mechanism + working; a flat gap = the answer isn't buying sharper early credit). Cheap: + driver-side tensor ops on tensors already in ``train_data``. + """ + mask = token_mask.bool() + if int(mask.sum()) < 2: + return {} + # 0-based index of each response token within its sample's response, scaled to [0,1) + resp_len = mask.sum(dim=1, keepdim=True).clamp(min=1) + rel = (torch.cumsum(mask.long(), dim=1) - 1).float() / resp_len.float() + names = ["early", "mid", "late"] if n_bins == 3 else [f"bin{i}" for i in range(n_bins)] + out: dict[str, float] = {} + for i in range(n_bins): + lo, hi = i / n_bins, (i + 1) / n_bins + upper = (rel < hi) if i < n_bins - 1 else (rel <= 1.0) + bmask = mask & (rel >= lo) & upper + n = int(bmask.sum()) + if n < 2: + continue + v = values[bmask].float() + r = returns[bmask].float() + var_r = r.var(unbiased=False) + ev = (1.0 - (r - v).var(unbiased=False) / var_r).item() if var_r > 1e-8 else 0.0 + out[f"critic/ev_{names[i]}"] = ev + out[f"critic/abs_err_{names[i]}"] = (r - v).abs().mean().item() + out[f"critic/mean_v_{names[i]}"] = v.mean().item() + out[f"critic/n_tokens_{names[i]}"] = float(n) + return out + + # =============================================================================== # Training & Validation # =============================================================================== @@ -1755,9 +1844,41 @@ def ppo_train( memory_tracker.snapshot_start_of_stage("Value inference", dir()) print("▶ Computing values...", flush=True) with timer.time("value_inference"): + # Privileged (answer-conditioned) critic: the value model scores an + # answer-augmented sequence [prompt(+gold), response] (built at the + # message level, response tokens verbatim) and the response-position + # values are mapped back into the original layout. Policy/GAE/value + # workers are untouched. See nemo_rl/algorithms/privileged_critic.py. + privileged_critic_cfg = master_config.value.get("privileged_critic") + if privileged_critic_cfg is not None and not privileged_critic_cfg.get( + "enabled" + ): + privileged_critic_cfg = None + critic_batch = None + value_model.prepare_for_inference() - values = value_model.get_values(train_data) - train_data["values"] = values["values"].squeeze(-1) + if privileged_critic_cfg is not None: + critic_batch = build_privileged_value_inputs( + repeated_batch, + tokenizer, + privileged_critic_cfg, + make_seq_len_divisible_by=master_config.policy[ + "make_sequence_length_divisible_by" + ], + ) + vals_aug = value_model.get_values(critic_batch)["values"].squeeze( + -1 + ) + # keep aug-layout old values around for the (optional) value clip + critic_batch["values"] = vals_aug + train_data["values"] = remap_by_response_mask( + vals_aug, + critic_batch["token_mask"], + train_data["token_mask"], + ) + else: + values = value_model.get_values(train_data) + train_data["values"] = values["values"].squeeze(-1) value_model.finish_inference() print( @@ -1869,6 +1990,15 @@ def ppo_train( train_data["advantages"] = advantages if returns is not None: train_data["returns"] = returns + # Privileged critic trains on the answer-augmented sequence: + # scatter the (original-layout) GAE returns onto the augmented + # response positions. Same response tokens => exact mapping. + if critic_batch is not None: + critic_batch["returns"] = remap_by_response_mask( + returns, + train_data["token_mask"], + critic_batch["token_mask"], + ) # PPO: Multiple training steps per rollout memory_tracker.snapshot_start_of_stage("Policy train", dir()) @@ -1884,8 +2014,16 @@ def ppo_train( with timer.time("value_training"): print("▶ Training value...", flush=True) + # Privileged critic: train on the answer-augmented batch + # (returns already scattered onto its response positions above; + # values=aug old-values for the optional value clip). + if critic_batch is not None: + critic_batch["sample_mask"] = train_data["sample_mask"] + value_train_batch = critic_batch + else: + value_train_batch = train_data value_results = value_model.train( - train_data, + value_train_batch, value_loss_fn, timer=timer, ) @@ -2006,6 +2144,17 @@ def ppo_train( # Extract critic metrics from value training results if value_results is not None: metrics.update(_compute_critic_metrics(value_results)) + # Positional critic-quality diagnostic (early/mid/late tokens): how + # well V predicts the return along the trajectory, and — for the + # privileged critic — where the golden answer sharpens it. + if "values" in train_data and "returns" in train_data: + metrics.update( + _positional_value_metrics( + train_data["values"], + train_data["returns"], + train_data["token_mask"], + ) + ) metrics.update( { "reward": rewards.numpy(), diff --git a/nemo_rl/algorithms/privileged_critic.py b/nemo_rl/algorithms/privileged_critic.py new file mode 100644 index 00000000000..9343f5fa76f --- /dev/null +++ b/nemo_rl/algorithms/privileged_critic.py @@ -0,0 +1,239 @@ +# Copyright (c) 2025, 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. +"""Privileged (answer-conditioned) critic for PPO. + +An *asymmetric actor-critic*: the value model additionally sees privileged +information the policy never sees — the golden/reference answer — so its +per-token value estimates (and therefore the GAE advantages) are sharper. The +policy stays blind and the privileged conditioning is training-only. See +``privileged_critic_proposal.md`` for the rationale, the unbiasedness argument +(``V(h, z)`` form; the answer is action-independent and prompt-determined), and +the safety rules. + +Design (kept deliberately small so GAE, the value workers and the policy are all +untouched): + +* The critic scores an **answer-augmented** sequence ``x' = [prompt(+gold), + response]`` that is constructed at the **message level** and re-rendered with + the chat template, so it stays a well-formed conversation the value model (a + chat model + value head) can process. Crucially, the **RESPONSE token-ids are + the verbatim generated tokens** — we never re-tokenize the response, so the + value at each response position corresponds to exactly the state the policy + produced. Only the prompt region differs (it now contains the answer). +* Values / returns are moved between the augmented layout ``x'`` and the original + layout ``x = [prompt, response]`` by the response mask. Because the response + tokens are identical, the two masks select the same tokens with equal per-row + counts, so the mapping is an exact masked scatter (asserted). + +SAFETY (enforced by construction, per the theory): +* The value model is a **separate worker** from the policy (PPO always builds one + via ``init_value``) — no shared parameters, so the answer cannot leak into the + policy through shared features. +* The answer only ever enters the critic's advantage **magnitude** (a baseline), + never the policy path or the reward — the outcome reward alone sets direction. + +Prompt construction is the delicate part: do NOT splice raw answer tokens between +the prompt and response token-ids (that yields a malformed conversation and +out-of-distribution value estimates). Instead fold the answer into a prompt turn +and re-render with the chat template, then append the verbatim response tokens. +""" + +from __future__ import annotations + +from typing import Any, Optional + +import torch + +from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + +# The answer is framed as a clearly-delimited grader note (not the solver's turn), +# so it is well-formed and less trivially string-matchable than a bare "\boxed{}". +DEFAULT_TEMPLATE = ( + "\n\n[Reference answer, provided to the grader only and NOT visible to the " + "solver — use it to judge whether the assistant's solution is on track: {answer}]" +) + + +def _truncate_answer( + gold: str, tokenizer: Any, max_answer_tokens: Optional[int] +) -> str: + """Cap the reference answer to ``max_answer_tokens`` (keeps sequences bounded + for full-solution exposure; a no-op for short final answers).""" + if not gold or not max_answer_tokens or max_answer_tokens <= 0: + return gold or "" + ids = tokenizer.encode(gold, add_special_tokens=False) + if len(ids) <= max_answer_tokens: + return gold + return tokenizer.decode(ids[:max_answer_tokens]) + + +def _inject_answer( + prompt_msgs: list[dict], gold: str, placement: str, template: str +) -> list[dict]: + """Fold the golden answer into the (text) prompt messages as a well-formed turn. + + placement: ``user_suffix`` (default, always supported) | ``user_prefix`` | + ``system``. + """ + note = template.format(answer=gold) + msgs = [{"role": m["role"], "content": m["content"]} for m in prompt_msgs] + if placement == "system": + return [{"role": "system", "content": note.lstrip("\n")}] + msgs + # find the last user turn to attach the note to + user_idxs = [i for i, m in enumerate(msgs) if m["role"] == "user"] + if not user_idxs: # no user turn to attach to -> fall back to a system note + return [{"role": "system", "content": note.lstrip("\n")}] + msgs + i = user_idxs[-1] + if placement == "user_prefix": + msgs[i]["content"] = note.lstrip("\n") + "\n\n" + msgs[i]["content"] + else: # user_suffix (default) + msgs[i]["content"] = msgs[i]["content"] + note + return msgs + + +def build_privileged_value_inputs( + repeated_batch: BatchedDataDict, + tokenizer: Any, + pcfg: dict[str, Any], + make_seq_len_divisible_by: int = 1, +) -> BatchedDataDict: + """Build the answer-augmented critic input batch, row-aligned with ``train_data``. + + Returns a ``BatchedDataDict`` with ``input_ids`` / ``input_lengths`` / + ``token_mask``, where ``token_mask`` marks exactly the verbatim response tokens. + + ``setup()`` raises the value model's sequence budget by ``max_answer_tokens`` (+ + margin) when this is enabled, so answer-augmented sequences fit its packing bins. + """ + placement = pcfg.get("placement", "user_suffix") + template = pcfg.get("template", DEFAULT_TEMPLATE) + max_answer_tokens = pcfg.get("max_answer_tokens", 256) + + message_logs = repeated_batch["message_log"] + env_infos = repeated_batch.get("extra_env_info", None) + if env_infos is None: + env_infos = [None] * len(message_logs) + + critic_message_logs: list[list[dict]] = [] + for msgs, info in zip(message_logs, env_infos): + # Split into prompt (leading non-assistant turns) and response (the + # assistant-generated turn[s]). Single-turn RLVR: the rollout is + # [user, assistant, environment-feedback]; the trailing environment turn is + # masked out AND comes after the response, so it can't affect the causal value + # at response positions — drop it. Reject genuinely interleaved multi-turn + # (an environment turn wedged BETWEEN assistant turns). + assistant_idxs = [i for i, m in enumerate(msgs) if m["role"] == "assistant"] + assert assistant_idxs, ( + "privileged critic: no assistant/response message found in the rollout" + ) + first_a, last_a = assistant_idxs[0], assistant_idxs[-1] + assert all(m["role"] == "assistant" for m in msgs[first_a : last_a + 1]), ( + "privileged critic currently supports single-turn rollouts (an environment " + "turn is interleaved between assistant turns — multi-turn not supported)." + ) + prompt_msgs = msgs[:first_a] + response_msgs = msgs[first_a : last_a + 1] + assert not any(("images" in m or "videos" in m) for m in msgs), ( + "privileged critic supports text-only rollouts." + ) + + gold = "" + if info is not None: + gold = _truncate_answer( + info.get("ground_truth", "") or "", tokenizer, max_answer_tokens + ) + + # Fold the answer into the prompt, render to a string with the chat template, + # then tokenize — the SAME two-step path the data processor uses to build the + # original prompt (apply_chat_template(tokenize=False) -> str; tokenizer() -> + # ids). Matching it keeps the critic's prompt format consistent with the + # policy's and avoids tokenizer-specific quirks of tokenize=True. + aug_prompt_text_msgs = _inject_answer(prompt_msgs, gold, placement, template) + rendered = tokenizer.apply_chat_template( + aug_prompt_text_msgs, + tokenize=False, + add_generation_prompt=True, + add_special_tokens=False, + ) + prompt_ids = tokenizer( + rendered, return_tensors="pt", add_special_tokens=False + )["input_ids"][0].to(dtype=torch.long) + + # ...then append the VERBATIM generated response token-ids. + critic_msgs: list[dict] = [ + { + "role": "user", + "content": "", # token_ids provided; content is unused by the flattener + "token_ids": prompt_ids, + "token_loss_mask": torch.zeros_like(prompt_ids), + } + ] + for rm in response_msgs: + rid = torch.as_tensor(rm["token_ids"], dtype=torch.long).flatten() + critic_msgs.append( + { + "role": "assistant", + "content": "", + "token_ids": rid, + "token_loss_mask": torch.ones_like(rid), + } + ) + critic_message_logs.append(critic_msgs) + + flat, input_lengths = batched_message_log_to_flat_message( + critic_message_logs, + pad_value_dict={"token_ids": tokenizer.pad_token_id}, + make_sequence_length_divisible_by=make_seq_len_divisible_by, + ) + return BatchedDataDict( + { + "input_ids": flat["token_ids"], + "input_lengths": input_lengths, + "token_mask": flat["token_loss_mask"], + } + ) + + +def remap_by_response_mask( + src: torch.Tensor, src_mask: torch.Tensor, dst_mask: torch.Tensor +) -> torch.Tensor: + """Move per-token values between two layouts that mark the SAME response tokens. + + ``src[src_mask] -> dst[dst_mask]``; ``dst`` is zero elsewhere (GAE's + carry-forward masking ignores those non-response positions). Requires equal + per-row response counts — asserted, since a mismatch means the response tokens + were not preserved verbatim (a construction bug). + """ + dev = dst_mask.device + src = src.to(dev) + src_mask_b = src_mask.to(dev).bool() + dst_mask_b = dst_mask.bool() + + src_counts = src_mask_b.sum(dim=-1) + dst_counts = dst_mask_b.sum(dim=-1) + if not torch.equal(src_counts, dst_counts): + raise AssertionError( + "privileged critic: per-row response-token counts differ between the " + "augmented and original layouts — the response tokens were not preserved " + "verbatim.\n" + f" augmented per-row counts: {src_counts.tolist()}\n" + f" original per-row counts: {dst_counts.tolist()}" + ) + + dst = torch.zeros( + dst_mask_b.shape[0], dst_mask_b.shape[1], dtype=src.dtype, device=dev + ) + dst[dst_mask_b] = src[src_mask_b] + return dst diff --git a/tests/unit/algorithms/test_privileged_critic.py b/tests/unit/algorithms/test_privileged_critic.py new file mode 100644 index 00000000000..fbbeaa3c697 --- /dev/null +++ b/tests/unit/algorithms/test_privileged_critic.py @@ -0,0 +1,189 @@ +# Copyright (c) 2025, 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. +"""Unit tests for the privileged (answer-conditioned) critic helpers.""" + +import pytest +import torch + +from nemo_rl.algorithms.privileged_critic import ( + build_privileged_value_inputs, + remap_by_response_mask, +) + + +# --------------------------------------------------------------------------- # +# remap_by_response_mask: the exact scatter/gather between the augmented and +# original layouts. This is the one piece of index logic that must be correct. +# --------------------------------------------------------------------------- # +def test_remap_moves_response_values_and_zeros_elsewhere(): + B, S_dst, S_src = 2, 7, 9 + dst_mask = torch.zeros(B, S_dst) + dst_mask[0, 3:6] = 1 # response len 3 + dst_mask[1, 4:6] = 1 # response len 2 + # augmented layout: response pushed right by the (longer) answer-augmented prompt + src_mask = torch.zeros(B, S_src) + src_mask[0, 6:9] = 1 # response len 3 + src_mask[1, 7:9] = 1 # response len 2 + src = torch.zeros(B, S_src) + src[0, 6:9] = torch.tensor([1.0, 2.0, 3.0]) + src[1, 7:9] = torch.tensor([4.0, 5.0]) + + dst = remap_by_response_mask(src, src_mask, dst_mask) + + assert dst.shape == (B, S_dst) + assert torch.equal(dst[0, 3:6], torch.tensor([1.0, 2.0, 3.0])) + assert torch.equal(dst[1, 4:6], torch.tensor([4.0, 5.0])) + # everything outside the response mask is exactly zero (GAE masks these anyway) + assert dst[dst_mask == 0].abs().sum() == 0 + + +def test_remap_is_invertible_roundtrip(): + B, S_dst, S_src = 3, 6, 8 + dst_mask = torch.zeros(B, S_dst) + src_mask = torch.zeros(B, S_src) + for b, (a, n) in enumerate([(2, 3), (1, 2), (3, 2)]): + dst_mask[b, a : a + n] = 1 + src_mask[b, S_src - n : S_src] = 1 # response is the tail in the aug layout + src = torch.randn(B, S_src) * src_mask + + orig = remap_by_response_mask(src, src_mask, dst_mask) + back = remap_by_response_mask(orig, dst_mask, src_mask) + assert torch.equal(back * src_mask, src * src_mask) + + +def test_remap_rejects_mismatched_response_counts(): + # A per-row count mismatch means the response tokens were NOT preserved verbatim. + dst_mask = torch.zeros(1, 5) + dst_mask[0, 2:5] = 1 # count 3 + src_mask = torch.zeros(1, 5) + src_mask[0, 3:5] = 1 # count 2 + with pytest.raises(AssertionError, match="response-token counts differ"): + remap_by_response_mask(torch.zeros(1, 5), src_mask, dst_mask) + + +# --------------------------------------------------------------------------- # +# build_privileged_value_inputs: the careful message-level prompt construction. +# The invariant that matters most: the RESPONSE tokens are byte-identical to the +# generated tokens, and the answer only lengthens the PROMPT region. +# --------------------------------------------------------------------------- # +class _FakeTokenizer: + """Mirrors the HF pattern the code uses: apply_chat_template(tokenize=False) + renders to a STRING, then __call__ tokenizes it. Prompt ids land in [200,249] + so they never collide with the test response/env ids (50..91).""" + + pad_token_id = 0 + + def encode(self, text, add_special_tokens=False): + return [ord(c) % 50 + 200 for c in text] + + def decode(self, ids): + return "".join(chr((int(i) - 200) % 50 + 65) for i in ids) + + def apply_chat_template( + self, + messages, + tokenize=False, + add_generation_prompt=False, + add_special_tokens=False, + ): + s = "".join(f"<{m['role']}>{m['content']}" for m in messages) + if add_generation_prompt: + s += "" + return s + + def __call__(self, text, return_tensors=None, add_special_tokens=False): + return {"input_ids": torch.tensor([self.encode(text)], dtype=torch.long)} + + +def _make_repeated_batch(response_ids, gold): + msgs = [ + {"role": "user", "content": "What is 2+2?", "token_ids": torch.tensor([50, 51, 52])}, + {"role": "assistant", "content": "resp", "token_ids": torch.tensor(response_ids)}, + ] + return {"message_log": [msgs], "extra_env_info": [{"ground_truth": gold}]} + + +def test_build_preserves_response_tokens_verbatim(): + tok = _FakeTokenizer() + response_ids = [77, 78, 79, 80] + batch = _make_repeated_batch(response_ids, gold="42") + pcfg = {"enabled": True, "placement": "user_suffix", "max_answer_tokens": 256} + + critic = build_privileged_value_inputs(batch, tok, pcfg) + + ids = critic["input_ids"][0] + mask = critic["token_mask"][0].bool() + # response region == the verbatim generated tokens, in order + assert ids[mask].tolist() == response_ids + # response count matches + assert int(mask.sum()) == len(response_ids) + # answer lives strictly in the prompt (mask==0) region + assert mask[: -len(response_ids)].sum() == 0 + + +def test_build_drops_trailing_environment_turn(): + # The math rollout's message_log is [user, assistant, environment-feedback]. + # The trailing environment turn is masked and comes AFTER the response, so it must + # be dropped from the critic input (not rejected, not included). + tok = _FakeTokenizer() + response_ids = [77, 78, 79] + msgs = [ + {"role": "user", "content": "What is 2+2?", "token_ids": torch.tensor([50, 51, 52])}, + {"role": "assistant", "content": "resp", "token_ids": torch.tensor(response_ids)}, + {"role": "environment", "content": "feedback", "token_ids": torch.tensor([90, 91])}, + ] + batch = {"message_log": [msgs], "extra_env_info": [{"ground_truth": "42"}]} + critic = build_privileged_value_inputs(batch, tok, {"enabled": True}) + + ids = critic["input_ids"][0] + mask = critic["token_mask"][0].bool() + assert ids[mask].tolist() == response_ids # response verbatim + assert int(mask.sum()) == len(response_ids) # count matches train_data's mask + assert 90 not in ids.tolist() and 91 not in ids.tolist() # env turn dropped entirely + + +def test_build_answer_lengthens_only_the_prompt(): + tok = _FakeTokenizer() + response_ids = [77, 78, 79, 80] + pcfg = {"enabled": True, "placement": "user_suffix", "max_answer_tokens": 256} + + with_answer = build_privileged_value_inputs( + _make_repeated_batch(response_ids, gold="the answer is 42 with reasoning"), + tok, + pcfg, + ) + no_answer = build_privileged_value_inputs( + _make_repeated_batch(response_ids, gold=""), tok, pcfg + ) + # a longer reference answer only grows the prompt; response count is unchanged + assert with_answer["input_lengths"][0] > no_answer["input_lengths"][0] + assert int(with_answer["token_mask"][0].sum()) == len(response_ids) + assert int(no_answer["token_mask"][0].sum()) == len(response_ids) + + +def test_build_truncates_long_answer(): + tok = _FakeTokenizer() + response_ids = [77, 78] + long_gold = "x" * 500 + short = build_privileged_value_inputs( + _make_repeated_batch(response_ids, long_gold), + tok, + {"enabled": True, "max_answer_tokens": 8}, + ) + untrunc = build_privileged_value_inputs( + _make_repeated_batch(response_ids, long_gold), + tok, + {"enabled": True, "max_answer_tokens": 500}, + ) + assert short["input_lengths"][0] < untrunc["input_lengths"][0] From d93381a3bc0d02cab12087569b487094b41a1a5f Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 16 Jul 2026 16:16:15 -0700 Subject: [PATCH 33/69] fix ray launching on cmh Signed-off-by: Jiaqi Zeng --- ray.sub | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/ray.sub b/ray.sub index b575ba3276b..a5b099d32ce 100644 --- a/ray.sub +++ b/ray.sub @@ -333,11 +333,21 @@ for node in $nodes; do # Try multiple methods to get IP address - ENHANCED VERSION v2.0 echo "[DEBUG] Resolving hostname: $node using enhanced resolution methods" ip_address="" - + + # Method 0: Ask the node itself. DNS records on this cluster can lag node + # re-addressing by several minutes after boot, so any DNS-based answer at + # job start may be stale (jobs 21344/21345 hung on ray start because the + # head connected to a stale DNS IP). The node's own view is authoritative. + echo "[DEBUG] Method 0: srun hostname -I on the node" + ip_address=$(timeout 60 srun --overlap -N1 -n1 -w "$node" hostname -I 2>/dev/null | awk '{ print $1 }' || true) + echo "[DEBUG] srun result: '$ip_address'" + # Method 1: Try host command - echo "[DEBUG] Method 1: host command" - ip_address=$(host $node 2>/dev/null | awk '/has address/ { print $4 }' | head -1 || true) - echo "[DEBUG] host result: '$ip_address'" + if [[ -z "$ip_address" ]]; then + echo "[DEBUG] Method 1: host command" + ip_address=$(host $node 2>/dev/null | awk '/has address/ { print $4 }' | head -1 || true) + echo "[DEBUG] host result: '$ip_address'" + fi # Method 2: If host fails, try getent if [[ -z "$ip_address" ]]; then From dada509dc067dfe36840cf54b8d7e90b13713197 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Mon, 20 Jul 2026 15:59:51 -0700 Subject: [PATCH 34/69] fix genrm index Signed-off-by: Jiaqi Zeng --- nemo_rl/algorithms/async_utils/__init__.py | 8 +- .../async_utils/trajectory_collector.py | 106 ++++++++++++++++ nemo_rl/algorithms/grpo.py | 16 ++- nemo_rl/algorithms/ppo.py | 16 ++- tests/unit/algorithms/test_ng_task_index.py | 120 ++++++++++++++++++ 5 files changed, 263 insertions(+), 3 deletions(-) create mode 100644 tests/unit/algorithms/test_ng_task_index.py diff --git a/nemo_rl/algorithms/async_utils/__init__.py b/nemo_rl/algorithms/async_utils/__init__.py index 3ff2774f3ae..6d76f5915d2 100644 --- a/nemo_rl/algorithms/async_utils/__init__.py +++ b/nemo_rl/algorithms/async_utils/__init__.py @@ -13,9 +13,15 @@ # limitations under the License. from nemo_rl.algorithms.async_utils.replay_buffer import ReplayBuffer -from nemo_rl.algorithms.async_utils.trajectory_collector import AsyncTrajectoryCollector +from nemo_rl.algorithms.async_utils.trajectory_collector import ( + AsyncTrajectoryCollector, + compute_resume_ng_task_index, + save_rollouts_state, +) __all__ = [ "ReplayBuffer", "AsyncTrajectoryCollector", + "compute_resume_ng_task_index", + "save_rollouts_state", ] diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index 3cd127105ca..81840b1b5db 100644 --- a/nemo_rl/algorithms/async_utils/trajectory_collector.py +++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py @@ -50,6 +50,72 @@ # NRL_ASYNC_VERBOSE=1 additionally enables the O(n) full-list dumps. _ASYNC_VERBOSE = os.environ.get("NRL_ASYNC_VERBOSE", "0").lower() in ("1", "true") +# NeMo-Gym cohort grouping. The GenRM verifier buffers rollouts into cohorts keyed +# by ``_ng_task_index`` (falling back to a content hash of the prompt when it is +# absent). Giving every prompt group a distinct index keeps duplicate / identical- +# input prompts — including the same prompt sampled with reasoning on and off, which +# share an input hash — from landing in one cohort and overflowing +# ``num_rollouts_per_prompt`` (an AssertionError that 500s the gym and crashes the run). +_NG_TASK_INDEX_KEY = "_ng_task_index" +_NEXT_NG_TASK_INDEX_KEY = "next_ng_task_index" +_ROLLOUTS_STATE_FILENAME = "rollouts.pt" + + +def _stamp_ng_task_index( + repeated_batch: BatchedDataDict[DatumSpec], task_index: int +) -> None: + """Stamp every NeMo-Gym row in a prompt group with its cohort task index. + + Rewrites ``extra_env_info`` in place with shallow-copied rows so the index + rides along to the gym ``/verify`` request (via the aliased ``task_index`` + field) without mutating any dict shared by ``repeat_interleave``. + """ + stamped_rows = [] + for row in repeated_batch["extra_env_info"]: + if not isinstance(row, dict): + raise TypeError( + f"Expected NeMo-Gym extra_env_info row to be a dict, got {type(row).__name__}" + ) + stamped_rows.append({**row, _NG_TASK_INDEX_KEY: task_index}) + repeated_batch["extra_env_info"] = stamped_rows + + +def compute_resume_ng_task_index( + last_checkpoint_path: Optional[str], + replay_buffer_state: Optional[dict[str, Any]], +) -> int: + """Recover the next NeMo-Gym cohort index when resuming from a checkpoint. + + Prefer the collector's saved counter (``rollouts.pt``). Also advance past the + largest ``_ng_task_index`` still live in the restored replay buffer, so resumed + rollouts never reuse an index that a buffered (in-flight) cohort still holds. + Returns 0 for a fresh (non-resumed) run. + """ + next_ng_task_index = 0 + if last_checkpoint_path is not None: + rollouts_path = os.path.join(last_checkpoint_path, _ROLLOUTS_STATE_FILENAME) + if os.path.exists(rollouts_path): + # weights_only=False: a trusted same-job dict of Python ints. + rollouts_state = torch.load(rollouts_path, weights_only=False) + next_ng_task_index = int( + (rollouts_state or {}).get(_NEXT_NG_TASK_INDEX_KEY, 0) + ) + if replay_buffer_state is not None: + live_indices = [ + int(trajectory[_NG_TASK_INDEX_KEY]) + for trajectory in replay_buffer_state.get("trajectories", []) + if trajectory.get(_NG_TASK_INDEX_KEY) is not None + ] + if live_indices: + next_ng_task_index = max(next_ng_task_index, 1 + max(live_indices)) + return next_ng_task_index + + +def save_rollouts_state(trajectory_collector: Any, checkpoint_path: str) -> None: + """Persist the collector's NeMo-Gym cohort counter alongside the checkpoint.""" + rollouts_state = ray.get(trajectory_collector.get_rollouts_state.remote()) + torch.save(rollouts_state, os.path.join(checkpoint_path, _ROLLOUTS_STATE_FILENAME)) + @ray.remote # pragma: no cover class AsyncTrajectoryCollector: @@ -66,6 +132,7 @@ def __init__( teacher_worker_groups: Optional[dict[str, Any]] = None, alias_to_group_alias: Optional[dict[str, str]] = None, on_policy_distillation_cfg: Optional[dict[str, Any]] = None, + next_ng_task_index: int = 0, ): self.policy_generation = policy_generation self.tokenizer = tokenizer @@ -175,6 +242,9 @@ def __init__( self._completed_per_target: dict[int, int] = {} self._spawning_targets: set[int] = set() self._counter_lock: _threading.Lock = _threading.Lock() + # Monotonic NeMo-Gym cohort index; reserved per prompt group under + # _counter_lock in _process_batch. Seeded from checkpoint on resume. + self._next_ng_task_index: int = int(next_ng_task_index) # Timer for efficiency metrics self._efficiency_timer = ThreadSafeTimer(context={"worker": "collector"}) @@ -461,6 +531,19 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None: f"prompts (need {trajectories_needed} more trajectories)" ) + # Reserve a contiguous block of globally-unique NeMo-Gym cohort + # indices for this batch — one per prompt group about to be spawned. + # Reserved under _counter_lock so concurrently-processed target + # batches never receive overlapping ranges. Only reserve for the + # NeMo-Gym path; other envs do not use cohort grouping. + from nemo_rl.algorithms.grpo import _should_use_nemo_gym + + prompt_group_base_task_index = None + if _should_use_nemo_gym(self.master_config): + with self._counter_lock: + prompt_group_base_task_index = self._next_ng_task_index + self._next_ng_task_index += num_prompts_to_generate + # Generate only the prompt groups needed for this target. While the # spawn loop is open, workers may finish before later workers start, # so reservation release is deferred until spawning closes. @@ -489,6 +572,11 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None: repeated_batch = single_prompt_batch.repeat_interleave( num_generations ) + prompt_group_task_index = ( + prompt_group_base_task_index + prompt_idx + if prompt_group_base_task_index is not None + else None + ) worker = _threading.Thread( target=self._run_prompt_group_worker, @@ -497,6 +585,7 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None: generation_weight_version, target_weight, prompt_idx, + prompt_group_task_index, ), daemon=True, ) @@ -680,6 +769,11 @@ def get_dataloader_state(self) -> dict: return self.dataloader.state_dict() return {} + def get_rollouts_state(self) -> dict[str, int]: + """Collector-side rollout state for checkpointing (NeMo-Gym cohort counter).""" + with self._counter_lock: + return {_NEXT_NG_TASK_INDEX_KEY: self._next_ng_task_index} + def get_efficiency_metrics(self) -> dict[str, float]: """Return accumulated efficiency metrics (sum of durations per category). @@ -843,6 +937,7 @@ def _run_prompt_group_worker( generation_weight_version: int, target_weight_version: int, prompt_idx: int, + prompt_group_task_index: Optional[int] = None, ) -> None: worker_start = time.perf_counter() try: @@ -853,6 +948,13 @@ def _run_prompt_group_worker( run_async_nemo_gym_rollout, ) + # Stamp every rollout in this prompt group with its cohort index so + # the GenRM verifier keys the whole group into a single cohort of + # num_generations instead of colliding with an identically-worded + # prompt elsewhere in the step (which overflows the cohort assert). + if prompt_group_task_index is not None and "extra_env_info" in repeated_batch: + _stamp_ng_task_index(repeated_batch, prompt_group_task_index) + # Run rollout for this prompt group # Async engine supports concurrent generation; avoid locking # Check if we should use nemo_gym (similar to synchronous GRPO) @@ -930,6 +1032,10 @@ def _run_prompt_group_worker( "rollout_metrics": rollout_metrics, "timestamp": time.time(), } + # Record the cohort index on the group so the counter can be resumed + # (one past the max live index) after a checkpoint restore. + if prompt_group_task_index is not None: + trajectory_group[_NG_TASK_INDEX_KEY] = prompt_group_task_index # Use exponential backoff when buffer is full try: diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 44d4047730a..60e4d928118 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -3620,7 +3620,12 @@ def async_grpo_train( ) # Import async utilities only when needed - from nemo_rl.algorithms.async_utils import AsyncTrajectoryCollector, ReplayBuffer + from nemo_rl.algorithms.async_utils import ( + AsyncTrajectoryCollector, + ReplayBuffer, + compute_resume_ng_task_index, + save_rollouts_state, + ) timer = Timer(context={"worker": "driver"}) training_wall_start = time.perf_counter() @@ -3715,6 +3720,7 @@ def async_grpo_train( ) last_checkpoint_path = checkpointer.get_latest_checkpoint_path() + replay_buffer_state = None if last_checkpoint_path is not None: replay_buffer_path = os.path.join(last_checkpoint_path, "replay_buffer.pt") if os.path.exists(replay_buffer_path): @@ -3737,6 +3743,12 @@ def async_grpo_train( "Starting with an empty replay buffer." ) + # Resume the NeMo-Gym cohort index counter so post-resume rollouts never + # reuse a _ng_task_index still held by a buffered/in-flight cohort. + next_ng_task_index = compute_resume_ng_task_index( + last_checkpoint_path, replay_buffer_state + ) + _tc_py_exec = get_actor_python_env( "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector" ) @@ -3772,6 +3784,7 @@ def async_grpo_train( teacher_worker_groups=teacher_worker_groups, alias_to_group_alias=alias_to_group_alias, on_policy_distillation_cfg=opd_module._opd_cfg(master_config), + next_ng_task_index=next_ng_task_index, ) # Start trajectory collection in background @@ -4483,6 +4496,7 @@ def async_grpo_train( "✅ Saved replay buffer with " f"{len(replay_buffer_state['trajectories'])} trajectories" ) + save_rollouts_state(trajectory_collector, checkpoint_path) checkpointer.finalize_checkpoint(checkpoint_path) # Record last-successful-checkpoint time/step for external diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index a7cf6b39cbe..bf1aba54e66 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -2448,7 +2448,12 @@ def async_ppo_train( ) # Import async utilities only when needed (heavy Ray actors). - from nemo_rl.algorithms.async_utils import AsyncTrajectoryCollector, ReplayBuffer + from nemo_rl.algorithms.async_utils import ( + AsyncTrajectoryCollector, + ReplayBuffer, + compute_resume_ng_task_index, + save_rollouts_state, + ) timer = Timer(context={"worker": "driver"}) training_wall_start = time.perf_counter() @@ -2563,6 +2568,7 @@ def _sample_max_age(s: int) -> int: ) last_checkpoint_path = checkpointer.get_latest_checkpoint_path() + replay_buffer_state = None if last_checkpoint_path is not None: replay_buffer_path = os.path.join(last_checkpoint_path, "replay_buffer.pt") if os.path.exists(replay_buffer_path): @@ -2585,6 +2591,12 @@ def _sample_max_age(s: int) -> int: "Starting with an empty replay buffer." ) + # Resume the NeMo-Gym cohort index counter so post-resume rollouts never + # reuse a _ng_task_index still held by a buffered/in-flight cohort. + next_ng_task_index = compute_resume_ng_task_index( + last_checkpoint_path, replay_buffer_state + ) + _tc_py_exec = get_actor_python_env( "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector" ) @@ -2615,6 +2627,7 @@ def _sample_max_age(s: int) -> int: master_config=master_config, replay_buffer=replay_buffer, start_step=step, + next_ng_task_index=next_ng_task_index, ) collection_task = trajectory_collector.start_collection.remote(dataloader) # noqa: F841 @@ -3263,6 +3276,7 @@ def _sample_max_age(s: int) -> int: replay_buffer_state, os.path.join(checkpoint_path, "replay_buffer.pt"), ) + save_rollouts_state(trajectory_collector, checkpoint_path) checkpointer.finalize_checkpoint(checkpoint_path) _write_latest_checkpoint_status( checkpointer, last_checkpoint_step=step + 1 diff --git a/tests/unit/algorithms/test_ng_task_index.py b/tests/unit/algorithms/test_ng_task_index.py new file mode 100644 index 00000000000..0629fd6e962 --- /dev/null +++ b/tests/unit/algorithms/test_ng_task_index.py @@ -0,0 +1,120 @@ +# 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. + +"""Unit tests for the NeMo-Gym cohort-index passthrough (``_ng_task_index``). + +These cover the pure helpers added for the GenRM cohort fix: the resume-counter +recovery math (``compute_resume_ng_task_index``) and the per-prompt-group row +stamping (``_stamp_ng_task_index``). Giving every prompt group a globally-unique +``_ng_task_index`` keeps duplicate / identical-input prompts from colliding on the +GenRM cohort key and overflowing ``num_rollouts_per_prompt``. +""" + +import pytest +import torch + +from nemo_rl.algorithms.async_utils import compute_resume_ng_task_index +from nemo_rl.algorithms.async_utils.trajectory_collector import ( + _NEXT_NG_TASK_INDEX_KEY, + _NG_TASK_INDEX_KEY, + _ROLLOUTS_STATE_FILENAME, + _stamp_ng_task_index, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +def _write_rollouts_state(ckpt_dir, next_ng_task_index): + torch.save( + {_NEXT_NG_TASK_INDEX_KEY: next_ng_task_index}, + ckpt_dir / _ROLLOUTS_STATE_FILENAME, + ) + + +def _buffer_state(indices): + """A replay-buffer ``state_dict``-like object with the given per-group indices. + + ``None`` entries model trajectories lacking the key (pre-fix checkpoints). + """ + trajectories = [] + for idx in indices: + traj = {"batch": None, "rollout_metrics": {}} + if idx is not None: + traj[_NG_TASK_INDEX_KEY] = idx + trajectories.append(traj) + return {"trajectories": trajectories} + + +# ---- compute_resume_ng_task_index --------------------------------------- + + +def test_resume_fresh_run_is_zero(): + assert compute_resume_ng_task_index(None, None) == 0 + + +def test_resume_reads_saved_counter(tmp_path): + _write_rollouts_state(tmp_path, 42) + assert compute_resume_ng_task_index(str(tmp_path), None) == 42 + + +def test_resume_missing_rollouts_file_uses_buffer_high_water(tmp_path): + # Old checkpoint: no rollouts.pt, but buffered groups carry indices. + assert compute_resume_ng_task_index(str(tmp_path), _buffer_state([3, 7, 5])) == 8 + + +def test_resume_takes_max_of_counter_and_buffer(tmp_path): + _write_rollouts_state(tmp_path, 42) + # Buffer high-water (100) exceeds the saved counter -> 1 + 100. + assert compute_resume_ng_task_index(str(tmp_path), _buffer_state([100, 10])) == 101 + _write_rollouts_state(tmp_path, 200) + # Saved counter dominates the buffer high-water. + assert compute_resume_ng_task_index(str(tmp_path), _buffer_state([100, 10])) == 200 + + +def test_resume_ignores_trajectories_without_index(tmp_path): + # Pre-fix buffer (no _ng_task_index on any group) -> start at 0. + assert compute_resume_ng_task_index(str(tmp_path), _buffer_state([None, None])) == 0 + + +def test_resume_handles_zero_index(tmp_path): + # index 0 must be honored (guarded by `is not None`, not truthiness). + assert compute_resume_ng_task_index(str(tmp_path), _buffer_state([0])) == 1 + + +# ---- _stamp_ng_task_index ------------------------------------------------ + + +def test_stamp_sets_same_index_on_all_rows_and_preserves_fields(): + batch = BatchedDataDict( + {"extra_env_info": [{"a": 1, _NG_TASK_INDEX_KEY: 999}, {"a": 2}]} + ) + _stamp_ng_task_index(batch, 5) + rows = batch["extra_env_info"] + # All rollouts of one prompt group share the same cohort index (overwriting + # any non-unique index that came in on the training data), other fields kept. + assert [r[_NG_TASK_INDEX_KEY] for r in rows] == [5, 5] + assert [r["a"] for r in rows] == [1, 2] + + +def test_stamp_does_not_mutate_original_row_dicts(): + original = {"a": 1} + batch = BatchedDataDict({"extra_env_info": [original]}) + _stamp_ng_task_index(batch, 7) + assert _NG_TASK_INDEX_KEY not in original # shallow-copied, not mutated in place + assert batch["extra_env_info"][0][_NG_TASK_INDEX_KEY] == 7 + + +def test_stamp_raises_on_non_dict_row(): + batch = BatchedDataDict({"extra_env_info": [{"a": 1}, "not-a-dict"]}) + with pytest.raises(TypeError): + _stamp_ng_task_index(batch, 5) From 9be0f8fa8bb5fb937dc0097349578d287762c800 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 31 Jul 2026 14:39:54 -0700 Subject: [PATCH 35/69] chore: ignore local Claude Code state and test scratch artifacts Co-Authored-By: Claude Opus 5 --- .gitignore | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/.gitignore b/.gitignore index 1d9441977f1..80a6f6454dd 100644 --- a/.gitignore +++ b/.gitignore @@ -47,3 +47,16 @@ code_snapshots*/ # Claude Code review memory (local only) .claude/review-memory/ + +# Local Claude Code / MCP state +.claude/ +!.claude/review-memory/ +.mcp.json + +# Local test + scratch artifacts +tests/unit/unit_results.json +tests/unit/unit_results/ +classifier.model +known_elements.txt +output.txt +perm*.txt From 1a6c298af246b44dcd292cabd1e58e65c219ea43 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 31 Jul 2026 14:40:16 -0700 Subject: [PATCH 36/69] feat(critic): offline rollout collection for the decoupled critic pipeline Stage A of the decoupled critic pipeline: drive the Gym rollout path with a frozen policy and persist per-prompt groups (post-penalty rewards, message logs, extra_env_info) as shards, so critic pretraining can consume them offline instead of paying generation cost inside the PPO loop. - nemo_rl/algorithms/rollout_collection.py: collection loop + shard writer - examples/nemo_gym/run_swe_rollout_collection.py: SWE entrypoint - tests/unit/algorithms/test_rollout_collection.py Co-Authored-By: Claude Opus 5 --- .../nemo_gym/run_swe_rollout_collection.py | 197 +++++ nemo_rl/algorithms/rollout_collection.py | 717 ++++++++++++++++++ .../algorithms/test_rollout_collection.py | 228 ++++++ 3 files changed, 1142 insertions(+) create mode 100644 examples/nemo_gym/run_swe_rollout_collection.py create mode 100644 nemo_rl/algorithms/rollout_collection.py create mode 100644 tests/unit/algorithms/test_rollout_collection.py diff --git a/examples/nemo_gym/run_swe_rollout_collection.py b/examples/nemo_gym/run_swe_rollout_collection.py new file mode 100644 index 00000000000..60229efab5e --- /dev/null +++ b/examples/nemo_gym/run_swe_rollout_collection.py @@ -0,0 +1,197 @@ +# 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. +"""Generation-only rollout collection with NeMo-Gym (stage A, decoupled PPO). + +Reuses the production PPO config so shards match what the coupled run's critic +warmup would consume, but spins up ONLY the vLLM engines + NeMo-Gym — no +policy/value workers, no refit NCCL group, no replay buffer. Job shape and +sharding knobs come from a ``collection:`` config block (see +``nemo_rl.algorithms.rollout_collection.resolve_collection_config``); launch +via scripts/swe/ppo/collect_rollouts.sh (1-node SLURM array tasks). +""" + +import argparse +import os +import pprint +import time +from concurrent.futures import ThreadPoolExecutor + +from omegaconf import OmegaConf + +from nemo_rl.algorithms.grpo import _should_use_nemo_gym +from nemo_rl.algorithms.ppo import MasterConfig +from nemo_rl.algorithms.rollout_collection import ( + collect_rollouts, + resolve_collection_config, + spinup_nemo_gym, +) +from nemo_rl.algorithms.utils import get_tokenizer +from nemo_rl.data.utils import setup_response_data +from nemo_rl.distributed.virtual_cluster import RayVirtualCluster, init_ray +from nemo_rl.environments.nemo_gym import setup_nemo_gym_config +from nemo_rl.models.generation import configure_generation_config +from nemo_rl.models.generation.vllm import VllmGeneration +from nemo_rl.utils.config import ( + load_config, + parse_hydra_overrides, + register_omegaconf_resolvers, +) + + +def parse_args() -> tuple[argparse.Namespace, list[str]]: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Generation-only SWE rollout collection with NeMo-Gym" + ) + parser.add_argument( + "--config", type=str, required=True, help="Path to YAML config file" + ) + args, overrides = parser.parse_known_args() + return args, overrides + + +def main() -> None: + """Main entry point.""" + register_omegaconf_resolvers() + args, overrides = parse_args() + + config = load_config(args.config) + print(f"Loaded configuration from: {args.config}") + if overrides: + print(f"Overrides: {overrides}") + config = parse_hydra_overrides(config, overrides) + config = OmegaConf.to_container(config, resolve=True) + config = MasterConfig(**config) + + collection = resolve_collection_config( + getattr(config, "collection", None), config.ppo + ) + + tokenizer = get_tokenizer(config.policy["tokenizer"]) + assert config.policy["generation"] is not None, ( + "A generation config is required for rollout collection" + ) + # is_eval=True => vllm_cfg.load_format="auto": the engine must load REAL + # pi_0 weights from policy.model_name. There is no policy worker and no + # refit in this job; the training default ("dummy") would generate from + # random weights. + config.policy["generation"] = configure_generation_config( + config.policy["generation"], tokenizer, is_eval=True + ) + setup_nemo_gym_config(config, tokenizer) + assert _should_use_nemo_gym(config), ( + "Rollout collection requires the NeMo-Gym path " + "(env.should_use_nemo_gym=true with an async, HTTP-exposed vLLM engine)." + ) + + print("\n▶ Setting up data...") + train_dataset, _ = setup_response_data(tokenizer, config.data, env_configs=None) + assert not isinstance(train_dataset, dict), ( + "use_multiple_dataloader is not supported for rollout collection" + ) + + print("Final config:") + pprint.pprint(config) + print(f"Collection config: {collection}") + + init_ray() + + # ------------------------------------------------------------------ + # Inference cluster over ALL nodes of this job (generation-only shape). + # ------------------------------------------------------------------ + cluster_config = config.cluster + num_nodes = cluster_config["num_nodes"] + gpus_per_node = cluster_config["gpus_per_node"] + cluster = RayVirtualCluster( + name="rollout_collection_cluster", + bundle_ct_per_node_list=[gpus_per_node] * num_nodes, + use_gpus=True, + num_gpus_per_node=gpus_per_node, + max_colocated_worker_groups=1, + port_range_low=cluster_config.get("master_port_range_low"), + port_range_high=cluster_config.get("master_port_range_high"), + ) + print(f" ✓ Ray cluster initialized: {num_nodes} nodes x {gpus_per_node} GPUs") + + # ------------------------------------------------------------------ + # Deferred vLLM init -> reserve server ports -> spin up NeMo-Gym while the + # model loads (same overlap pattern as ppo.setup()). + # ------------------------------------------------------------------ + generation_config = config.policy["generation"] + generation_config["model_name"] = config.policy["model_name"] + generation_config["vllm_kwargs"]["hf_overrides"] = config.policy.get( + "hf_config_overrides", {} + ) + + setup_start = time.perf_counter() + print(" ⚡ Deferred vLLM load: reserving ports for overlapped NeMo-Gym init") + policy_generation = VllmGeneration( + cluster=cluster, config=generation_config, defer_model_load=True + ) + print( + f" ✓ Reserved {len(policy_generation.dp_openai_server_base_urls)} vLLM " + f"server URLs: {policy_generation.dp_openai_server_base_urls}" + ) + + def _load_vllm(): + policy_generation.load_and_start() + policy_generation.finish_generation() + + def _init_gym(): + return spinup_nemo_gym( + config, + policy_generation.dp_openai_server_base_urls, + generation_config["model_name"], + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + vllm_future = executor.submit(_load_vllm) + gym_future = executor.submit(_init_gym) + vllm_future.result() + nemo_gym = gym_future.result() + print(f" ✓ vLLM + NeMo-Gym ready in {time.perf_counter() - setup_start:.1f}s") + + # ------------------------------------------------------------------ + # Collect. + # ------------------------------------------------------------------ + task_to_env = {"nemo_gym": nemo_gym} + summary = collect_rollouts( + policy_generation=policy_generation, + tokenizer=tokenizer, + task_to_env=task_to_env, + master_config=config, + dataset=train_dataset, + collection=collection, + ) + + print("🛑 Shutting down generation workers...") + try: + policy_generation.shutdown() + except Exception as e: + print(f"⚠️ vLLM shutdown failed (non-fatal): {e}") + + if summary.get("aborted"): + raise SystemExit( + "Collection aborted after repeated consecutive failures; " + "see logs above." + ) + if os.environ.get("SLURM_ARRAY_TASK_ID") is not None and summary["remaining"] > 0: + print( + f"ℹ️ {summary['remaining']} assigned groups still missing " + "(walltime/failures) — resubmit the same array to resume." + ) + + +if __name__ == "__main__": + main() diff --git a/nemo_rl/algorithms/rollout_collection.py b/nemo_rl/algorithms/rollout_collection.py new file mode 100644 index 00000000000..3049dc67649 --- /dev/null +++ b/nemo_rl/algorithms/rollout_collection.py @@ -0,0 +1,717 @@ +# 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. +"""Standalone rollout collection (stage A of the decoupled PPO pipeline). + +Generates frozen-policy (pi_0) rollouts with a generation-only job — no policy +or value workers, no weight refits, no replay buffer — and banks each finished +prompt group as one file on disk: + + /shard_/group_.pt + +Each group file is the same trajectory dict the async PPO collector pushes to +the replay buffer (``{"batch": final_batch, "rollout_metrics": ...}``), plus +input-side fields the NeMo-Gym rollout path drops (``extra_env_info``, ``idx``, +``task_name`` are grafted back into ``batch``) and provenance metadata. The +critic-pretraining stage (:mod:`nemo_rl.algorithms.critic_pretrain`) consumes +these files. + +Sharding is coordination-free: shard ``k`` of ``num_shards`` owns dataset +indices with ``idx % num_shards == k``; resume simply skips indices whose group +file already exists, so a killed job loses only its in-flight episodes. + +Concurrency is admission-controlled at the SAMPLE level: each of a group's +``gens_per_prompt`` episodes is its own rollout call, bounded by +``collection.max_inflight_samples``, so when one episode finishes the freed +slot is immediately backfilled by the next pending sample (from any group). +Completed samples are reassembled into whole-group files; a group is banked +all-or-nothing, exactly as before. +""" + +import hashlib +import json +import math +import os +import subprocess +import threading +import time +from pathlib import Path +from typing import Any, Optional + +import torch + +from nemo_rl.data.collate_fn import rl_collate_fn +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + +SHARD_FORMAT_VERSION = 1 +_GROUP_PREFIX = "group_" +_GROUP_SUFFIX = ".pt" + + +# =============================================================================== +# Pure helpers (unit-tested, no heavy deps) +# =============================================================================== +def assigned_indices( + dataset_len: int, + shard_id: int, + num_shards: int, + index_start: int = 0, + index_end: Optional[int] = None, + max_groups: Optional[int] = None, +) -> list[int]: + """Dataset indices owned by this shard (strided split over [start, end)). + + Strided (``idx % num_shards == shard_id``) rather than contiguous so every + shard samples the full curriculum mix. + """ + assert 0 <= shard_id < num_shards, ( + f"shard_id must be in [0, {num_shards}), got {shard_id}" + ) + end = dataset_len if index_end is None else min(index_end, dataset_len) + assert index_start >= 0 and index_start <= end, ( + f"invalid index range [{index_start}, {end})" + ) + indices = [i for i in range(index_start, end) if i % num_shards == shard_id] + if max_groups is not None: + indices = indices[:max_groups] + return indices + + +def group_filename(dataset_idx: int) -> str: + """Group file name for a dataset index (fixed-width so listings sort).""" + return f"{_GROUP_PREFIX}{dataset_idx:08d}{_GROUP_SUFFIX}" + + +def parse_group_index(filename: str) -> Optional[int]: + """Inverse of :func:`group_filename`; None for non-group files.""" + name = os.path.basename(filename) + if not (name.startswith(_GROUP_PREFIX) and name.endswith(_GROUP_SUFFIX)): + return None + stem = name[len(_GROUP_PREFIX) : -len(_GROUP_SUFFIX)] + if not stem.isdigit(): + return None + return int(stem) + + +def existing_group_indices(shard_dir: str | Path) -> set[int]: + """Dataset indices with a completed group file in ``shard_dir``. + + Only finalized files count: in-progress writes use a ``.tmp*`` suffix and + are atomically renamed on completion, so a crash never leaves a partial + file that would be skipped on resume. + """ + shard_dir = Path(shard_dir) + if not shard_dir.is_dir(): + return set() + out = set() + for name in os.listdir(shard_dir): + idx = parse_group_index(name) + if idx is not None: + out.add(idx) + return out + + +def existing_group_indices_all(out_dir: str | Path) -> set[int]: + """Union of completed group indices across ALL shard dirs under ``out_dir``. + + Resume scans the whole output tree rather than just this task's shard dir, + so ``num_shards`` can be changed freely between submissions: a group banked + under any previous sharding layout is never regenerated (``dataset_idx`` is + globally unique and encoded in the filename, regardless of which shard dir + holds it). + """ + out_dir = Path(out_dir) + done: set[int] = set() + if not out_dir.is_dir(): + return done + done |= existing_group_indices(out_dir) + for sub in out_dir.glob("shard_*"): + if sub.is_dir(): + done |= existing_group_indices(sub) + return done + + +def write_group_atomic(shard_dir: str | Path, dataset_idx: int, payload: dict) -> Path: + """torch.save ``payload`` to a tmp file, then atomically rename into place. + + The tmp file is fsynced before the rename so a node/kernel crash cannot + leave a truncated file under the FINAL name (which resume would treat as + done and stage B would fail to load). + """ + shard_dir = Path(shard_dir) + final_path = shard_dir / group_filename(dataset_idx) + tmp_path = shard_dir / f"{group_filename(dataset_idx)}.tmp.{os.getpid()}" + with open(tmp_path, "wb") as f: + torch.save(payload, f) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, final_path) + return final_path + + +def load_group(path: str | Path) -> dict: + """Load a group file written by :func:`write_group_atomic`. + + weights_only=False: trajectories are pickled BatchedDataDict/dicts written + by this same pipeline (trusted local artifact), not plain tensors. + """ + return torch.load(path, weights_only=False, map_location="cpu") + + +def build_group_payload( + dataset_idx: int, + input_batch: BatchedDataDict, + final_batch: BatchedDataDict, + rollout_metrics: dict[str, Any], +) -> dict[str, Any]: + """Assemble the on-disk group payload. + + The NeMo-Gym rollout path returns a ``final_batch`` WITHOUT + ``extra_env_info`` / ``idx`` / ``task_name`` (unlike the generic multi-turn + path); graft them back from the input batch so downstream consumers see the + same batch shape the PPO driver loops build, and so answer-conditioned + (privileged) critics can reach ``extra_env_info`` at train time. + """ + for key in ("extra_env_info", "idx", "task_name"): + if key not in final_batch and key in input_batch: + final_batch[key] = input_batch[key] + # _rowidx is a per-call scratch field run_async_nemo_gym_rollout writes into + # extra_env_info rows (always 0 for single-sample calls); drop it so stored + # rows match the pre-rollout inputs regardless of call batching. + for row in final_batch.get("extra_env_info") or []: + if isinstance(row, dict): + row.pop("_rowidx", None) + return { + "format_version": SHARD_FORMAT_VERSION, + "dataset_idx": dataset_idx, + "batch": final_batch, + "rollout_metrics": dict(rollout_metrics), + "timestamp": time.time(), + } + + +def _sha256(text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def _best_effort_git_commit() -> str: + try: + return ( + subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=os.path.dirname(os.path.abspath(__file__)), + capture_output=True, + text=True, + timeout=10, + ).stdout.strip() + or "unknown" + ) + except Exception: + return "unknown" + + +def build_shard_meta( + master_config: Any, + collection: dict[str, Any], + tokenizer: Any, + dataset_len: int, +) -> dict[str, Any]: + """Provenance metadata written once per shard dir (``meta.json``). + + Downstream stages assert on the model/tokenizer identity recorded here — + shards are token-id level and are invalidated by any tokenizer, chat + template, max-length, or dataset change. + """ + data_train = master_config.data.get("train") + if isinstance(data_train, list): + data_train = data_train[0] if data_train else {} + chat_template = getattr(tokenizer, "chat_template", None) or "" + return { + "format_version": SHARD_FORMAT_VERSION, + "model_name": master_config.policy["model_name"], + "tokenizer_name_or_path": getattr(tokenizer, "name_or_path", "unknown"), + "chat_template_sha256": _sha256(chat_template), + "max_total_sequence_length": master_config.policy[ + "max_total_sequence_length" + ], + "dataset_path": (data_train or {}).get("data_path"), + "dataset_len": dataset_len, + "gens_per_prompt": collection["gens_per_prompt"], + "shard_id": collection["shard_id"], + "num_shards": collection["num_shards"], + "index_start": collection["index_start"], + "index_end": collection["index_end"], + "git_commit": _best_effort_git_commit(), + "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + } + + +def resolve_collection_config( + raw: Optional[dict[str, Any]], ppo_config: dict[str, Any] +) -> dict[str, Any]: + """Fill defaults for the ``collection:`` config block.""" + cfg = dict(raw or {}) + assert cfg.get("out_dir"), ( + "collection.out_dir is required (pass ++collection.out_dir=)" + ) + cfg.setdefault("shard_id", 0) + cfg.setdefault("num_shards", 1) + cfg.setdefault("gens_per_prompt", ppo_config["num_generations_per_prompt"]) + cfg.setdefault("index_start", 0) + cfg.setdefault("index_end", None) + cfg.setdefault("max_groups", None) + cfg.setdefault("log_every", 5) + cfg.setdefault("max_consecutive_failures", 10) + for key in ("shard_id", "num_shards", "gens_per_prompt"): + cfg[key] = int(cfg[key]) + # Sample-level admission: concurrency is bounded per SAMPLE so a group's + # stragglers never idle the engine. The legacy group-level knob converts. + if cfg.get("max_inflight_samples") is None: + legacy_groups = cfg.get("max_inflight_groups") + cfg["max_inflight_samples"] = ( + int(legacy_groups) if legacy_groups is not None else 3 + ) * cfg["gens_per_prompt"] + cfg["max_inflight_samples"] = int(cfg["max_inflight_samples"]) + assert cfg["max_inflight_samples"] >= 1 + return cfg + + +# =============================================================================== +# NeMo-Gym spinup (mirrors the inline _spinup_nemo_gym in ppo.setup(), which is +# not importable standalone; kept behaviorally identical) +# =============================================================================== +def spinup_nemo_gym(master_config: Any, base_urls: list[str], model_name: str): + """Spin up the NeMo-Gym actor against the given vLLM server URLs.""" + import ray + from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + + from nemo_rl.distributed.ray_actor_environment_registry import ( + get_actor_python_env, + ) + from nemo_rl.environments.nemo_gym import ( + NemoGym, + NemoGymConfig, + get_nemo_gym_uv_cache_dir, + get_nemo_gym_venv_dir, + ) + from nemo_rl.utils.venvs import create_local_venv_on_each_node + + nemo_gym_py_exec = get_actor_python_env("nemo_rl.environments.nemo_gym.NemoGym") + if nemo_gym_py_exec.startswith("uv"): + nemo_gym_py_exec = create_local_venv_on_each_node( + nemo_gym_py_exec, "nemo_rl.environments.nemo_gym.NemoGym" + ) + nemo_gym_dict = dict(master_config.env["nemo_gym"]) + invalid_tool_call_patterns = nemo_gym_dict.pop("invalid_tool_call_patterns", None) + thinking_tags = nemo_gym_dict.pop("thinking_tags", None) + uv_cache_dir = get_nemo_gym_uv_cache_dir() + if uv_cache_dir is not None: + nemo_gym_dict.setdefault("uv_cache_dir", uv_cache_dir) + uv_venv_dir = get_nemo_gym_venv_dir() + if uv_venv_dir is not None: + nemo_gym_dict.setdefault("uv_venv_dir", uv_venv_dir) + nemo_gym_cfg = NemoGymConfig( + model_name=model_name, + base_urls=base_urls, + invalid_tool_call_patterns=invalid_tool_call_patterns, + thinking_tags=thinking_tags, + require_routed_experts=False, + initial_global_config_dict=nemo_gym_dict, + ) + nemo_gym_opts: dict[str, Any] = {} + if master_config.env.get("nemo_gym", {}).get("num_gpu_nodes", 0): + nemo_gym_opts["scheduling_strategy"] = NodeAffinitySchedulingStrategy( + node_id=ray.get_runtime_context().get_node_id(), + soft=True, + ) + nemo_gym_opts["runtime_env"] = { + "py_executable": nemo_gym_py_exec, + "env_vars": { + **os.environ, + "VIRTUAL_ENV": nemo_gym_py_exec, + "UV_PROJECT_ENVIRONMENT": nemo_gym_py_exec, + }, + } + actor = NemoGym.options(**nemo_gym_opts).remote(nemo_gym_cfg) + ray.get(actor._spinup.remote()) + return actor + + +# =============================================================================== +# Collection loop +# =============================================================================== +def _run_rollout_batch( + policy_generation: Any, + input_batch: BatchedDataDict, + tokenizer: Any, + task_to_env: dict[str, Any], + master_config: Any, +) -> tuple[BatchedDataDict, dict[str, Any]]: + """Run one batch (here: a single sample) through the NeMo-Gym rollout path. + + Mirrors AsyncTrajectoryCollector._run_prompt_group_worker's gym branch: + stop tokens are cleared on a copied generation config so this path is safe + by construction (run_async_nemo_gym_rollout asserts they are unset). + """ + from nemo_rl.experience.rollouts import ( + get_nemo_gym_thinking_tags, + run_async_nemo_gym_rollout, + ) + + generation_config = { + **master_config.policy["generation"], + "stop_token_ids": None, + "stop_strings": None, + } + result = run_async_nemo_gym_rollout( + policy_generation=policy_generation, + input_batch=input_batch, + tokenizer=tokenizer, + task_to_env=task_to_env, + max_seq_len=master_config.policy["max_total_sequence_length"], + generation_config=generation_config, + max_rollout_turns=None, + greedy=False, + reward_penalty_config=master_config.reward_penalties, + thinking_tags=get_nemo_gym_thinking_tags(master_config.env), + ) + return result.final_batch.to("cpu"), result.rollout_metrics + + +def _aggregate_sample_metrics(metrics_list: list[dict[str, Any]]) -> dict[str, Any]: + """Mean-aggregate numeric per-sample rollout metrics into one group dict. + + Sample-level admission produces one rollout_metrics dict per episode; the + group payload keeps a single dict (as a whole-group call would), so numeric + fields are averaged and non-numeric fields dropped. Non-finite values are + skipped too: single-sample calls report NaN for */stddev-style statistics, + which would otherwise poison the group mean. + """ + merged: dict[str, list[float]] = {} + for m in metrics_list: + for k, v in (m or {}).items(): + if isinstance(v, (bool, int, float)) and math.isfinite(float(v)): + merged.setdefault(k, []).append(float(v)) + out: dict[str, Any] = {k: sum(v) / len(v) for k, v in merged.items()} + out["aggregated_from_samples"] = float(len(metrics_list)) + return out + + +def assemble_group_payload( + dataset_idx: int, + input_batch: BatchedDataDict, + sample_batches: list[BatchedDataDict], + sample_metrics: list[dict[str, Any]], +) -> dict[str, Any]: + """Reassemble per-sample rollout results into one group payload. + + ``sample_batches`` must be in generation order; concatenating them yields + the same batch content a single whole-group rollout call produces + (semantically identical for all consumers; byte layout may differ, e.g. + ``from_batches`` sorts keys). + """ + final_batch = BatchedDataDict.from_batches(sample_batches) + return build_group_payload( + dataset_idx, input_batch, final_batch, _aggregate_sample_metrics(sample_metrics) + ) + + +def _vllm_engine_stats_line(policy_generation: Any, interval_s: float) -> str: + """One-line engine summary from the in-process vLLM metrics logger. + + The async worker samples cumulative ``vllm:generation_tokens`` plus + running/waiting/KV gauges every ``interval_s`` (see + ``enable_vllm_metrics_logger``); tokens/s over the window is the counter + delta divided by the window span. Metrics are cleared after reading so each + heartbeat reports a fresh window. Best-effort: any failure returns "". + """ + try: + m = policy_generation.get_vllm_logger_metrics() + parts = [] + tputs = [] + for series in (m.get("generation_tokens") or {}).values(): + if len(series) >= 2: + tputs.append((series[-1] - series[0]) / ((len(series) - 1) * interval_s)) + if tputs: + parts.append(f"{sum(tputs):.0f} gen tok/s") + running = [v for s in (m.get("inflight_batch_sizes") or {}).values() for v in s] + if running: + parts.append(f"running {sum(running) / len(running):.0f} reqs") + kv = [v[-1] for v in (m.get("kv_cache_usage_perc") or {}).values() if v] + if kv: + parts.append(f"kv {100 * max(kv):.1f}%") + policy_generation.clear_vllm_logger_metrics() + return " | ".join(parts) + except Exception: + return "" + + +def collect_rollouts( + policy_generation: Any, + tokenizer: Any, + task_to_env: dict[str, Any], + master_config: Any, + dataset: Any, + collection: dict[str, Any], +) -> dict[str, Any]: + """Generate this shard's assigned prompt groups and write them to disk. + + Admission control is per SAMPLE: one worker thread per in-flight episode + (bounded by ``collection.max_inflight_samples``), each blocking in its own + single-sample ``run_async_nemo_gym_rollout`` call. When an episode + finishes, the freed slot is immediately backfilled by the next pending + sample from any group — a slow straggler holds one slot, not its whole + group's worth. Completed samples are reassembled (in generation order) + into the atomic per-group file; a group is banked all-or-nothing, and any + sample failure fails the whole group (its remaining samples are skipped + and partials discarded; resume regenerates it). + + Returns a stats dict (also written to ``shard_summary.json``). + """ + out_dir = Path(collection["out_dir"]) + shard_dir = out_dir / f"shard_{collection['shard_id']:03d}" + shard_dir.mkdir(parents=True, exist_ok=True) + + meta_path = shard_dir / "meta.json" + meta = build_shard_meta(master_config, collection, tokenizer, len(dataset)) + if not meta_path.exists(): + with open(meta_path, "w") as f: + json.dump(meta, f, indent=2) + print(f"📁 Shard dir: {shard_dir}") + + todo = assigned_indices( + len(dataset), + collection["shard_id"], + collection["num_shards"], + index_start=collection["index_start"], + index_end=collection["index_end"], + max_groups=collection["max_groups"], + ) + # Scan ALL shard dirs (not just ours): num_shards may differ from a prior + # submission, and a group banked under any layout must not be regenerated. + existing = existing_group_indices_all(out_dir) + skipped = [i for i in todo if i in existing] + pending = [i for i in todo if i not in existing] + print( + f"🎯 Assigned {len(todo)} groups " + f"(resume: {len(skipped)} already done, {len(pending)} to generate) " + f"x {collection['gens_per_prompt']} gens/prompt" + ) + + gens_per_prompt = collection["gens_per_prompt"] + max_inflight_samples = collection["max_inflight_samples"] + log_every = max(1, int(collection["log_every"])) + inflight = threading.Semaphore(max_inflight_samples) + state_lock = threading.Lock() + manifest_path = shard_dir / "manifest.jsonl" + # Per-group assembly state, keyed by dataset_idx. Bounded: samples are + # submitted group-by-group, so at most ~max_inflight_samples/gens + 1 + # groups are open at once. Failed groups keep only a tombstone in + # `failed_groups` (partials are dropped immediately). + groups_state: dict[int, dict[str, Any]] = {} + failed_groups: set[int] = set() + stats = { + "completed": 0, + "failed": 0, + "samples": 0, + "reward_sum": 0.0, + "consecutive_failures": 0, + "inflight_samples": 0, + } + start_time = time.perf_counter() + abort = threading.Event() + + def _fail_group(dataset_idx: int, err: Exception, where: str) -> None: + """First failure fails the whole group; partials are discarded.""" + with state_lock: + if dataset_idx in failed_groups: + return + failed_groups.add(dataset_idx) + groups_state.pop(dataset_idx, None) + stats["failed"] += 1 + stats["consecutive_failures"] += 1 + consecutive = stats["consecutive_failures"] + print(f"❌ group {dataset_idx} failed ({where}): {err}", flush=True) + if consecutive >= collection["max_consecutive_failures"]: + print( + f"🛑 {consecutive} consecutive group failures — aborting " + "collection (systemic problem, e.g. dead engine or gym).", + flush=True, + ) + abort.set() + + def _sample_worker( + dataset_idx: int, sample_idx: int, sample_batch: BatchedDataDict + ) -> None: + try: + final_one, metrics_one = _run_rollout_batch( + policy_generation, sample_batch, tokenizer, task_to_env, master_config + ) + complete_group = None + with state_lock: + group = groups_state.get(dataset_idx) + if group is None: + return # group already failed; discard this sample + group["parts"][sample_idx] = final_one + group["metrics"][sample_idx] = metrics_one + if len(group["parts"]) == group["expected"]: + groups_state.pop(dataset_idx) + complete_group = group + if complete_group is None: + return + ordered = [ + complete_group["parts"][j] for j in range(complete_group["expected"]) + ] + ordered_metrics = [ + complete_group["metrics"][j] for j in range(complete_group["expected"]) + ] + payload = assemble_group_payload( + dataset_idx, complete_group["input_batch"], ordered, ordered_metrics + ) + write_group_atomic(shard_dir, dataset_idx, payload) + # From here the group is banked on disk; a bookkeeping error must + # not mark it failed (that would double-count it and bump the + # consecutive-failure abort counter for a successful group). + try: + rewards = payload["batch"]["total_reward"] + record = { + "dataset_idx": dataset_idx, + "num_samples": int(payload["batch"].size), + "reward_mean": float(rewards.float().mean()), + "truncated_frac": float( + payload["batch"]["truncated"].float().mean() + ), + "seconds": round( + time.perf_counter() - complete_group["start"], 1 + ), + } + with state_lock: + stats["completed"] += 1 + stats["samples"] += record["num_samples"] + stats["reward_sum"] += ( + record["reward_mean"] * record["num_samples"] + ) + stats["consecutive_failures"] = 0 + done = stats["completed"] + inflight_now = stats["inflight_samples"] + with open(manifest_path, "a") as f: + f.write(json.dumps(record) + "\n") + if done % log_every == 0 or done == len(pending): + elapsed = time.perf_counter() - start_time + rate = stats["samples"] / max(elapsed, 1e-6) * 3600 + engine_line = _vllm_engine_stats_line( + policy_generation, + master_config.policy["generation"]["vllm_cfg"].get( + "vllm_metrics_logger_interval", 0.5 + ), + ) + print( + f"✅ [{done}/{len(pending)}] group {dataset_idx}: " + f"reward={record['reward_mean']:.3f} " + f"({record['seconds']}s) | {rate:.0f} samples/h | " + f"{inflight_now} samples in flight" + + (f" | vLLM: {engine_line}" if engine_line else ""), + flush=True, + ) + except Exception as e: + print( + f"⚠️ post-write bookkeeping failed for group {dataset_idx} " + f"(group file IS banked): {e}", + flush=True, + ) + except Exception as e: + import traceback + + traceback.print_exc() + _fail_group(dataset_idx, e, "rollout") + finally: + with state_lock: + stats["inflight_samples"] -= 1 + inflight.release() + + threads: set[threading.Thread] = set() + for dataset_idx in pending: + if abort.is_set(): + break + # Build the group's input batch in the submit loop (cheap CPU work); + # repeat_interleave deep-copies list fields, so each sample slice owns + # its rows (run_async_nemo_gym_rollout mutates extra_env_info in place). + # A deterministically-bad dataset row must not wedge the shard: count it + # as a group failure and move on instead of crashing the whole job. + try: + datum = dataset[dataset_idx] + single = rl_collate_fn([datum]) + repeated_batch = single.repeat_interleave(gens_per_prompt) + except Exception as e: + _fail_group(dataset_idx, e, "input batch") + continue + with state_lock: + groups_state[dataset_idx] = { + "parts": {}, + "metrics": {}, + "expected": gens_per_prompt, + "input_batch": repeated_batch, + "start": time.perf_counter(), + } + for sample_idx in range(gens_per_prompt): + if abort.is_set(): + break + with state_lock: + if dataset_idx in failed_groups: + break # a sibling sample failed; skip the rest of the group + sample_batch = repeated_batch.slice(sample_idx, sample_idx + 1) + inflight.acquire() + with state_lock: + dead = abort.is_set() or dataset_idx in failed_groups + if not dead: + stats["inflight_samples"] += 1 + if dead: + inflight.release() + break + t = threading.Thread( + target=_sample_worker, + args=(dataset_idx, sample_idx, sample_batch), + daemon=True, + ) + threads.add(t) + t.start() + # Prune finished thread objects so the set stays bounded on long runs. + if len(threads) > 4 * max_inflight_samples: + for t in [t for t in threads if not t.is_alive()]: + threads.discard(t) + + for t in list(threads): + t.join() + + elapsed = time.perf_counter() - start_time + summary = { + "assigned": len(todo), + "skipped_existing": len(skipped), + "completed": stats["completed"], + "failed": stats["failed"], + "samples": stats["samples"], + "mean_reward": ( + stats["reward_sum"] / stats["samples"] if stats["samples"] else None + ), + "elapsed_s": round(elapsed, 1), + "samples_per_hour": ( + round(stats["samples"] / elapsed * 3600, 1) if elapsed > 0 else None + ), + "aborted": abort.is_set(), + "remaining": len(pending) - stats["completed"], + } + with open(shard_dir / "shard_summary.json", "w") as f: + json.dump(summary, f, indent=2) + print(f"🏁 Collection summary: {json.dumps(summary)}") + return summary diff --git a/tests/unit/algorithms/test_rollout_collection.py b/tests/unit/algorithms/test_rollout_collection.py new file mode 100644 index 00000000000..cfebbd53a82 --- /dev/null +++ b/tests/unit/algorithms/test_rollout_collection.py @@ -0,0 +1,228 @@ +# 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. + +import pytest +import torch + +from nemo_rl.algorithms.rollout_collection import ( + _aggregate_sample_metrics, + assemble_group_payload, + assigned_indices, + build_group_payload, + existing_group_indices, + existing_group_indices_all, + group_filename, + load_group, + parse_group_index, + resolve_collection_config, + write_group_atomic, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +def test_assigned_indices_partition(): + """Shards partition the index range: disjoint and complete.""" + n, num_shards = 103, 7 + seen = [] + for k in range(num_shards): + idx = assigned_indices(n, k, num_shards) + assert all(i % num_shards == k for i in idx) + seen.extend(idx) + assert sorted(seen) == list(range(n)) + + +def test_assigned_indices_range_and_cap(): + """index_start/index_end clamp the range; max_groups caps the count.""" + idx = assigned_indices(100, 0, 2, index_start=10, index_end=20) + assert idx == [10, 12, 14, 16, 18] + assert assigned_indices(100, 0, 1, max_groups=3) == [0, 1, 2] + # index_end beyond dataset length is clamped + assert assigned_indices(5, 0, 1, index_end=100) == [0, 1, 2, 3, 4] + + +def test_group_filename_roundtrip(): + """parse_group_index inverts group_filename and rejects non-group files.""" + assert parse_group_index(group_filename(42)) == 42 + assert parse_group_index(group_filename(12345678)) == 12345678 + assert parse_group_index("group_00000042.pt.tmp.123") is None + assert parse_group_index("meta.json") is None + assert parse_group_index("group_abc.pt") is None + + +def test_existing_group_indices_ignores_tmp(tmp_path): + """Only finalized group files count toward resume state.""" + write_group_atomic(tmp_path, 3, {"x": 1}) + write_group_atomic(tmp_path, 7, {"x": 2}) + torch.save({"x": 3}, tmp_path / "group_00000009.pt.tmp.999") # crashed write + (tmp_path / "meta.json").write_text("{}") + assert existing_group_indices(tmp_path) == {3, 7} + assert existing_group_indices(tmp_path / "does_not_exist") == set() + + +def test_existing_group_indices_all_spans_shard_layouts(tmp_path): + """Resume unions across all shard dirs so num_shards can be rescaled.""" + (tmp_path / "shard_000").mkdir() + (tmp_path / "shard_001").mkdir() + (tmp_path / "shard_042").mkdir() + write_group_atomic(tmp_path / "shard_000", 0, {"x": 1}) + write_group_atomic(tmp_path / "shard_001", 5, {"x": 1}) + write_group_atomic(tmp_path / "shard_042", 9, {"x": 1}) + assert existing_group_indices_all(tmp_path) == {0, 5, 9} + assert existing_group_indices_all(tmp_path / "missing") == set() + + +def test_write_load_group_roundtrip(tmp_path): + """Payloads round-trip through the atomic writer, tensors intact.""" + batch = BatchedDataDict( + { + "total_reward": torch.tensor([0.0, 1.0]), + "message_log": [[{"role": "assistant"}], [{"role": "assistant"}]], + } + ) + payload = {"format_version": 1, "dataset_idx": 5, "batch": batch} + path = write_group_atomic(tmp_path, 5, payload) + loaded = load_group(path) + assert loaded["dataset_idx"] == 5 + torch.testing.assert_close( + loaded["batch"]["total_reward"], batch["total_reward"] + ) + + +def test_build_group_payload_grafts_input_keys(): + """extra_env_info/idx/task_name are grafted from the input batch (the + NeMo-Gym rollout path drops them from final_batch).""" + input_batch = BatchedDataDict( + { + "extra_env_info": [{"gt": "a"}, {"gt": "b"}], + "idx": [4, 4], + "task_name": ["nemo_gym", "nemo_gym"], + "loss_multiplier": torch.ones(2), + } + ) + final_batch = BatchedDataDict( + { + "total_reward": torch.tensor([1.0, 0.0]), + "loss_multiplier": torch.ones(2), + } + ) + payload = build_group_payload(4, input_batch, final_batch, {"m": 1}) + assert payload["batch"]["extra_env_info"] == [{"gt": "a"}, {"gt": "b"}] + assert payload["batch"]["idx"] == [4, 4] + assert payload["batch"]["task_name"] == ["nemo_gym", "nemo_gym"] + assert payload["dataset_idx"] == 4 + assert payload["rollout_metrics"] == {"m": 1} + + +def test_build_group_payload_strips_rowidx(): + """The rollout's per-call _rowidx scratch field is not persisted.""" + input_batch = BatchedDataDict( + { + "extra_env_info": [{"gt": "a", "_rowidx": 0}, {"gt": "b", "_rowidx": 0}], + "loss_multiplier": torch.ones(2), + } + ) + final_batch = BatchedDataDict({"total_reward": torch.zeros(2)}) + payload = build_group_payload(9, input_batch, final_batch, {}) + assert payload["batch"]["extra_env_info"] == [{"gt": "a"}, {"gt": "b"}] + + +def test_resolve_collection_config_defaults_and_required(): + """Defaults are filled from the ppo block; out_dir is mandatory.""" + ppo_cfg = {"num_generations_per_prompt": 16} + cfg = resolve_collection_config({"out_dir": "/tmp/x"}, ppo_cfg) + assert cfg["gens_per_prompt"] == 16 + assert cfg["num_shards"] == 1 and cfg["shard_id"] == 0 + assert cfg["max_inflight_samples"] == 48 # 3 groups-equivalent x 16 + cfg = resolve_collection_config( + {"out_dir": "/tmp/x", "gens_per_prompt": "4", "shard_id": "2", + "num_shards": "8"}, + ppo_cfg, + ) + assert cfg["gens_per_prompt"] == 4 and cfg["shard_id"] == 2 + assert cfg["max_inflight_samples"] == 12 # 3 x gens_per_prompt=4 + with pytest.raises(AssertionError): + resolve_collection_config({}, ppo_cfg) + + +def test_resolve_collection_config_inflight_knobs(): + """max_inflight_samples wins; legacy max_inflight_groups converts.""" + ppo_cfg = {"num_generations_per_prompt": 16} + cfg = resolve_collection_config( + {"out_dir": "/tmp/x", "max_inflight_samples": "64"}, ppo_cfg + ) + assert cfg["max_inflight_samples"] == 64 + cfg = resolve_collection_config( + {"out_dir": "/tmp/x", "max_inflight_groups": 2, "gens_per_prompt": 8}, + ppo_cfg, + ) + assert cfg["max_inflight_samples"] == 16 + # explicit samples takes precedence over the legacy knob + cfg = resolve_collection_config( + {"out_dir": "/tmp/x", "max_inflight_groups": 2, "max_inflight_samples": 5}, + ppo_cfg, + ) + assert cfg["max_inflight_samples"] == 5 + + +def test_aggregate_sample_metrics(): + """Numeric fields are averaged; non-numeric/non-finite dropped.""" + agg = _aggregate_sample_metrics( + [ + {"reward": 1.0, "turns": 4, "note": "text", "hit_max": True, + "len_stddev": float("nan")}, + {"reward": 0.0, "turns": 6, "hit_max": False, + "len_stddev": float("nan")}, + ] + ) + assert agg["reward"] == 0.5 + assert agg["turns"] == 5.0 + assert agg["hit_max"] == 0.5 + assert "note" not in agg + assert "len_stddev" not in agg # single-sample NaN stats must not survive + assert agg["aggregated_from_samples"] == 2.0 + + +def test_assemble_group_payload_orders_and_concats(): + """Per-sample results reassemble into a whole-group-shaped payload.""" + def sample_batch(reward, truncated): + return BatchedDataDict( + { + "message_log": [[{"role": "assistant"}]], + "total_reward": torch.tensor([reward]), + "truncated": torch.tensor([truncated]), + "loss_multiplier": torch.ones(1), + } + ) + + input_batch = BatchedDataDict( + { + "extra_env_info": [{"k": 1}, {"k": 2}], + "idx": [7, 7], + "task_name": ["nemo_gym", "nemo_gym"], + "loss_multiplier": torch.ones(2), + } + ) + payload = assemble_group_payload( + 7, + input_batch, + [sample_batch(1.0, False), sample_batch(0.0, True)], + [{"reward": 1.0}, {"reward": 0.0}], + ) + assert payload["dataset_idx"] == 7 + assert payload["batch"].size == 2 + torch.testing.assert_close( + payload["batch"]["total_reward"], torch.tensor([1.0, 0.0]) + ) + assert payload["batch"]["extra_env_info"] == [{"k": 1}, {"k": 2}] + assert payload["rollout_metrics"]["reward"] == 0.5 From 47f8a99eada3ec16c313922b9ca9c46f36ace164 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 31 Jul 2026 14:40:16 -0700 Subject: [PATCH 37/69] feat(critic): offline critic pretraining on collected rollout shards Stage B: value-only pretraining over the shards written by stage A, gated on held-out explained variance before the critic is handed to PPO. - frozen train-file order (critic_pretrain_files.json) so resume replays the identical stream even if new shards appear later - multi-epoch stream: epoch e is shuffled with Random(seed + e), so epoch 0 reproduces the original one-epoch order exactly and raising num_epochs EXTENDS a frozen stream rather than rewriting it; mismatch raises instead of silently training on a different sequence - a finished run relaunched unchanged now says so and returns, instead of falling through the train loop having done nothing - held-out value dumps (format_version 2) carry per-token strings for a bounded, contrastful subset of samples (dump_token_samples) to power the token-level heatmap without paying a decode per token everywhere - terminal value/reward AUC + positional diagnostics for the gate Co-Authored-By: Claude Opus 5 --- examples/nemo_gym/run_critic_pretrain.py | 75 ++ nemo_rl/algorithms/critic_pretrain.py | 887 ++++++++++++++++++ tests/unit/algorithms/test_critic_pretrain.py | 194 ++++ 3 files changed, 1156 insertions(+) create mode 100644 examples/nemo_gym/run_critic_pretrain.py create mode 100644 nemo_rl/algorithms/critic_pretrain.py create mode 100644 tests/unit/algorithms/test_critic_pretrain.py diff --git a/examples/nemo_gym/run_critic_pretrain.py b/examples/nemo_gym/run_critic_pretrain.py new file mode 100644 index 00000000000..da1817430c0 --- /dev/null +++ b/examples/nemo_gym/run_critic_pretrain.py @@ -0,0 +1,75 @@ +# 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. +"""Offline PPO critic pretraining on stored rollout shards (stage B). + +Reuses the production PPO config (value model, adv_estimator, value_loss_fn) +but initializes ONLY the value model — no policy, no vLLM, no gym. Rollout +shards come from examples/nemo_gym/run_swe_rollout_collection.py; knobs live in +a ``critic_pretrain:`` config block (see +``nemo_rl.algorithms.critic_pretrain.resolve_critic_pretrain_config``). +Launch via scripts/swe/ppo/critic_pretrain.sh. +""" + +import argparse +import pprint + +from omegaconf import OmegaConf + +from nemo_rl.algorithms.critic_pretrain import critic_pretrain +from nemo_rl.algorithms.ppo import MasterConfig +from nemo_rl.algorithms.utils import get_tokenizer +from nemo_rl.distributed.virtual_cluster import init_ray +from nemo_rl.utils.config import ( + load_config, + parse_hydra_overrides, + register_omegaconf_resolvers, +) + + +def parse_args() -> tuple[argparse.Namespace, list[str]]: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Offline PPO critic pretraining on stored rollouts" + ) + parser.add_argument( + "--config", type=str, required=True, help="Path to YAML config file" + ) + args, overrides = parser.parse_known_args() + return args, overrides + + +def main() -> None: + """Main entry point.""" + register_omegaconf_resolvers() + args, overrides = parse_args() + + config = load_config(args.config) + print(f"Loaded configuration from: {args.config}") + if overrides: + print(f"Overrides: {overrides}") + config = parse_hydra_overrides(config, overrides) + config = OmegaConf.to_container(config, resolve=True) + config = MasterConfig(**config) + + print("Final config:") + pprint.pprint(config) + + tokenizer = get_tokenizer(config.policy["tokenizer"]) + + init_ray() + critic_pretrain(config, tokenizer) + + +if __name__ == "__main__": + main() diff --git a/nemo_rl/algorithms/critic_pretrain.py b/nemo_rl/algorithms/critic_pretrain.py new file mode 100644 index 00000000000..e535d57dc44 --- /dev/null +++ b/nemo_rl/algorithms/critic_pretrain.py @@ -0,0 +1,887 @@ +# 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. +"""Offline critic pretraining on stored rollouts (stage B, decoupled PPO). + +Trains ONLY the PPO value model on prompt-group shards written by +:mod:`nemo_rl.algorithms.rollout_collection` — no policy worker, no generation +engines, no gym. Each step mirrors the async PPO loop's critic path exactly +(reward/mask processing -> fresh value forward -> GAE returns -> value train), +so with the production SWE config (``gae_lambda_value=1``, ``gae_gamma=1``, +KL=0) one offline epoch is the same optimization as the online critic warmup. + +Extras the online warmup cannot provide: + * a held-out shard split (``dataset_idx % heldout_mod == 0``) with + explained-variance / calibration / terminal-AUC eval on unseen rollouts; + * checkpoints in the standard layout with ONLY a ``value/`` dir — the PPO + resume path explicitly tolerates a missing ``policy/``, which is how + stage C warm-starts from these checkpoints (scripts/swe/ppo/ + prep_warm_start.sh). + +The train-file ORDER is frozen in the first checkpoint +(``critic_pretrain_files.json``) so resume replays the identical stream even if +new shards appear later. ``critic_pretrain.num_epochs`` (default 1) sets how +many passes over the train split the stream contains; epoch ``e`` is shuffled +with ``Random(seed + e)``, so epoch 0 is byte-identical to the original +one-epoch order and raising ``num_epochs`` on a finished run EXTENDS its frozen +stream (resume continues) rather than rewriting it. +""" + +import json +import os +import random +import time +from pathlib import Path +from typing import Any, Optional + +import numpy as np +import torch + +from nemo_rl.algorithms.rollout_collection import load_group, parse_group_index +from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + +FILE_LIST_NAME = "critic_pretrain_files.json" + + +# =============================================================================== +# Pure helpers (unit-tested, no heavy deps) +# =============================================================================== +def list_group_files(shards_dir: str | Path) -> list[Path]: + """All group files under ``shards_dir`` (searches shard_*/ and the dir itself).""" + shards_dir = Path(shards_dir) + files = sorted(shards_dir.glob("shard_*/group_*.pt")) + sorted( + shards_dir.glob("group_*.pt") + ) + return [f for f in files if parse_group_index(f.name) is not None] + + +def split_heldout( + files: list[Path], heldout_mod: int +) -> tuple[list[Path], list[Path]]: + """Deterministic train/held-out split by dataset index. + + ``dataset_idx % heldout_mod == 0`` goes to held-out; ``heldout_mod <= 0`` + disables the split (everything trains). + """ + if heldout_mod <= 0: + return list(files), [] + train, heldout = [], [] + for f in files: + idx = parse_group_index(f.name) + (heldout if idx % heldout_mod == 0 else train).append(f) + return train, heldout + + +def build_epoch_stream( + base_files: list[Path], num_epochs: int, seed: int +) -> list[Path]: + """Concatenate ``num_epochs`` independent shuffles of ``base_files``. + + Epoch ``e`` is shuffled with ``Random(seed + e)``, which makes two + properties hold and both matter: + + * epoch 0 reproduces the original single-epoch order EXACTLY (that order + was ``Random(seed).shuffle(train_files)``), so raising ``num_epochs`` on + an existing run extends the stream instead of rewriting it; + * any prefix of the stream is a pure function of (base set, seed, + epoch index), so resume replays the consumed prefix identically. + + Each epoch is a fresh permutation rather than a repeat of the same order, so + the model does not see the same batch composition twice. + """ + stream: list[Path] = [] + for e in range(num_epochs): + epoch = list(base_files) + random.Random(seed + e).shuffle(epoch) + stream += epoch + return stream + + +def terminal_value_reward_auc( + values: torch.Tensor, + rewards: torch.Tensor, + token_mask: torch.Tensor, + positive_threshold: float = 0.5, +) -> float: + """AUC of the LAST response token's value as a predictor of success. + + Rank-based (Mann-Whitney) AUC with tie correction; returns nan when the + batch has a single outcome class. This is the "end-verification" critic + quality signal from the privileged-critic analyses. + """ + mask = token_mask.bool() + has_response = mask.any(dim=1) + if int(has_response.sum()) < 2: + return float("nan") + last_idx = mask.shape[1] - 1 - mask.fliplr().float().argmax(dim=1) + v = values[has_response, last_idx[has_response]].float() + y = (rewards[has_response].float() >= positive_threshold).float() + n_pos, n_neg = int(y.sum()), int((1 - y).sum()) + if n_pos == 0 or n_neg == 0: + return float("nan") + order = torch.argsort(v) + ranks = torch.empty_like(v) + ranks[order] = torch.arange(1, v.numel() + 1, dtype=v.dtype) + # midranks for ties + for val in torch.unique(v): + tie = v == val + if int(tie.sum()) > 1: + ranks[tie] = ranks[tie].mean() + auc = (ranks[y.bool()].sum().item() - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg) + return float(auc) + + +def verify_shard_meta( + shards_dir: str | Path, master_config: Any, tokenizer: Any +) -> None: + """Assert stored-shard provenance matches this run's model/tokenizer. + + Shards are token-id level: a different base model, tokenizer/chat template, + or max sequence length silently invalidates them. Checks every shard + meta.json written by stage A; missing meta files only warn (older shards). + """ + from nemo_rl.algorithms.rollout_collection import _sha256 + + metas = sorted(Path(shards_dir).glob("shard_*/meta.json")) + if not metas: + print(f"⚠️ No shard meta.json found under {shards_dir}; skipping provenance check.") + return + expected = { + "model_name": master_config.policy["model_name"], + "chat_template_sha256": _sha256(getattr(tokenizer, "chat_template", None) or ""), + "max_total_sequence_length": master_config.policy["max_total_sequence_length"], + } + for meta_path in metas: + with open(meta_path) as f: + meta = json.load(f) + for key, want in expected.items(): + got = meta.get(key) + assert got == want, ( + f"Shard provenance mismatch in {meta_path}: {key}={got!r} but this " + f"run expects {want!r}. Shards are token-id level and are only " + "valid for the exact model/tokenizer/max-length they were " + "generated with." + ) + print(f" ✓ Shard provenance verified ({len(metas)} shard meta files)") + + +def resolve_critic_pretrain_config( + raw: Optional[dict[str, Any]], ppo_config: dict[str, Any] +) -> dict[str, Any]: + """Fill defaults for the ``critic_pretrain:`` config block.""" + cfg = dict(raw or {}) + assert cfg.get("shards_dir"), ( + "critic_pretrain.shards_dir is required " + "(pass ++critic_pretrain.shards_dir=)" + ) + cfg.setdefault("groups_per_step", ppo_config["num_prompts_per_step"]) + cfg.setdefault("heldout_mod", 16) + cfg.setdefault("eval_period", 10) + cfg.setdefault("heldout_max_groups", cfg["groups_per_step"]) + cfg.setdefault("max_steps", None) + # Passes over the train split. 1 = the original one-epoch behaviour; epoch e + # is a fresh permutation seeded with (seed + e), so raising this on a + # finished run EXTENDS its frozen stream rather than rewriting it. + cfg.setdefault("num_epochs", 1) + cfg.setdefault("seed", ppo_config.get("seed", 42)) + # Eval/dump mode: no training — load a specific checkpoint, score the + # held-out groups, and dump per-token values aligned to message spans for + # offline value-vs-behavior analysis. + cfg.setdefault("eval_only", False) + cfg.setdefault("eval_checkpoint_path", None) + cfg.setdefault("dump_dir", None) # default: /value_dumps + cfg.setdefault("dump_text_groups", 8) # decode message text for first N groups + # per-token strings (for the token-level HTML heatmap) are ~35k/sample, so + # store them only for a bounded, contrastful subset of samples per text group + cfg.setdefault("dump_token_samples", 4) + for key in ( + "groups_per_step", + "heldout_mod", + "eval_period", + "heldout_max_groups", + "num_epochs", + ): + cfg[key] = int(cfg[key]) + assert cfg["num_epochs"] >= 1, ( + f"critic_pretrain.num_epochs must be >= 1, got {cfg['num_epochs']}" + ) + return cfg + + +def message_spans(message_log: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Per-message (role, start, len) spans within a sample's flattened tokens. + + Mirrors ``batched_message_log_to_flat_message``'s concatenation order, so + span positions index directly into the flat per-token value/mask tensors. + """ + spans = [] + pos = 0 + for m in message_log: + n = len(m["token_ids"]) + spans.append({"role": m["role"], "start": pos, "len": n}) + pos += n + return spans + + +# =============================================================================== +# Batch construction (mirrors the async PPO loop's steps 2-3 minus logprobs) +# =============================================================================== +def build_value_train_data( + groups: list[dict[str, Any]], + tokenizer: Any, + master_config: Any, +) -> tuple[BatchedDataDict, BatchedDataDict]: + """Assemble (train_data, repeated_batch) from loaded group payloads. + + Follows async_ppo_train's reward-processing + inline loss-mask block + verbatim (overlong filtering, env-flagged sample masking, unmask ALL + assistant messages) so offline critic batches are bit-identical in shape + and masking to what the online warmup trains on. Policy/reference logprobs + are deliberately absent: the critic path never consumes them (KL-in-reward + handles logprobs=None) and computing them is the warmup's main train-node + waste. + """ + per_prompt_batches = [g["batch"] for g in groups] + repeated_batch = BatchedDataDict.from_batches(per_prompt_batches) + + use_overlong_filtering = master_config.ppo["overlong_filtering"] + if use_overlong_filtering: + loss_multiplier = repeated_batch["loss_multiplier"].clone() + truncated = repeated_batch["truncated"] + if isinstance(truncated, list): + truncated = torch.tensor(truncated, dtype=torch.bool) + loss_multiplier[truncated] = 0 + repeated_batch["loss_multiplier"] = loss_multiplier + + if "mask_sample" in repeated_batch: + loss_multiplier = repeated_batch["loss_multiplier"].clone() + mask_sample = repeated_batch["mask_sample"] + if isinstance(mask_sample, list): + mask_sample = torch.tensor(mask_sample, dtype=torch.bool) + loss_multiplier[mask_sample.bool()] = 0 + repeated_batch["loss_multiplier"] = loss_multiplier + + # PPO's inline loss-mask setup: unmask all assistant messages. + for message_log in repeated_batch["message_log"]: + for message in message_log: + if message["role"] == "assistant": + message["token_loss_mask"] = torch.ones_like(message["token_ids"]) + else: + message["token_loss_mask"] = torch.zeros_like(message["token_ids"]) + if "generation_logprobs" not in message: + message["generation_logprobs"] = torch.zeros_like( + message["token_ids"], dtype=torch.float32 + ) + + flat_messages, input_lengths = batched_message_log_to_flat_message( + repeated_batch["message_log"], + pad_value_dict={"token_ids": tokenizer.pad_token_id}, + make_sequence_length_divisible_by=master_config.policy[ + "make_sequence_length_divisible_by" + ], + ) + + train_data = BatchedDataDict( + { + "input_ids": flat_messages["token_ids"], + "input_lengths": input_lengths, + "rewards": repeated_batch["total_reward"], + "token_mask": flat_messages["token_loss_mask"], + "sample_mask": repeated_batch["loss_multiplier"], + } + ) + train_data.to("cpu") + return train_data, repeated_batch + + +# =============================================================================== +# Value forward + returns (shared by train and held-out eval) +# =============================================================================== +def _forward_values_and_returns( + value_model: Any, + adv_estimator: Any, + train_data: BatchedDataDict, + repeated_batch: BatchedDataDict, + tokenizer: Any, + master_config: Any, +) -> Optional[BatchedDataDict]: + """Populate train_data['values'/'returns'] in place. + + Returns the privileged critic batch (with its own values/returns/token_mask) + when the privileged critic is enabled, else None. + + Mirrors async_ppo_train steps 3 (value inference, incl. the privileged + answer-conditioned remap) and 6 (GAE returns; logprobs=None is valid for + the critic path since KL-in-reward is the only consumer of logprobs). + """ + from nemo_rl.algorithms.grpo import extract_initial_prompt_messages + from nemo_rl.algorithms.privileged_critic import ( + build_privileged_value_inputs, + remap_by_response_mask, + ) + + privileged_critic_cfg = master_config.value.get("privileged_critic") + if privileged_critic_cfg is not None and not privileged_critic_cfg.get("enabled"): + privileged_critic_cfg = None + critic_batch = None + + value_model.prepare_for_inference() + if privileged_critic_cfg is not None: + critic_batch = build_privileged_value_inputs( + repeated_batch, + tokenizer, + privileged_critic_cfg, + make_seq_len_divisible_by=master_config.policy[ + "make_sequence_length_divisible_by" + ], + ) + vals_aug = value_model.get_values(critic_batch)["values"].squeeze(-1) + critic_batch["values"] = vals_aug + train_data["values"] = remap_by_response_mask( + vals_aug, + critic_batch["token_mask"], + train_data["token_mask"], + ) + else: + train_data["values"] = value_model.get_values(train_data)["values"].squeeze( + -1 + ) + value_model.finish_inference() + + initial_prompt_message_logs = extract_initial_prompt_messages( + repeated_batch["message_log"], + repeated_batch["length"], + ) + prompt_batched_flat, _ = batched_message_log_to_flat_message( + initial_prompt_message_logs, + pad_value_dict={"token_ids": tokenizer.pad_token_id}, + ) + advantages, returns = adv_estimator.compute_advantage( + prompt_ids=prompt_batched_flat["token_ids"], + rewards=train_data["rewards"], + mask=train_data["token_mask"], + values=train_data["values"], + reference_logprobs=None, + logprobs=None, + ) + del advantages # critic pretraining has no actor; only returns are used + train_data["returns"] = returns + if critic_batch is not None: + critic_batch["returns"] = remap_by_response_mask( + returns, + train_data["token_mask"], + critic_batch["token_mask"], + ) + critic_batch["sample_mask"] = train_data["sample_mask"] + return critic_batch + + +def _heldout_metrics( + value_model: Any, + adv_estimator: Any, + heldout_files: list[Path], + tokenizer: Any, + master_config: Any, +) -> dict[str, float]: + """Critic quality on held-out rollouts: EV, positional EV/ECE, terminal AUC.""" + from nemo_rl.algorithms.ppo import _positional_value_metrics + + groups = [load_group(p) for p in heldout_files] + train_data, repeated_batch = build_value_train_data( + groups, tokenizer, master_config + ) + _forward_values_and_returns( + value_model, adv_estimator, train_data, repeated_batch, tokenizer, + master_config, + ) + values, returns = train_data["values"], train_data["returns"] + mask = train_data["token_mask"].bool() + metrics: dict[str, float] = {} + if int(mask.sum()) >= 2: + v, r = values[mask].float(), returns[mask].float() + var_r = r.var(unbiased=False) + metrics["critic/explained_var"] = ( + (1.0 - (r - v).var(unbiased=False) / var_r).item() + if var_r > 1e-8 + else 0.0 + ) + metrics["critic/mse"] = ((r - v) ** 2).mean().item() + metrics.update( + _positional_value_metrics(values, returns, train_data["token_mask"]) + ) + metrics["critic/terminal_auc"] = terminal_value_reward_auc( + values, train_data["rewards"], train_data["token_mask"] + ) + metrics["reward"] = train_data["rewards"].float().mean().item() + metrics["num_heldout_samples"] = float(train_data["input_ids"].shape[0]) + return metrics + + +def _dump_heldout_values( + value_model: Any, + adv_estimator: Any, + heldout_files: list[Path], + tokenizer: Any, + master_config: Any, + dump_dir: Path, + dump_text_groups: int, + dump_token_samples: int = 4, +) -> None: + """Score held-out groups with the loaded critic and dump per-token values. + + Values/returns are packed over response tokens; per-message spans (with + decoded text for the first ``dump_text_groups`` groups) let offline + analysis align value movements to agent/tool behavior in the trajectory. + """ + dump_dir.mkdir(parents=True, exist_ok=True) + for gi, path in enumerate(heldout_files): + g = load_group(path) + train_data, repeated_batch = build_value_train_data( + [g], tokenizer, master_config + ) + _forward_values_and_returns( + value_model, adv_estimator, train_data, repeated_batch, tokenizer, + master_config, + ) + mask = train_data["token_mask"].bool() + coords = mask.nonzero(as_tuple=False) + with_text = gi < dump_text_groups + # per-token strings power the token-level HTML heatmap but cost a decode + # per token (~35k/sample), so carry them only for a bounded, contrastful + # subset: successes first, then fails, capped at dump_token_samples. + render_samples = [] + if with_text and dump_token_samples > 0: + rew = train_data["rewards"].float().tolist() + succ = [i for i in range(len(rew)) if rew[i] > 0.5] + fail = [i for i in range(len(rew)) if rew[i] <= 0.5] + half = max(1, dump_token_samples // 2) + render_samples = (succ[:half] + fail[: dump_token_samples - len(succ[:half])]) + render_samples = sorted(render_samples[:dump_token_samples]) + render_set = set(render_samples) + samples_msgs = [] + for si, ml in enumerate(repeated_batch["message_log"]): + spans = message_spans(ml) + if with_text: + want_toks = si in render_set + for m, s in zip(ml, spans): + s["text"] = tokenizer.decode(m["token_ids"]) + if want_toks and m["role"] == "assistant": + s["toks"] = [ + tokenizer.decode([int(t)]) for t in m["token_ids"] + ] + samples_msgs.append(spans) + payload = { + "format_version": 2, + "dataset_idx": g["dataset_idx"], + "source_file": str(path), + "rewards": train_data["rewards"].float().cpu(), + "sample_mask": train_data["sample_mask"].float().cpu(), + "token_sample_index": coords[:, 0].to(torch.int32), + "token_position": coords[:, 1].to(torch.int32), + "values": train_data["values"][mask].to(torch.float16).cpu(), + "returns": train_data["returns"][mask].to(torch.float16).cpu(), + "messages": samples_msgs, + "has_text": with_text, + "render_samples": render_samples, + } + out = dump_dir / f"valuedump_{g['dataset_idx']:08d}.pt" + torch.save(payload, out) + if (gi + 1) % 10 == 0 or gi + 1 == len(heldout_files): + print(f" 💾 dumped {gi + 1}/{len(heldout_files)} groups", flush=True) + + +# =============================================================================== +# Main entry point +# =============================================================================== +def critic_pretrain(master_config: Any, tokenizer: Any) -> None: + """Set up the value model and run offline critic pretraining. + + Heavy setup (Ray cluster, Megatron value workers, checkpointing) lives here + rather than in a separate setup() so the driver stays thin; the module-level + helpers above stay importable without Ray/Megatron for unit tests. + """ + from pathlib import Path as _Path + + from nemo_rl.algorithms.loss.loss_functions import MseValueLossFn + from nemo_rl.algorithms.ppo import ( + _compute_critic_metrics, + _create_advantage_estimator, + _positional_value_metrics, + _resolve_resume_optimizer_path, + ) + from nemo_rl.distributed.virtual_cluster import RayVirtualCluster + from nemo_rl.models.value.lm_value import Value + from nemo_rl.utils.checkpoint import CheckpointManager + from nemo_rl.utils.logger import Logger + + cp_config = resolve_critic_pretrain_config( + getattr(master_config, "critic_pretrain", None), master_config.ppo + ) + value_config = master_config.value + cluster_config = master_config.cluster + + # Known, inherent divergence from the online warmup: the seq-level + # train/inference logprob-error masking (ppo.seq_logprob_error_threshold) + # needs policy-engine logprobs, which a value-only job cannot compute. The + # online loop zeroes sample_mask for badly mismatched sequences; offline + # those sequences stay in the critic loss. + if master_config.ppo.get("seq_logprob_error_threshold") is not None: + print( + "⚠️ ppo.seq_logprob_error_threshold is set, but offline critic " + "pretraining cannot apply seq-logprob-error masking (no policy " + "worker). Sequences the online warmup would mask are trained on." + ) + + # Privileged critic scores [prompt + answer + response]: raise the value + # model's sequence/packing budgets exactly as ppo.setup() does, so + # answer-augmented near-max-length samples fit the packing bins. + _privileged_critic = value_config.get("privileged_critic") + if _privileged_critic is not None and _privileged_critic.get("enabled"): + _needed = ( + master_config.policy["max_total_sequence_length"] + + int(_privileged_critic.get("max_answer_tokens", 256) or 0) + + 128 # grader-note template + chat re-render slack + ) + if value_config["max_total_sequence_length"] < _needed: + print( + " ↑ privileged critic: raising value.max_total_sequence_length " + f"{value_config['max_total_sequence_length']} -> {_needed}", + flush=True, + ) + value_config["max_total_sequence_length"] = _needed + for _bcfg_key in ("sequence_packing", "dynamic_batching"): + _bcfg = value_config.get(_bcfg_key) or {} + if not _bcfg.get("enabled"): + continue + for _tok_key in ("train_mb_tokens", "logprob_mb_tokens"): + if _bcfg.get(_tok_key) is not None and _bcfg[_tok_key] < _needed: + print( + f" ↑ privileged critic: raising value.{_bcfg_key}." + f"{_tok_key} {_bcfg[_tok_key]} -> {_needed}", + flush=True, + ) + _bcfg[_tok_key] = _needed + + logger = Logger(master_config.logger) + logger.log_hyperparams(master_config.model_dump()) + + checkpointer = CheckpointManager(master_config.checkpointing) + last_checkpoint_path = checkpointer.get_latest_checkpoint_path() + save_state = checkpointer.load_training_info(last_checkpoint_path) or { + "total_steps": 0, + "groups_consumed": 0, + "consumed_samples": 0, + } + step = int(save_state["total_steps"]) + + # ------------------------------------------------------------------ + # Shard discovery + frozen multi-epoch order (replayed exactly on resume). + # ------------------------------------------------------------------ + all_files = list_group_files(cp_config["shards_dir"]) + assert all_files, f"No group files found under {cp_config['shards_dir']}" + verify_shard_meta(cp_config["shards_dir"], master_config, tokenizer) + num_epochs = cp_config["num_epochs"] + frozen = None + if last_checkpoint_path is not None: + file_list_path = os.path.join(last_checkpoint_path, FILE_LIST_NAME) + if os.path.exists(file_list_path): + with open(file_list_path) as f: + frozen = json.load(f) + if frozen is not None: + heldout_files = [_Path(p) for p in frozen["heldout"]] + frozen_train = [_Path(p) for p in frozen["train"]] + # The frozen stream is `num_epochs_then` shuffles of a base set; recover + # the base in its pre-shuffle order (list_group_files sorts by path) and + # regenerate the stream for the num_epochs asked for NOW. + base_train = [_Path(p) for p in sorted({str(p) for p in frozen_train})] + train_files = build_epoch_stream(base_train, num_epochs, cp_config["seed"]) + n_frozen = len(frozen_train) + # Fail loud rather than train on a different sequence than the + # checkpoint recorded: the regenerated stream MUST reproduce the frozen + # one as a prefix, otherwise resume would silently replay other data. + if len(train_files) < n_frozen or train_files[:n_frozen] != frozen_train: + raise ValueError( + f"Cannot reconcile critic_pretrain.num_epochs={num_epochs} with the " + f"frozen stream in {last_checkpoint_path}: the regenerated order does " + f"not reproduce its {n_frozen} entries as a prefix, so resuming would " + "train on a different sequence than the checkpoint recorded. Check " + "that critic_pretrain.seed / heldout_mod / shards_dir are unchanged " + f"(seed={cp_config['seed']}, heldout_mod={cp_config['heldout_mod']}), " + f"and that num_epochs is not below the {n_frozen // max(len(base_train), 1)} " + "epoch(s) already frozen." + ) + if len(train_files) > n_frozen: + print( + f"↻ Extending the frozen stream: {n_frozen} -> {len(train_files)} " + f"groups ({num_epochs} epochs x {len(base_train)} train groups)." + ) + newly_seen = len(all_files) - len(base_train) - len(heldout_files) + if newly_seen > 0: + print( + f"ℹ️ {newly_seen} group files appeared after the file list was " + "frozen; they are ignored this run (frozen-dataset semantics)." + ) + else: + base_train, heldout_files = split_heldout( + all_files, cp_config["heldout_mod"] + ) + train_files = build_epoch_stream(base_train, num_epochs, cp_config["seed"]) + missing = [ + p for p in set(train_files) | set(heldout_files) if not os.path.exists(p) + ] + assert not missing, f"{len(missing)} frozen group files are missing, e.g. {missing[:3]}" + + groups_per_step = cp_config["groups_per_step"] + # drop-last within the whole multi-epoch stream + planned_steps = len(train_files) // groups_per_step + if cp_config["max_steps"] is not None: + planned_steps = min(planned_steps, int(cp_config["max_steps"])) + assert planned_steps > 0, ( + f"Not enough train groups ({len(train_files)}) for one step of " + f"{groups_per_step} groups." + ) + print( + f"📚 {len(base_train)} train groups x {num_epochs} epoch(s) = " + f"{len(train_files)}, {len(heldout_files)} held-out groups " + f"-> {planned_steps} steps of {groups_per_step} groups (resuming at {step})" + ) + # A finished run relaunched unchanged used to fall straight through the + # train loop and exit having done nothing. Say so instead of exiting silently. + if step >= planned_steps: + print( + f"\n✅ Nothing to do: {step} steps already completed and this config " + f"plans {planned_steps} (num_epochs={num_epochs}, " + f"max_steps={cp_config['max_steps']}). Raise " + "++critic_pretrain.num_epochs to train further.\n" + ) + return + + # Scheduler budget: one tick per train() call, one call per step. + if value_config.get("megatron_cfg", {}).get("enabled", False): + value_config["megatron_cfg"]["train_iters"] = planned_steps + + # ------------------------------------------------------------------ + # Cluster + value model (mirrors ppo.setup()'s init_value resume probe). + # ------------------------------------------------------------------ + cluster = RayVirtualCluster( + name="critic_pretrain_cluster", + bundle_ct_per_node_list=[cluster_config["gpus_per_node"]] + * cluster_config["num_nodes"], + use_gpus=True, + num_gpus_per_node=cluster_config["gpus_per_node"], + max_colocated_worker_groups=1, + port_range_low=cluster_config.get("master_port_range_low"), + port_range_high=cluster_config.get("master_port_range_high"), + ) + print( + f" ✓ Ray cluster: {cluster_config['num_nodes']} nodes x " + f"{cluster_config['gpus_per_node']} GPUs (value model only)" + ) + + eval_only = bool(cp_config.get("eval_only")) + if eval_only and cp_config.get("eval_checkpoint_path"): + # Eval/dump mode scores with an EXPLICIT checkpoint (e.g. .../step_10), + # independent of this dir's latest; no optimizer needed. + _value_weights = _Path(cp_config["eval_checkpoint_path"]) / "value" / "weights" + assert _value_weights.exists(), ( + f"eval_checkpoint_path has no value/weights: {_value_weights}" + ) + value_weights_path = _value_weights + value_optimizer_path = None + print(f" ✓ Eval mode: loading critic from {value_weights_path}") + elif last_checkpoint_path: + _value_weights = _Path(last_checkpoint_path) / "value" / "weights" + _value_optim = _Path(last_checkpoint_path) / "value" / "optimizer" + value_weights_path = _value_weights if _value_weights.exists() else None + value_optimizer_path = _resolve_resume_optimizer_path( + _value_optim, value_weights_path, value_config + ) + if value_weights_path is not None: + print(f" ✓ Resuming value model from: {value_weights_path}") + else: + value_weights_path = None + value_optimizer_path = None + + value_model = Value( + cluster=cluster, + config=value_config, + tokenizer=tokenizer, + name_prefix="lm_value", + weights_path=value_weights_path, + optimizer_path=value_optimizer_path, + init_optimizer=not eval_only, + ) + value_model.finish_training() # block init, offload until first use + print(" ✓ Value model initialized") + + value_loss_fn = MseValueLossFn(master_config.value_loss_fn) + adv_estimator = _create_advantage_estimator(master_config) + + if eval_only: + dump_dir = _Path( + cp_config.get("dump_dir") + or os.path.join(master_config.checkpointing["checkpoint_dir"], "value_dumps") + ) + print( + f"🔍 Eval-only: scoring {len(heldout_files)} held-out groups -> {dump_dir}" + ) + val_metrics = _heldout_metrics( + value_model, adv_estimator, heldout_files, tokenizer, master_config + ) + logger.log_metrics(val_metrics, 0, prefix="validation") + print(" heldout metrics:", {k: round(v, 4) for k, v in val_metrics.items()}) + _dump_heldout_values( + value_model, + adv_estimator, + heldout_files, + tokenizer, + master_config, + dump_dir, + int(cp_config["dump_text_groups"]), + int(cp_config["dump_token_samples"]), + ) + with open(dump_dir / "summary.json", "w") as f: + json.dump( + { + "checkpoint": str(value_weights_path), + "num_groups": len(heldout_files), + "metrics": val_metrics, + }, + f, + indent=2, + ) + print(f"🏁 Eval dump complete: {dump_dir}") + return + + expected_gbs = value_config["train_global_batch_size"] + save_period = master_config.checkpointing["save_period"] + checkpointing_enabled = master_config.checkpointing["enabled"] + eval_period = cp_config["eval_period"] + heldout_eval_files = heldout_files[: cp_config["heldout_max_groups"]] + + # ------------------------------------------------------------------ + # Train loop: one pass over the frozen file order. + # ------------------------------------------------------------------ + while step < planned_steps: + step_start = time.perf_counter() + print(f"\n{'=' * 25} Critic step {step + 1}/{planned_steps} {'=' * 25}") + + step_files = train_files[step * groups_per_step : (step + 1) * groups_per_step] + groups = [load_group(p) for p in step_files] + train_data, repeated_batch = build_value_train_data( + groups, tokenizer, master_config + ) + if train_data["input_ids"].shape[0] != expected_gbs: + raise ValueError( + f"Step batch has {train_data['input_ids'].shape[0]} samples but " + f"value.train_global_batch_size={expected_gbs}. Override " + "value.train_global_batch_size (and critic_pretrain." + "groups_per_step) to match groups_per_step * gens_per_prompt " + "of the stored shards." + ) + + print("▶ Computing values...") + critic_batch = _forward_values_and_returns( + value_model, adv_estimator, train_data, repeated_batch, tokenizer, + master_config, + ) + + print("▶ Training critic...") + value_model.prepare_for_training() + value_train_batch = critic_batch if critic_batch is not None else train_data + value_results = value_model.train( + value_train_batch, value_loss_fn + ) + value_model.finish_training() + + # ---- Metrics ---- + metrics = _compute_critic_metrics(value_results) + # critic/loss and critic/grad_norm come back as numpy arrays; the async + # loop scalarizes ndarray metrics before printing/logging — mirror that. + for k, v in metrics.items(): + if isinstance(v, (np.ndarray, list)): + metrics[k] = np.sum(v).item() + metrics.update( + _positional_value_metrics( + train_data["values"], train_data["returns"], train_data["token_mask"] + ) + ) + metrics["reward"] = train_data["rewards"].float().mean().item() + metrics["num_samples"] = float(train_data["input_ids"].shape[0]) + metrics["total_step_time"] = time.perf_counter() - step_start + logger.log_metrics(metrics, step + 1, prefix="train") + print( + f" step {step + 1}: loss={metrics.get('critic/loss'):.6f} " + f"ev={metrics.get('critic/explained_var'):.4f} " + f"reward={metrics['reward']:.3f} " + f"({metrics['total_step_time']:.1f}s)" + ) + + # ---- Held-out eval ---- + is_last_step = step + 1 == planned_steps + if heldout_eval_files and ( + (eval_period > 0 and (step + 1) % eval_period == 0) or is_last_step + ): + print("🔍 Held-out eval...") + val_metrics = _heldout_metrics( + value_model, + adv_estimator, + heldout_eval_files, + tokenizer, + master_config, + ) + logger.log_metrics(val_metrics, step + 1, prefix="validation") + print( + f" heldout: ev={val_metrics.get('critic/explained_var', float('nan')):.4f} " + f"terminal_auc={val_metrics.get('critic/terminal_auc', float('nan')):.4f}" + ) + + # ---- Checkpoint (value/ only — stage C's warm-start seed layout) ---- + step += 1 + save_state["total_steps"] = step + save_state["groups_consumed"] = step * groups_per_step + save_state["consumed_samples"] = save_state.get("consumed_samples", 0) + int( + train_data["input_ids"].shape[0] + ) + if checkpointing_enabled and ( + is_last_step or step % save_period == 0 + ): + print(f"💾 Saving checkpoint for step {step}...") + checkpoint_path = checkpointer.init_tmp_checkpoint( + step, save_state, master_config + ) + value_model.prepare_for_training() + value_model.save_checkpoint( + weights_path=os.path.join(checkpoint_path, "value", "weights"), + optimizer_path=os.path.join(checkpoint_path, "value", "optimizer"), + tokenizer_path=os.path.join(checkpoint_path, "value", "tokenizer"), + checkpointing_cfg=master_config.checkpointing, + ) + value_model.finish_training() + with open(os.path.join(checkpoint_path, FILE_LIST_NAME), "w") as f: + json.dump( + { + "train": [str(p) for p in train_files], + "heldout": [str(p) for p in heldout_files], + }, + f, + ) + checkpointer.finalize_checkpoint(checkpoint_path) + print(f" ✓ Checkpoint saved: step_{step}") + + print( + f"\n🏁 Critic pretraining complete: {step} steps, " + f"{save_state['consumed_samples']} samples. Latest checkpoint: " + f"{checkpointer.get_latest_checkpoint_path()}" + ) diff --git a/tests/unit/algorithms/test_critic_pretrain.py b/tests/unit/algorithms/test_critic_pretrain.py new file mode 100644 index 00000000000..cabcee97dfe --- /dev/null +++ b/tests/unit/algorithms/test_critic_pretrain.py @@ -0,0 +1,194 @@ +# 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. + +import math +from pathlib import Path +from types import SimpleNamespace + +import torch + +from nemo_rl.algorithms.critic_pretrain import ( + build_value_train_data, + split_heldout, + terminal_value_reward_auc, +) + + +def _paths(indices): + return [Path(f"shard_000/group_{i:08d}.pt") for i in indices] + + +def test_split_heldout_partition(): + """idx % mod == 0 goes to held-out; mod<=0 disables the split.""" + files = _paths(range(33)) + train, heldout = split_heldout(files, 16) + assert {int(p.name[6:14]) for p in heldout} == {0, 16, 32} + assert len(train) + len(heldout) == len(files) + assert set(train).isdisjoint(heldout) + + train_all, heldout_none = split_heldout(files, 0) + assert train_all == files and heldout_none == [] + + +def test_terminal_auc_perfect_and_inverted(): + """Terminal-token value ranking success perfectly => AUC 1; inverted => 0.""" + B, S = 4, 5 + mask = torch.ones(B, S) + rewards = torch.tensor([1.0, 0.0, 1.0, 0.0]) + values = torch.zeros(B, S) + values[:, -1] = torch.tensor([0.9, 0.1, 0.8, 0.2]) # separates classes + assert terminal_value_reward_auc(values, rewards, mask) == 1.0 + values[:, -1] = torch.tensor([0.1, 0.9, 0.2, 0.8]) # inverted + assert terminal_value_reward_auc(values, rewards, mask) == 0.0 + + +def test_terminal_auc_single_class_and_ties(): + """Single outcome class => nan; fully tied scores => 0.5.""" + B, S = 3, 4 + mask = torch.ones(B, S) + assert math.isnan( + terminal_value_reward_auc(torch.rand(B, S), torch.ones(B), mask) + ) + values = torch.full((4, S), 0.5) + rewards = torch.tensor([1.0, 0.0, 1.0, 0.0]) + auc = terminal_value_reward_auc(values, rewards, torch.ones(4, S)) + assert abs(auc - 0.5) < 1e-6 + + +def test_terminal_auc_respects_mask(): + """The scored token is the LAST masked position, not the last column.""" + B, S = 2, 6 + mask = torch.zeros(B, S) + mask[0, 1:3] = 1 # response ends at position 2 + mask[1, 1:5] = 1 # response ends at position 4 + values = torch.zeros(B, S) + values[0, 2] = 0.9 # positive sample's terminal value + values[1, 4] = 0.1 + values[:, -1] = torch.tensor([0.0, 1.0]) # decoys at the last column + rewards = torch.tensor([1.0, 0.0]) + assert terminal_value_reward_auc(values, rewards, mask) == 1.0 + + +def _make_group(rewards, truncated, mask_sample, prompt_len=3, resp_len=4): + """Build a group payload shaped like a stored stage-A shard.""" + n = len(rewards) + message_log = [] + for _ in range(n): + message_log.append( + [ + {"role": "user", "token_ids": torch.arange(prompt_len)}, + {"role": "assistant", "token_ids": torch.arange(resp_len)}, + ] + ) + batch = { + "message_log": message_log, + "length": torch.full((n,), prompt_len), + "total_reward": torch.tensor(rewards, dtype=torch.float32), + "loss_multiplier": torch.ones(n), + "truncated": torch.tensor(truncated, dtype=torch.bool), + "mask_sample": torch.tensor(mask_sample, dtype=torch.bool), + "idx": [0] * n, + "task_name": ["nemo_gym"] * n, + "extra_env_info": [{} for _ in range(n)], + } + return {"format_version": 1, "dataset_idx": 0, "batch": batch} + + +def _master_config(overlong_filtering=True): + return SimpleNamespace( + ppo={"overlong_filtering": overlong_filtering}, + policy={"make_sequence_length_divisible_by": 1}, + ) + + +def test_build_value_train_data_masks_and_shapes(): + """Assistant-only token mask; truncated/env-flagged samples get + sample_mask=0 (matching the async PPO loop's processing).""" + tokenizer = SimpleNamespace(pad_token_id=0) + g1 = _make_group([1.0, 0.0], truncated=[False, True], mask_sample=[False, False]) + g2 = _make_group([0.0, 1.0], truncated=[False, False], mask_sample=[True, False]) + train_data, repeated_batch = build_value_train_data( + [g1, g2], tokenizer, _master_config() + ) + + assert train_data["input_ids"].shape[0] == 4 + assert train_data["input_ids"].shape == train_data["token_mask"].shape + # 3 prompt tokens masked out, 4 assistant tokens unmasked, per sample + assert torch.equal( + train_data["token_mask"].sum(dim=1), torch.full((4,), 4.0) + ) + # sample 1 truncated (overlong filtering), sample 2 env-flagged + torch.testing.assert_close( + train_data["sample_mask"], torch.tensor([1.0, 0.0, 0.0, 1.0]) + ) + torch.testing.assert_close( + train_data["rewards"], torch.tensor([1.0, 0.0, 0.0, 1.0]) + ) + assert repeated_batch["message_log"][0][1]["token_loss_mask"].sum() == 4 + + +def test_build_value_train_data_no_overlong_filtering(): + """With overlong_filtering off, truncated samples keep sample_mask=1.""" + tokenizer = SimpleNamespace(pad_token_id=0) + g = _make_group([1.0, 0.0], truncated=[True, True], mask_sample=[False, False]) + train_data, _ = build_value_train_data( + [g], tokenizer, _master_config(overlong_filtering=False) + ) + torch.testing.assert_close(train_data["sample_mask"], torch.ones(2)) + + +def test_offline_returns_are_reward_to_go(): + """The stage-B invariant: with gae_lambda_value=1, gae_gamma=1, KL off, the + critic's regression targets equal the terminal reward broadcast over + response tokens — independent of the values fed in (so offline pretraining + on stored rollouts is exactly the online warmup's optimization).""" + from nemo_rl.algorithms.advantage_estimator import ( + GeneralizedAdvantageEstimator, + ) + + estimator = GeneralizedAdvantageEstimator( + { + "name": "gae", + "gae_lambda": 1.0, + "gae_gamma": 1, + "normalize_advantages": True, + "gae_lambda_value": 1.0, + "gae_lambda_policy": 1, + "length_adaptive_alpha": 1.5, + }, + SimpleNamespace( + use_kl_in_reward=False, + reference_policy_kl_penalty=0.0, + reference_policy_kl_type="low_var_kl", + ), + ) + B, S = 3, 6 + mask = torch.zeros(B, S) + mask[0, 2:6] = 1 + mask[1, 1:4] = 1 + mask[2, 3:5] = 1 + rewards = torch.tensor([1.0, 0.0, 0.5]) + values = torch.randn(B, S) # returns must NOT depend on these + + _, returns = estimator.compute_advantage( + prompt_ids=torch.zeros(B, 2, dtype=torch.long), + rewards=rewards, + mask=mask, + values=values, + reference_logprobs=None, + logprobs=None, # the critic-pretrain calling pattern (no policy) + ) + + expected = rewards.unsqueeze(1) * mask + torch.testing.assert_close(returns * mask, expected) From 4fb64f8be7dbbe2b6d30e6e7284722ca7b42b696 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 31 Jul 2026 14:41:12 -0700 Subject: [PATCH 38/69] feat(ppo): support the privileged critic in async PPO The answer-conditioned critic was gated to synchronous PPO by an assertion in setup(), because the async value stage never built the answer-augmented batch. Build it there instead: the buffered trajectories carry message_log + extra_env_info, so async_ppo_train can follow the same driver-side pattern as the sync loop. - value inference: score [prompt(+gold), response] and remap the response-position values back into the policy layout - GAE returns are scattered onto the augmented response positions (same response tokens => exact mapping) so the critic trains on what it scored - value training consumes the augmented batch, carrying sample_mask over Co-Authored-By: Claude Opus 5 --- nemo_rl/algorithms/ppo.py | 70 ++++++++++++++++++++++++++++++--------- 1 file changed, 55 insertions(+), 15 deletions(-) diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 5f8f69a1d0a..3bd25e58164 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -368,19 +368,11 @@ def setup( policy_config = master_config.policy value_config = master_config.value - # Privileged (answer-conditioned) critic is currently wired for synchronous PPO - # only; fail loudly rather than silently falling back to a blind critic on the - # async path (which does not build the answer-augmented value batch). + # Privileged (answer-conditioned) critic — supported in BOTH the sync (ppo_train) + # and async (async_ppo_train) loops; both build the answer-augmented value batch + # at train time from message_log + extra_env_info. _privileged_critic = value_config.get("privileged_critic") if _privileged_critic is not None and _privileged_critic.get("enabled"): - assert not ( - "async_ppo" in master_config.ppo - and master_config.ppo["async_ppo"].get("enabled") - ), ( - "value.privileged_critic is supported only for synchronous PPO for now " - "(the async_ppo value stage does not build the answer-augmented batch). " - "Disable async_ppo or value.privileged_critic.enabled." - ) # The critic scores [prompt + answer + response], which is longer than the # policy's [prompt + response] by up to max_answer_tokens (+ grader-note # template overhead). Give the VALUE model that much headroom so every @@ -3178,10 +3170,41 @@ def _sample_max_age(s: int) -> int: # end below). Load value only. print("▶ Computing values...") with timer.time("value_inference"): + # Privileged (answer-conditioned) critic — same driver-side pattern + # as the sync loop: score [prompt(+gold), response] and remap the + # response-position values back into the policy layout. The buffered + # trajectories carry message_log + extra_env_info, so the augmented + # batch is built at train time exactly as in sync. + privileged_critic_cfg = master_config.value.get("privileged_critic") + if privileged_critic_cfg is not None and not privileged_critic_cfg.get( + "enabled" + ): + privileged_critic_cfg = None + critic_batch = None + value_model.prepare_for_inference() - train_data["values"] = value_model.get_values(train_data)[ - "values" - ].squeeze(-1) + if privileged_critic_cfg is not None: + critic_batch = build_privileged_value_inputs( + repeated_batch, + tokenizer, + privileged_critic_cfg, + make_seq_len_divisible_by=master_config.policy[ + "make_sequence_length_divisible_by" + ], + ) + vals_aug = value_model.get_values(critic_batch)["values"].squeeze( + -1 + ) + critic_batch["values"] = vals_aug + train_data["values"] = remap_by_response_mask( + vals_aug, + critic_batch["token_mask"], + train_data["token_mask"], + ) + else: + train_data["values"] = value_model.get_values(train_data)[ + "values" + ].squeeze(-1) value_model.finish_inference() # ---- 4. Policy / reference logprobs (policy on GPU, then off) ---- @@ -3260,6 +3283,15 @@ def _sample_max_age(s: int) -> int: train_data["advantages"] = advantages if returns is not None: train_data["returns"] = returns + # Privileged critic trains on the answer-augmented sequence: + # scatter the (policy-layout) GAE returns onto the augmented + # response positions. Same response tokens => exact mapping. + if critic_batch is not None: + critic_batch["returns"] = remap_by_response_mask( + returns, + train_data["token_mask"], + critic_batch["token_mask"], + ) # ---- 7. ppo_epochs inner loop (critic, then actor) ---- # Each epoch: value on GPU -> train -> off. Then, once past critic @@ -3276,8 +3308,16 @@ def _sample_max_age(s: int) -> int: with timer.time("value_training_prep"): value_model.prepare_for_training() with timer.time("value_training"): + # Privileged critic: train on the answer-augmented batch + # (returns already scattered onto its response positions above; + # values=aug old-values for the optional value clip). + if critic_batch is not None: + critic_batch["sample_mask"] = train_data["sample_mask"] + value_train_batch = critic_batch + else: + value_train_batch = train_data value_results = value_model.train( - train_data, + value_train_batch, value_loss_fn, timer=timer, ) From f750f02e887affd20d8514a86830b2bd2804faf4 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 31 Jul 2026 14:41:25 -0700 Subject: [PATCH 39/69] feat(critic): calibration diagnostics and a fail-loud privilege precondition MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Calibration: EV tells you the critic ranks prefixes well, not that V is literally P(success | prefix). Add ECE (token-weighted |mean(V) - mean(R)| over 10 predicted-value bins) and signed bias per early/mid/late bucket, and wire the positional metrics into the async loop too. Miscalibration distorts advantage MAGNITUDES even when EV looks fine. Precondition: a missing/empty ground_truth silently degrades a sample to a blind critic (the grader note renders with an empty answer). Tolerate stragglers with a warning, but raise when the WHOLE batch lacks privilege — that means extra_env_info did not survive the data path and the run would measure a blind critic while labelled privileged. Also updates the positional-metrics docstring with the priv-vs-blind 7B DAPO finding: the answer acts as a verifier (wins at LATE tokens) rather than a forecaster of the policy's future behavior. Co-Authored-By: Claude Opus 5 --- nemo_rl/algorithms/ppo.py | 58 ++++++++++++++++--- nemo_rl/algorithms/privileged_critic.py | 23 ++++++++ .../unit/algorithms/test_privileged_critic.py | 21 +++++-- 3 files changed, 88 insertions(+), 14 deletions(-) diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 3bd25e58164..e649fcc3280 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -1456,20 +1456,47 @@ def _compute_critic_metrics(value_results: dict[str, Any]) -> dict[str, Any]: return critic_metrics +def _calibration_ece( + v: torch.Tensor, r: torch.Tensor, n_conf_bins: int = 10 +) -> float: + """Expected Calibration Error of V as an estimate of P(success). + + With outcome returns in [0, 1], a well-calibrated critic satisfies + ``mean(R | V ≈ p) == p``. Tokens are binned by predicted value (clamped to + [0, 1]); ECE is the token-weighted mean |mean(V) - mean(R)| over bins. A critic + can have decent explained variance yet be systematically over/under-confident — + that distorts advantage MAGNITUDES even when the ranking is right. + """ + vc = v.clamp(0.0, 1.0) + n_total = v.numel() + ece = 0.0 + for i in range(n_conf_bins): + lo, hi = i / n_conf_bins, (i + 1) / n_conf_bins + m = (vc >= lo) & ((vc < hi) if i < n_conf_bins - 1 else (vc <= 1.0)) + n = int(m.sum()) + if n == 0: + continue + ece += (n / n_total) * abs((v[m].mean() - r[m].mean()).item()) + return ece + + def _positional_value_metrics( values: torch.Tensor, returns: torch.Tensor, token_mask: torch.Tensor, n_bins: int = 3, ) -> dict[str, float]: - """Explained variance of V(s_t) vs the GAE return, bucketed by RELATIVE position - within each response (early / mid / late thirds). - - Answers "where does the critic do its job well?" — expected to rise early->late, - since near the end the outcome is nearly determined. Comparing a privileged - (answer-conditioned) run against a blind one, the golden answer should lift the - EARLY/MID buckets most (a large early-bucket gap = the privileged-critic mechanism - working; a flat gap = the answer isn't buying sharper early credit). Cheap: + """Critic-quality diagnostics bucketed by RELATIVE position within each response + (early / mid / late thirds): explained variance, calibration (ECE), and signed + bias of V(s_t) vs the GAE return. + + EV is expected to rise early->late (the outcome is nearly determined near the + end). Empirically (priv vs blind 7B DAPO runs), an answer-conditioned critic + matches the blind one at EARLY tokens and wins mostly at LATE tokens — the + privilege acts as a *verifier* (matching written content against the gold + answer), not a forecaster of the policy's future behavior. ECE/bias add the + calibration axis: V should literally be P(success | prefix), and miscalibration + distorts advantage magnitudes even when EV/ranking looks fine. Cheap: driver-side tensor ops on tensors already in ``train_data``. """ mask = token_mask.bool() @@ -1495,6 +1522,10 @@ def _positional_value_metrics( out[f"critic/abs_err_{names[i]}"] = (r - v).abs().mean().item() out[f"critic/mean_v_{names[i]}"] = v.mean().item() out[f"critic/n_tokens_{names[i]}"] = float(n) + # Calibration: ECE (magnitude of miscalibration across confidence bins) and + # signed bias (direction: >0 = overconfident/optimistic, <0 = pessimistic). + out[f"critic/ece_{names[i]}"] = _calibration_ece(v, r) + out[f"critic/bias_{names[i]}"] = (v - r).mean().item() return out @@ -3474,6 +3505,17 @@ def _sample_max_age(s: int) -> int: metrics.update(train_results["all_mb_metrics"]) if value_results is not None: metrics.update(_compute_critic_metrics(value_results)) + # Positional critic-quality diagnostic (early/mid/late tokens): how + # well V predicts the return along the trajectory, and — for the + # privileged critic — where the golden answer sharpens it. + if "values" in train_data and "returns" in train_data: + metrics.update( + _positional_value_metrics( + train_data["values"], + train_data["returns"], + train_data["token_mask"], + ) + ) for k, v in metrics.items(): if k in {"probs_ratio_min", "probs_ratio_clamped_min"}: diff --git a/nemo_rl/algorithms/privileged_critic.py b/nemo_rl/algorithms/privileged_critic.py index 9343f5fa76f..436a0c598fb 100644 --- a/nemo_rl/algorithms/privileged_critic.py +++ b/nemo_rl/algorithms/privileged_critic.py @@ -192,6 +192,29 @@ def build_privileged_value_inputs( ) critic_message_logs.append(critic_msgs) + # A missing/empty ground truth silently degrades the sample to a blind critic + # (the grader note renders with an empty answer). Tolerate stragglers but fail + # loudly if the WHOLE batch lacks privilege — that means extra_env_info didn't + # survive the data path (e.g. a buffer/collector change) and the run would + # silently measure a blind critic while labelled privileged. + n_empty = sum( + 1 + for info in env_infos + if not (info or {}).get("ground_truth") + ) + if n_empty == len(message_logs): + raise AssertionError( + "privileged critic: ground_truth missing for EVERY sample in the batch " + "(extra_env_info absent or empty) — the privilege would be silently " + "inert. Check that the data path carries extra_env_info." + ) + if n_empty: + print( + f"⚠️ privileged critic: {n_empty}/{len(message_logs)} samples have no " + "ground_truth — those critic inputs are effectively blind.", + flush=True, + ) + flat, input_lengths = batched_message_log_to_flat_message( critic_message_logs, pad_value_dict={"token_ids": tokenizer.pad_token_id}, diff --git a/tests/unit/algorithms/test_privileged_critic.py b/tests/unit/algorithms/test_privileged_critic.py index fbbeaa3c697..5b5a85eca18 100644 --- a/tests/unit/algorithms/test_privileged_critic.py +++ b/tests/unit/algorithms/test_privileged_critic.py @@ -158,18 +158,27 @@ def test_build_answer_lengthens_only_the_prompt(): response_ids = [77, 78, 79, 80] pcfg = {"enabled": True, "placement": "user_suffix", "max_answer_tokens": 256} - with_answer = build_privileged_value_inputs( + long_answer = build_privileged_value_inputs( _make_repeated_batch(response_ids, gold="the answer is 42 with reasoning"), tok, pcfg, ) - no_answer = build_privileged_value_inputs( - _make_repeated_batch(response_ids, gold=""), tok, pcfg + short_answer = build_privileged_value_inputs( + _make_repeated_batch(response_ids, gold="4"), tok, pcfg ) # a longer reference answer only grows the prompt; response count is unchanged - assert with_answer["input_lengths"][0] > no_answer["input_lengths"][0] - assert int(with_answer["token_mask"][0].sum()) == len(response_ids) - assert int(no_answer["token_mask"][0].sum()) == len(response_ids) + assert long_answer["input_lengths"][0] > short_answer["input_lengths"][0] + assert int(long_answer["token_mask"][0].sum()) == len(response_ids) + assert int(short_answer["token_mask"][0].sum()) == len(response_ids) + + +def test_build_raises_when_privilege_missing_for_whole_batch(): + # gold absent for EVERY sample => the "privileged" critic would silently be blind; + # the builder must fail loudly instead. + tok = _FakeTokenizer() + batch = _make_repeated_batch([77, 78], gold="") + with pytest.raises(AssertionError, match="ground_truth missing for EVERY sample"): + build_privileged_value_inputs(batch, tok, {"enabled": True}) def test_build_truncates_long_answer(): From 2519a03863134f31d13236ee8e8179595422e84b Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 31 Jul 2026 14:41:44 -0700 Subject: [PATCH 40/69] fix: crashes and false preconditions hit in the SWE / 1B runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - ppo.setup: a fabricated warm-start seed (a step_0 checkpoint holding only a pretrained value/ dir) carries no train_dataloader.pt; start the dataloader fresh with a warning instead of failing, mirroring how a missing policy/ falls back to base weights. - ppo.setup: segment topology only applies to multi-node clusters. In the single-node carve, train and inference SHARE the one node, so train_nodes + inference_nodes double-counted the roles and demanded 2 alive nodes on a 1-node cluster. NVLink-domain placement is meaningless there. - nemo_gym: when the agent fails before its first model request, the result has neither output items nor input messages and apply_chat_template([]) threw a bare IndexError that masked the diagnosis. Raise something actionable with the truncated gym result instead. - MegatronConfig.mtp_num_layers: type as int | None. For hybrid (mamba) models None is the only full off-switch — hybrid_model.py gates the MTP block on `is not None`, so 0 still enters it and asserts > 0. Co-Authored-By: Claude Opus 5 --- nemo_rl/algorithms/ppo.py | 21 +++++++++++++++++++-- nemo_rl/environments/nemo_gym.py | 14 ++++++++++++++ nemo_rl/models/policy/__init__.py | 6 ++++-- 3 files changed, 37 insertions(+), 4 deletions(-) diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index e649fcc3280..3d9d15b5f7f 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -483,7 +483,19 @@ def setup( num_workers=data_config["num_workers"], ) if last_checkpoint_path is not None: - load_dataloader_state(dataloader, last_checkpoint_path, data_config) + # Fabricated warm-start seeds (a step_0 checkpoint holding only a + # pretrained value/ dir — see scripts/swe/ppo/prep_warm_start.sh) carry + # no dataloader state; start the dataloader fresh instead of failing, + # mirroring how a missing policy/ falls back to base weights below. + if os.path.exists(os.path.join(last_checkpoint_path, "train_dataloader.pt")): + load_dataloader_state(dataloader, last_checkpoint_path, data_config) + else: + print( + f" ⚠ No train_dataloader.pt in checkpoint {last_checkpoint_path} " + "(e.g. a fabricated warm-start seed); starting dataloader from " + "the beginning.", + flush=True, + ) print(f" ✓ Training dataloader loaded with {len(dataset)} samples", flush=True) @@ -654,7 +666,12 @@ def setup( # across multiple nodes, so the inference cluster below gets no # topology constraints and falls back to default PACK placement. node_resource_constraints = None - if segment_size is not None: + # Segment topology only applies to multi-node clusters: in the single-node + # carve case train and inference SHARE the one node (GPUs split between + # them), so counting train_nodes + inference_nodes would double-count the + # roles and demand 2 alive nodes on a 1-node cluster. NVLink-domain + # placement is meaningless for a single node anyway. + if segment_size is not None and total_nodes > 1: topology = get_ray_cluster_topology() num_alive_nodes = len(topology) required_nodes = train_nodes + inference_nodes diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index b3f8dcbfbd1..9c7975edfd0 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -439,6 +439,20 @@ def _postprocess_nemo_gym_to_nemo_rl_result( if not nemo_rl_message_log: input_messages = nemo_gym_result["responses_create_params"]["input"] + if not input_messages: + # The agent failed before issuing any model request (e.g. its + # in-container setup command crashed), so the gym result carries + # neither output items nor input messages. Raise something + # actionable instead of letting apply_chat_template([]) throw a + # bare IndexError that masks the diagnosis. + raise ValueError( + "NeMo Gym returned a result with no generation data AND an " + "empty input message list. This usually means the agent " + "failed before its first model request (e.g. the SWE agent's " + "in-container setup command crashed on this node — check the " + "gym server logs around this time). Gym result (truncated): " + f"{str(nemo_gym_result)[:2000]}" + ) prompt_token_ids = tokenizer.apply_chat_template( input_messages, tokenize=True ) diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 8f2448e4e76..674b0a12838 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -387,8 +387,10 @@ class MegatronConfig(TypedDict): # Number of tokens per chunk when computing fused linear logprobs. # Smaller values reduce peak memory further but may decrease throughput. fused_linear_logprobs_chunk_size: NotRequired[int] - # When mtp_num_layers=0, Multi-Token Prediction is disabled. - mtp_num_layers: NotRequired[int] + # None disables Multi-Token Prediction entirely. For hybrid (mamba) models + # None is the only full off-switch: hybrid_model.py gates the MTP block on + # `mtp_num_layers is not None`, so 0 still enters it and asserts > 0. + mtp_num_layers: NotRequired[int | None] # MTP loss weight added to the main next-token loss (0.0 disables the MTP loss contribution). mtp_loss_scaling_factor: NotRequired[float] # When True, repeat a single MTP layer mtp_num_layers times instead of using distinct layers. From b1912bed3e9d212035956c1dc181bd101bba9732 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 31 Jul 2026 14:41:53 -0700 Subject: [PATCH 41/69] feat(logging): log mtp_metrics from both PPO loops Mirrors the existing moe_metrics passthrough so Multi-Token Prediction loss terms are visible in wandb under mtp/*, in sync ppo_train and async_ppo_train. Co-Authored-By: Claude Opus 5 --- nemo_rl/algorithms/ppo.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 3d9d15b5f7f..c52adf885f2 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -2332,6 +2332,13 @@ def ppo_train( for k, v in train_results["moe_metrics"].items() } ) + if "mtp_metrics" in train_results: + metrics.update( + { + f"mtp/{k}": v + for k, v in train_results["mtp_metrics"].items() + } + ) # Extract critic metrics from value training results if value_results is not None: @@ -3519,6 +3526,13 @@ def _sample_max_age(s: int) -> int: for k, v in train_results["moe_metrics"].items() } ) + if "mtp_metrics" in train_results: + metrics.update( + { + f"mtp/{k}": v + for k, v in train_results["mtp_metrics"].items() + } + ) metrics.update(train_results["all_mb_metrics"]) if value_results is not None: metrics.update(_compute_critic_metrics(value_results)) From 4be6d570ae11d9a9de1274b7813d28ebc87454b1 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 31 Jul 2026 14:42:11 -0700 Subject: [PATCH 42/69] config: SWE PPO run config + long-sequence memory headroom for 1B megatron MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - ppo_nano_v3_5_swe_cmh.yaml: the SWE PPO run config. Sets vllm disable_custom_all_reduce: jobs 37606/38016 each lost a generation engine ~15 min in, and the NCCL flight recorder showed all 4 TP ranks with their next collective enqueued-but-never-started — all four streams parked in the custom-allreduce kernel ahead of it (symmetric cross-GPU deadlock, different node each time). Costs a few % decode latency on TP-4. Also drops agent concurrency 3x -> 2x prompts*generations. - ppo_math_1B_megatron.yaml: activation_checkpointing + defer_fp32_logits. At long packed sequences the actor forward materializes full fp32 LM-head logits (~11GB over a 152k vocab) on top of unrecomputed activations, and the first policy-training step OOMs in the grad all-reduce. Memory-shape only. Co-Authored-By: Claude Opus 5 --- examples/configs/ppo_math_1B_megatron.yaml | 7 + examples/configs/ppo_nano_v3_5_swe_cmh.yaml | 698 ++++++++++++++++++++ 2 files changed, 705 insertions(+) create mode 100644 examples/configs/ppo_nano_v3_5_swe_cmh.yaml diff --git a/examples/configs/ppo_math_1B_megatron.yaml b/examples/configs/ppo_math_1B_megatron.yaml index 9ed2ce2d661..c145b1276cc 100644 --- a/examples/configs/ppo_math_1B_megatron.yaml +++ b/examples/configs/ppo_math_1B_megatron.yaml @@ -11,6 +11,13 @@ policy: enabled: false megatron_cfg: enabled: true + # Required headroom for the actor train step at long packed sequences (e.g. + # 7B math @ 18432 tokens): the actor forward materializes full LM-head + # logits (~11GB fp32 over a 152k vocab) on top of unrecomputed activations; + # without these, the first policy-training step OOMs in the grad all-reduce + # (NCCL "Cuda failure 2 'out of memory'"). Memory-shape only — same math. + activation_checkpointing: true + defer_fp32_logits: true make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} optimizer: null scheduler: null diff --git a/examples/configs/ppo_nano_v3_5_swe_cmh.yaml b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml new file mode 100644 index 00000000000..b3a4459718f --- /dev/null +++ b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml @@ -0,0 +1,698 @@ +# ============================================================================= +# PPO Ultra E2E — SWE Config for GB200 NVL72 (PPO/VAPO port of the GRPO config) +# ============================================================================= +# Duplicated from grpo_nano_v3_5_swe_hsg.yaml and adapted for PPO with a GAE +# critic (run with examples/nemo_gym/run_ppo_nemo_gym.py). Differences vs the +# GRPO parent: +# - `grpo:` section replaced by `ppo:` (GAE adv_estimator, ppo_epochs, +# critic warmup via policy_training_start_step, async_ppo instead of +# async_grpo). max_val_samples / val_batch_size MUST stay null — the +# NeMo-Gym PPO runner rejects preset values and sets both to +# len(val_dataset). +# - new `value:` block — Megatron-backend value model initialized from the +# same nano checkpoint (fresh scalar value head), mirroring the policy's +# parallelism/MoE settings. NOTE: the megatron value worker has only been +# validated on GPT-style dense models; the nano hybrid (mamba+MoE) critic +# is experimental. +# - new `value_loss_fn:` block (clipped MSE, VAPO-style scale=1.0/clip=0.5 +# from the validated 7B math PPO run). +# - loss_fn.force_on_policy_ratio false: PPO trains on the real prev/current +# logprob ratio (GRPO nano forced ratio=1 and leaned on TIS only). +# - grpo's invalid_tool_call_advantage / malformed_thinking_advantage are +# DROPPED: message-level advantage-overwrite penalties are rejected on the +# PPO path (they would break GAE advantage/return consistency). The +# reward-zeroing `reward_penalties` block is the PPO-compatible path; left +# all-off (default) for parity with the GRPO baseline, which also ran with +# them off (its top-level penalize_* keys are inert in this checkout). +# +# Static algorithm, loss, and environment settings live here. Per-run +# overrides (model path, data paths, parallelism, batch sizes) are applied by +# the launch script (scripts/swe/ppo/nano_on_main.sh) or Hydra CLI. +# ============================================================================= + +# ============================================================================= +# Cluster — overridden by launch script for parallelism changes +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 128 + segment_size: 16 + +# ============================================================================= +# Checkpointing +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/ppo_ultra_e2e" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 3 + ft_keep_latest_k: 1 + ft_save_period: 1 + save_optimizer: true + checkpoint_must_save_by: "00:03:30:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# PPO Algorithm (replaces the parent's grpo: section) +# ============================================================================= +ppo: + num_prompts_per_step: 16 + num_generations_per_prompt: 32 + max_rollout_turns: 1 + max_num_epochs: 4 + max_num_steps: 1000000 + # One optimizer pass per rollout batch (PPO-epochs>1 reuses the batch). + ppo_epochs: 1 + # Critic-only warmup: policy training is skipped for the first N steps while + # the (fresh-head) value model trains on rollouts from the frozen policy. + # Overridden by the launch script. + policy_training_start_step: 50 + val_period: 10000 + val_at_start: false + val_at_end: false + overlong_filtering: false + # NeMo-Gym principle: the validation set is used verbatim (no hidden pre/post + # processing). run_ppo_nemo_gym.py rejects a preset max_val_samples and sets + # both to len(val_dataset), so they must start null. + max_val_samples: null + val_batch_size: null + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + # loss_fn.reference_policy_kl_penalty is 0 (SWE baseline), so skip the + # reference-policy logprob forward. The launch script flips this to false + # automatically when a nonzero KL penalty is requested. + skip_reference_policy_logprobs_calculation: true + seq_logprob_error_threshold: 2 + + # GAE with VAPO decoupled lambdas (defaults match the validated 7B math PPO + # run: lambda_value=1.0, length-adaptive lambda_policy with alpha=0.05). + adv_estimator: + name: "gae" + gae_lambda: 1.0 + gae_gamma: 1 + normalize_advantages: true + gae_lambda_value: 1.0 + gae_lambda_policy: 1 + # Length-adaptive lambda_policy = 1 - 1/(alpha*l). 0 = disabled. + length_adaptive_alpha: null + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + # NOTE: grpo's penalize_invalid_tool_call / invalid_tool_call_advantage / + # penalize_malformed_thinking / malformed_thinking_advantage are intentionally + # absent: setup() raises NotImplementedError if the *_advantage keys are set + # on the PPO path (advantage overwrites break GAE consistency). + + async_ppo: + enabled: true + # lag-1: at most one policy version off-policy (recommended/validated). + # During critic warmup the frozen actor makes any-age trajectories + # on-policy for free (warmup_max_trajectory_age_steps; null => same value). + max_trajectory_age_steps: 1 + warmup_max_trajectory_age_steps: 8 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + # Drop the incomplete (survivorship-biased) frontier target on resume and + # regenerate it fresh — unbiased batches at the cost of a one-target + # generation bubble per resume. + drop_incomplete_targets_on_restore: false + # Throttle collector/replay-buffer per-rollout heartbeat prints. + log_every: 500 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + # Defaults from the validated 7B math PPO run (symmetric 0.2 clip + dual-clip + # c=3), not the GRPO-DAPO asymmetric 0.2/0.28. Overridable from the script. + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: 3 + use_on_policy_kl_approximation: true + # Async PPO requires the off-policy importance-sampling correction. + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + # PPO trains on the real prev/current-logprob ratio; forcing ratio=1 (as the + # GRPO nano config does) would disable PPO clipping entirely. + force_on_policy_ratio: false + use_kl_in_reward: false + +# ============================================================================= +# Value Loss (MSE with PPO-style clipping) +# ============================================================================= +value_loss_fn: + scale: 1.0 + cliprange: 0.5 + +# ============================================================================= +# Policy +# ============================================================================= +policy: + model_name: "/lustre/fsw/portfolios/llmservice/users/wdai/megatron-lm-ultra/checkpoints/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-288k-nano-loss-032026/iter_0007000/hf" + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 512 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 131072 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 32 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + use_gloo_process_groups: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + # community_import.py reads megatron_cfg["gradient_accumulation_fusion"] directly + # (no default) on the HF->megatron conversion path; absent here it raises KeyError. + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled + mtp_loss_scaling_factor: 0.3 + mtp_use_repeated_layer: true + # null disables MTP entirely (matches the upstream mtp=0 run's behavior). This + # checkout's hybrid model gates MTP on `mtp_num_layers is not None` + # (hybrid_model.py:544): 0 still enters the MTP block and asserts > 0, and 1 runs + # MTP forward on only the last-stage ranks (suspected cause of the EXPERT_MODEL_ + # PARALLEL NCCL collective timeout). null makes that check False -> MTP is skipped + # on all ranks, no partial-rank EP collective. Requires the schema change making + # MegatronConfig.mtp_num_layers accept None (nemo_rl/models/policy/__init__.py). + mtp_num_layers: 5 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 3.0e-6 + min_lr: 3.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 3e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + # FP8 — overridden by precision recipe in launch script + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo + # were dropped from this checkout's megatron-bridge CheckpointConfig (submodule + # 554c7b9); leaving them in raises InstantiationException "Unexpected config keys". + # Removed for the current-repo reproduction. + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + # Must be a multiple of minimum_pad_factor = cp_size*2*tp_size (see + # nemo_rl/models/megatron/data.py). With CP=16 + TP=4 + sequence_parallel that is + # 128; the old value (=TP=4) made get_logprobs assert and abort every training + # step. Formula matches examples/nemo_gym/grpo_nanov3.yaml (TP*CP*2). + make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 11001 + port_range_high: 15000 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 1 + gpu_memory_utilization: 0.8 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + # reasoning_parser_plugin lives at vllm_cfg top level in this checkout + # (vllm_worker_async reads vllm_cfg["reasoning_parser_plugin"]); passing it + # inside http_server_serving_chat_kwargs makes OpenAIServingChat raise + # "unexpected keyword argument 'reasoning_parser_plugin'". Path updated to this + # repo's location (the parser moved out of nemo_rl/utils/). + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + + vllm_kwargs: + attention_backend: FLASH_ATTN + moe_backend: triton + mamba_ssm_cache_dtype: "float32" + # Route TP allreduces through NCCL instead of vLLM's custom CUDA-IPC + # spin kernel. Jobs 37606/38016 each lost a generation engine ~15 min in: + # the NCCL flight recorder showed all 4 TP ranks with their next + # collective enqueued-but-never-started, i.e. all four streams parked in + # the custom-allreduce kernel ahead of it (symmetric cross-GPU deadlock, + # different node each time). Costs a few % decode latency on TP-4. + disable_custom_all_reduce: true + compilation_config: + cudagraph_mode: PIECEWISE + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 64 + +# ============================================================================= +# Value Model (GAE critic) — Megatron backend, initialized from the policy +# checkpoint with a fresh scalar value head. Parallelism/MoE settings mirror +# the policy block; the launch script overrides them together. +# ============================================================================= +value: + model_name: ${policy.model_name} + tokenizer: + name: ${value.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: ${policy.train_global_batch_size} + train_micro_batch_size: ${policy.train_micro_batch_size} + logprob_batch_size: ${policy.logprob_batch_size} + max_total_sequence_length: ${policy.max_total_sequence_length} + precision: "bfloat16" + + reward_model_cfg: + enabled: false # only used for the DTensor V2 value worker + reward_model_type: "regression" + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 32 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${value.precision} + sequence_parallel: true + + # MoE — mirror the policy + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + use_gloo_process_groups: false + + # Compute — mirror the policy + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — fully disabled for the critic. The value model replaces the LM + # head with a scalar head, so MTP next-token heads are dead weight; null + # skips the MTP block on all ranks (no partial-rank EP collective). + mtp_loss_scaling_factor: 0.0 + mtp_use_repeated_layer: true + mtp_num_layers: null + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + # Critic LR from the validated 7B math PPO run (fresh value head trains + # faster than the pretrained actor). Overridden by the launch script. + lr: 6.0e-6 + min_lr: 6.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${value.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${value.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${value.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 6.0e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + + env_vars: null + + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + + make_sequence_length_divisible_by: ${mul:${mul:${value.megatron_cfg.tensor_model_parallel_size}, ${value.megatron_cfg.context_parallel_size}}, 2} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + use_multiple_dataloader: false + train: + data_path: null # Set by launch script + validation: + data_path: null # Set by launch script + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + SWE Agents +# ============================================================================= +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + # Reuse pre-built gym venvs. With false, all 3 gym servers (policy_model, + # swe_agents_{train,val}) rebuild the editable nemo-gym package in parallel and + # contend on the same uv distribution-cache lock on Lustre; policy_model then + # times out (default 300s) and NemoGym spinup fails. Mirrors the qwen-30b SWE run. + skip_venv_if_present: true + num_gpu_nodes: 0 + port_range_low: 15001 + port_range_high: 20000 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml + swe_agents_train: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}} + swebench_agent_timeout: 3600 + run_with_mixed_prompts: true + dataset_path: ${data.train.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swerebench/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/nv_internal/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/r2e_gym/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/sweb.eval.arm64.{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swebench/swe-bench.eval.arm64.{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/swebenchpro_ots/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swegym/sweb.eval.arm64.{instance_id}.sif" + swe_agents_val: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}} + swebench_agent_timeout: 3600 + dataset_path: ${data.validation.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter} + use_absolute_ip: true + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "logs" + num_val_samples_to_print: 0 + wandb_enabled: true + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "ultra-v3-swe-e2e" + name: "ppo-ultra-e2e" + tensorboard: {} + mlflow: + experiment_name: "ppo-ultra-e2e" + run_name: "ppo-ultra-e2e" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels / Token IDs / penalize_* — carried over from the GRPO config +# for parity. NOTE: these top-level keys are INERT in this checkout (nothing +# reads them; the GRPO baseline also ran with them having no effect). The +# PPO-compatible penalty mechanism is the structured `reward_penalties:` block +# (reward-zeroing, pre-GAE) — left at its all-off default to match the GRPO +# baseline's effective behavior. +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + + +# Reward-zeroing penalties applied to NeMo-Gym rollout results. +reward_penalties: + penalize_duplicated_reasoning: true + penalize_empty_final_answer: true + penalize_unwanted_tokens: true + penalize_malformed_think_tag: true + # Optional model/tokenizer-specific token IDs. Add token IDs that should + # be penalized when emitted by the model. Unwanted token IDs must be + # specified explicitly; think-tag IDs are inferred when each tag encodes + # to one token. Example: {unwanted: [2], think_open: 12, think_close: 13} + token_ids: {unwanted: [2], think_open: 12, think_close: 13} \ No newline at end of file From 380346bf78eec5c191703b8199233edb0ad6b407 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 31 Jul 2026 14:42:11 -0700 Subject: [PATCH 43/69] config: run-config tuning for the ultra and 1B math sweeps - grpo_ultra_256n4g_bf16.yaml: alltoall MoE dispatcher, lr 3e-6 -> 4e-6, sequence-length divisor accounts for context parallelism, generation lengths follow policy.max_total_sequence_length, wandb on - grpo/ppo_ultra_256n4g_bf16_pivotonly.yaml: pivot-only variants - ppo_math_1B.yaml: rollout dumps on (every 5 steps), seq logprob error threshold 2, drop incomplete async targets on restore, segment_size 2 Co-Authored-By: Claude Opus 5 --- examples/configs/grpo_ultra_256n4g_bf16.yaml | 16 +- .../grpo_ultra_256n4g_bf16_pivotonly.yaml | 529 +++++++++++++ examples/configs/ppo_math_1B.yaml | 9 +- .../ppo_ultra_256n4g_bf16_pivotonly.yaml | 697 ++++++++++++++++++ 4 files changed, 1239 insertions(+), 12 deletions(-) create mode 100644 examples/configs/grpo_ultra_256n4g_bf16_pivotonly.yaml create mode 100644 examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml diff --git a/examples/configs/grpo_ultra_256n4g_bf16.yaml b/examples/configs/grpo_ultra_256n4g_bf16.yaml index b9c3c69bc67..c167ac41002 100644 --- a/examples/configs/grpo_ultra_256n4g_bf16.yaml +++ b/examples/configs/grpo_ultra_256n4g_bf16.yaml @@ -187,7 +187,7 @@ policy: moe_router_enable_expert_bias: true moe_permute_fusion: true moe_enable_deepep: false - moe_token_dispatcher_type: "flex" + moe_token_dispatcher_type: "alltoall" moe_flex_dispatcher_backend: "hybridep" moe_hybridep_num_sms: 32 moe_aux_loss_coeff: 0.0 @@ -216,8 +216,8 @@ policy: optimizer: optimizer: "adam" - lr: 3.0e-6 - min_lr: 3.0e-6 + lr: 4.0e-6 + min_lr: 4.0e-6 weight_decay: 0.0 bf16: true fp16: false @@ -244,7 +244,7 @@ policy: lr_decay_style: "constant" lr_decay_iters: null lr_warmup_iters: 10 - lr_warmup_init: 3e-7 + lr_warmup_init: 4e-7 override_opt_param_scheduler: true distributed_data_parallel_config: @@ -294,7 +294,7 @@ policy: sequence_length_round: 64 fuse_loss: true - make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2} max_grad_norm: 1.0 optimizer: null scheduler: null @@ -306,7 +306,7 @@ policy: port_range_low: 3000 port_range_high: 4999 backend: "vllm" - max_new_tokens: 49152 + max_new_tokens: ${policy.max_total_sequence_length} temperature: 1.0 top_p: 1.0 top_k: null @@ -323,7 +323,7 @@ policy: pipeline_parallel_size: 1 expert_parallel_size: 8 gpu_memory_utilization: 0.85 - max_model_len: 49152 + max_model_len: ${policy.max_total_sequence_length} enforce_eager: false use_deep_gemm: false num_last_layers_in_bf16: 0 @@ -622,7 +622,7 @@ env: logger: log_dir: "/lustre/fsw/portfolios/llmservice/users/ansubramania/ultra_v3/runs/logs" num_val_samples_to_print: 0 - wandb_enabled: false + wandb_enabled: true tensorboard_enabled: false mlflow_enabled: false monitor_gpus: true diff --git a/examples/configs/grpo_ultra_256n4g_bf16_pivotonly.yaml b/examples/configs/grpo_ultra_256n4g_bf16_pivotonly.yaml new file mode 100644 index 00000000000..d3454a15a64 --- /dev/null +++ b/examples/configs/grpo_ultra_256n4g_bf16_pivotonly.yaml @@ -0,0 +1,529 @@ +# ============================================================================= +# GRPO Ultra V3 — 256-node GB200 NVL72 Config (bf16) +# ============================================================================= +# Config for GRPO training on 256 nodes × 4 GPUs/node (1024 GPUs). +# Full batch sizes and sequence lengths for convergence runs. +# +# Node allocation (256 total, set via launch script env vars): +# - Training: 64 nodes (256 GPUs) — 4 segments of 16 +# - vLLM: 182 nodes (728 GPUs) — 91 instances at TP=8 EP=8 (2 nodes each) +# - Gym/Judge: 10 nodes ( 40 GPUs) — judges scaled for production throughput +# +# Generation-heavy split follows the SuperV3 production ratio (~25/71/4). +# +# Training parallelism (256 GPUs = 64 nodes): +# - TP: 8 +# - EP: 64 +# - CP: 8 +# - PP: 1 +# - SP: true +# +# vLLM parallelism (bf16, TP=8, EP=8): +# NeMo Gym requires async_engine=true, but vLLM DP+EP (EP > TP) requires +# async_engine=false (see https://github.com/NVIDIA-NeMo/RL/issues/1101). +# This forces EP <= TP. With EP=8 (=TP), vllm_dp_size = 8/8 = 1 so +# async_engine=true works. +# ============================================================================= + +# ============================================================================= +# Cluster — overridden by launch script +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 256 + segment_size: 16 + +# ============================================================================= +# Checkpointing +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/grpo_ultra_v3" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_optimizer: true + save_period: 10 + ft_keep_latest_k: 1 + ft_save_period: 1 + checkpoint_must_save_by: "00:03:30:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# GRPO Algorithm +# ============================================================================= +grpo: + num_prompts_per_step: 256 + num_generations_per_prompt: 16 + num_val_generations_per_prompt: 2 + max_rollout_turns: 1 + max_num_epochs: 1 + max_num_steps: 1000000 + normalize_rewards: true + use_leave_one_out_baseline: true + advantage_clip_low: -50 + advantage_clip_high: 50 + val_period: -1 + val_at_start: false + val_at_end: false + overlong_filtering: false + max_val_samples: null + val_batch_size: 256 + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + penalize_invalid_tool_call: true + invalid_tool_call_advantage: -5.0 + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + + use_best_at_k: false + best_at_k_k: 8 + best_at_k_m: 1000 + + use_combined_training: false + combined_training_weight_mode: "auto" + combined_training_best_at_k_weight: 0.2 + combined_training_pass_at_1_weight: 1.0 + + dynamic_sampling_oversample_ratio: 1.0 + seq_logprob_error_threshold: 2 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: null + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + force_on_policy_ratio: true + use_kl_in_reward: false + +# ============================================================================= +# Policy +# ============================================================================= +policy: + model_name: "/lustre/fsw/portfolios/llmservice/users/soumyes/sft-runs/eval_and_sleep/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-mar13-blend-512k-filt-1e-5/iter_0001900/hf" + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 4096 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 65536 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + # EP=64: 512 experts / 64 = 8 experts per EP rank + # All-to-all spans 64 GPUs (16 nodes), fits within one NVLink domain + expert_model_parallel_size: 64 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled + mtp_loss_scaling_factor: 0.3 + mtp_use_repeated_layer: true + mtp_num_layers: 5 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 4.0e-6 + min_lr: 4.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 4e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + # FP8 — disabled for bf16 runs. Enable for mxfp8 validation. + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + use_gloo_process_groups: false + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo + # were dropped: this container's megatron-bridge CheckpointConfig no longer has + # them and raises InstantiationException "Unexpected config keys" (see + # ppo_nano_v3_5_swe_cmh.yaml). + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 3000 + port_range_high: 4999 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. + # Same memory as EP=1 but uses all-to-all for expert routing. + # EP > TP blocked by https://github.com/NVIDIA-NeMo/RL/issues/1101. + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 8 + gpu_memory_utilization: 0.85 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + # nano_v3 is a CUSTOM reasoning parser; its plugin file is registered at vllm_cfg + # top level (read by vllm_worker_async.py:481 -> ReasoningParserManager.import_reasoning_parser). + # It must NOT sit inside http_server_serving_chat_kwargs: that dict is splatted into + # OpenAIServingChat(**kwargs), which has no reasoning_parser_plugin arg -> TypeError. + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + + vllm_kwargs: + attention_backend: FLASH_ATTN + moe_backend: triton + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_mode: PIECEWISE + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 182 # Overridden by launch script + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + use_multiple_dataloader: false + train: + data_path: null # Set by launch script + validation: + data_path: null # Set by launch script + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + Judge Models +# ============================================================================= +env: + should_use_nemo_gym: true + # true: skip expensive train_data_step*.jsonl (recommended for large Gym runs); false: write full jsonl. + should_log_nemo_gym_responses: true + nemo_gym: + nemo_gym_log_dir: "logs/nemo_gym" + skip_venv_if_present: true + port_range_low: 5000 + port_range_high: 5999 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/math_with_judge/configs/math_with_judge.yaml + - resources_servers/code_gen/configs/code_gen.yaml + - resources_servers/instruction_following/configs/instruction_following.yaml + - resources_servers/calendar/configs/calendar.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml + - resources_servers/terminus_judge/configs/terminus_judge_string_only.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/search_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/toolcall_schema_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/swe_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/droid_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/structured_outputs/configs/structured_outputs_json_yaml_xml_v1.yaml + - resources_servers/structured_outputs/configs/structured_outputs_v3.yaml + - resources_servers/format_verification/configs/freeform_formatting.yaml + - resources_servers/format_verification/configs/citation_format.yaml + + # nl2bash / General Judge: TP=4 on GB200 192GB + nl2bash_judge_model: + responses_api_models: + local_vllm_model: + entrypoint: app.py + model: null # Set by launch script + return_token_id_information: false + uses_reasoning_parser: false + debug: true + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + # The gym `local_vllm_model` venv has no deep_gemm (nemo_rl's policy venv + # does, but this stock-vLLM judge venv doesn't). On an FP8 model, vLLM's + # deep_gemm warmup then hard-fails. Disable it -> fall back to vLLM's other + # FP8 kernels, consistent with the policy running use_deep_gemm=false. + VLLM_USE_DEEP_GEMM: "0" + + vllm_serve_kwargs: + attention_backend: FLASH_ATTN + tensor_parallel_size: 2 + data_parallel_size: 4 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + enable_expert_parallel: true + enable_auto_tool_choice: true + tool_call_parser: hermes + gpu_memory_utilization: 0.85 + max_model_len: 131072 + max_num_seqs: 512 + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + math_with_judge: + resources_servers: + math_with_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + should_use_judge: true + + code_gen: + resources_servers: + code_gen: + num_processes: 2048 + unit_test_timeout_secs: 10 + debug: false + + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "/lustre/fsw/portfolios/llmservice/users/ansubramania/ultra_v3/runs/logs" + num_val_samples_to_print: 0 + wandb_enabled: false + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "grpo-ultra-v3" + name: "grpo-ultra-v3-256n" + tensorboard: {} + mlflow: + experiment_name: "grpo-ultra-v3" + run_name: "grpo-ultra-v3-256n" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + +# ============================================================================= +# Token IDs (model-specific, used by token-based penalties) +# ============================================================================= +token_ids: + eos: 2 # + think_open: 12 # + think_close: 13 # + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +penalize_duplicated_reasoning: true # reasoning content == final answer +penalize_empty_final_answer: true # last message output has empty content +penalize_eos_token: true # eos token appears in generation +penalize_malformed_think_tag: true # /<\/think> count != 1 per turn + +# Reward-zeroing penalties applied to NeMo-Gym rollout results. +reward_penalties: + penalize_duplicated_reasoning: true + penalize_empty_final_answer: true + penalize_unwanted_tokens: true + penalize_malformed_think_tag: true + # Optional model/tokenizer-specific token IDs. Add token IDs that should + # be penalized when emitted by the model. Unwanted token IDs must be + # specified explicitly; think-tag IDs are inferred when each tag encodes + # to one token. Example: {unwanted: [2], think_open: 12, think_close: 13} + token_ids: {unwanted: [2], think_open: 12, think_close: 13} \ No newline at end of file diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index 8da488861c8..dbf48e44f56 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -25,8 +25,8 @@ ppo: # critic values / GAE advantages / logprobs (plus decoded text and per-sample # context) to {logger.log_dir}/ppo_rollout_dump_step{N}.pt. Tensors reflect # the state entering the PPO-epoch loop (before any update this step). - log_rollout_dump: false - rollout_dump_period: 1 # dump every N PPO steps (when log_rollout_dump is true) + log_rollout_dump: true + rollout_dump_period: 5 # dump every N PPO steps (when log_rollout_dump is true) reward_shaping: enabled: true @@ -53,7 +53,7 @@ ppo: source_max: 1.0 target_min: -1.0 # DAPO: scale rewards to [-1, 1] target_max: 1.0 - seq_logprob_error_threshold: null + seq_logprob_error_threshold: 2 async_ppo: enabled: false # Set to true for async PPO (requires non-colocated vLLM generation) @@ -74,7 +74,7 @@ ppo: # short rollouts => a dip in mean_gen_tokens_per_sample + reward spike after # each resume). true drops it and regenerates fresh (unbiased, one-target # bubble at resume). See AsyncPPOConfig in nemo_rl/algorithms/ppo.py. - drop_incomplete_targets_on_restore: false + drop_incomplete_targets_on_restore: true # Heartbeat frequency for the collector/replay-buffer per-rollout progress # prints (they flood the log at large num_prompts_per_step). Log every Nth # event (plus the final one per target); 0 silences them; unset => every event. @@ -486,3 +486,4 @@ logger: cluster: gpus_per_node: 1 num_nodes: 1 + segment_size: 2 diff --git a/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml b/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml new file mode 100644 index 00000000000..41f15529cb1 --- /dev/null +++ b/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml @@ -0,0 +1,697 @@ +# ============================================================================= +# PPO Ultra V3 — Multienv Pivot-Only Config (bf16), GB200 NVL72 +# ============================================================================= +# PPO/VAPO port of grpo_ultra_256n4g_bf16_pivotonly.yaml (multienv pivot data: +# code_gen, instruction_following, calendar, tool-use pivots, structured +# outputs, format verification). Run with examples/nemo_gym/run_ppo_nemo_gym.py. +# The GRPO->PPO transformation mirrors ppo_nano_v3_5_swe_cmh.yaml (the +# validated SWE PPO port): +# - `grpo:` section replaced by `ppo:` (GAE adv_estimator, ppo_epochs, +# critic warmup via policy_training_start_step, async_ppo instead of +# async_grpo). max_val_samples / val_batch_size MUST stay null — the +# NeMo-Gym PPO runner rejects preset values and sets both to +# len(val_dataset). GRPO-only keys are dropped (not in PPOConfig): +# normalize_rewards, use_leave_one_out_baseline, advantage_clip_low/high, +# num_val_generations_per_prompt, use_best_at_k*, use_combined_training*, +# dynamic_sampling_oversample_ratio. +# - grpo's penalize_invalid_tool_call / invalid_tool_call_advantage (-5.0) +# are DROPPED: ppo setup() raises NotImplementedError on *_advantage keys +# (message-level advantage overwrites break GAE advantage/return +# consistency). NOTE this means PPO does NOT reproduce the GRPO baseline's +# invalid-tool-call advantage penalty; the PPO-compatible mechanism is the +# reward-zeroing `reward_penalties:` block below (carried over unchanged), +# which has no invalid-tool-call entry. +# - new `value:` block — Megatron-backend value model initialized from the +# policy checkpoint (fresh scalar value head), mirroring the policy's +# parallelism/MoE settings. The nano hybrid (mamba+MoE) critic is +# experimental (megatron value worker validated on GPT-style dense only). +# - new `value_loss_fn:` block (clipped MSE, scale=1.0 / cliprange=0.5 from +# the validated 7B math PPO run). +# - loss_fn: force_on_policy_ratio false (PPO trains on the real +# prev/current logprob ratio; GRPO forced ratio=1 and leaned on TIS only), +# ratio_clip_c 3 (dual-clip). +# - policy/value megatron checkpoint blocks omit the fully_parallel_* keys +# (present in the GRPO parent): this container's megatron-bridge +# CheckpointConfig dropped them and raises InstantiationException +# "Unexpected config keys" (see ppo_nano_v3_5_swe_cmh.yaml). +# - vllm_kwargs adds disable_custom_all_reduce: true — validated fix for the +# GB200 TP-4 custom-allreduce generation-engine deadlock seen on the SWE +# PPO runs (same container/hardware). +# +# Static algorithm, loss, and environment settings live here. Per-run +# overrides (model path, data paths, parallelism, batch sizes) are applied by +# the launch script (scripts/multienv/nano_pivot_ppo.sh) or Hydra CLI. +# ============================================================================= + +# ============================================================================= +# Cluster — overridden by launch script +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 256 + segment_size: 16 + +# ============================================================================= +# Checkpointing +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/ppo_ultra_v3_pivotonly" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_optimizer: true + save_period: 10 + ft_keep_latest_k: 1 + ft_save_period: 1 + checkpoint_must_save_by: "00:03:30:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# PPO Algorithm (replaces the parent's grpo: section) +# ============================================================================= +ppo: + num_prompts_per_step: 256 + num_generations_per_prompt: 16 + max_rollout_turns: 1 + max_num_epochs: 1 + max_num_steps: 1000000 + # One optimizer pass per rollout batch (PPO-epochs>1 reuses the batch). + ppo_epochs: 1 + # Critic-only warmup: policy training is skipped for the first N steps while + # the (fresh-head) value model trains on rollouts from the frozen policy. + # Overridden by the launch script. + policy_training_start_step: 100 + val_period: 10000 + val_at_start: false + val_at_end: false + overlong_filtering: false + # NeMo-Gym principle: the validation set is used verbatim (no hidden pre/post + # processing). run_ppo_nemo_gym.py rejects a preset max_val_samples and sets + # both to len(val_dataset), so they must start null (the GRPO parent's + # val_batch_size: 256 must NOT carry over). + max_val_samples: null + val_batch_size: null + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + # loss_fn.reference_policy_kl_penalty is 0 (GRPO pivot baseline), so skip the + # reference-policy logprob forward. The launch script flips this to false + # automatically when a nonzero KL penalty is requested. + skip_reference_policy_logprobs_calculation: true + seq_logprob_error_threshold: 2 + + # GAE with VAPO decoupled lambdas (defaults match the validated 7B math PPO + # run; the launch script sets length_adaptive_alpha). + adv_estimator: + name: "gae" + gae_lambda: 1.0 + gae_gamma: 1 + normalize_advantages: true + gae_lambda_value: 1.0 + gae_lambda_policy: 1 + # Length-adaptive lambda_policy = 1 - 1/(alpha*l). 0 = disabled. + length_adaptive_alpha: null + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + # NOTE: grpo's penalize_invalid_tool_call / invalid_tool_call_advantage are + # intentionally absent: setup() raises NotImplementedError if the *_advantage + # keys are set on the PPO path (advantage overwrites break GAE consistency). + + async_ppo: + enabled: true + # lag-1: at most one policy version off-policy (recommended/validated). + # During critic warmup the frozen actor makes any-age trajectories + # on-policy for free (warmup_max_trajectory_age_steps; null => same value). + max_trajectory_age_steps: 1 + warmup_max_trajectory_age_steps: 8 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + # Drop the incomplete (survivorship-biased) frontier target on resume and + # regenerate it fresh — unbiased batches at the cost of a one-target + # generation bubble per resume. + drop_incomplete_targets_on_restore: false + # Throttle collector/replay-buffer per-rollout heartbeat prints. + log_every: 500 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + # clip_min/max carried from the GRPO pivot baseline (DAPO 0.2/0.28); the + # launch script overrides clip_max to the validated PPO recipe's 0.2. + # Dual-clip c=3 from the validated 7B math PPO run. + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: 3 + use_on_policy_kl_approximation: true + # Async PPO requires the off-policy importance-sampling correction. + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + # PPO trains on the real prev/current-logprob ratio; forcing ratio=1 (as the + # GRPO parent does) would disable PPO clipping entirely. + force_on_policy_ratio: false + use_kl_in_reward: false + +# ============================================================================= +# Value Loss (MSE with PPO-style clipping) +# ============================================================================= +value_loss_fn: + scale: 1.0 + cliprange: 0.5 + +# ============================================================================= +# Policy +# ============================================================================= +policy: + model_name: "/lustre/fsw/portfolios/llmservice/users/soumyes/sft-runs/eval_and_sleep/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-mar13-blend-512k-filt-1e-5/iter_0001900/hf" + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 4096 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 65536 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + # EP=64: 512 experts / 64 = 8 experts per EP rank + # All-to-all spans 64 GPUs (16 nodes), fits within one NVLink domain + expert_model_parallel_size: 64 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled + mtp_loss_scaling_factor: 0.3 + mtp_use_repeated_layer: true + mtp_num_layers: 5 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 4.0e-6 + min_lr: 4.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 4e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + # FP8 — disabled for bf16 runs. Enable for mxfp8 validation. + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + use_gloo_process_groups: false + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + # NOTE: the GRPO parent's fully_parallel_{save,load}_process_group / + # fully_parallel_load_exchange_algo are dropped: this container's + # megatron-bridge CheckpointConfig no longer has them and raises + # InstantiationException "Unexpected config keys" (see the SWE PPO + # config, ppo_nano_v3_5_swe_cmh.yaml). + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 3000 + port_range_high: 4999 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. + # Same memory as EP=1 but uses all-to-all for expert routing. + # EP > TP blocked by https://github.com/NVIDIA-NeMo/RL/issues/1101. + # The launch script overrides TP and EP together (TP=EP=4 by default) to + # keep the EP<=TP constraint when it shrinks TP. + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 8 + gpu_memory_utilization: 0.85 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + # nano_v3 is a CUSTOM reasoning parser; its plugin file is registered at vllm_cfg + # top level (read by vllm_worker_async.py -> ReasoningParserManager.import_reasoning_parser). + # It must NOT sit inside http_server_serving_chat_kwargs: that dict is splatted into + # OpenAIServingChat(**kwargs), which has no reasoning_parser_plugin arg -> TypeError. + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + + vllm_kwargs: + attention_backend: FLASH_ATTN + moe_backend: triton + mamba_ssm_cache_dtype: "float32" + # Route TP allreduces through NCCL instead of vLLM's custom CUDA-IPC + # spin kernel. On the SWE PPO runs (same container/GB200 nodes), TP-4 + # engines deadlocked in the custom-allreduce kernel (all 4 TP ranks with + # the next collective enqueued-but-never-started, symmetric cross-GPU + # deadlock). Costs a few % decode latency on TP-4. + disable_custom_all_reduce: true + compilation_config: + cudagraph_mode: PIECEWISE + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 182 # Overridden by launch script + +# ============================================================================= +# Value Model (GAE critic) — Megatron backend, initialized from the policy +# checkpoint with a fresh scalar value head. Parallelism/MoE settings mirror +# the policy block; the launch script overrides them together. +# ============================================================================= +value: + model_name: ${policy.model_name} + tokenizer: + name: ${value.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: ${policy.train_global_batch_size} + train_micro_batch_size: ${policy.train_micro_batch_size} + logprob_batch_size: ${policy.logprob_batch_size} + max_total_sequence_length: ${policy.max_total_sequence_length} + precision: "bfloat16" + + reward_model_cfg: + enabled: false # only used for the DTensor V2 value worker + reward_model_type: "regression" + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 64 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${value.precision} + sequence_parallel: true + + # MoE — mirror the policy + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + use_gloo_process_groups: false + + # Compute — mirror the policy + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — fully disabled for the critic. The value model replaces the LM + # head with a scalar head, so MTP next-token heads are dead weight; null + # skips the MTP block on all ranks (no partial-rank EP collective). + mtp_loss_scaling_factor: 0.0 + mtp_use_repeated_layer: true + mtp_num_layers: null + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + # Critic LR (fresh value head trains faster than the pretrained actor). + # Overridden by the launch script. + lr: 6.0e-6 + min_lr: 6.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${value.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${value.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${value.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 6.0e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + + env_vars: null + + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + + make_sequence_length_divisible_by: ${mul:${mul:${value.megatron_cfg.tensor_model_parallel_size}, ${value.megatron_cfg.context_parallel_size}}, 2} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + use_multiple_dataloader: false + train: + data_path: null # Set by launch script + validation: + data_path: null # Set by launch script + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym multienv pivot servers (copied verbatim from the +# GRPO parent grpo_ultra_256n4g_bf16_pivotonly.yaml) +# ============================================================================= +env: + should_use_nemo_gym: true + # true: skip expensive train_data_step*.jsonl (recommended for large Gym runs); false: write full jsonl. + should_log_nemo_gym_responses: true + nemo_gym: + nemo_gym_log_dir: "logs/nemo_gym" + skip_venv_if_present: true + port_range_low: 5000 + port_range_high: 5999 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/math_with_judge/configs/math_with_judge.yaml + - resources_servers/code_gen/configs/code_gen.yaml + - resources_servers/instruction_following/configs/instruction_following.yaml + - resources_servers/calendar/configs/calendar.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml + - resources_servers/terminus_judge/configs/terminus_judge_string_only.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/search_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/toolcall_schema_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/swe_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/droid_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/structured_outputs/configs/structured_outputs_json_yaml_xml_v1.yaml + - resources_servers/structured_outputs/configs/structured_outputs_v3.yaml + - resources_servers/format_verification/configs/freeform_formatting.yaml + - resources_servers/format_verification/configs/citation_format.yaml + + code_gen: + resources_servers: + code_gen: + num_processes: 2048 + unit_test_timeout_secs: 10 + debug: false + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "logs" + num_val_samples_to_print: 0 + # The GRPO parent shipped wandb_enabled: false (and the launch script does + # not override it) — enabled here so W&B logging actually works. + wandb_enabled: true + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "ppo-ultra-v3-pivotonly" + name: "ppo-ultra-v3-pivotonly-256n" + tensorboard: {} + mlflow: + experiment_name: "ppo-ultra-v3-pivotonly" + run_name: "ppo-ultra-v3-pivotonly-256n" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels — carried over from the GRPO parent for parity. The parent's +# other top-level keys (token_ids, penalize_*) are inert in this checkout and +# were dropped on the PPO path (matching ppo_nano_v3_5_swe_cmh.yaml). +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + +# Reward-zeroing penalties applied to NeMo-Gym rollout results (the +# PPO-compatible penalty path: applied to rewards pre-GAE, unlike grpo's +# advantage-overwrite penalties). Carried over unchanged from the GRPO parent. +reward_penalties: + penalize_duplicated_reasoning: true + penalize_empty_final_answer: true + penalize_unwanted_tokens: true + penalize_malformed_think_tag: true + # Optional model/tokenizer-specific token IDs. Add token IDs that should + # be penalized when emitted by the model. Unwanted token IDs must be + # specified explicitly; think-tag IDs are inferred when each tag encodes + # to one token. Example: {unwanted: [2], think_open: 12, think_close: 13} + token_ids: {unwanted: [2], think_open: 12, think_close: 13} From ac59649592d08879bef33385064d20a0fa432a80 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 31 Jul 2026 14:53:14 -0700 Subject: [PATCH 44/69] chore: bump Gym submodule (rollout /run retry + sandbox DNS fix) 67f7955 -> 5613d41: - upstream jiaqiz/ppo-dev (incl. responses_converter tool_choice fix #1925) - fix(rollout_collection): retry /run on transient connection failures - fix(swe_agents): re-bind /etc/resolv.conf into the apptainer sandbox Co-Authored-By: Claude Opus 5 --- 3rdparty/Gym-workspace/Gym | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index 67f7955902a..5613d417271 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit 67f7955902aa7a54678e22fa680858dc1fe7e613 +Subproject commit 5613d417271338ea522ff34b5d63f995ba21ebef From 1fa70a8f9626d38dbbd78ebc09b4b3fb5d93ad8d Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Tue, 4 Aug 2026 20:29:00 -0700 Subject: [PATCH 45/69] ppo: let the critic train for more epochs than the actor PPO's inner loop drove the critic and the actor off a single `ppo.ppo_epochs` count, so fitting the critic harder meant also pushing the policy further off the data that produced its advantages. Add `ppo.critic_ppo_epochs` (null/absent => coupled to ppo_epochs, the previous behavior; must be >= ppo_epochs). The surplus runs as critic-only passes before the shared critic/actor loop, which is left byte-identical -- the only change to existing code is hoisting the value_train_batch selection out of it. Ordering is free because every pass consumes the same returns/advantages, computed once per step and frozen before any update, so the critic never observes the actor's within-step updates. prepare_for_training/finish_training bracket the extra passes as a whole, so they cost one additional param+grad+optimizer onload cycle per step rather than one per pass. Each Megatron worker ticks its LR scheduler once per train() call, so the value model's train_iters budget now scales by critic_ppo_epochs instead of ppo_epochs; otherwise its decay schedule would end before training does. Note that lr_warmup_iters is counted in those same ticks, so a critic running N epochs warms up N times faster in PPO-step terms. Also add `ppo.log_post_update_critic_metrics` (default false). critic/* metrics come from a training pass's forward and therefore describe the critic BEFORE that pass's update; when enabled, one forward-only pass (eval_mode, so it skips optimizer.step and scheduler.step) re-scores the same batch after the final update and emits critic/explained_var_post_update and critic/loss_post_update. It runs while the value model is still resident and the policy is not, so it adds no co-residency -- just the forward. _resolve_critic_ppo_epochs also asserts ppo_epochs >= 1, which sync ppo_train never checked; ppo_epochs: 0 previously produced a run that trained nothing. Co-Authored-By: Claude Opus 5 --- docs/guides/ppo.md | 31 ++- examples/configs/ppo_math_1B.yaml | 11 +- examples/configs/ppo_nano_v3_5_swe_cmh.yaml | 11 + .../ppo_ultra_256n4g_bf16_pivotonly.yaml | 11 + nemo_rl/algorithms/ppo.py | 191 ++++++++++++++++-- tests/unit/algorithms/test_ppo.py | 111 +++++++++- .../ppo_math_1B_megatron.yaml | 11 +- 7 files changed, 341 insertions(+), 36 deletions(-) diff --git a/docs/guides/ppo.md b/docs/guides/ppo.md index 6d4ff7ae745..a38040ca03b 100644 --- a/docs/guides/ppo.md +++ b/docs/guides/ppo.md @@ -48,7 +48,7 @@ PPO also supports **non-colocated generation**, where the generation engine runs PPO supports an **asynchronous** training mode (`ppo.async_ppo.enabled: true`) that mirrors [async GRPO](grpo.md): a background `AsyncTrajectoryCollector` continuously generates trajectories into a replay buffer while the driver trains, so generation and training overlap instead of running in lockstep. Async PPO requires non-colocated vLLM generation (`policy.generation.colocated.enabled: false`, `vllm_cfg.async_engine: true`) and importance-sampling correction (`loss_fn.use_importance_sampling_correction: true`) — all enforced by asserts in `async_ppo_train`. Enable it via `examples/run_ppo.py`, which dispatches to `async_ppo_train` when `ppo.async_ppo.enabled` is set. -Each async step: sample a fixed batch from the replay buffer → compute fresh critic **values** at train time → compute fresh policy/reference logprobs → compute GAE advantages → run the `ppo_epochs` inner loop (critic then actor) → perform a **single** weight refit to the generation engine and bump the replay-buffer weight version. Computing values at train time (rather than stashing them when the trajectory was generated) keeps the critic side as fresh as possible. +Each async step: sample a fixed batch from the replay buffer → compute fresh critic **values** at train time → compute fresh policy/reference logprobs → compute GAE advantages → run the `ppo_epochs` inner loop (critic then actor), preceded by any extra `critic_ppo_epochs` critic-only passes → perform a **single** weight refit to the generation engine and bump the replay-buffer weight version. Computing values at train time (rather than stashing them when the trajectory was generated) keeps the critic side as fresh as possible. **Critic-staleness caveat.** Off-policy trajectories are corrected on the *actor* side by importance sampling (`exp(prev_logprobs - generation_logprobs)`), exactly as in async GRPO. PPO's GAE, however, recursively bootstraps value estimates across each trajectory, so bias from a stale trajectory's actions/rewards compounds along the sequence in a way GRPO's memoryless reward-only advantage does not, and there is **no** corresponding critic-side correction. This bias is *bounded* (not eliminated) by keeping `ppo.async_ppo.max_trajectory_age_steps: 1` (at most one policy version stale), which is the recommended and validated setting. Values above 1 are permitted but only warned about; a rigorous fix (folding importance ratios into the GAE recursion, V-trace style) is future work. @@ -176,7 +176,7 @@ The PPO training loop, [ppo_train](../../nemo_rl/algorithms/ppo.py), follows thi 6. **Value training**: the critic is updated first (critic-before-actor, following [veRL](https://arxiv.org/abs/2412.09613)) 7. **Policy training**: the actor is updated with the clipped surrogate objective -Steps 6–7 repeat `ppo_epochs` times per rollout before generating new responses. +Steps 6–7 repeat `ppo_epochs` times per rollout before generating new responses. Step 6 additionally runs `critic_ppo_epochs - ppo_epochs` extra critic-only passes up front. ### Multiple Training Steps per Rollout @@ -184,10 +184,29 @@ Unlike GRPO, which performs one training update per rollout, PPO can perform mul ```yaml ppo: - ppo_epochs: 4 # Train 4 times on each rollout batch + ppo_epochs: 4 # Train the actor 4 times on each rollout batch ``` -Each step trains both the critic and the actor on the same advantage estimates computed from the initial rollout. +Every pass trains on the same advantage estimates, computed once from the initial rollout. + +#### Training the critic for more epochs + +The critic often benefits from more passes over a rollout batch than the actor: extra actor epochs push the policy further off the data that produced the advantages, while extra critic epochs just fit a regression target harder. Set `critic_ppo_epochs` to give the critic its own epoch count: + +```yaml +ppo: + ppo_epochs: 1 # one actor update per rollout batch + critic_ppo_epochs: 4 # four critic updates on the same batch +``` + +`critic_ppo_epochs: null` (or omitting it) keeps the critic coupled to `ppo_epochs`, which is the default behavior. It must be `>= ppo_epochs`: the critic already trains once per shared epoch, so the setting only adds passes. + +The extra critic-only passes run before the shared critic/actor loop, which is otherwise unchanged. This is safe because every pass consumes the same returns and advantages, computed once per step before any update — the critic never sees the actor's within-step updates, so the ordering does not change the result. For Megatron backends the LR-schedule budget (`train_iters`) is derived per model from its own epoch count, so a longer critic loop does not truncate the value model's decay schedule. + +Two caveats when reading the plots: + +- The critic's LR scheduler ticks once per pass, so with `critic_ppo_epochs: 4` it advances 4x faster per PPO step than the actor's. `lr_warmup_iters` on the value optimizer is counted in those ticks — multiply it by the same factor to keep warmup spanning the same number of PPO steps. +- `critic/explained_var` is measured by a training pass's forward, so it reflects the critic **before** that pass's update. Setting `ppo.log_post_update_critic_metrics: true` adds `critic/explained_var_post_update` and `critic/loss_post_update`, which re-measure the same batch **after** every update that step. The gap between them is how much the critic actually gained on the batch it just fit, and the post-update value is the fairer comparison across different `critic_ppo_epochs` settings. It costs one extra forward-only critic pass over the batch per step, so it defaults to `false`. ### Critic Warmup @@ -235,6 +254,7 @@ ppo: max_num_epochs: 100000 max_num_steps: 100000 ppo_epochs: 4 + critic_ppo_epochs: null policy_training_start_step: 0 val_period: 20 val_at_start: true @@ -280,7 +300,8 @@ value_loss_fn: ``` **PPO-specific parameters:** -- **`ppo.ppo_epochs`**: Number of training updates per rollout batch +- **`ppo.ppo_epochs`**: Number of actor updates per rollout batch +- **`ppo.critic_ppo_epochs`**: Number of critic updates per rollout batch (`null` = same as `ppo.ppo_epochs`) - **`ppo.policy_training_start_step`**: Number of critic-only warmup steps before policy training begins - **`ppo.adv_estimator.name`**: Set to `"gae"` for GAE advantage estimation (PPO default) - **`ppo.adv_estimator.gae_lambda`**: GAE $\lambda$ parameter (bias-variance tradeoff, typically 0.95) diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index dbf48e44f56..45db28bd7ae 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -8,7 +8,16 @@ ppo: max_rollout_turns: 1 max_num_epochs: 100000 max_num_steps: 100000 - ppo_epochs: 4 + ppo_epochs: 4 # actor passes over each rollout batch + # Critic passes over each rollout batch. null = same as ppo_epochs (coupled, + # the historical behavior); must be >= ppo_epochs. The surplus runs as extra + # critic-only passes, fitting the critic harder without over-training the actor. + critic_ppo_epochs: null + # Re-score the batch with a forward-only critic pass after the final update, + # adding critic/explained_var_post_update + critic/loss_post_update (the + # unsuffixed critic/* are measured BEFORE the last update). Costs one extra + # critic forward per step. + log_post_update_critic_metrics: false policy_training_start_step: 0 # number of PPO steps of critic-only warmup before policy training begins val_period: 20 val_at_start: true diff --git a/examples/configs/ppo_nano_v3_5_swe_cmh.yaml b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml index b3a4459718f..76fe092cf23 100644 --- a/examples/configs/ppo_nano_v3_5_swe_cmh.yaml +++ b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml @@ -65,7 +65,18 @@ ppo: max_num_epochs: 4 max_num_steps: 1000000 # One optimizer pass per rollout batch (PPO-epochs>1 reuses the batch). + # This is the ACTOR's epoch count. ppo_epochs: 1 + # Critic (value) passes over each rollout batch. null = same as ppo_epochs + # (coupled, the historical behavior); must be >= ppo_epochs. The surplus runs + # as extra critic-only passes, fitting the critic harder without + # over-training the actor. + critic_ppo_epochs: null + # Re-score the batch with a forward-only critic pass after the final update, + # adding critic/explained_var_post_update + critic/loss_post_update (the + # unsuffixed critic/* are measured BEFORE the last update). Costs one extra + # critic forward per step. + log_post_update_critic_metrics: false # Critic-only warmup: policy training is skipped for the first N steps while # the (fresh-head) value model trains on rollouts from the frozen policy. # Overridden by the launch script. diff --git a/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml b/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml index 41f15529cb1..804bb05091c 100644 --- a/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml +++ b/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml @@ -78,7 +78,18 @@ ppo: max_num_epochs: 1 max_num_steps: 1000000 # One optimizer pass per rollout batch (PPO-epochs>1 reuses the batch). + # This is the ACTOR's epoch count. ppo_epochs: 1 + # Critic (value) passes over each rollout batch. null = same as ppo_epochs + # (coupled, the historical behavior); must be >= ppo_epochs. The surplus runs + # as extra critic-only passes, fitting the critic harder without + # over-training the actor. + critic_ppo_epochs: null + # Re-score the batch with a forward-only critic pass after the final update, + # adding critic/explained_var_post_update + critic/loss_post_update (the + # unsuffixed critic/* are measured BEFORE the last update). Costs one extra + # critic forward per step. + log_post_update_critic_metrics: false # Critic-only warmup: policy training is skipped for the first N steps while # the (fresh-head) value model trains on rollouts from the frozen policy. # Overridden by the launch script. diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index c52adf885f2..3763e342a15 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -221,7 +221,21 @@ class PPOConfig(TypedDict): # When using dynamic sampling, generation prompt batch size will equal # num_prompts_per_step * batch_multiplier batch_multiplier: NotRequired[float] + # Number of actor (policy) passes over each rollout batch. ppo_epochs: int + # Number of critic (value) passes over each rollout batch. None/absent keeps + # the critic coupled to `ppo_epochs` (the historical behavior); set it higher + # (never lower) to fit the critic harder on the same rollout without + # over-training the actor. The surplus runs as critic-only passes before the + # shared inner loop; safe to reorder because every pass consumes the same + # returns/advantages, frozen before any update this step. + critic_ppo_epochs: NotRequired[Optional[int]] + # Re-score the rollout batch with a forward-only pass after the final critic + # update, emitting critic/explained_var_post_update and + # critic/loss_post_update. critic/* without the suffix are measured by a + # training pass, i.e. BEFORE that pass's update. Costs one extra critic + # forward over the batch per step, so this defaults to off. + log_post_update_critic_metrics: NotRequired[bool] reward_shaping: RewardShapingConfig reward_scaling: RewardScalingConfig # By default advantages are calculated on CPU. Setting this flag to true leverages GPU for their computation. @@ -274,6 +288,27 @@ def _default_ppo_save_state() -> PPOSaveState: } +def _resolve_critic_ppo_epochs(ppo_config: PPOConfig) -> int: + """Number of critic (value) passes over each rollout batch. + + ``ppo.critic_ppo_epochs`` lets the critic train longer than the actor on the + same rollout. Unset/None keeps the two coupled, which is the historical + behavior. + """ + ppo_epochs = ppo_config["ppo_epochs"] + assert ppo_epochs >= 1, f"ppo.ppo_epochs must be >= 1 (got {ppo_epochs})." + critic_ppo_epochs = ppo_config.get("critic_ppo_epochs") + if critic_ppo_epochs is None: + return ppo_epochs + # The critic trains once per shared epoch, so it can only be given MORE + # epochs than the actor, never fewer. + assert critic_ppo_epochs >= ppo_epochs, ( + f"ppo.critic_ppo_epochs ({critic_ppo_epochs}) must be >= ppo.ppo_epochs " + f"({ppo_epochs})." + ) + return critic_ppo_epochs + + class PPOLoggerConfig(LoggerConfig): num_val_samples_to_print: int # number of val samples to print to stdout @@ -761,11 +796,14 @@ def setup( # train_iters is the total scheduler-tick budget. Each Megatron worker # ticks once per train() call (matching upstream main's per-rollout - # convention), and PPO calls each worker's train() `ppo_epochs` times - # per outer step. So total ticks = (outer steps) * ppo_epochs. + # convention), and PPO calls each worker's train() once per inner epoch per + # outer step. So total ticks = (outer steps) * (that model's epochs) -- the + # critic may run more than the actor (ppo.critic_ppo_epochs), so it gets its + # own budget, else its LR schedule would finish decaying early. # Scale train_iters accordingly so the configured warmup/decay horizon # matches the actual scheduler-step count. ppo_epochs = ppo_config["ppo_epochs"] + critic_ppo_epochs = _resolve_critic_ppo_epochs(ppo_config) if policy_config.get("megatron_cfg", {}).get("enabled", False): total_train_iters = ( min( @@ -782,7 +820,7 @@ def setup( ppo_config["max_num_steps"], ppo_config["max_num_epochs"] * len(dataloader), ) - * ppo_epochs + * critic_ppo_epochs ) value_config["megatron_cfg"]["train_iters"] = total_train_iters @@ -1700,7 +1738,8 @@ def ppo_train( Based on the grpo_train loop with PPO-specific modifications: - Value model inference and training (actor-critic) - GAE advantage estimation with value bootstrap - - Multiple training steps per rollout (ppo_epochs) + - Multiple training steps per rollout (ppo_epochs), plus optional extra + critic-only passes (critic_ppo_epochs) - Configurable policy training start epoch """ timer = Timer() @@ -1737,6 +1776,14 @@ def ppo_train( current_epoch = ppo_save_state["current_epoch"] max_num_epochs = master_config.ppo["max_num_epochs"] ppo_epochs = master_config.ppo["ppo_epochs"] + # Total critic epochs (ppo.critic_ppo_epochs; defaults to ppo_epochs). Any + # surplus over ppo_epochs runs as extra critic-only passes below. + critic_ppo_epochs = _resolve_critic_ppo_epochs(master_config.ppo) + # Optional forward-only pass after the final critic update (costs one extra + # critic forward per step); see log_post_update_critic_metrics in PPOConfig. + log_post_update_critic_metrics = master_config.ppo.get( + "log_post_update_critic_metrics", False + ) # Number of PPO steps to train only the critic before starting policy # training. Despite the legacy name, this is compared against total_steps # (not current_epoch) to match veRL's critic_warmup semantics. @@ -2194,6 +2241,41 @@ def ppo_train( # PPO: Multiple training steps per rollout memory_tracker.snapshot_start_of_stage("Policy train", dir()) + + # Privileged critic: train on the answer-augmented batch + # (returns already scattered onto its response positions above; + # values=aug old-values for the optional value clip). + if critic_batch is not None: + critic_batch["sample_mask"] = train_data["sample_mask"] + value_train_batch = critic_batch + else: + value_train_batch = train_data + + # Forward-only pass after the final critic update, for the + # post-update critic metrics (see below). + post_value_results = None + + # Extra critic-only passes when ppo.critic_ppo_epochs exceeds + # ppo_epochs. They run before the shared loop below so that loop + # -- and the GPU state it leaves behind -- is untouched. Ordering + # is free: every pass consumes the same returns/advantages, frozen + # above for the whole step. + if critic_ppo_epochs > ppo_epochs: + print( + f"▶ Training value ({critic_ppo_epochs - ppo_epochs} extra " + f"critic-only epochs)...", + flush=True, + ) + with timer.time("value_training_prep"): + value_model.prepare_for_training() + for _ in range(critic_ppo_epochs - ppo_epochs): + with timer.time("value_training"): + value_model.train( + value_train_batch, value_loss_fn, timer=timer + ) + with timer.time("value_training"): + value_model.finish_training() + for step in range(ppo_epochs): print( f"▶ Step {step + 1}/{ppo_epochs}...", @@ -2206,19 +2288,22 @@ def ppo_train( with timer.time("value_training"): print("▶ Training value...", flush=True) - # Privileged critic: train on the answer-augmented batch - # (returns already scattered onto its response positions above; - # values=aug old-values for the optional value clip). - if critic_batch is not None: - critic_batch["sample_mask"] = train_data["sample_mask"] - value_train_batch = critic_batch - else: - value_train_batch = train_data value_results = value_model.train( value_train_batch, value_loss_fn, timer=timer, ) + # After the LAST critic update: one forward-only pass to + # score the updated critic on the same batch. Done here, + # while the value model is still on GPU and the policy is + # not, so it adds no co-residency. + if log_post_update_critic_metrics and step == ppo_epochs - 1: + post_value_results = value_model.train( + value_train_batch, + value_loss_fn, + eval_mode=True, + timer=timer, + ) value_model.finish_training() @@ -2343,6 +2428,19 @@ def ppo_train( # Extract critic metrics from value training results if value_results is not None: metrics.update(_compute_critic_metrics(value_results)) + # critic/explained_var above is measured by the last training + # pass's forward, i.e. BEFORE that pass's update. This is the + # same quantity re-measured AFTER every update this step, so + # the pair shows how much the critic actually gained on the + # batch it just fit. + if post_value_results is not None: + post_metrics = _compute_critic_metrics(post_value_results) + metrics["critic/explained_var_post_update"] = post_metrics[ + "critic/explained_var" + ] + metrics["critic/loss_post_update"] = np.mean( + post_metrics["critic/loss"] + ).item() # Positional critic-quality diagnostic (early/mid/late tokens): how # well V predicts the return along the trajectory, and — for the # privileged critic — where the golden answer sharpens it. @@ -2734,6 +2832,7 @@ def async_ppo_train( 3. computes fresh policy/reference logprobs, 4. computes GAE advantages/returns using those values, 5. runs the ``ppo_epochs`` inner loop (critic train, then actor train), + preceded by any extra ``critic_ppo_epochs`` critic-only passes, 6. performs a single weight refit to the generation engine and bumps the replay-buffer weight version. @@ -2827,6 +2926,14 @@ def async_ppo_train( total_valid_tokens = ppo_save_state.get("total_valid_tokens", 0) max_num_steps = master_config.ppo["max_num_steps"] ppo_epochs = master_config.ppo["ppo_epochs"] + # Total critic epochs (ppo.critic_ppo_epochs; defaults to ppo_epochs). Any + # surplus over ppo_epochs runs as extra critic-only passes below. + critic_ppo_epochs = _resolve_critic_ppo_epochs(master_config.ppo) + # Optional forward-only pass after the final critic update (costs one extra + # critic forward per step); see log_post_update_critic_metrics in PPOConfig. + log_post_update_critic_metrics = master_config.ppo.get( + "log_post_update_critic_metrics", False + ) val_period = master_config.ppo["val_period"] val_at_start = master_config.ppo["val_at_start"] val_at_end = master_config.ppo["val_at_end"] @@ -3358,24 +3465,58 @@ def _sample_max_age(s: int) -> int: is_policy_training_step = step >= policy_training_start_step train_results = None value_results = None + # Forward-only pass after the final critic update, for the + # post-update critic metrics (see the sync loop). + post_value_results = None + + # Privileged critic: train on the answer-augmented batch + # (returns already scattered onto its response positions above; + # values=aug old-values for the optional value clip). + if critic_batch is not None: + critic_batch["sample_mask"] = train_data["sample_mask"] + value_train_batch = critic_batch + else: + value_train_batch = train_data + + # Extra critic-only passes when ppo.critic_ppo_epochs exceeds + # ppo_epochs. They run before the shared loop below so that loop + # -- and the GPU state it leaves behind -- is untouched. Ordering + # is free: every pass consumes the same returns/advantages, frozen + # above for the whole step. + if critic_ppo_epochs > ppo_epochs: + print( + f"▶ {critic_ppo_epochs - ppo_epochs} extra critic-only " + "epochs..." + ) + with timer.time("value_training_prep"): + value_model.prepare_for_training() + for _ in range(critic_ppo_epochs - ppo_epochs): + with timer.time("value_training"): + value_model.train( + value_train_batch, value_loss_fn, timer=timer + ) + with timer.time("value_training"): + value_model.finish_training() + for epoch in range(ppo_epochs): print(f"▶ PPO epoch {epoch + 1}/{ppo_epochs}...") with timer.time("value_training_prep"): value_model.prepare_for_training() with timer.time("value_training"): - # Privileged critic: train on the answer-augmented batch - # (returns already scattered onto its response positions above; - # values=aug old-values for the optional value clip). - if critic_batch is not None: - critic_batch["sample_mask"] = train_data["sample_mask"] - value_train_batch = critic_batch - else: - value_train_batch = train_data value_results = value_model.train( value_train_batch, value_loss_fn, timer=timer, ) + # After the LAST critic update: forward-only pass to score + # the updated critic (value still on GPU, policy not). + if log_post_update_critic_metrics and epoch == ppo_epochs - 1: + post_value_results = value_model.train( + value_train_batch, + value_loss_fn, + eval_mode=True, + timer=timer, + ) value_model.finish_training() if is_policy_training_step: @@ -3536,6 +3677,16 @@ def _sample_max_age(s: int) -> int: metrics.update(train_results["all_mb_metrics"]) if value_results is not None: metrics.update(_compute_critic_metrics(value_results)) + # Same quantity re-measured AFTER every update this step; see + # the sync loop. + if post_value_results is not None: + post_metrics = _compute_critic_metrics(post_value_results) + metrics["critic/explained_var_post_update"] = post_metrics[ + "critic/explained_var" + ] + metrics["critic/loss_post_update"] = np.mean( + post_metrics["critic/loss"] + ).item() # Positional critic-quality diagnostic (early/mid/late tokens): how # well V predicts the return along the trajectory, and — for the # privileged critic — where the golden answer sharpens it. diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 4e593bbd471..2402d6df121 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -836,17 +836,12 @@ def test_ppo_noncolocated_rejects_sglang_backend(): setup(master_config, tokenizer, dataset, None) -@pytest.mark.parametrize("colocated", [True, False]) -def test_ppo_setup_cluster_split_matches_colocation_mode(monkeypatch, colocated): - """train_cluster/inference_cluster identity and worker-group budget by mode. +def _patch_ppo_setup_deps(monkeypatch): + """Stub out every heavy dependency of setup() so it runs end to end in-process. - Colocated: train_cluster is inference_cluster (single shared pool of - generation+policy+value). Non-colocated: they are distinct clusters, and - train_cluster must budget for 2 co-timesharing worker groups (policy + - value) since generation now lives on its own inference_cluster. + Returns the DummyCluster class; its ``instances`` list records the clusters + setup() created, in order. """ - from unittest.mock import MagicMock - from nemo_rl.algorithms import ppo as ppo_mod class DummyLogger: @@ -935,6 +930,23 @@ def init_collective(self, *_args, **_kwargs): ppo_mod, "VllmGeneration", lambda *_a, **_k: DummyVllmGeneration() ) monkeypatch.setattr(ppo_mod.ray, "get", lambda x: x) + return DummyCluster + + +@pytest.mark.parametrize("colocated", [True, False]) +def test_ppo_setup_cluster_split_matches_colocation_mode(monkeypatch, colocated): + """train_cluster/inference_cluster identity and worker-group budget by mode. + + Colocated: train_cluster is inference_cluster (single shared pool of + generation+policy+value). Non-colocated: they are distinct clusters, and + train_cluster must budget for 2 co-timesharing worker groups (policy + + value) since generation now lives on its own inference_cluster. + """ + from unittest.mock import MagicMock + + from nemo_rl.algorithms import ppo as ppo_mod + + _patch_ppo_setup_deps(monkeypatch) master_config = _build_ppo_master_config() if colocated: @@ -967,6 +979,52 @@ def init_collective(self, *_args, **_kwargs): assert inference_cluster.max_colocated_worker_groups == 1 +@pytest.mark.parametrize( + "critic_ppo_epochs,expected_value_iters", + [ + pytest.param(None, 20, id="coupled_defaults_to_ppo_epochs"), + pytest.param(4, 40, id="decoupled_critic_gets_its_own_budget"), + ], +) +def test_ppo_setup_megatron_train_iters_per_model_epochs( + monkeypatch, critic_ppo_epochs, expected_value_iters +): + """Each Megatron model's LR-schedule budget follows its own epoch count. + + A Megatron worker ticks its scheduler once per train() call, so total ticks + are (outer steps) x (that model's inner epochs). With a longer critic loop + the value model must get a proportionally longer train_iters, or its decay + schedule ends before training does. + """ + from unittest.mock import MagicMock + + from nemo_rl.algorithms import ppo as ppo_mod + + _patch_ppo_setup_deps(monkeypatch) + + master_config = _build_ppo_master_config() + master_config.policy["megatron_cfg"] = {"enabled": True} + # context_parallel_size is read by setup()'s value-config validation before + # it ever reaches the train_iters computation. + master_config.value["megatron_cfg"] = {"enabled": True, "context_parallel_size": 1} + master_config.ppo["ppo_epochs"] = 2 + master_config.ppo["critic_ppo_epochs"] = critic_ppo_epochs + # outer_steps = min(max_num_steps, max_num_epochs * len(dataloader)); the + # stubbed dataloader has length 1. + master_config.ppo["max_num_steps"] = 10 + master_config.ppo["max_num_epochs"] = 10 + + tokenizer = MagicMock() + dataset = MagicMock() + dataset.__len__ = MagicMock(return_value=10) + + ppo_mod.setup(master_config, tokenizer, dataset, None) + + # 10 outer steps x 2 actor epochs; value uses critic_ppo_epochs (or 2 when null). + assert master_config.policy["megatron_cfg"]["train_iters"] == 20 + assert master_config.value["megatron_cfg"]["train_iters"] == expected_value_iters + + # =============================================================================== # Async PPO entry guards (async_ppo_train) # =============================================================================== @@ -1062,6 +1120,41 @@ def test_async_ppo_requires_positive_ppo_epochs(): _call_async_ppo_train(master_config) +# --------------------------------------------------------------------------- +# Decoupled critic epochs (ppo.critic_ppo_epochs) +# --------------------------------------------------------------------------- +def test_critic_ppo_epochs_defaults_to_ppo_epochs(): + """Unset or null keeps the critic coupled to the actor (legacy behavior).""" + from nemo_rl.algorithms.ppo import _resolve_critic_ppo_epochs + + assert _resolve_critic_ppo_epochs({"ppo_epochs": 3}) == 3 + assert _resolve_critic_ppo_epochs({"ppo_epochs": 3, "critic_ppo_epochs": None}) == 3 + + +def test_critic_ppo_epochs_decouples_from_ppo_epochs(): + from nemo_rl.algorithms.ppo import _resolve_critic_ppo_epochs + + assert _resolve_critic_ppo_epochs({"ppo_epochs": 1, "critic_ppo_epochs": 4}) == 4 + + +@pytest.mark.parametrize("bad", [0, -1, 1]) +def test_critic_ppo_epochs_rejects_fewer_than_actor_epochs(bad): + """The critic trains once per shared epoch, so it can only be given more.""" + from nemo_rl.algorithms.ppo import _resolve_critic_ppo_epochs + + with pytest.raises(AssertionError, match="must be >= ppo.ppo_epochs"): + _resolve_critic_ppo_epochs({"ppo_epochs": 2, "critic_ppo_epochs": bad}) + + +def test_resolve_critic_ppo_epochs_rejects_zero_ppo_epochs(): + """Guards sync ppo_train, which has no ppo_epochs assert of its own; 0 would + otherwise train neither model while still consuming the whole dataset.""" + from nemo_rl.algorithms.ppo import _resolve_critic_ppo_epochs + + with pytest.raises(AssertionError, match="ppo_epochs must be >= 1"): + _resolve_critic_ppo_epochs({"ppo_epochs": 0}) + + # NeMo-Gym IS supported for async PPO: the AsyncTrajectoryCollector runs the gym # rollout internally (checks _should_use_nemo_gym, passes master_config.reward_penalties), # so async_ppo_train no longer guards it out. Validate-side dispatch is covered by diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 66c2284e040..d5c5df045e9 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -8,7 +8,16 @@ ppo: max_rollout_turns: 1 max_num_epochs: 100000 max_num_steps: 100000 - ppo_epochs: 4 + ppo_epochs: 4 # actor passes over each rollout batch + # Critic passes over each rollout batch. null = same as ppo_epochs (coupled, + # the historical behavior); must be >= ppo_epochs. The surplus runs as extra + # critic-only passes, fitting the critic harder without over-training the actor. + critic_ppo_epochs: null + # Re-score the batch with a forward-only critic pass after the final update, + # adding critic/explained_var_post_update + critic/loss_post_update (the + # unsuffixed critic/* are measured BEFORE the last update). Costs one extra + # critic forward per step. + log_post_update_critic_metrics: false policy_training_start_step: 0 # number of PPO steps of critic-only warmup before policy training begins val_period: 20 val_at_start: true From d8ec18929915919e804d4ff76d76dabab7c8860f Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Tue, 4 Aug 2026 20:39:33 -0700 Subject: [PATCH 46/69] examples: add PPO nano v3.5 SWE config for the hsg cluster PPO/VAPO port of grpo_nano_v3_5_swe_hsg.yaml for GB200 NVL72, run via examples/nemo_gym/run_ppo_nemo_gym.py. Mirrors the existing cmh variant, with the apptainer SWE image paths pointing at this cluster's image store and moe_hybridep_num_sms omitted (Megatron in the NCCL-2.30.4 container deprecates it in favour of moe_flex_dispatcher_num_sms, which this checkout does not plumb, and errors out when both deprecated knobs are set). Enables the per-token rollout dump every 5 steps and the post-update critic metrics for critic-quality debugging. Co-Authored-By: Claude Opus 5 --- examples/configs/ppo_nano_v3_5_swe_hsg.yaml | 724 ++++++++++++++++++++ 1 file changed, 724 insertions(+) create mode 100644 examples/configs/ppo_nano_v3_5_swe_hsg.yaml diff --git a/examples/configs/ppo_nano_v3_5_swe_hsg.yaml b/examples/configs/ppo_nano_v3_5_swe_hsg.yaml new file mode 100644 index 00000000000..a7b9ba63fe8 --- /dev/null +++ b/examples/configs/ppo_nano_v3_5_swe_hsg.yaml @@ -0,0 +1,724 @@ +# ============================================================================= +# PPO Ultra E2E — SWE Config for GB200 NVL72 (PPO/VAPO port of the GRPO config) +# ============================================================================= +# Duplicated from grpo_nano_v3_5_swe_hsg.yaml and adapted for PPO with a GAE +# critic (run with examples/nemo_gym/run_ppo_nemo_gym.py). Differences vs the +# GRPO parent: +# - `grpo:` section replaced by `ppo:` (GAE adv_estimator, ppo_epochs, +# critic warmup via policy_training_start_step, async_ppo instead of +# async_grpo). max_val_samples / val_batch_size MUST stay null — the +# NeMo-Gym PPO runner rejects preset values and sets both to +# len(val_dataset). +# - new `value:` block — Megatron-backend value model initialized from the +# same nano checkpoint (fresh scalar value head), mirroring the policy's +# parallelism/MoE settings. NOTE: the megatron value worker has only been +# validated on GPT-style dense models; the nano hybrid (mamba+MoE) critic +# is experimental. +# - new `value_loss_fn:` block (clipped MSE, VAPO-style scale=1.0/clip=0.5 +# from the validated 7B math PPO run). +# - loss_fn.force_on_policy_ratio false: PPO trains on the real prev/current +# logprob ratio (GRPO nano forced ratio=1 and leaned on TIS only). +# - grpo's invalid_tool_call_advantage / malformed_thinking_advantage are +# DROPPED: message-level advantage-overwrite penalties are rejected on the +# PPO path (they would break GAE advantage/return consistency). The +# reward-zeroing `reward_penalties` block is the PPO-compatible path; left +# all-off (default) for parity with the GRPO baseline, which also ran with +# them off (its top-level penalize_* keys are inert in this checkout). +# +# Static algorithm, loss, and environment settings live here. Per-run +# overrides (model path, data paths, parallelism, batch sizes) are applied by +# the launch script (scripts/swe/ppo/nano_on_main.sh) or Hydra CLI. +# ============================================================================= + +# ============================================================================= +# Cluster — overridden by launch script for parallelism changes +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 128 + segment_size: 16 + +# ============================================================================= +# Checkpointing +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/ppo_ultra_e2e" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 3 + ft_keep_latest_k: 1 + ft_save_period: 1 + save_optimizer: true + checkpoint_must_save_by: "00:03:30:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# PPO Algorithm (replaces the parent's grpo: section) +# ============================================================================= +ppo: + num_prompts_per_step: 16 + num_generations_per_prompt: 32 + max_rollout_turns: 1 + max_num_epochs: 4 + max_num_steps: 1000000 + # One optimizer pass per rollout batch (PPO-epochs>1 reuses the batch). + # This is the ACTOR's epoch count. + ppo_epochs: 1 + # Critic (value) passes over each rollout batch. null = same as ppo_epochs + # (coupled, the historical behavior); must be >= ppo_epochs. The surplus runs + # as extra critic-only passes, fitting the critic harder without + # over-training the actor. + critic_ppo_epochs: null + # Re-score the batch with a forward-only critic pass after the final update, + # adding critic/explained_var_post_update + critic/loss_post_update (the + # unsuffixed critic/* are measured BEFORE the last update). Costs one extra + # critic forward per step. + log_post_update_critic_metrics: true + # Critic-only warmup: policy training is skipped for the first N steps while + # the (fresh-head) value model trains on rollouts from the frozen policy. + # Overridden by the launch script. + policy_training_start_step: 50 + val_period: 10000 + val_at_start: false + val_at_end: false + overlong_filtering: false + # NeMo-Gym principle: the validation set is used verbatim (no hidden pre/post + # processing). run_ppo_nemo_gym.py rejects a preset max_val_samples and sets + # both to len(val_dataset), so they must start null. + max_val_samples: null + val_batch_size: null + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + # loss_fn.reference_policy_kl_penalty is 0 (SWE baseline), so skip the + # reference-policy logprob forward. The launch script flips this to false + # automatically when a nonzero KL penalty is requested. + skip_reference_policy_logprobs_calculation: true + seq_logprob_error_threshold: 2 + # Per-token rollout dump for offline analysis: writes packed token ids / + # critic values / GAE advantages / logprobs (plus decoded text and per-sample + # context) to {logger.log_dir}/ppo_rollout_dump_step{N}.pt. Tensors reflect + # the state entering the PPO-epoch loop (before any update this step). + log_rollout_dump: true + rollout_dump_period: 5 # dump every N PPO steps (when log_rollout_dump is true) + + # GAE with VAPO decoupled lambdas (defaults match the validated 7B math PPO + # run: lambda_value=1.0, length-adaptive lambda_policy with alpha=0.05). + adv_estimator: + name: "gae" + gae_lambda: 1.0 + gae_gamma: 1 + normalize_advantages: true + gae_lambda_value: 1.0 + gae_lambda_policy: 1 + # Length-adaptive lambda_policy = 1 - 1/(alpha*l). 0 = disabled. + length_adaptive_alpha: null + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + # NOTE: grpo's penalize_invalid_tool_call / invalid_tool_call_advantage / + # penalize_malformed_thinking / malformed_thinking_advantage are intentionally + # absent: setup() raises NotImplementedError if the *_advantage keys are set + # on the PPO path (advantage overwrites break GAE consistency). + + async_ppo: + enabled: true + # lag-1: at most one policy version off-policy (recommended/validated). + # During critic warmup the frozen actor makes any-age trajectories + # on-policy for free (warmup_max_trajectory_age_steps; null => same value). + max_trajectory_age_steps: 1 + warmup_max_trajectory_age_steps: 8 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + # Drop the incomplete (survivorship-biased) frontier target on resume and + # regenerate it fresh — unbiased batches at the cost of a one-target + # generation bubble per resume. + drop_incomplete_targets_on_restore: false + # Throttle collector/replay-buffer per-rollout heartbeat prints. + log_every: 500 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + # Defaults from the validated 7B math PPO run (symmetric 0.2 clip + dual-clip + # c=3), not the GRPO-DAPO asymmetric 0.2/0.28. Overridable from the script. + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: 3 + use_on_policy_kl_approximation: true + # Async PPO requires the off-policy importance-sampling correction. + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + # PPO trains on the real prev/current-logprob ratio; forcing ratio=1 (as the + # GRPO nano config does) would disable PPO clipping entirely. + force_on_policy_ratio: false + use_kl_in_reward: false + +# ============================================================================= +# Value Loss (MSE with PPO-style clipping) +# ============================================================================= +value_loss_fn: + scale: 1.0 + cliprange: 0.5 + +# ============================================================================= +# Policy +# ============================================================================= +policy: + model_name: "/lustre/fsw/portfolios/llmservice/users/wdai/megatron-lm-ultra/checkpoints/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-288k-nano-loss-032026/iter_0007000/hf" + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 512 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 131072 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 32 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_flex_dispatcher_backend: "hybridep" + # moe_hybridep_num_sms removed: Megatron in the NCCL-2.30.4 container deprecates + # both moe_hybridep_num_sms and moe_deepep_num_sms in favour of + # moe_flex_dispatcher_num_sms and errors out when both deprecated knobs are set + # ("Conflicting deprecated SM-count knobs"). nemo_rl on this branch does not plumb + # the new knob, so drop this and take Megatron's default SM count. + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + use_gloo_process_groups: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + # community_import.py reads megatron_cfg["gradient_accumulation_fusion"] directly + # (no default) on the HF->megatron conversion path; absent here it raises KeyError. + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled + mtp_loss_scaling_factor: 0.3 + mtp_use_repeated_layer: true + # null disables MTP entirely (matches the upstream mtp=0 run's behavior). This + # checkout's hybrid model gates MTP on `mtp_num_layers is not None` + # (hybrid_model.py:544): 0 still enters the MTP block and asserts > 0, and 1 runs + # MTP forward on only the last-stage ranks (suspected cause of the EXPERT_MODEL_ + # PARALLEL NCCL collective timeout). null makes that check False -> MTP is skipped + # on all ranks, no partial-rank EP collective. Requires the schema change making + # MegatronConfig.mtp_num_layers accept None (nemo_rl/models/policy/__init__.py). + mtp_num_layers: 5 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 3.0e-6 + min_lr: 3.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 3e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + # FP8 — overridden by precision recipe in launch script + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo + # were dropped from this checkout's megatron-bridge CheckpointConfig (submodule + # 554c7b9); leaving them in raises InstantiationException "Unexpected config keys". + # Removed for the current-repo reproduction. + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + # Must be a multiple of minimum_pad_factor = cp_size*2*tp_size (see + # nemo_rl/models/megatron/data.py). With CP=16 + TP=4 + sequence_parallel that is + # 128; the old value (=TP=4) made get_logprobs assert and abort every training + # step. Formula matches examples/nemo_gym/grpo_nanov3.yaml (TP*CP*2). + make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 11001 + port_range_high: 15000 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 1 + gpu_memory_utilization: 0.8 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + # reasoning_parser_plugin lives at vllm_cfg top level in this checkout + # (vllm_worker_async reads vllm_cfg["reasoning_parser_plugin"]); passing it + # inside http_server_serving_chat_kwargs makes OpenAIServingChat raise + # "unexpected keyword argument 'reasoning_parser_plugin'". Path updated to this + # repo's location (the parser moved out of nemo_rl/utils/). + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + + vllm_kwargs: + attention_backend: FLASH_ATTN + moe_backend: triton + mamba_ssm_cache_dtype: "float32" + # Route TP allreduces through NCCL instead of vLLM's custom CUDA-IPC + # spin kernel. Jobs 37606/38016 each lost a generation engine ~15 min in: + # the NCCL flight recorder showed all 4 TP ranks with their next + # collective enqueued-but-never-started, i.e. all four streams parked in + # the custom-allreduce kernel ahead of it (symmetric cross-GPU deadlock, + # different node each time). Costs a few % decode latency on TP-4. + disable_custom_all_reduce: true + compilation_config: + cudagraph_mode: PIECEWISE + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 64 + +# ============================================================================= +# Value Model (GAE critic) — Megatron backend, initialized from the policy +# checkpoint with a fresh scalar value head. Parallelism/MoE settings mirror +# the policy block; the launch script overrides them together. +# ============================================================================= +value: + model_name: ${policy.model_name} + tokenizer: + name: ${value.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: ${policy.train_global_batch_size} + train_micro_batch_size: ${policy.train_micro_batch_size} + logprob_batch_size: ${policy.logprob_batch_size} + max_total_sequence_length: ${policy.max_total_sequence_length} + precision: "bfloat16" + + reward_model_cfg: + enabled: false # only used for the DTensor V2 value worker + reward_model_type: "regression" + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 32 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${value.precision} + sequence_parallel: true + + # MoE — mirror the policy + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_flex_dispatcher_backend: "hybridep" + # moe_hybridep_num_sms removed: Megatron in the NCCL-2.30.4 container deprecates + # both moe_hybridep_num_sms and moe_deepep_num_sms in favour of + # moe_flex_dispatcher_num_sms and errors out when both deprecated knobs are set + # ("Conflicting deprecated SM-count knobs"). nemo_rl on this branch does not plumb + # the new knob, so drop this and take Megatron's default SM count. + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + use_gloo_process_groups: false + + # Compute — mirror the policy + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — fully disabled for the critic. The value model replaces the LM + # head with a scalar head, so MTP next-token heads are dead weight; null + # skips the MTP block on all ranks (no partial-rank EP collective). + mtp_loss_scaling_factor: 0.0 + mtp_use_repeated_layer: true + mtp_num_layers: null + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + # Critic LR from the validated 7B math PPO run (fresh value head trains + # faster than the pretrained actor). Overridden by the launch script. + lr: 6.0e-6 + min_lr: 6.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${value.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${value.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${value.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 6.0e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + + env_vars: null + + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + + make_sequence_length_divisible_by: ${mul:${mul:${value.megatron_cfg.tensor_model_parallel_size}, ${value.megatron_cfg.context_parallel_size}}, 2} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + use_multiple_dataloader: false + train: + data_path: null # Set by launch script + validation: + data_path: null # Set by launch script + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + SWE Agents +# ============================================================================= +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + # Reuse pre-built gym venvs. With false, all 3 gym servers (policy_model, + # swe_agents_{train,val}) rebuild the editable nemo-gym package in parallel and + # contend on the same uv distribution-cache lock on Lustre; policy_model then + # times out (default 300s) and NemoGym spinup fails. Mirrors the qwen-30b SWE run. + skip_venv_if_present: true + num_gpu_nodes: 0 + port_range_low: 15001 + port_range_high: 20000 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml + swe_agents_train: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}} + swebench_agent_timeout: 2700 + swebench_tests_timeout: 900 + run_with_mixed_prompts: true + dataset_path: ${data.train.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swerebench/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/nv_internal/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/r2e_gym/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swegym/sweb.eval.arm64.{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swebench/swe-bench.eval.arm64.{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/mercor/swebenchpro_ots/{instance_id}.sif" + - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/mercor/{instance_id}.sif" + swe_agents_val: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}} + swebench_agent_timeout: 3600 + dataset_path: ${data.validation.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter} + use_absolute_ip: true + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "logs" + num_val_samples_to_print: 0 + wandb_enabled: true + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "ultra-v3-swe-e2e" + name: "ppo-ultra-e2e" + tensorboard: {} + mlflow: + experiment_name: "ppo-ultra-e2e" + run_name: "ppo-ultra-e2e" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels / Token IDs / penalize_* — carried over from the GRPO config +# for parity. NOTE: these top-level keys are INERT in this checkout (nothing +# reads them; the GRPO baseline also ran with them having no effect). The +# PPO-compatible penalty mechanism is the structured `reward_penalties:` block +# (reward-zeroing, pre-GAE) — left at its all-off default to match the GRPO +# baseline's effective behavior. +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + + +# Reward-zeroing penalties applied to NeMo-Gym rollout results. +reward_penalties: + penalize_duplicated_reasoning: true + penalize_empty_final_answer: true + penalize_unwanted_tokens: true + penalize_malformed_think_tag: true + # Optional model/tokenizer-specific token IDs. Add token IDs that should + # be penalized when emitted by the model. Unwanted token IDs must be + # specified explicitly; think-tag IDs are inferred when each tag encodes + # to one token. Example: {unwanted: [2], think_open: 12, think_close: 13} + token_ids: {unwanted: [2], think_open: 12, think_close: 13} \ No newline at end of file From eb028cee5d66b0fd85873a09220d3bc63d0c79d2 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Tue, 4 Aug 2026 20:40:43 -0700 Subject: [PATCH 47/69] data: import decord lazily in load_media_from_message decord is only reachable on the audio fallback path, but the top-level import made it a hard dependency of multimodal_utils, which every entrypoint pulls in via batched_data_dict. In containers whose prebaked venvs omit decord that killed the driver with ModuleNotFoundError before any text-only run could start. Move the import to its use site, matching what nemo_rl/data/datasets/utils.py already does. Co-Authored-By: Claude Opus 5 --- nemo_rl/data/multimodal_utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 6a16e18fec8..3b62ca5be96 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -20,7 +20,6 @@ from io import BytesIO from typing import Any, Optional, Union -import decord import requests import torch from PIL import Image @@ -358,7 +357,12 @@ def load_media_from_message( ) except (RuntimeError, FileNotFoundError, OSError) as e: logger.warning("Audio loading failed. Fall back to decord.") - # use decord + # Imported lazily (same as nemo_rl/data/datasets/utils.py) so that + # text-only runs work in containers whose prebaked venvs omit + # decord -- a top-level import here reaches every entrypoint via + # batched_data_dict and kills the driver with ModuleNotFoundError. + import decord + loaded_audio = decord.AudioReader( aud, sample_rate=multimodal_load_kwargs["audio"]["sampling_rate"], From 4c1626bf7b2c10264616ecd5937ab8fab027dacf Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Tue, 4 Aug 2026 20:41:36 -0700 Subject: [PATCH 48/69] ppo: load critic weights only when resuming from a warm-start seed A warm-start seed fabricated by prep_warm_start.sh carries critic weights from a stage-B run whose expert-parallel layout may differ from the resuming run's. Megatron's distributed optimizer state is sharded by (expert-)DP group index and cannot reshard across a changed EP, so the load died with "Missing key in checkpoint state_dict: chained_1.optimizer.distributed.dp_group_idx_N...". Model weights do reshard, so take those and rebuild Adam plus the LR schedule from scratch (value_lr_warmup_iters re-warms it). Gated on the warm_start_provenance.txt marker that only the fabricated seed carries, so ordinary step_N resumes keep restoring the optimizer as before. Co-Authored-By: Claude Opus 5 --- nemo_rl/algorithms/ppo.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 3763e342a15..aa4badeae7b 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -853,6 +853,25 @@ def init_value(): value_optimizer_path = _resolve_resume_optimizer_path( _value_optim, value_weights_path, value_config ) + # Warm-start seed (prep_warm_start.sh): the critic weights come from a + # stage-B run whose expert-parallel layout may differ from this run's. + # Megatron's distributed optimizer state is sharded by (expert-)DP group + # index, so it cannot reshard across a changed EP -- the load dies with + # "Missing key in checkpoint state_dict: chained_1.optimizer.distributed. + # dp_group_idx_N...". Model weights DO reshard, so take those and rebuild + # Adam + the LR schedule fresh (value_lr_warmup_iters re-warms it). + # Self-limiting: only the fabricated seed carries the provenance file, so + # ordinary step_N resumes keep restoring the optimizer normally. + if ( + value_optimizer_path is not None + and (Path(last_checkpoint_path) / "warm_start_provenance.txt").exists() + ): + print( + " ⚠ Warm-start seed detected: loading critic weights only " + "(fresh optimizer + LR schedule).", + flush=True, + ) + value_optimizer_path = None if value_weights_path is None: print( f" ⚠ Value weights not found in checkpoint {last_checkpoint_path} " From 0003f1b036a00df7fcd35c42d6f8f0a3680f7a0d Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Tue, 4 Aug 2026 20:42:10 -0700 Subject: [PATCH 49/69] examples: cap SWE test execution separately from agent time (cmh) Split the single 3600s swebench_agent_timeout into 2700s of agent time plus a dedicated 900s swebench_tests_timeout, leaving the 3600s total unchanged so a long test run can no longer consume the whole budget. Co-Authored-By: Claude Opus 5 --- examples/configs/ppo_nano_v3_5_swe_cmh.yaml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/configs/ppo_nano_v3_5_swe_cmh.yaml b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml index 76fe092cf23..f4fb6d88ae2 100644 --- a/examples/configs/ppo_nano_v3_5_swe_cmh.yaml +++ b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml @@ -636,7 +636,8 @@ env: swe_agents: agent_max_turns: 200 concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}} - swebench_agent_timeout: 3600 + swebench_agent_timeout: 2700 + swebench_tests_timeout: 900 run_with_mixed_prompts: true dataset_path: ${data.train.data_path} apptainer_memory_limit_mb: 65536 From e7ac37bc42118528e4a8f6d90636f80f4c125e11 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Wed, 5 Aug 2026 12:08:35 -0700 Subject: [PATCH 50/69] ppo: report critic/explained_var from the pre-update values With critic_ppo_epochs > 1 the loss-derived explained_var came from the last training pass's forward -- a critic that had already fit this batch several times -- so it read as batch-fit rather than critic quality. On run awk40z9y (critic_ppo_epochs=4) it sat at ~0.93 while the positional ev_early/mid/late diagnostics, computed from the rollout-time values, showed the critic's actual pre-update EV of ~0.33. Overwrite critic/explained_var at both metric sites with a pooled EV computed driver-side from train_data values/returns -- the exact tensors GAE consumed, i.e. the critic before ANY update this step, regardless of epoch count or gbs-vs-rollout microbatching. Masked like the value loss (token * sample) so it stays directly comparable to critic/explained_var_post_update, which remains the after-all-updates end of the bracket. critic/loss and critic/grad_norm still describe the last training pass. No extra compute: the tensors are already on the driver. Co-Authored-By: Claude Fable 5 --- docs/guides/ppo.md | 2 +- examples/configs/ppo_math_1B.yaml | 7 +- examples/configs/ppo_nano_v3_5_swe_cmh.yaml | 7 +- examples/configs/ppo_nano_v3_5_swe_hsg.yaml | 7 +- .../ppo_ultra_256n4g_bf16_pivotonly.yaml | 7 +- nemo_rl/algorithms/ppo.py | 70 ++++++++++++++++--- tests/unit/algorithms/test_ppo.py | 42 +++++++++++ .../ppo_math_1B_megatron.yaml | 7 +- 8 files changed, 124 insertions(+), 25 deletions(-) diff --git a/docs/guides/ppo.md b/docs/guides/ppo.md index a38040ca03b..31f0ed743fa 100644 --- a/docs/guides/ppo.md +++ b/docs/guides/ppo.md @@ -206,7 +206,7 @@ The extra critic-only passes run before the shared critic/actor loop, which is o Two caveats when reading the plots: - The critic's LR scheduler ticks once per pass, so with `critic_ppo_epochs: 4` it advances 4x faster per PPO step than the actor's. `lr_warmup_iters` on the value optimizer is counted in those ticks — multiply it by the same factor to keep warmup spanning the same number of PPO steps. -- `critic/explained_var` is measured by a training pass's forward, so it reflects the critic **before** that pass's update. Setting `ppo.log_post_update_critic_metrics: true` adds `critic/explained_var_post_update` and `critic/loss_post_update`, which re-measure the same batch **after** every update that step. The gap between them is how much the critic actually gained on the batch it just fit, and the post-update value is the fairer comparison across different `critic_ppo_epochs` settings. It costs one extra forward-only critic pass over the batch per step, so it defaults to `false`. +- `critic/explained_var` is computed from the rollout-time values — the exact tensors GAE consumed — so it describes the critic **before any update that step**, regardless of `critic_ppo_epochs` (it will sit near the positional `critic/ev_early|mid|late` diagnostics, which use the same values). Setting `ppo.log_post_update_critic_metrics: true` adds `critic/explained_var_post_update` and `critic/loss_post_update`, which re-score the same batch **after** every update that step, at the cost of one extra forward-only critic pass per step. A large post-vs-pre gap that does not lift the pre-update curve over training means the critic is memorizing each batch rather than generalizing. `critic/loss` and `critic/grad_norm` come from the last training pass. ### Critic Warmup diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index 45db28bd7ae..dd73b615d08 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -14,9 +14,10 @@ ppo: # critic-only passes, fitting the critic harder without over-training the actor. critic_ppo_epochs: null # Re-score the batch with a forward-only critic pass after the final update, - # adding critic/explained_var_post_update + critic/loss_post_update (the - # unsuffixed critic/* are measured BEFORE the last update). Costs one extra - # critic forward per step. + # adding critic/explained_var_post_update + critic/loss_post_update. + # critic/explained_var is always the PRE-update EV (from the rollout-time + # values GAE consumed); critic/loss is from the last training pass. Costs one + # extra critic forward per step. log_post_update_critic_metrics: false policy_training_start_step: 0 # number of PPO steps of critic-only warmup before policy training begins val_period: 20 diff --git a/examples/configs/ppo_nano_v3_5_swe_cmh.yaml b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml index f4fb6d88ae2..3bff312d166 100644 --- a/examples/configs/ppo_nano_v3_5_swe_cmh.yaml +++ b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml @@ -73,9 +73,10 @@ ppo: # over-training the actor. critic_ppo_epochs: null # Re-score the batch with a forward-only critic pass after the final update, - # adding critic/explained_var_post_update + critic/loss_post_update (the - # unsuffixed critic/* are measured BEFORE the last update). Costs one extra - # critic forward per step. + # adding critic/explained_var_post_update + critic/loss_post_update. + # critic/explained_var is always the PRE-update EV (from the rollout-time + # values GAE consumed); critic/loss is from the last training pass. Costs one + # extra critic forward per step. log_post_update_critic_metrics: false # Critic-only warmup: policy training is skipped for the first N steps while # the (fresh-head) value model trains on rollouts from the frozen policy. diff --git a/examples/configs/ppo_nano_v3_5_swe_hsg.yaml b/examples/configs/ppo_nano_v3_5_swe_hsg.yaml index a7b9ba63fe8..aff7e41f8dd 100644 --- a/examples/configs/ppo_nano_v3_5_swe_hsg.yaml +++ b/examples/configs/ppo_nano_v3_5_swe_hsg.yaml @@ -73,9 +73,10 @@ ppo: # over-training the actor. critic_ppo_epochs: null # Re-score the batch with a forward-only critic pass after the final update, - # adding critic/explained_var_post_update + critic/loss_post_update (the - # unsuffixed critic/* are measured BEFORE the last update). Costs one extra - # critic forward per step. + # adding critic/explained_var_post_update + critic/loss_post_update. + # critic/explained_var is always the PRE-update EV (from the rollout-time + # values GAE consumed); critic/loss is from the last training pass. Costs one + # extra critic forward per step. log_post_update_critic_metrics: true # Critic-only warmup: policy training is skipped for the first N steps while # the (fresh-head) value model trains on rollouts from the frozen policy. diff --git a/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml b/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml index 804bb05091c..5e8a3a32441 100644 --- a/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml +++ b/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml @@ -86,9 +86,10 @@ ppo: # over-training the actor. critic_ppo_epochs: null # Re-score the batch with a forward-only critic pass after the final update, - # adding critic/explained_var_post_update + critic/loss_post_update (the - # unsuffixed critic/* are measured BEFORE the last update). Costs one extra - # critic forward per step. + # adding critic/explained_var_post_update + critic/loss_post_update. + # critic/explained_var is always the PRE-update EV (from the rollout-time + # values GAE consumed); critic/loss is from the last training pass. Costs one + # extra critic forward per step. log_post_update_critic_metrics: false # Critic-only warmup: policy training is skipped for the first N steps while # the (fresh-head) value model trains on rollouts from the frozen policy. diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index aa4badeae7b..2ade166867b 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -232,8 +232,9 @@ class PPOConfig(TypedDict): critic_ppo_epochs: NotRequired[Optional[int]] # Re-score the rollout batch with a forward-only pass after the final critic # update, emitting critic/explained_var_post_update and - # critic/loss_post_update. critic/* without the suffix are measured by a - # training pass, i.e. BEFORE that pass's update. Costs one extra critic + # critic/loss_post_update. critic/explained_var is always the PRE-update EV, + # computed from the rollout-time values GAE consumed (critic/loss and + # critic/grad_norm come from the last training pass). Costs one extra critic # forward over the batch per step, so this defaults to off. log_post_update_critic_metrics: NotRequired[bool] reward_shaping: RewardShapingConfig @@ -1520,6 +1521,12 @@ def _compute_critic_metrics(value_results: dict[str, Any]) -> dict[str, Any]: # Compute explained variance from sufficient statistics: # EV = 1 - Var(returns - values) / Var(returns) + # NOTE: these statistics come from a TRAINING pass's forward, so with + # critic_ppo_epochs > 1 they see a critic already fitted by the earlier + # passes. ppo_train/async_ppo_train overwrite critic/explained_var with the + # pre-update EV from the rollout-time values (_pooled_explained_var); this + # loss-derived one only survives as critic/explained_var_post_update, where + # the post-update timing is the point. r_mean = critic_metrics.get("critic/returns_mean", 0) v_mean = critic_metrics.get("critic/values_mean", 0) r_sq = critic_metrics.get("critic/returns_sq_mean", 0) @@ -1530,6 +1537,31 @@ def _compute_critic_metrics(value_results: dict[str, Any]) -> dict[str, Any]: return critic_metrics +def _pooled_explained_var( + values: torch.Tensor, + returns: torch.Tensor, + token_mask: torch.Tensor, + sample_mask: torch.Tensor, +) -> float: + """Pooled EV = 1 - Var(returns - values) / Var(returns) over valid tokens. + + Computed driver-side from the rollout-time values -- the exact tensors GAE + consumed -- so it describes the critic BEFORE any update this step, + regardless of critic_ppo_epochs or gbs-vs-rollout microbatching. Masked like + the value loss (token_mask * sample_mask) so it is directly comparable to + critic/explained_var_post_update. + """ + mask = (token_mask * sample_mask.unsqueeze(-1)).bool() + if int(mask.sum()) < 2: + return 0.0 + r = returns[mask].float() + v = values[mask].float() + var_returns = r.var(unbiased=False) + if var_returns <= 1e-8: + return 0.0 + return (1.0 - (r - v).var(unbiased=False) / var_returns).item() + + def _calibration_ece( v: torch.Tensor, r: torch.Tensor, n_conf_bins: int = 10 ) -> float: @@ -2447,11 +2479,22 @@ def ppo_train( # Extract critic metrics from value training results if value_results is not None: metrics.update(_compute_critic_metrics(value_results)) - # critic/explained_var above is measured by the last training - # pass's forward, i.e. BEFORE that pass's update. This is the - # same quantity re-measured AFTER every update this step, so - # the pair shows how much the critic actually gained on the - # batch it just fit. + # critic/explained_var must describe the critic BEFORE any + # update this step (the values GAE consumed). The loss-derived + # EV from _compute_critic_metrics comes from the last training + # pass's forward, which with critic_ppo_epochs > 1 has already + # fit this batch -- overwrite it with the pre-update EV. + # critic/loss and critic/grad_norm still describe the last + # training pass. + if "values" in train_data and "returns" in train_data: + metrics["critic/explained_var"] = _pooled_explained_var( + train_data["values"], + train_data["returns"], + train_data["token_mask"], + train_data["sample_mask"], + ) + # The other end of the bracket: the same batch re-scored AFTER + # every update this step (ppo.log_post_update_critic_metrics). if post_value_results is not None: post_metrics = _compute_critic_metrics(post_value_results) metrics["critic/explained_var_post_update"] = post_metrics[ @@ -3696,8 +3739,17 @@ def _sample_max_age(s: int) -> int: metrics.update(train_results["all_mb_metrics"]) if value_results is not None: metrics.update(_compute_critic_metrics(value_results)) - # Same quantity re-measured AFTER every update this step; see - # the sync loop. + # Overwrite critic/explained_var with the pre-update EV (the + # values GAE consumed); see the sync loop for why. + if "values" in train_data and "returns" in train_data: + metrics["critic/explained_var"] = _pooled_explained_var( + train_data["values"], + train_data["returns"], + train_data["token_mask"], + train_data["sample_mask"], + ) + # The other end of the bracket: re-scored AFTER every update + # this step (ppo.log_post_update_critic_metrics). if post_value_results is not None: post_metrics = _compute_critic_metrics(post_value_results) metrics["critic/explained_var_post_update"] = post_metrics[ diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 2402d6df121..3aa58a94901 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -1155,6 +1155,48 @@ def test_resolve_critic_ppo_epochs_rejects_zero_ppo_epochs(): _resolve_critic_ppo_epochs({"ppo_epochs": 0}) +def test_pooled_explained_var_pre_update(): + """Pre-update EV: perfect prediction -> 1, mean prediction -> 0, and both + masks are respected (values at masked positions must not leak in).""" + from nemo_rl.algorithms.ppo import _pooled_explained_var + + token_mask = torch.tensor([[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]]) + sample_mask = torch.tensor([1.0, 1.0]) + returns = torch.tensor([[1.0, 1.0, 1.0, 9.0], [0.0, 0.0, 9.0, 9.0]]) + + # Perfect prediction on valid tokens; garbage on masked tokens is ignored. + values = returns.clone() + values[0, 3] = 7.0 + values[1, 2] = 7.0 + assert _pooled_explained_var(values, returns, token_mask, sample_mask) == ( + pytest.approx(1.0) + ) + + # Constant prediction at the valid-token mean explains nothing. + mean_pred = torch.full_like(returns, 0.6) # mean of [1,1,1,0,0] + assert _pooled_explained_var(mean_pred, returns, token_mask, sample_mask) == ( + pytest.approx(0.0, abs=1e-6) + ) + + # sample_mask knocks out the second sequence; the survivor has constant + # returns (zero variance), which reports 0.0 rather than dividing by ~0. + solo = torch.tensor([1.0, 0.0]) + assert _pooled_explained_var(returns.clone(), returns, token_mask, solo) == 0.0 + + # Matches the sufficient-statistics form used by the loss-side metrics. + mask = (token_mask * sample_mask.unsqueeze(-1)).bool() + noisy = returns + 0.25 * torch.tensor( + [[1.0, -1.0, 0.5, 0.0], [-0.5, 1.0, 0.0, 0.0]] + ) + r, v = returns[mask], noisy[mask] + r_mean, v_mean = r.mean(), v.mean() + r_sq, res_sq = (r**2).mean(), ((r - v) ** 2).mean() + ev_suff = 1.0 - (res_sq - (r_mean - v_mean) ** 2) / (r_sq - r_mean**2) + assert _pooled_explained_var(noisy, returns, token_mask, sample_mask) == ( + pytest.approx(ev_suff.item(), abs=1e-6) + ) + + # NeMo-Gym IS supported for async PPO: the AsyncTrajectoryCollector runs the gym # rollout internally (checks _should_use_nemo_gym, passes master_config.reward_penalties), # so async_ppo_train no longer guards it out. Validate-side dispatch is covered by diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index d5c5df045e9..bc89b731cb7 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -14,9 +14,10 @@ ppo: # critic-only passes, fitting the critic harder without over-training the actor. critic_ppo_epochs: null # Re-score the batch with a forward-only critic pass after the final update, - # adding critic/explained_var_post_update + critic/loss_post_update (the - # unsuffixed critic/* are measured BEFORE the last update). Costs one extra - # critic forward per step. + # adding critic/explained_var_post_update + critic/loss_post_update. + # critic/explained_var is always the PRE-update EV (from the rollout-time + # values GAE consumed); critic/loss is from the last training pass. Costs one + # extra critic forward per step. log_post_update_critic_metrics: false policy_training_start_step: 0 # number of PPO steps of critic-only warmup before policy training begins val_period: 20 From afc21546b35a569993662dc62d0ca4b37e15b6c8 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 7 Aug 2026 17:43:43 +0000 Subject: [PATCH 51/69] Log pre-whitening advantage statistics from GAE normalize_advantages rescales advantages to unit std every step, so the observable advantage spread is 1.0 by construction and carries no signal. The pre-whitening spread does: at lambda=1 the advantage is R - V(s_t), so its std is the critic's residual scale and should shrink as the critic improves. GeneralizedAdvantageEstimator now records mean/std/abs_mean/max_abs (plus whiten_gain = 1/std) over valid tokens immediately before the whitening branch, exposed via last_metrics following the OPDAdvantageEstimator convention. Both ppo_train and async_ppo_train merge it into their metrics. Diagnostic only: no change to advantages, returns, or the loss. Co-Authored-By: Claude Opus 5 --- nemo_rl/algorithms/advantage_estimator.py | 39 +++++++++++++++++++++++ nemo_rl/algorithms/ppo.py | 13 ++++++++ 2 files changed, 52 insertions(+) diff --git a/nemo_rl/algorithms/advantage_estimator.py b/nemo_rl/algorithms/advantage_estimator.py index 5a73902f329..0a9999fa669 100644 --- a/nemo_rl/algorithms/advantage_estimator.py +++ b/nemo_rl/algorithms/advantage_estimator.py @@ -305,6 +305,41 @@ def __init__(self, estimator_config: dict, loss_config: ClippedPGLossConfig): self.kl_coef = loss_config.reference_policy_kl_penalty self.kl_type = loss_config.reference_policy_kl_type + # Populated by compute_advantage; surfaced by ppo.py into rollout metrics. + self.last_metrics: dict[str, float] = {} + + def _compute_raw_advantage_metrics( + self, + advantages: torch.Tensor, + mask: torch.Tensor, + ) -> dict[str, float]: + """Stats on the PRE-whitening advantages, over valid tokens. + + ``normalize_advantages`` rescales advantages to unit std every step, so + the post-whitening spread is 1.0 by construction and carries no + information. The pre-whitening spread does: with lambda=1 the advantage + is ``R - V(s_t)``, so ``adv_raw/std`` is the critic's residual scale and + should SHRINK as the critic improves. A flat ``adv_raw/std`` alongside a + flat ``critic/explained_var`` is independent confirmation that the critic + is not learning. ``whiten_gain`` is the amplification whitening applies + (1/std); a large and growing value means the loss is being scaled up to + compensate for a shrinking signal. + """ + m = mask.bool() + if int(m.sum()) < 2: + return {} + a = advantages[m].float() + std = a.std(unbiased=False) + metrics = { + "adv_raw/mean": a.mean().item(), + "adv_raw/std": std.item(), + "adv_raw/abs_mean": a.abs().mean().item(), + "adv_raw/max_abs": a.abs().max().item(), + } + if self.normalize_advantages: + metrics["adv_raw/whiten_gain"] = (1.0 / (std + 1e-8)).item() + return metrics + def _reward_whiten( self, rewards: torch.Tensor, @@ -452,6 +487,10 @@ def compute_advantage( mask, ) + # Capture the pre-whitening advantage scale (critic-quality diagnostic) + # BEFORE normalize_advantages pins the std to 1.0. + self.last_metrics = self._compute_raw_advantage_metrics(advantages, mask) + # Whiten advantages (optional) and zero out masked positions (always) if self.normalize_advantages: advantages = self._reward_whiten(advantages, mask) diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 2ade166867b..8c4117004ab 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -2454,6 +2454,12 @@ def ppo_train( ) memory_tracker.snapshot_start_of_stage("Metrics", dir()) + # Pre-whitening advantage scale (see + # GeneralizedAdvantageEstimator._compute_raw_advantage_metrics). + # normalize_advantages pins the post-whitening std to 1.0, so this + # is the only place the critic's residual scale is observable. + if getattr(adv_estimator, "last_metrics", None): + metrics.update(adv_estimator.last_metrics) if train_results is not None: metrics = { **metrics, @@ -3699,6 +3705,13 @@ def _sample_max_age(s: int) -> int: flat_advantages, flat_token_mask.bool() ) + # Pre-whitening advantage scale (see + # GeneralizedAdvantageEstimator._compute_raw_advantage_metrics). + # normalize_advantages pins the post-whitening std to 1.0, so this + # is the only place the critic's residual scale is observable. + if getattr(adv_estimator, "last_metrics", None): + metrics.update(adv_estimator.last_metrics) + metrics.update( { "reward": rewards.numpy(), From b5eb09e570a9fa78551fbd8bcca90e85d54482ca Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 7 Aug 2026 17:44:37 +0000 Subject: [PATCH 52/69] Record pre-whitening advantage scale in the PPO rollout dump The dump stores train_data["advantages"], which is post-whitening, so offline analysis cannot recover the critic's residual scale. Reconstructing it from values/reward instead requires knowing lambda_policy, which the payload does not record. Whitening is affine, so the two scalars suffice: adv_raw = advantages * adv_raw_std + adv_raw_mean Stored as scalars rather than a second packed tensor -- same information, ~13MB/dump cheaper. Threaded from the estimator's last_metrics at both dump sites; absent for estimators that do not provide it. Co-Authored-By: Claude Opus 5 --- nemo_rl/algorithms/ppo.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 8c4117004ab..e3e08fa20fb 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -1657,6 +1657,7 @@ def _build_ppo_rollout_dump_payload( prompt_lengths: torch.Tensor, content: list[str], repeated_batch: BatchedDataDict[DatumSpec], + adv_raw_metrics: dict[str, float] | None = None, ) -> dict[str, Any]: """Build the packed per-token rollout dump payload for torch.save. @@ -1742,6 +1743,18 @@ def _pack_float(field: str) -> torch.Tensor: if "reference_policy_logprobs" in train_data: payload["reference_policy_logprobs"] = _pack_float("reference_policy_logprobs") + # ``advantages`` above are POST-whitening (normalize_advantages pins their std + # to 1.0). Whitening is affine, so the pre-whitening advantages -- the critic's + # actual residual scale, R - V(s_t) at lambda=1 -- are recoverable exactly: + # adv_raw = advantages * adv_raw_std + adv_raw_mean + # Stored as two scalars rather than a second packed tensor (same information, + # ~13MB/dump cheaper). + if adv_raw_metrics: + for key in ("mean", "std"): + value = adv_raw_metrics.get(f"adv_raw/{key}") + if value is not None: + payload[f"adv_raw_{key}"] = float(value) + if num_generations_per_prompt > 0: payload["prompt_group_index"] = ( sample_indices // num_generations_per_prompt @@ -2285,6 +2298,9 @@ def ppo_train( prompt_lengths=repeated_batch["length"], content=flat_messages["content"], repeated_batch=repeated_batch, + adv_raw_metrics=getattr( + adv_estimator, "last_metrics", None + ), ) torch.save(rollout_dump_payload, rollout_dump_path) print(f" 📝 Dumped rollout data to {rollout_dump_path}") @@ -3966,6 +3982,7 @@ def _sample_max_age(s: int) -> int: prompt_lengths=repeated_batch["length"], content=flat_messages_content, repeated_batch=repeated_batch, + adv_raw_metrics=getattr(adv_estimator, "last_metrics", None), ) torch.save(rollout_dump_payload, rollout_dump_path) print(f" 📝 Dumped rollout data to {rollout_dump_path}") From 2f5ff50814574249ccee89db5db7a28c63d391b1 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 7 Aug 2026 17:04:27 -0700 Subject: [PATCH 53/69] feat(ppo): turn-level critic and turn-level GAE for agentic RL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an opt-in turn MDP for long agentic rollouts, selected with ppo.adv_estimator.name=turn_gae. One assistant message is one action: d_k = r_k + g*V(s_{k+1}) - V(s_k), A_k = d_k + g*lam*A_{k+1}, G_k = A_k + V(s_k) The motivation is not that the token-level critic is noisy within a turn (measured: only 3.8% of V's variance is intra-turn, and at lambda=1 token- and turn-level baselines are near-identical). It is that lambda is structurally unusable at the token level: its horizon is 1/(1-lambda) TOKENS, so on a ~45k-token SWE rollout length_adaptive_alpha is forced to drive lambda to 1-1.5e-5 and GAE degenerates to A_t = R - V(s_t) — a pure baseline with no temporal credit assignment at all. Over ~92 turns, lambda=0.97 is a 33-turn horizon: a usable knob. At lambda<1 the advantage is built from TD increments, in which any constant per-trajectory offset cancels — and the critic's measured weakness on this workload IS the constant part (it barely reads task difficulty from the prompt), while what survives, d_K = R - V(s_K), is its strongest signal. Where V(s_k) is read: the value workers right-shift the value head, so values[t] = V(state before token t). V(s_k) is therefore the FIRST token of assistant message k, which is already inside the response mask. That is what keeps the change small — no model, worker, sequence-packing or context-parallel changes are needed. Critic supervision is anchor-only: the value model trains on its own batch whose token_mask is the anchor mask (the same mechanism the privileged critic already uses). Since process_global_batch derives global_valid_toks from token_mask, MseValueLossFn becomes an equal-weighted mean over turns for free — today the longest decile of turns owns ~43% of the critic loss. This costs no information: at lambda_value=1 every target inside a trajectory is the same number, R. gamma / lambda_value / lambda_policy are required config keys with no silent defaults, so a sweep cannot appear to have run when it did not. Verification: 25 CPU unit tests (span construction incl. consecutive and empty assistant messages, ragged batches, all fail-loud paths, turn GAE vs a hand-computed reference, the gamma=lambda=1 => G_k == R identity, gather/scatter round trips, per-turn KL aggregation, and that the estimator's returns are covered by the critic batch's mask); plus an offline replay of the production estimator over stored SWE value dumps (8 held-out groups, 13,330 turns, 3.24M response tokens) checking anchor placement, the MC-return identity, within-turn constancy of advantages and the returns/mask bijection on every group. The default path is unchanged: with name=gae or raw_reward this is inert. Co-Authored-By: Claude Opus 5 --- examples/configs/ppo_nano_v3_5_swe_cmh.yaml | 30 ++ nemo_rl/algorithms/advantage_estimator.py | 155 +++++++ nemo_rl/algorithms/critic_pretrain.py | 89 +++- nemo_rl/algorithms/ppo.py | 163 ++++++- nemo_rl/algorithms/turn_level.py | 435 ++++++++++++++++++ tests/unit/algorithms/test_turn_level.py | 463 ++++++++++++++++++++ 6 files changed, 1306 insertions(+), 29 deletions(-) create mode 100644 nemo_rl/algorithms/turn_level.py create mode 100644 tests/unit/algorithms/test_turn_level.py diff --git a/examples/configs/ppo_nano_v3_5_swe_cmh.yaml b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml index b3a4459718f..43d7ee9c706 100644 --- a/examples/configs/ppo_nano_v3_5_swe_cmh.yaml +++ b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml @@ -93,6 +93,20 @@ ppo: # GAE with VAPO decoupled lambdas (defaults match the validated 7B math PPO # run: lambda_value=1.0, length-adaptive lambda_policy with alpha=0.05). + # + # name: "gae" -> token-level GAE (the historical path, unchanged). + # name: "turn_gae" -> TURN-level GAE: one assistant message is one action, + # V(s_k) is read at that turn's first token (the value head is right-shifted, + # so that position has seen the whole preceding observation and none of the + # action), the advantage is constant across the turn's tokens, and the critic + # is supervised at one anchor per turn. Only the turn_gae_* keys apply. + # + # Why it exists: at the token level lambda has effective horizon 1/(1-lambda) + # TOKENS, so on a ~45k-token SWE rollout anything below 1 severs the terminal + # reward — which is why length_adaptive_alpha pushes lambda to 1 - 1.5e-5 and + # GAE degenerates to the pure baseline A_t = R - V(s_t) with no temporal + # credit assignment at all. Over ~92 turns, lambda=0.97 is a 33-turn horizon. + # See research/ppo/turn_level_critic_plan.md. adv_estimator: name: "gae" gae_lambda: 1.0 @@ -103,6 +117,22 @@ ppo: # Length-adaptive lambda_policy = 1 - 1/(alpha*l). 0 = disabled. length_adaptive_alpha: null + # --- turn-level GAE (used only when name: "turn_gae") --- + # All three lambdas/gamma are required in that mode and are never silently + # defaulted, so a sweep cannot appear to have run when it did not. + turn_gae_gamma: 1.0 + # lambda_value=1.0 keeps the critic's target Monte-Carlo (G_k == R): the same + # target it regresses on today, just at ~92 anchors instead of ~45k tokens. + # Lower it only deliberately — bootstrapped targets feed the critic's own + # (still weak) estimates back into its own supervision. + turn_gae_lambda_value: 1.0 + # lambda_policy is THE knob. 1.0 reproduces today's A_k = R - V(s_k); + # 0.99 ~ 100-turn horizon, 0.97 ~ 33 turns, 0.95 ~ 20 turns. + turn_gae_lambda_policy: 1.0 + # The critic is supervised at ONE anchor per turn, every turn weighted + # equally (under token weighting the longest decile of turns owns ~43% of + # the critic loss). + reward_shaping: enabled: false overlong_buffer_length: 128 diff --git a/nemo_rl/algorithms/advantage_estimator.py b/nemo_rl/algorithms/advantage_estimator.py index 5a73902f329..f07571cba7a 100644 --- a/nemo_rl/algorithms/advantage_estimator.py +++ b/nemo_rl/algorithms/advantage_estimator.py @@ -20,6 +20,7 @@ - ReinforcePlusPlusAdvantageEstimator: Reinforce++ with optional baseline subtraction (minus_baseline) and KL penalty in reward - RawRewardAdvantageEstimator: Raw reward as advantage with optional batch normalization (no baseline, no value model) - GeneralizedAdvantageEstimator: Generalized Advantage Estimation (GAE) with temporal bootstrapping +- TurnLevelGeneralizedAdvantageEstimator: GAE over agent TURNS instead of tokens - OPDAdvantageEstimator: Multi-Teacher On-Policy Distillation (MOPD) token-level distillation advantages Reference papers: - ProRLv2: https://developer.nvidia.com/blog/scaling-llm-reinforcement-learning-with-prolonged-training-using-prorl-v2/ @@ -511,6 +512,160 @@ def _compute_gae( return advantages, returns +class TurnLevelGeneralizedAdvantageEstimator: + """GAE over agent turns instead of tokens (see nemo_rl.algorithms.turn_level). + + The turn MDP treats one assistant message as one action: + + δ_k = r_k + γ V(s_{k+1}) - V(s_k), A_k = δ_k + γλ A_{k+1}, G_k = A_k + V(s_k) + + with ``V(s_k)`` read at the FIRST token of assistant message k (the value + head is right-shifted, so that position sees the whole preceding observation + and none of the action) and ``V(s_{K+1}) = 0``. + + Why not just run the token-level estimator: at the token level λ has + effective horizon 1/(1-λ) TOKENS, so on a 45k-token rollout any λ that is not + ≈1 severs the terminal reward entirely — which is exactly why the production + config lands at λ = 1 - 1.5e-5 and GAE collapses to the pure baseline + ``A_t = R - V(s_t)`` with no temporal credit assignment at all. Over ~92 + turns, λ=0.97 is a 33-turn horizon: a usable knob. + + A second, structural benefit: at λ<1 the advantage is built from TD + increments, in which any constant per-trajectory offset cancels identically. + The critic's measured weakness on this workload is precisely the constant + part (it cannot read task difficulty from the prompt); the part it is good at + is the terminal ``δ_K = R - V(s_K)``, which is what survives. + + Outputs: + advantages: ``[B, S]``, constant across the tokens of one assistant + message (the standard treatment of a multi-token action). + returns: ``[B, S]``, ``G_k`` placed at the turn ANCHOR only, zero + elsewhere — paired with ``token_mask = anchor_mask`` in the critic's + batch so the value loss is an equal-weighted mean over decision + points instead of a token-count-weighted one. + + Requires ``turn_spans`` (a :class:`~nemo_rl.algorithms.turn_level.TurnSpans`) + in ``compute_advantage``; the caller builds it once per step from the batch's + message logs. + """ + + def __init__(self, estimator_config: dict, loss_config: ClippedPGLossConfig): + # No silent defaults: a λ that quietly falls back to 1.0 would make a + # sweep look like it ran when it did not. + for key in ("turn_gae_gamma", "turn_gae_lambda_value", "turn_gae_lambda_policy"): + if estimator_config.get(key) is None: + raise ValueError( + f"adv_estimator.{key} must be set explicitly when " + "adv_estimator.name='turn_gae' (no default is assumed). " + "See research/ppo/turn_level_critic_plan.md." + ) + self.gamma = float(estimator_config["turn_gae_gamma"]) + self.lambda_value = float(estimator_config["turn_gae_lambda_value"]) + self.lambda_policy = float(estimator_config["turn_gae_lambda_policy"]) + self.normalize_advantages = estimator_config["normalize_advantages"] + + self.use_kl_in_reward = loss_config.use_kl_in_reward + self.kl_coef = loss_config.reference_policy_kl_penalty + self.kl_type = loss_config.reference_policy_kl_type + + self.last_metrics: dict[str, float] = {} + + def compute_advantage( + self, + prompt_ids, + rewards, + mask, + values, + turn_spans=None, + reference_logprobs=None, + logprobs=None, + sample_mask=None, + **kwargs, + ): + """Compute turn-level GAE advantages and critic targets. + + Args: + prompt_ids: unused (kept for interface parity). + rewards: ``[B]`` terminal reward per sample. + mask: ``[B, S]`` response token mask. + values: ``[B, S]`` per-token values from a fresh critic forward. + turn_spans: :class:`TurnSpans` for this batch (required). + reference_logprobs / logprobs: ``[B, S]``, only used when + ``use_kl_in_reward`` is on; the per-token penalty is summed into + its turn's reward. + sample_mask: ``[B]``, used for metrics only. + + Returns: + ``(advantages, returns)``, both ``[B, S]``. + """ + from nemo_rl.algorithms.turn_level import ( + build_turn_rewards, + gather_turn_values, + scatter_turns_to_anchors, + scatter_turns_to_tokens, + turn_gae, + turn_level_metrics, + ) + + if turn_spans is None: + raise ValueError( + "TurnLevelGeneralizedAdvantageEstimator requires turn_spans; " + "build it with nemo_rl.algorithms.turn_level.build_turn_spans() " + "from the batch's message logs." + ) + + seq_len = mask.shape[1] + turn_values = gather_turn_values(values, turn_spans) + + token_penalty = None + if ( + self.use_kl_in_reward + and self.kl_coef > 0 + and logprobs is not None + and reference_logprobs is not None + ): + kl = calculate_kl(logprobs, reference_logprobs, self.kl_type) + token_penalty = -self.kl_coef * kl * mask + turn_rewards = build_turn_rewards(rewards, turn_spans, token_penalty) + + # Decoupled λ (VAPO-style): critic targets and policy advantages may use + # different horizons. Skip the second pass when they agree. + _, turn_returns = turn_gae( + turn_values, + turn_rewards, + turn_spans.turn_valid, + self.gamma, + self.lambda_value, + ) + if self.lambda_policy == self.lambda_value: + turn_advantages = turn_returns - turn_values + else: + turn_advantages, _ = turn_gae( + turn_values, + turn_rewards, + turn_spans.turn_valid, + self.gamma, + self.lambda_policy, + ) + + advantages = scatter_turns_to_tokens(turn_advantages, turn_spans, seq_len) + # Anchor layout, matched by token_mask=anchor_mask in the critic's batch + # (build_turn_value_batch). The two MUST agree: returns placed anywhere + # the critic's mask does not cover are silently discarded. + returns = scatter_turns_to_anchors(turn_returns, turn_spans, seq_len) + + if self.normalize_advantages: + adv_mean = masked_mean(advantages, mask) + adv_var = masked_var(advantages, mask, adv_mean) + advantages = (advantages - adv_mean) * torch.rsqrt(adv_var + 1e-8) + advantages = torch.masked_fill(advantages, ~(mask.bool()), 0) + + self.last_metrics = turn_level_metrics( + turn_values, turn_advantages, turn_spans, sample_mask + ) + return advantages, returns + + class OPDAdvantageEstimator: """Multi-Teacher On-Policy Distillation (MOPD) advantage estimator (arXiv:2601.02780). diff --git a/nemo_rl/algorithms/critic_pretrain.py b/nemo_rl/algorithms/critic_pretrain.py index e535d57dc44..016c9fe297e 100644 --- a/nemo_rl/algorithms/critic_pretrain.py +++ b/nemo_rl/algorithms/critic_pretrain.py @@ -20,6 +20,13 @@ so with the production SWE config (``gae_lambda_value=1``, ``gae_gamma=1``, KL=0) one offline epoch is the same optimization as the online critic warmup. +Turn-level mode (``ppo.adv_estimator.name=turn_gae``) changes what a "position" +means here: the critic is supervised at ONE anchor per assistant turn (the turn's +first token, where the right-shifted value head reads ``V(s_k)``) with turn-level +GAE returns, instead of at all ~45k response tokens. Stage C must run with the +same setting — a token-level critic and a turn-level one are not interchangeable +warm starts. See research/ppo/turn_level_critic_plan.md. + Extras the online warmup cannot provide: * a held-out shard split (``dataset_idx % heldout_mod == 0``) with explained-variance / calibration / terminal-AUC eval on unseen rollouts; @@ -315,17 +322,22 @@ def _forward_values_and_returns( repeated_batch: BatchedDataDict, tokenizer: Any, master_config: Any, -) -> Optional[BatchedDataDict]: +) -> tuple[Optional[BatchedDataDict], Optional[Any]]: """Populate train_data['values'/'returns'] in place. - Returns the privileged critic batch (with its own values/returns/token_mask) - when the privileged critic is enabled, else None. + Returns ``(critic_batch, turn_spans)``: the batch the critic should actually + train on when that differs from ``train_data`` (the privileged critic's + answer-augmented batch, or the turn-level anchor batch) else None, and the + turn structure (None on the token-level path) so callers can score metrics + at the positions the critic is actually supervised at. Mirrors async_ppo_train steps 3 (value inference, incl. the privileged answer-conditioned remap) and 6 (GAE returns; logprobs=None is valid for the critic path since KL-in-reward is the only consumer of logprobs). """ from nemo_rl.algorithms.grpo import extract_initial_prompt_messages + from nemo_rl.algorithms.ppo import build_turn_spans_for_batch + from nemo_rl.algorithms.turn_level import build_turn_value_batch from nemo_rl.algorithms.privileged_critic import ( build_privileged_value_inputs, remap_by_response_mask, @@ -336,6 +348,11 @@ def _forward_values_and_returns( privileged_critic_cfg = None critic_batch = None + # Turn structure (None on the token-level path). Stage B must build this the + # same way stage C does, or the pretrained critic is supervised at positions + # PPO never reads. + turn_spans = build_turn_spans_for_batch(master_config, repeated_batch, train_data) + value_model.prepare_for_inference() if privileged_critic_cfg is not None: critic_batch = build_privileged_value_inputs( @@ -367,24 +384,32 @@ def _forward_values_and_returns( initial_prompt_message_logs, pad_value_dict={"token_ids": tokenizer.pad_token_id}, ) - advantages, returns = adv_estimator.compute_advantage( + adv_kwargs = dict( prompt_ids=prompt_batched_flat["token_ids"], rewards=train_data["rewards"], mask=train_data["token_mask"], values=train_data["values"], reference_logprobs=None, logprobs=None, + sample_mask=train_data["sample_mask"], ) + if turn_spans is not None: + adv_kwargs["turn_spans"] = turn_spans + advantages, returns = adv_estimator.compute_advantage(**adv_kwargs) del advantages # critic pretraining has no actor; only returns are used train_data["returns"] = returns - if critic_batch is not None: + if turn_spans is not None: + # One supervised position per turn, equally weighted (swapping token_mask + # for the anchor mask is what makes MseValueLossFn a per-turn mean). + critic_batch = build_turn_value_batch(train_data, turn_spans) + elif critic_batch is not None: critic_batch["returns"] = remap_by_response_mask( returns, train_data["token_mask"], critic_batch["token_mask"], ) critic_batch["sample_mask"] = train_data["sample_mask"] - return critic_batch + return critic_batch, turn_spans def _heldout_metrics( @@ -394,19 +419,29 @@ def _heldout_metrics( tokenizer: Any, master_config: Any, ) -> dict[str, float]: - """Critic quality on held-out rollouts: EV, positional EV/ECE, terminal AUC.""" + """Critic quality on held-out rollouts: EV, positional EV/ECE, terminal AUC. + + On the turn-level path every metric is scored at the positions the critic is + actually supervised at (turn anchors) — scoring an anchor-layout return + tensor over the full response mask would average each real target against + ~270 structural zeros — and the per-turn metrics from the estimator are + merged in. + """ from nemo_rl.algorithms.ppo import _positional_value_metrics groups = [load_group(p) for p in heldout_files] train_data, repeated_batch = build_value_train_data( groups, tokenizer, master_config ) - _forward_values_and_returns( + _, turn_spans = _forward_values_and_returns( value_model, adv_estimator, train_data, repeated_batch, tokenizer, master_config, ) values, returns = train_data["values"], train_data["returns"] - mask = train_data["token_mask"].bool() + scored_mask = ( + turn_spans.anchor_mask if turn_spans is not None else train_data["token_mask"] + ) + mask = scored_mask.bool() metrics: dict[str, float] = {} if int(mask.sum()) >= 2: v, r = values[mask].float(), returns[mask].float() @@ -417,12 +452,13 @@ def _heldout_metrics( else 0.0 ) metrics["critic/mse"] = ((r - v) ** 2).mean().item() - metrics.update( - _positional_value_metrics(values, returns, train_data["token_mask"]) - ) + metrics.update(_positional_value_metrics(values, returns, scored_mask)) + # Scored on `scored_mask`: in turn mode the last RESPONSE token carries an + # untrained value, while the last anchor is the supervised V(s_K). metrics["critic/terminal_auc"] = terminal_value_reward_auc( - values, train_data["rewards"], train_data["token_mask"] + values, train_data["rewards"], scored_mask ) + metrics.update(getattr(adv_estimator, "last_metrics", {}) or {}) metrics["reward"] = train_data["rewards"].float().mean().item() metrics["num_heldout_samples"] = float(train_data["input_ids"].shape[0]) return metrics @@ -443,6 +479,13 @@ def _dump_heldout_values( Values/returns are packed over response tokens; per-message spans (with decoded text for the first ``dump_text_groups`` groups) let offline analysis align value movements to agent/tool behavior in the trajectory. + + In turn-level mode ``returns`` is an ANCHOR-layout tensor: it is the turn + return at each turn's first token and structurally 0 at the other ~270 + tokens of the turn. Averaging it over all stored tokens is meaningless, so + the payload carries ``credit_level`` and a per-token ``is_anchor`` flag + (format_version 3) and consumers must filter on it. Values are per-token in + both modes. """ dump_dir.mkdir(parents=True, exist_ok=True) for gi, path in enumerate(heldout_files): @@ -450,7 +493,7 @@ def _dump_heldout_values( train_data, repeated_batch = build_value_train_data( [g], tokenizer, master_config ) - _forward_values_and_returns( + _, turn_spans = _forward_values_and_returns( value_model, adv_estimator, train_data, repeated_batch, tokenizer, master_config, ) @@ -481,10 +524,15 @@ def _dump_heldout_values( tokenizer.decode([int(t)]) for t in m["token_ids"] ] samples_msgs.append(spans) + anchor_mask = turn_spans.anchor_mask if turn_spans is not None else None payload = { - "format_version": 2, + "format_version": 3, "dataset_idx": g["dataset_idx"], "source_file": str(path), + # "token": returns are per-token. "turn": returns are the turn + # return at anchors and structurally 0 elsewhere — filter on + # is_anchor before averaging or computing EV. + "credit_level": "turn" if anchor_mask is not None else "token", "rewards": train_data["rewards"].float().cpu(), "sample_mask": train_data["sample_mask"].float().cpu(), "token_sample_index": coords[:, 0].to(torch.int32), @@ -495,6 +543,8 @@ def _dump_heldout_values( "has_text": with_text, "render_samples": render_samples, } + if anchor_mask is not None: + payload["is_anchor"] = anchor_mask[mask].bool().cpu() out = dump_dir / f"valuedump_{g['dataset_idx']:08d}.pt" torch.save(payload, out) if (gi + 1) % 10 == 0 or gi + 1 == len(heldout_files): @@ -792,7 +842,7 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: ) print("▶ Computing values...") - critic_batch = _forward_values_and_returns( + critic_batch, turn_spans = _forward_values_and_returns( value_model, adv_estimator, train_data, repeated_batch, tokenizer, master_config, ) @@ -814,9 +864,14 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: metrics[k] = np.sum(v).item() metrics.update( _positional_value_metrics( - train_data["values"], train_data["returns"], train_data["token_mask"] + train_data["values"], + train_data["returns"], + turn_spans.anchor_mask + if turn_spans is not None + else train_data["token_mask"], ) ) + metrics.update(getattr(adv_estimator, "last_metrics", {}) or {}) metrics["reward"] = train_data["rewards"].float().mean().item() metrics["num_samples"] = float(train_data["input_ids"].shape[0]) metrics["total_step_time"] = time.perf_counter() - step_start diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index c52adf885f2..295806135a0 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -31,6 +31,7 @@ from nemo_rl.algorithms.advantage_estimator import ( GeneralizedAdvantageEstimator, RawRewardAdvantageEstimator, + TurnLevelGeneralizedAdvantageEstimator, ) from nemo_rl.algorithms.grpo import ( RewardPenaltyConfig, @@ -182,9 +183,9 @@ class AsyncPPOConfig(TypedDict): class AdvEstimatorConfig(TypedDict): - """Configuration for PPO advantage estimator (GAE or raw_reward).""" + """Configuration for PPO advantage estimator (GAE, turn-level GAE, or raw_reward).""" - name: str # "gae" or "raw_reward" + name: str # "gae", "turn_gae", or "raw_reward" # GAE-specific (only used when name="gae") gae_lambda: NotRequired[float] gae_gamma: NotRequired[float] @@ -194,6 +195,13 @@ class AdvEstimatorConfig(TypedDict): gae_lambda_policy: NotRequired[Optional[float]] # Length-adaptive λ_policy = 1 - 1/(α·l). 0 = disabled. length_adaptive_alpha: NotRequired[float] + # Turn-level GAE (only used when name="turn_gae"). One assistant message is + # one action; λ is measured in TURNS, where a value below 1 is finally + # usable (at the token level a 45k-token rollout forces λ ≈ 1). All three + # are required when name="turn_gae" — no silent defaults. + turn_gae_gamma: NotRequired[Optional[float]] + turn_gae_lambda_value: NotRequired[Optional[float]] + turn_gae_lambda_policy: NotRequired[Optional[float]] class PPOConfig(TypedDict): @@ -1394,15 +1402,16 @@ def _create_advantage_estimator(master_config: MasterConfig): """Create and return an advantage estimator based on configuration. PPO's training loop consumes a `(advantages, returns)` pair from a - value-model-based estimator, so only `gae` and `raw_reward` are supported - here. Group-relative estimators like GRPO / Reinforce++ are not compatible - with PPO's loop and live in `grpo.py`. + value-model-based estimator, so only `gae`, `turn_gae` and `raw_reward` are + supported here. Group-relative estimators like GRPO / Reinforce++ are not + compatible with PPO's loop and live in `grpo.py`. Args: master_config: The master configuration dictionary. Returns: - A `GeneralizedAdvantageEstimator` or `RawRewardAdvantageEstimator` instance. + A `GeneralizedAdvantageEstimator`, `TurnLevelGeneralizedAdvantageEstimator` + or `RawRewardAdvantageEstimator` instance. Raises: ValueError: If the advantage estimator name is not recognized. @@ -1418,18 +1427,84 @@ def _create_advantage_estimator(master_config: MasterConfig): gae_lambda = adv_estimator_config["gae_lambda"] gae_gamma = adv_estimator_config["gae_gamma"] print(f" ✓ Using GAE advantage estimator (λ={gae_lambda}, γ={gae_gamma})") + elif adv_estimator_name == "turn_gae": + _raise_if_turn_gae_unsupported(master_config) + adv_estimator = TurnLevelGeneralizedAdvantageEstimator( + adv_estimator_config, loss_config + ) + print( + " ✓ Using TURN-level GAE advantage estimator " + f"(γ={adv_estimator.gamma}, λ_value={adv_estimator.lambda_value}, " + f"λ_policy={adv_estimator.lambda_policy}); the critic is supervised " + "at one anchor per assistant turn" + ) elif adv_estimator_name == "raw_reward": adv_estimator = RawRewardAdvantageEstimator(adv_estimator_config, loss_config) print(" ✓ Using raw reward advantage estimator (no value model, no baselines)") else: raise ValueError( f"Invalid adv_estimator name for PPO: {adv_estimator_name!r}. " - f"PPO only supports 'gae' or 'raw_reward'." + f"PPO only supports 'gae', 'turn_gae' or 'raw_reward'." ) return adv_estimator +def _raise_if_turn_gae_unsupported(master_config: MasterConfig) -> None: + """Reject turn-level GAE combined with features that assume a token-level critic.""" + # getattr: this factory is shared with grpo_sync, whose configs have no + # `value:` section at all. + value_config = getattr(master_config, "value", None) or {} + privileged_critic = value_config.get("privileged_critic") + if privileged_critic is not None and privileged_critic.get("enabled"): + raise NotImplementedError( + "adv_estimator.name='turn_gae' with value.privileged_critic.enabled " + "is not supported: the privileged critic scores an answer-augmented " + "sequence whose token layout differs from the policy's, so turn " + "anchors computed on the policy layout would land on the wrong " + "positions. Disable one of the two." + ) + + +def build_turn_spans_for_batch( + master_config: MasterConfig, + repeated_batch: BatchedDataDict[DatumSpec], + train_data: BatchedDataDict[Any], +) -> Optional[Any]: + """Locate the batch's agent turns, or None when running token-level GAE. + + Shared by both PPO loops and by offline critic pretraining so all three see + identical turn structure. Validates on the way out rather than letting a + misaligned anchor silently attribute a turn's credit to the wrong tokens. + """ + if master_config.ppo["adv_estimator"]["name"] != "turn_gae": + return None + + from nemo_rl.algorithms.turn_level import build_turn_spans, validate_turn_spans + + spans = build_turn_spans( + repeated_batch["message_log"], + seq_len=train_data["token_mask"].shape[1], + mask_dtype=train_data["token_mask"].dtype, + ) + validate_turn_spans(spans, train_data["token_mask"], train_data["sample_mask"]) + return spans + + +def _value_metric_mask( + train_data: BatchedDataDict[Any], turn_spans: Optional[Any] +) -> torch.Tensor: + """Positions the critic is actually supervised at, for value diagnostics. + + Token-level runs supervise every response token; turn-level anchor runs + supervise one position per turn, and scoring them over the full response + mask would average a real target against ~270 structural zeros. + """ + if turn_spans is None: + return train_data["token_mask"] + return turn_spans.anchor_mask + + def _compute_critic_metrics(value_results: dict[str, Any]) -> dict[str, Any]: """Build the ``critic/*`` metrics dict from a value model's train results. @@ -1568,6 +1643,7 @@ def _build_ppo_rollout_dump_payload( prompt_lengths: torch.Tensor, content: list[str], repeated_batch: BatchedDataDict[DatumSpec], + turn_spans: Optional[Any] = None, ) -> dict[str, Any]: """Build the packed per-token rollout dump payload for torch.save. @@ -1587,6 +1663,12 @@ def _build_ppo_rollout_dump_payload( content: Per-sample decoded conversation text. repeated_batch: Rollout batch; optional per-sample metadata (``task_name``, ``idx``, ``truncated``) is copied when present. + turn_spans: Turn structure when running turn-level GAE, else None. In + that mode ``returns`` is anchor-layout — the turn return at each + turn's first token, structurally 0 at the other ~270 tokens of the + turn — so the payload records ``credit_level`` and a per-token + ``is_anchor`` flag, and consumers MUST filter on it before averaging + returns or computing explained variance. Returns: Payload dict of packed per-token tensors, per-sample arrays, and @@ -1616,7 +1698,10 @@ def _pack_float(field: str) -> torch.Tensor: sample_indices = torch.arange(batch_size, dtype=torch.int32) payload: dict[str, Any] = { - "format_version": 1, + "format_version": 2, + # "token": returns are per-token. "turn": returns are the turn return at + # anchors and structurally 0 elsewhere — filter on is_anchor first. + "credit_level": "turn" if turn_spans is not None else "token", "description": ( "Packed response-token PPO rollout dump. Per-token tensors are " "flattened over token_mask; entry i belongs to sample " @@ -1650,6 +1735,10 @@ def _pack_float(field: str) -> torch.Tensor: } if "returns" in train_data: payload["returns"] = _pack_float("returns") + if turn_spans is not None: + payload["is_anchor"] = ( + turn_spans.anchor_mask.detach().bool().cpu()[token_mask] + ) if "reference_policy_logprobs" in train_data: payload["reference_policy_logprobs"] = _pack_float("reference_policy_logprobs") @@ -2012,6 +2101,13 @@ def ppo_train( metrics_logging_data["content"] = flat_messages["content"] + # Turn structure for turn-level credit assignment (None on + # the token-level path). Built here, before any GPU work, so + # a malformed batch fails before the expensive forwards. + turn_spans = build_turn_spans_for_batch( + master_config, repeated_batch, train_data + ) + memory_tracker.snapshot_start_of_stage("Value inference", dir()) print("▶ Computing values...", flush=True) with timer.time("value_inference"): @@ -2148,9 +2244,12 @@ def ppo_train( mask=train_data["token_mask"], reference_logprobs=train_data.get("reference_policy_logprobs"), logprobs=train_data["prev_logprobs"], + sample_mask=train_data["sample_mask"], ) if "values" in train_data: adv_kwargs["values"] = train_data["values"] + if turn_spans is not None: + adv_kwargs["turn_spans"] = turn_spans result = adv_estimator.compute_advantage(**adv_kwargs) if isinstance(result, tuple): advantages, returns = result @@ -2161,10 +2260,20 @@ def ppo_train( train_data["advantages"] = advantages if returns is not None: train_data["returns"] = returns + # Turn-level: the critic trains on one anchor per turn, + # so it needs its own batch (same sequences, anchor mask). + if turn_spans is not None: + from nemo_rl.algorithms.turn_level import ( + build_turn_value_batch, + ) + + critic_batch = build_turn_value_batch( + train_data, turn_spans + ) # Privileged critic trains on the answer-augmented sequence: # scatter the (original-layout) GAE returns onto the augmented # response positions. Same response tokens => exact mapping. - if critic_batch is not None: + elif critic_batch is not None: critic_batch["returns"] = remap_by_response_mask( returns, train_data["token_mask"], @@ -2187,6 +2296,7 @@ def ppo_train( prompt_lengths=repeated_batch["length"], content=flat_messages["content"], repeated_batch=repeated_batch, + turn_spans=turn_spans, ) torch.save(rollout_dump_payload, rollout_dump_path) print(f" 📝 Dumped rollout data to {rollout_dump_path}") @@ -2351,9 +2461,13 @@ def ppo_train( _positional_value_metrics( train_data["values"], train_data["returns"], - train_data["token_mask"], + _value_metric_mask(train_data, turn_spans), ) ) + # Turn-level counterparts (per-turn EV/bias, last-turn AUC, + # turn counts). The token-level versions above are dominated + # by the longest turns; these weight every decision equally. + metrics.update(getattr(adv_estimator, "last_metrics", {}) or {}) metrics.update( { "reward": rewards.numpy(), @@ -3220,6 +3334,13 @@ def _sample_max_age(s: int) -> int: train_data.update(extra_multimodal_data) train_data.to("cpu") + # Turn structure for turn-level credit assignment (None on + # the token-level path). Built here, before any GPU work, so + # a malformed batch fails before the expensive forwards. + turn_spans = build_turn_spans_for_batch( + master_config, repeated_batch, train_data + ) + # ---- 3. Value forward (critic on GPU, then offloaded) ---- # GPU state entering here: policy OFF, value OFF (see refit/step # end below). Load value only. @@ -3326,9 +3447,12 @@ def _sample_max_age(s: int) -> int: mask=train_data["token_mask"], reference_logprobs=train_data.get("reference_policy_logprobs"), logprobs=train_data["prev_logprobs"], + sample_mask=train_data["sample_mask"], ) if "values" in train_data: adv_kwargs["values"] = train_data["values"] + if turn_spans is not None: + adv_kwargs["turn_spans"] = turn_spans result = adv_estimator.compute_advantage(**adv_kwargs) if isinstance(result, tuple): advantages, returns = result @@ -3338,10 +3462,20 @@ def _sample_max_age(s: int) -> int: train_data["advantages"] = advantages if returns is not None: train_data["returns"] = returns + # Turn-level: the critic trains on one anchor per turn, + # so it needs its own batch (same sequences, anchor mask). + if turn_spans is not None: + from nemo_rl.algorithms.turn_level import ( + build_turn_value_batch, + ) + + critic_batch = build_turn_value_batch( + train_data, turn_spans + ) # Privileged critic trains on the answer-augmented sequence: # scatter the (policy-layout) GAE returns onto the augmented # response positions. Same response tokens => exact mapping. - if critic_batch is not None: + elif critic_batch is not None: critic_batch["returns"] = remap_by_response_mask( returns, train_data["token_mask"], @@ -3544,9 +3678,13 @@ def _sample_max_age(s: int) -> int: _positional_value_metrics( train_data["values"], train_data["returns"], - train_data["token_mask"], + _value_metric_mask(train_data, turn_spans), ) ) + # Turn-level counterparts (per-turn EV/bias, last-turn AUC, + # turn counts). The token-level versions above are dominated + # by the longest turns; these weight every decision equally. + metrics.update(getattr(adv_estimator, "last_metrics", {}) or {}) for k, v in metrics.items(): if k in {"probs_ratio_min", "probs_ratio_clamped_min"}: @@ -3731,6 +3869,7 @@ def _sample_max_age(s: int) -> int: prompt_lengths=repeated_batch["length"], content=flat_messages_content, repeated_batch=repeated_batch, + turn_spans=turn_spans, ) torch.save(rollout_dump_payload, rollout_dump_path) print(f" 📝 Dumped rollout data to {rollout_dump_path}") diff --git a/nemo_rl/algorithms/turn_level.py b/nemo_rl/algorithms/turn_level.py new file mode 100644 index 00000000000..35d6f66119a --- /dev/null +++ b/nemo_rl/algorithms/turn_level.py @@ -0,0 +1,435 @@ +# 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. +"""Turn-level MDP bookkeeping for agentic PPO (see research/ppo/turn_level_critic_plan.md). + +Long agentic rollouts (SWE: ~92 assistant turns, up to 200, over ~45k response +tokens) are a *turn* MDP wearing a token MDP's clothes: + + state s_k = the context before the k-th assistant message + action a_k = the whole k-th assistant message + reward r_k = 0 for k < K, r_K = R (environment-graded, terminal) + +Running GAE over tokens instead of turns is not merely wasteful, it removes the +λ knob entirely: at the token level λ has effective horizon 1/(1-λ) **tokens**, +so λ=0.95 reaches 20 tokens and the terminal reward cannot propagate. That is +why the production config ends up at λ = 1 - 1.5e-5 (``length_adaptive_alpha``), +where GAE degenerates to the pure baseline ``A_t = R - V(s_t)`` and nothing in +the objective assigns credit to a particular turn. Over ~92 turns, λ=0.97 is a +33-turn horizon — a usable bias/variance knob. + +Where V(s_k) lives +------------------ +The value workers right-shift the value head +(``megatron_value_worker._value_loss_prepare_fn``), so + + values[t] = V(state before token t) + +which means **V(s_k) is read at the FIRST token of assistant message k**: that +position has attended to the entire preceding observation and none of the +action. Anchors are therefore always inside the response mask, so nothing here +needs a gradient at a non-response position, and no model / worker / sequence +packing / context-parallel code changes. + +Nothing in this module is sharded: the driver builds these tensors, uses them to +scatter advantages and to assemble the critic's anchor batch, and throws them +away. Keeping them off the workers is what keeps the change small. +""" + +from dataclasses import dataclass +from typing import Any, Optional + +import torch + +# Roles that constitute an agent action. Everything else (user, tool, +# environment, system) is an observation and only ever contributes state. +_ACTION_ROLE = "assistant" + + +@dataclass +class TurnSpans: + """Turn structure of one flattened batch of message logs. + + Attributes: + anchor_mask: ``[B, S]``, 1 at the first token of each assistant message. + Always a subset of the response ``token_mask``. + turn_index: ``[B, S]`` int32, the turn ordinal ``k`` for every token of + assistant message ``k``; ``-1`` at observation/padding positions. + anchor_pos: ``[B, K_max]`` int64, token index of each turn's anchor. + Invalid slots hold 0 (a prompt position, never an anchor). + turn_valid: ``[B, K_max]`` bool, whether slot ``k`` is a real turn. + num_turns: ``[B]`` int64, assistant messages per sample. + turn_ntokens: ``[B, K_max]`` int64, tokens in each turn (0 if invalid). + """ + + anchor_mask: torch.Tensor + turn_index: torch.Tensor + anchor_pos: torch.Tensor + turn_valid: torch.Tensor + num_turns: torch.Tensor + turn_ntokens: torch.Tensor + + +def build_turn_spans( + message_log_batch: list[list[dict[str, Any]]], + seq_len: int, + mask_dtype: torch.dtype = torch.long, +) -> TurnSpans: + """Locate every assistant message in the flattened ``[B, seq_len]`` layout. + + Mirrors :func:`batched_message_log_to_flat_message`'s concatenation order + (messages laid out back to back, right padding), so the returned indices + address ``train_data["input_ids"]`` / ``values`` directly. + + Args: + message_log_batch: one message log per sample, each a list of messages + with ``role`` and ``token_ids``. + seq_len: padded sequence length of the flattened batch. + mask_dtype: dtype for ``anchor_mask`` (match the batch's ``token_mask``). + + Raises: + ValueError: if a sample's flattened length exceeds ``seq_len``, or an + assistant message starts at token 0 (there is always a prompt, and + index 0 carries no value after the right-shift). + """ + batch_size = len(message_log_batch) + per_sample_starts: list[torch.Tensor] = [] + per_sample_lens: list[torch.Tensor] = [] + turn_index = torch.full((batch_size, seq_len), -1, dtype=torch.int32) + + for i, message_log in enumerate(message_log_batch): + lengths = torch.tensor( + [len(m["token_ids"]) for m in message_log], dtype=torch.long + ) + total = int(lengths.sum()) + if total > seq_len: + raise ValueError( + f"Sample {i} flattens to {total} tokens but the batch is padded " + f"to {seq_len}. build_turn_spans must be called with the same " + "seq_len as the flattened batch." + ) + is_action = torch.tensor( + [m["role"] == _ACTION_ROLE for m in message_log], dtype=torch.bool + ) + starts_all = torch.cumsum(lengths, 0) - lengths + + # Non-empty assistant messages only: a zero-length message has no anchor + # and no tokens to carry an advantage. + keep = is_action & (lengths > 0) + starts = starts_all[keep] + lens = lengths[keep] + if starts.numel() > 0 and int(starts[0]) == 0: + raise ValueError( + f"Sample {i} starts with an assistant message at token 0. The " + "value head is right-shifted, so position 0 carries no value " + "and cannot anchor a turn; every sample must begin with a prompt." + ) + + # Turn ordinal per token, without a Python loop over messages: map each + # token to its message, then each message to its turn (-1 if not an action). + turn_of_msg = torch.cumsum(keep.long(), 0) - 1 + turn_of_msg[~keep] = -1 + if total > 0: + msg_of_tok = torch.repeat_interleave( + torch.arange(len(message_log), dtype=torch.long), lengths + ) + turn_index[i, :total] = turn_of_msg[msg_of_tok].to(torch.int32) + + per_sample_starts.append(starts) + per_sample_lens.append(lens) + + num_turns = torch.tensor([s.numel() for s in per_sample_starts], dtype=torch.long) + max_turns = int(num_turns.max()) if batch_size > 0 else 0 + max_turns = max(max_turns, 1) # keep a well-formed [B, K] even if no turns + + anchor_pos = torch.zeros((batch_size, max_turns), dtype=torch.long) + turn_valid = torch.zeros((batch_size, max_turns), dtype=torch.bool) + turn_ntokens = torch.zeros((batch_size, max_turns), dtype=torch.long) + for i, (starts, lens) in enumerate(zip(per_sample_starts, per_sample_lens)): + k = starts.numel() + if k == 0: + continue + anchor_pos[i, :k] = starts + turn_valid[i, :k] = True + turn_ntokens[i, :k] = lens + + anchor_mask = torch.zeros((batch_size, seq_len), dtype=mask_dtype) + anchor_mask.scatter_(1, anchor_pos, turn_valid.to(mask_dtype)) + # Slot 0 of an all-invalid row scattered a 0, but a valid row may also have + # written to index 0 only if it had an anchor there — ruled out above. + + return TurnSpans( + anchor_mask=anchor_mask, + turn_index=turn_index, + anchor_pos=anchor_pos, + turn_valid=turn_valid, + num_turns=num_turns, + turn_ntokens=turn_ntokens, + ) + + +def validate_turn_spans( + spans: TurnSpans, + token_mask: torch.Tensor, + sample_mask: Optional[torch.Tensor] = None, +) -> None: + """Fail loud on a turn structure that disagrees with the batch it describes. + + The invariant that matters is a BIJECTION between response tokens and turns: + every anchor is a response token, and every response token belongs to some + turn. Either direction failing means the turn structure and the flattened + batch were built from different things, and credit would be attributed to + the wrong tokens — silently. + + A sample with no assistant message is NOT an error: its ``token_mask`` row is + empty, so it contributes nothing to either loss. It is only reported, because + aborting a 512-sample step over one empty trajectory is a far worse failure + than the trajectory itself. + """ + anchors = spans.anchor_mask.bool() + resp = token_mask.bool() + + stray = int((anchors & ~resp).sum()) + if stray: + raise ValueError( + f"{stray} turn anchors fall outside the response mask. Anchors are " + "the first token of an assistant message and must be trainable " + "response positions; a mismatch means the turn structure and the " + "flattened batch disagree." + ) + + orphan = int((resp & (spans.turn_index.to(resp.device) < 0)).sum()) + if orphan: + raise ValueError( + f"{orphan} response tokens belong to no turn. Every trainable token " + "must sit inside an assistant message, or its policy gradient would " + "carry no advantage; the turn structure and the loss mask disagree." + ) + + if sample_mask is not None: + no_turns = (spans.num_turns == 0) & (sample_mask.detach().cpu() > 0) + n = int(no_turns.sum()) + if n: + bad = torch.nonzero(no_turns).flatten().tolist()[:5] + print( + f" ⚠️ {n} unmasked samples have no assistant message (e.g. " + f"indices {bad}); they contribute no tokens to either loss.", + flush=True, + ) + + +def gather_turn_values( + values: torch.Tensor, spans: TurnSpans +) -> torch.Tensor: + """``V(s_k)`` for every turn: ``[B, S] -> [B, K_max]`` (0 at invalid slots).""" + v = torch.gather(values, 1, spans.anchor_pos.to(values.device)) + return v * spans.turn_valid.to(device=values.device, dtype=values.dtype) + + +def build_turn_rewards( + rewards: torch.Tensor, + spans: TurnSpans, + token_level_penalty: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Per-turn reward ``[B, K_max]``: terminal ``R`` at the last turn. + + ``token_level_penalty`` (e.g. ``-kl_coef * KL`` per token, already signed) is + summed within each turn and added to that turn's reward, so a per-token + shaping signal survives the move to turn granularity. + """ + device = rewards.device + turn_valid = spans.turn_valid.to(device) + turn_rewards = torch.zeros( + turn_valid.shape, device=device, dtype=rewards.dtype + ) + + if token_level_penalty is not None: + # Sum the per-token penalty into its turn. index_add over a flattened + # [B*K] buffer keeps this a single vectorised op. + b, s = token_level_penalty.shape + k = turn_valid.shape[1] + ti = spans.turn_index.to(device).long() + valid = ti >= 0 + flat_idx = ( + torch.arange(b, device=device).unsqueeze(1) * k + ti.clamp(min=0) + )[valid] + acc = torch.zeros(b * k, device=device, dtype=rewards.dtype) + acc.index_add_(0, flat_idx, token_level_penalty[valid].to(rewards.dtype)) + turn_rewards = turn_rewards + acc.view(b, k) + + # Terminal reward on the last valid turn of each sample. + last = (spans.num_turns.to(device) - 1).clamp(min=0) + has_turns = spans.num_turns.to(device) > 0 + rows = torch.nonzero(has_turns).flatten() + if rows.numel(): + turn_rewards[rows, last[rows]] += rewards[rows] + + return turn_rewards * turn_valid.to(turn_rewards.dtype) + + +def turn_gae( + turn_values: torch.Tensor, + turn_rewards: torch.Tensor, + turn_valid: torch.Tensor, + gamma: float, + gae_lambda: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """GAE over turns. + + ``δ_k = r_k + γ V(s_{k+1}) - V(s_k)``, ``A_k = δ_k + γλ A_{k+1}``, with + ``V(s_{K+1}) = 0``. Truncated rollouts are still terminal in this MDP: the + environment grades the final state whether or not the agent ran out of + turns, so ``R`` is the realised episode return and bootstrapping 0 is right. + + Uses the same carry-forward masking as the token-level estimator (invalid + slots preserve the accumulators rather than zeroing them), so the trailing + padding of short samples is skipped instead of injecting phantom TD errors. + + Args: + turn_values: ``[B, K]`` V(s_k). + turn_rewards: ``[B, K]`` per-turn rewards. + turn_valid: ``[B, K]`` bool. + gamma: discount. + gae_lambda: scalar λ. Deliberately not per-sample: a length-adaptive λ is + the token-level pathology this estimator exists to remove. + + Returns: + ``(advantages, returns)``, each ``[B, K]``. ``returns = advantages + values``. + """ + dtype, device = turn_values.dtype, turn_values.device + num_turns = turn_values.shape[1] + lam = gae_lambda + + next_values = torch.zeros(turn_values.shape[0], device=device, dtype=dtype) + last_gae = torch.zeros_like(next_values) + out = torch.zeros_like(turn_values) + valid = turn_valid.to(device=device, dtype=dtype) + + for k in reversed(range(num_turns)): + delta = turn_rewards[:, k] + gamma * next_values - turn_values[:, k] + new_gae = delta + gamma * lam * last_gae + m = valid[:, k] + next_values = turn_values[:, k] * m + (1 - m) * next_values + last_gae = new_gae * m + (1 - m) * last_gae + out[:, k] = last_gae + + return out, out + turn_values + + +def scatter_turns_to_tokens( + turn_quantity: torch.Tensor, spans: TurnSpans, seq_len: int +) -> torch.Tensor: + """Broadcast a per-turn quantity ``[B, K]`` onto every token of its turn. + + This is what makes the policy update turn-level: all tokens of one assistant + message share one advantage, the standard treatment of a multi-token action. + """ + ti = spans.turn_index.to(turn_quantity.device) + valid = ti >= 0 + gathered = torch.gather(turn_quantity, 1, ti.clamp(min=0).long()[:, :seq_len]) + return gathered * valid[:, :seq_len].to(turn_quantity.dtype) + + +def scatter_turns_to_anchors( + turn_quantity: torch.Tensor, spans: TurnSpans, seq_len: int +) -> torch.Tensor: + """Place a per-turn quantity at its anchor token only; 0 everywhere else. + + Used for the critic's regression targets: paired with + ``token_mask = anchor_mask`` the value loss sees exactly one target per + decision point, equally weighted, instead of one per token. + """ + out = torch.zeros( + (turn_quantity.shape[0], seq_len), + device=turn_quantity.device, + dtype=turn_quantity.dtype, + ) + vals = turn_quantity * spans.turn_valid.to( + device=turn_quantity.device, dtype=turn_quantity.dtype + ) + out.scatter_(1, spans.anchor_pos.to(turn_quantity.device), vals) + return out + + +def build_turn_value_batch(train_data: Any, spans: TurnSpans) -> Any: + """The critic's anchor batch: same sequences, one supervised position per turn. + + Swapping only ``token_mask`` is enough to retarget the whole critic path: + ``process_global_batch`` derives ``global_valid_toks`` from ``token_mask``, + so :class:`MseValueLossFn` becomes an equal-weighted mean over turns with no + change to the loss, the value workers, sequence packing, or CP. This mirrors + how the privileged critic already trains on its own batch. + + ``returns`` must already be in anchor layout (see + :func:`scatter_turns_to_anchors`). + """ + from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + required = ("input_ids", "input_lengths", "sample_mask", "returns") + missing = [k for k in required if k not in train_data] + if missing: + raise ValueError( + f"build_turn_value_batch is missing {missing}; the turn returns must " + "be written to train_data['returns'] before the critic batch is built." + ) + keys = required + ("values",) # values: old values, for the PPO value clip + batch = BatchedDataDict({k: train_data[k] for k in keys if k in train_data}) + batch["token_mask"] = spans.anchor_mask.to(train_data["token_mask"].dtype) + # Carry multimodal side inputs through untouched (same rule the PPO loops + # use to assemble extra_multimodal_data). + batch.update(train_data.get_multimodal_dict(as_tensors=False)) + return batch + + +# =============================================================================== +# Metrics +# =============================================================================== +def turn_level_metrics( + turn_values: torch.Tensor, + turn_advantages: torch.Tensor, + spans: TurnSpans, + sample_mask: Optional[torch.Tensor] = None, +) -> dict[str, float]: + """Diagnostics that only exist at turn granularity. + + Deliberately does NOT re-derive per-position EV/bias or a terminal AUC: with + ``token_mask = anchor_mask`` the existing ``_positional_value_metrics`` and + ``terminal_value_reward_auc`` already bin by turn ordinal and read the last + turn's anchor, and produce bit-identical numbers. Only the quantities with no + token-level counterpart live here. + + ``turn_advantages`` is expected PRE-normalization (the estimator whitens the + token-level tensor afterwards), so the ``advantage/turn_*_prenorm`` keys show + the real scale of ``A_k`` — which collapses as λ_policy drops — rather than + the post-whitening ``advantages/std`` == 1 that PPO logs separately. + """ + valid = spans.turn_valid.to(turn_values.device) + if sample_mask is not None: + valid = valid & (sample_mask.to(turn_values.device) > 0).unsqueeze(1) + if int(valid.sum()) < 2: + return {} + + v = turn_values[valid].float() + a = turn_advantages[valid].float() + return { + "critic/turn_value_mean": v.mean().item(), + "critic/turn_value_std": v.std().item(), + "advantage/turn_abs_mean_prenorm": a.abs().mean().item(), + "advantage/turn_std_prenorm": a.std().item(), + "turn/num_turns_mean": spans.num_turns.float().mean().item(), + "turn/num_turns_max": float(spans.num_turns.max()), + "turn/tokens_per_turn_mean": ( + spans.turn_ntokens[spans.turn_valid].float().mean().item() + ), + "turn/total_turns": float(int(valid.sum())), + } diff --git a/tests/unit/algorithms/test_turn_level.py b/tests/unit/algorithms/test_turn_level.py new file mode 100644 index 00000000000..63d3bb0b723 --- /dev/null +++ b/tests/unit/algorithms/test_turn_level.py @@ -0,0 +1,463 @@ +# 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. +"""Unit tests for turn-level credit assignment (CPU only, no Ray/Megatron).""" + +import pytest +import torch + +from nemo_rl.algorithms.advantage_estimator import ( + TurnLevelGeneralizedAdvantageEstimator, +) +from nemo_rl.algorithms.turn_level import ( + build_turn_rewards, + build_turn_spans, + build_turn_value_batch, + gather_turn_values, + scatter_turns_to_anchors, + scatter_turns_to_tokens, + turn_gae, + turn_level_metrics, + validate_turn_spans, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +# ---------------------------------------------------------------- helpers +def _msg(role, n, start_id=0): + return {"role": role, "token_ids": torch.arange(start_id, start_id + n)} + + +def _agentic_log(turn_lens, obs_len=3, prompt_len=5): + """[prompt, (assistant, observation)*] — the NeMo-Gym alternating layout.""" + log = [_msg("user", prompt_len)] + for i, n in enumerate(turn_lens): + log.append(_msg("assistant", n)) + if i < len(turn_lens) - 1: + log.append(_msg("user", obs_len)) + return log + + +def _flat_len(log): + return sum(len(m["token_ids"]) for m in log) + + +def _token_mask_from(log, seq_len): + mask = torch.zeros(seq_len, dtype=torch.long) + pos = 0 + for m in log: + n = len(m["token_ids"]) + if m["role"] == "assistant": + mask[pos : pos + n] = 1 + pos += n + return mask + + +class _FakeLossCfg: + use_kl_in_reward = False + reference_policy_kl_penalty = 0.0 + reference_policy_kl_type = "low_var_kl" + + +def _estimator(**overrides): + cfg = { + "turn_gae_gamma": 1.0, + "turn_gae_lambda_value": 1.0, + "turn_gae_lambda_policy": 1.0, + "normalize_advantages": False, + } + cfg.update(overrides) + return TurnLevelGeneralizedAdvantageEstimator(cfg, _FakeLossCfg()) + + +# ---------------------------------------------------------------- spans +def test_spans_locate_every_assistant_message(): + log = _agentic_log([4, 2, 3], obs_len=3, prompt_len=5) + seq_len = _flat_len(log) + 7 # padding + spans = build_turn_spans([log], seq_len) + + # prompt 5 | asst 4 @5 | obs 3 | asst 2 @12 | obs 3 | asst 3 @17 + assert spans.num_turns.tolist() == [3] + assert spans.anchor_pos[0, :3].tolist() == [5, 12, 17] + assert spans.turn_ntokens[0, :3].tolist() == [4, 2, 3] + assert spans.turn_valid[0, :3].all() + assert spans.anchor_mask[0].sum() == 3 + assert spans.anchor_mask[0, [5, 12, 17]].tolist() == [1, 1, 1] + + ti = spans.turn_index[0] + assert ti[:5].tolist() == [-1] * 5 + assert ti[5:9].tolist() == [0] * 4 + assert ti[9:12].tolist() == [-1] * 3 + assert ti[12:14].tolist() == [1, 1] + assert ti[17:20].tolist() == [2, 2, 2] + assert ti[20:].tolist() == [-1] * (seq_len - 20) + + +def test_anchors_are_always_response_tokens(): + logs = [_agentic_log([4, 2, 3]), _agentic_log([7])] + seq_len = max(_flat_len(x) for x in logs) + spans = build_turn_spans(logs, seq_len) + token_mask = torch.stack([_token_mask_from(x, seq_len) for x in logs]) + # the property the whole design rests on: anchors ⊆ response mask + assert bool((spans.anchor_mask.bool() & ~token_mask.bool()).sum() == 0) + validate_turn_spans(spans, token_mask, torch.ones(2)) + + +def test_ragged_batch_is_left_aligned_and_padded(): + logs = [_agentic_log([2, 2, 2, 2]), _agentic_log([5])] + seq_len = max(_flat_len(x) for x in logs) + spans = build_turn_spans(logs, seq_len) + assert spans.anchor_pos.shape[1] == 4 + assert spans.num_turns.tolist() == [4, 1] + assert spans.turn_valid[1].tolist() == [True, False, False, False] + assert spans.turn_ntokens[1].tolist() == [5, 0, 0, 0] + + +def test_consecutive_assistant_messages_are_separate_turns(): + log = [_msg("user", 3), _msg("assistant", 2), _msg("assistant", 4)] + spans = build_turn_spans([log], _flat_len(log)) + assert spans.num_turns.tolist() == [2] + assert spans.anchor_pos[0, :2].tolist() == [3, 5] + + +def test_empty_assistant_message_is_not_a_turn(): + log = [_msg("user", 3), _msg("assistant", 0), _msg("assistant", 4)] + spans = build_turn_spans([log], _flat_len(log)) + assert spans.num_turns.tolist() == [1] + assert spans.anchor_pos[0, 0].item() == 3 + + +def test_assistant_at_index_zero_raises(): + # The value head is right-shifted, so position 0 carries no value. + with pytest.raises(ValueError, match="token 0"): + build_turn_spans([[_msg("assistant", 3)]], 3) + + +def test_seq_len_overflow_raises(): + log = _agentic_log([4, 4]) + with pytest.raises(ValueError, match="flattens to"): + build_turn_spans([log], _flat_len(log) - 1) + + +def test_sample_without_turns_is_reported_not_fatal(capsys): + """An empty trajectory contributes no tokens, so it must not kill the step.""" + log = [_msg("user", 4)] + spans = build_turn_spans([log], 4) + empty_mask = torch.zeros(1, 4, dtype=torch.long) + validate_turn_spans(spans, empty_mask, torch.ones(1)) + assert "no assistant message" in capsys.readouterr().out + # a sample that carries no gradient is not even worth reporting + validate_turn_spans(spans, empty_mask, torch.zeros(1)) + assert "no assistant message" not in capsys.readouterr().out + + +def test_validate_rejects_response_tokens_outside_any_turn(): + """A response token with no turn would carry a policy gradient with no advantage.""" + log = _agentic_log([3]) + seq_len = _flat_len(log) + spans = build_turn_spans([log], seq_len) + bad_mask = torch.ones(1, seq_len, dtype=torch.long) # claims the prompt is trainable + with pytest.raises(ValueError, match="belong to no turn"): + validate_turn_spans(spans, bad_mask, torch.ones(1)) + + +# ---------------------------------------------------------------- GAE math +def test_turn_gae_matches_hand_computation(): + v = torch.tensor([[0.2, 0.5, 0.4]]) + r = torch.tensor([[0.0, 0.0, 1.0]]) + valid = torch.ones(1, 3, dtype=torch.bool) + gamma, lam = 0.9, 0.5 + + # backwards by hand + d2 = 1.0 + gamma * 0.0 - 0.4 + a2 = d2 + d1 = 0.0 + gamma * 0.4 - 0.5 + a1 = d1 + gamma * lam * a2 + d0 = 0.0 + gamma * 0.5 - 0.2 + a0 = d0 + gamma * lam * a1 + + adv, ret = turn_gae(v, r, valid, gamma, lam) + assert torch.allclose(adv, torch.tensor([[a0, a1, a2]]), atol=1e-6) + assert torch.allclose(ret, adv + v, atol=1e-6) + + +def test_lambda_one_gamma_one_gives_monte_carlo_returns(): + """The property that makes turn-level stage-B comparable to the token-level run. + + With gamma=lambda=1 and a terminal-only reward, G_k == R for every turn, so + the critic regresses on the same Monte-Carlo target it does today — only at + ~200 anchors instead of ~45k tokens. + """ + torch.manual_seed(0) + v = torch.randn(4, 6) + valid = torch.ones(4, 6, dtype=torch.bool) + valid[1, 4:] = False + valid[2, 2:] = False + R = torch.tensor([1.0, 0.0, 1.0, 0.5]) + r = torch.zeros(4, 6) + last = valid.long().sum(1) - 1 + r[torch.arange(4), last] = R + + _, ret = turn_gae(v * valid, r, valid, 1.0, 1.0) + for i in range(4): + n = int(valid[i].sum()) + assert torch.allclose(ret[i, :n], R[i].expand(n), atol=1e-5) + + +def test_invalid_turns_do_not_leak_into_valid_ones(): + v = torch.tensor([[0.3, 0.7, 99.0]]) + r = torch.tensor([[0.0, 1.0, 0.0]]) + valid = torch.tensor([[True, True, False]]) + adv, _ = turn_gae(v, r, valid, 1.0, 1.0) + # turn 1 is terminal: A = R - V = 1 - 0.7; turn 0: A = 1 - 0.3 + assert torch.allclose(adv[0, :2], torch.tensor([0.7, 0.3]), atol=1e-6) + + +# ---------------------------------------------------------------- scatter +def test_gather_and_scatter_round_trip(): + log = _agentic_log([4, 2, 3]) + seq_len = _flat_len(log) + spans = build_turn_spans([log], seq_len) + values = torch.arange(seq_len, dtype=torch.float32).unsqueeze(0) + + tv = gather_turn_values(values, spans) + assert tv[0, :3].tolist() == [5.0, 12.0, 17.0] + + per_turn = torch.tensor([[10.0, 20.0, 30.0]]) + tok = scatter_turns_to_tokens(per_turn, spans, seq_len) + assert tok[0, 5:9].tolist() == [10.0] * 4 + assert tok[0, 9:12].tolist() == [0.0] * 3 # observation gets nothing + assert tok[0, 12:14].tolist() == [20.0] * 2 + assert tok[0, 17:20].tolist() == [30.0] * 3 + + anc = scatter_turns_to_anchors(per_turn, spans, seq_len) + assert anc.sum().item() == 60.0 + assert anc[0, [5, 12, 17]].tolist() == [10.0, 20.0, 30.0] + + +def test_turn_rewards_land_on_the_last_turn(): + logs = [_agentic_log([2, 2, 2]), _agentic_log([3])] + seq_len = max(_flat_len(x) for x in logs) + spans = build_turn_spans(logs, seq_len) + tr = build_turn_rewards(torch.tensor([1.0, 0.25]), spans) + assert tr[0].tolist()[:3] == [0.0, 0.0, 1.0] + assert tr[1, 0].item() == 0.25 + assert tr[1, 1:].sum().item() == 0.0 + + +def test_token_penalty_is_summed_into_its_turn(): + log = _agentic_log([4, 2]) + seq_len = _flat_len(log) + spans = build_turn_spans([log], seq_len) + penalty = torch.zeros(1, seq_len) + penalty[0, 5:9] = -0.1 # turn 0 (4 tokens) + penalty[0, 12:14] = -0.5 # turn 1 (2 tokens) + tr = build_turn_rewards(torch.tensor([1.0]), spans, penalty) + assert tr[0, 0].item() == pytest.approx(-0.4, abs=1e-6) + assert tr[0, 1].item() == pytest.approx(1.0 - 1.0, abs=1e-6) + + +# ---------------------------------------------------------------- estimator +def test_estimator_advantage_is_constant_within_a_turn(): + logs = [_agentic_log([4, 2, 3])] + seq_len = _flat_len(logs[0]) + spans = build_turn_spans(logs, seq_len) + token_mask = torch.stack([_token_mask_from(x, seq_len) for x in logs]) + torch.manual_seed(1) + values = torch.randn(1, seq_len) + + est = _estimator(turn_gae_lambda_policy=0.9) + adv, ret = est.compute_advantage( + prompt_ids=None, + rewards=torch.tensor([1.0]), + mask=token_mask, + values=values, + turn_spans=spans, + sample_mask=torch.ones(1), + ) + assert len(set(adv[0, 5:9].tolist())) == 1 + assert len(set(adv[0, 12:14].tolist())) == 1 + assert adv[0, 9:12].abs().sum().item() == 0.0 # observations carry none + # anchor supervision: returns live only at anchors + assert ret[0].nonzero().flatten().tolist() == [5, 12, 17] + + +def test_estimator_lambda_one_reproduces_the_monte_carlo_baseline(): + logs = [_agentic_log([4, 2, 3])] + seq_len = _flat_len(logs[0]) + spans = build_turn_spans(logs, seq_len) + token_mask = torch.stack([_token_mask_from(x, seq_len) for x in logs]) + values = torch.rand(1, seq_len) + R = 1.0 + + est = _estimator() + adv, ret = est.compute_advantage( + prompt_ids=None, + rewards=torch.tensor([R]), + mask=token_mask, + values=values, + turn_spans=spans, + sample_mask=torch.ones(1), + ) + # A_k = R - V(s_k) at gamma=lambda=1 + for anchor, lo, hi in ((5, 5, 9), (12, 12, 14), (17, 17, 20)): + expected = R - values[0, anchor] + assert torch.allclose(adv[0, lo:hi], expected.expand(hi - lo), atol=1e-5) + assert torch.allclose(ret[0, [5, 12, 17]], torch.full((3,), R), atol=1e-5) + + +def test_estimator_returns_are_covered_by_the_critic_mask(): + """Every turn return must land where the critic batch is actually supervised. + + The estimator's `returns` layout and `build_turn_value_batch`'s token_mask + are set in two different places; if they drift apart, targets outside the + mask are silently dropped and the critic trains on less than it looks like. + """ + logs = [_agentic_log([4, 2, 3]), _agentic_log([5])] + seq_len = max(_flat_len(x) for x in logs) + spans = build_turn_spans(logs, seq_len) + td = _train_data(logs, seq_len) + + est = _estimator(turn_gae_lambda_policy=0.9) + _, ret = est.compute_advantage( + prompt_ids=None, + rewards=torch.tensor([1.0, 0.0]), + mask=td["token_mask"], + values=td["values"], + turn_spans=spans, + sample_mask=td["sample_mask"], + ) + td["returns"] = ret + batch = build_turn_value_batch(td, spans) + covered = batch["token_mask"].bool() + assert bool(((ret != 0) & ~covered).sum() == 0) + # and the mask has exactly one position per real turn + assert int(covered.sum()) == int(spans.num_turns.sum()) + + +def test_estimator_normalizes_over_response_tokens_only(): + logs = [_agentic_log([4, 2, 3])] + seq_len = _flat_len(logs[0]) + spans = build_turn_spans(logs, seq_len) + token_mask = torch.stack([_token_mask_from(x, seq_len) for x in logs]) + est = _estimator(normalize_advantages=True, turn_gae_lambda_policy=0.5) + adv, _ = est.compute_advantage( + prompt_ids=None, + rewards=torch.tensor([1.0]), + mask=token_mask, + values=torch.rand(1, seq_len), + turn_spans=spans, + sample_mask=torch.ones(1), + ) + resp = adv[token_mask.bool()] + assert resp.mean().abs().item() < 1e-5 + assert abs(resp.std(unbiased=False).item() - 1.0) < 1e-3 + assert adv[~token_mask.bool()].abs().sum().item() == 0.0 + + +def test_missing_lambda_fails_loud(): + with pytest.raises(ValueError, match="must be set explicitly"): + TurnLevelGeneralizedAdvantageEstimator( + { + "turn_gae_gamma": 1.0, + "turn_gae_lambda_value": None, + "turn_gae_lambda_policy": 1.0, + "normalize_advantages": True, + }, + _FakeLossCfg(), + ) + + +def test_missing_turn_spans_fails_loud(): + est = _estimator() + with pytest.raises(ValueError, match="requires turn_spans"): + est.compute_advantage( + prompt_ids=None, + rewards=torch.tensor([1.0]), + mask=torch.ones(1, 4, dtype=torch.long), + values=torch.zeros(1, 4), + ) + + +# ---------------------------------------------------------------- critic batch +def _train_data(logs, seq_len): + token_mask = torch.stack([_token_mask_from(x, seq_len) for x in logs]) + return BatchedDataDict( + { + "input_ids": torch.zeros(len(logs), seq_len, dtype=torch.long), + "input_lengths": torch.tensor([_flat_len(x) for x in logs]), + "token_mask": token_mask, + "sample_mask": torch.ones(len(logs)), + "values": torch.rand(len(logs), seq_len), + "returns": torch.zeros(len(logs), seq_len), + # keys the policy needs but the critic must not inherit + "advantages": torch.zeros(len(logs), seq_len), + "prev_logprobs": torch.zeros(len(logs), seq_len), + } + ) + + +def test_critic_batch_swaps_in_the_anchor_mask(): + logs = [_agentic_log([4, 2, 3])] + seq_len = _flat_len(logs[0]) + spans = build_turn_spans(logs, seq_len) + td = _train_data(logs, seq_len) + + batch = build_turn_value_batch(td, spans) + # This swap is the whole mechanism: process_global_batch derives + # global_valid_toks from token_mask, so the value loss becomes a per-turn + # mean with no change to MseValueLossFn or the value workers. + assert batch["token_mask"].sum() == 3 + assert torch.equal(batch["token_mask"], spans.anchor_mask) + assert batch["token_mask"].dtype == td["token_mask"].dtype + # the policy-only tensors stay out of the critic's batch + assert "advantages" not in batch and "prev_logprobs" not in batch + for k in ("input_ids", "input_lengths", "sample_mask", "returns", "values"): + assert k in batch + # train_data itself is untouched (the policy still needs the full mask) + assert td["token_mask"].sum() == 9 + + +def test_critic_batch_requires_returns(): + logs = [_agentic_log([2, 2])] + seq_len = _flat_len(logs[0]) + spans = build_turn_spans(logs, seq_len) + td = _train_data(logs, seq_len) + del td["returns"] + with pytest.raises(ValueError, match="returns"): + build_turn_value_batch(td, spans) + + +# ---------------------------------------------------------------- metrics +def test_turn_metrics_are_reported(): + logs = [_agentic_log([2, 2, 2]), _agentic_log([3])] + seq_len = max(_flat_len(x) for x in logs) + spans = build_turn_spans(logs, seq_len) + tv = torch.tensor([[0.2, 0.3, 0.4], [0.5, 0.0, 0.0]]) + ta = torch.tensor([[1.0, 1.0, 1.0], [0.0, 0.0, 0.0]]) - tv + m = turn_level_metrics(tv, ta, spans, torch.ones(2)) + assert m["turn/total_turns"] == 4.0 + assert m["turn/num_turns_mean"] == pytest.approx(2.0) + assert m["turn/num_turns_max"] == 3.0 + assert "advantage/turn_abs_mean_prenorm" in m + # per-position EV / terminal AUC deliberately live in the existing + # _positional_value_metrics / terminal_value_reward_auc, not here + assert not any(k.startswith("critic/turn_ev") for k in m) + + +def test_turn_metrics_respect_the_sample_mask(): + logs = [_agentic_log([2, 2, 2]), _agentic_log([3])] + spans = build_turn_spans(logs, max(_flat_len(x) for x in logs)) + tv = torch.tensor([[0.2, 0.3, 0.4], [0.5, 0.0, 0.0]]) + m = turn_level_metrics(tv, torch.zeros_like(tv), spans, torch.tensor([1.0, 0.0])) + assert m["turn/total_turns"] == 3.0 From 76c49edad99c337bdd1bd5c59f61a4097a4bf1cf Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Tue, 11 Aug 2026 22:59:47 +0000 Subject: [PATCH 54/69] behaviour kl loss between pi_gen and pi_curr Signed-off-by: Jiaqi Zeng --- examples/configs/grpo_ultra_256n4g_bf16.yaml | 14 ++++- nemo_rl/algorithms/loss/loss_functions.py | 60 +++++++++++++++++++- 2 files changed, 72 insertions(+), 2 deletions(-) diff --git a/examples/configs/grpo_ultra_256n4g_bf16.yaml b/examples/configs/grpo_ultra_256n4g_bf16.yaml index c167ac41002..7b39729ca96 100644 --- a/examples/configs/grpo_ultra_256n4g_bf16.yaml +++ b/examples/configs/grpo_ultra_256n4g_bf16.yaml @@ -119,6 +119,14 @@ loss_fn: kl_input_clamp_value: null kl_output_clamp_value: null + # KL(pi_gen || pi_curr) penalty toward the behaviour (rollout) policy. + # Bounds cumulative drift across async steps; clamps must stay set (k3 + # blows up exponentially on bf16 logprob outliers). + behaviour_kl_penalty: 0.0 + behaviour_kl_type: "k3" + behaviour_kl_input_clamp_value: 20.0 + behaviour_kl_output_clamp_value: 10.0 + ratio_clip_min: 0.2 ratio_clip_max: 0.28 ratio_clip_c: null @@ -189,7 +197,11 @@ policy: moe_enable_deepep: false moe_token_dispatcher_type: "alltoall" moe_flex_dispatcher_backend: "hybridep" - moe_hybridep_num_sms: 32 + # moe_hybridep_num_sms removed: Megatron in the NCCL-2.30.4 container deprecates + # both moe_hybridep_num_sms and moe_deepep_num_sms in favour of + # moe_flex_dispatcher_num_sms and errors out when both deprecated knobs are set + # ("Conflicting deprecated SM-count knobs"). nemo_rl on this branch does not plumb + # the new knob, so drop this and take Megatron's default SM count. moe_aux_loss_coeff: 0.0 moe_shared_expert_overlap: false diff --git a/nemo_rl/algorithms/loss/loss_functions.py b/nemo_rl/algorithms/loss/loss_functions.py index 5b8b0bf9cd0..3d20b37489c 100755 --- a/nemo_rl/algorithms/loss/loss_functions.py +++ b/nemo_rl/algorithms/loss/loss_functions.py @@ -130,6 +130,19 @@ class ClippedPGLossConfig(BaseModel, extra="allow"): # If True, add KL penalty to reward instead of loss (used by Reinforce++) use_kl_in_reward: bool = False + # --- Behaviour-policy KL regularization --- + # Penalize KL(π_gen || π_curr) on rollout tokens, anchoring the policy to + # the behaviour (generation) policy. Bounds cumulative drift across async + # steps: clipping only limits per-update movement (prev_logprobs refreshes + # every update) and TIS only reweights the estimator after drift happened. + behaviour_kl_penalty: float = 0.0 + # Can be set to k1, k2, k3 (see reference_policy_kl_type) + behaviour_kl_type: str = "k3" + # Separate clamps from the reference-policy KL so enabling this term is + # safe even in configs that set kl_{input,output}_clamp_value to null. + behaviour_kl_input_clamp_value: Optional[float] = 20.0 + behaviour_kl_output_clamp_value: Optional[float] = 10.0 + # --- Importance sampling correction --- # Async GRPO requires importance sampling correction enabled # Set to true when async_grpo.enabled is true @@ -239,6 +252,13 @@ def __init__( self.reference_policy_kl_type = cfg.reference_policy_kl_type self.kl_input_clamp_value = cfg.kl_input_clamp_value self.kl_output_clamp_value = cfg.kl_output_clamp_value + self.behaviour_kl_penalty = cfg.behaviour_kl_penalty + self.behaviour_kl_type = cfg.behaviour_kl_type + self.behaviour_kl_input_clamp_value = cfg.behaviour_kl_input_clamp_value + self.behaviour_kl_output_clamp_value = cfg.behaviour_kl_output_clamp_value + assert self.behaviour_kl_type in ("k1", "k2", "k3"), ( + f"behaviour_kl_type must be 'k1', 'k2', or 'k3', got {self.behaviour_kl_type}" + ) self.use_importance_sampling_correction = cfg.use_importance_sampling_correction # Type of truncated importance sampling: "tis" | "icepop" | "seq-mask-tis" self.truncated_importance_sampling_type = cfg.truncated_importance_sampling_type @@ -333,6 +353,7 @@ def __init__( # Normalized like the gradient (loss_type-dependent). "loss": grad_normalizer, "kl_penalty": grad_normalizer, + "behaviour_kl": grad_normalizer, # Token-normalized diagnostics, independent of loss_type. "probs_ratio": MetricNormalizer.TOKENS, "probs_ratio_clamped": MetricNormalizer.TOKENS, @@ -512,6 +533,35 @@ def __call__( else: kl = torch.tensor(0.0) + # Behaviour-policy KL penalty: KL(π_gen || π_curr) estimated on rollout + # tokens. Samples come from π_gen (θ-independent), so unlike the + # reference-policy KL above, no importance-sampling weights are needed + # and the pathwise gradient through curr_logprobs is unbiased. + # curr_logprobs (not curr_logprobs_unfiltered) is used to match the + # other π_gen-facing terms (actor IS weights, approx_entropy). + if self.behaviour_kl_penalty != 0: + behaviour_kl = self.behaviour_kl_penalty * calculate_kl( + logprobs=generation_logprobs, + logprobs_reference=curr_logprobs, + kl_type=self.behaviour_kl_type, + input_clamp_value=self.behaviour_kl_input_clamp_value, + output_clamp_value=self.behaviour_kl_output_clamp_value, + ) + if self.loss_type == LossType.TOKEN_LEVEL: + behaviour_kl = masked_mean( + behaviour_kl, + mask, + global_normalization_factor=global_valid_toks, + ) + else: + behaviour_kl = masked_mean( + masked_mean(behaviour_kl, token_mask, dim=-1), + sample_mask, + global_normalization_factor=global_valid_seqs, + ) + else: + behaviour_kl = torch.tensor(0.0) + # Calculate clipped loss function if ppo ratio is enabled. if self.force_on_policy_ratio: # Force ratio to 1.0 for truly on-policy behavior @@ -728,7 +778,12 @@ def __call__( global_normalization_factor=correct_valid_toks, ) - loss = actor_loss + kl + self.positive_example_nll_weight * nll_loss + loss = ( + actor_loss + + kl + + behaviour_kl + + self.positive_example_nll_weight * nll_loss + ) with torch.no_grad(): probs_ratio = masked_mean( ratios.detach(), @@ -771,6 +826,9 @@ def __call__( "probs_ratio_clamped_min": probs_ratio_clamped_min, "probs_ratio_clamped_max": probs_ratio_clamped_max, "kl_penalty": kl.item() / self.reference_policy_kl_penalty if kl else 0, + "behaviour_kl": behaviour_kl.item() / self.behaviour_kl_penalty + if behaviour_kl + else 0, "token_mult_prob_error": mult_prob_error, "gen_kl_error": gen_kl_error, "policy_kl_error": policy_kl_error, From 97385994aedb863e527f450b59856c09c54aad9a Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 13 Aug 2026 12:49:41 -0700 Subject: [PATCH 55/69] build: re-lock to pick up boto3 required by the pinned Gym The Gym submodule declares boto3>=1.34 (ECS Fargate sandbox provider) and is a uv workspace member, so its dependencies resolve into the root uv.lock. The pin crossed that boundary at 11b1240f8 ("bump gym") and no bump since re-locked, leaving HEAD's lockfile without boto3/botocore/s3transfer. Every worker runs `uv run --locked --directory {git_root}` (nemo_rl/distributed/virtual_cluster.py), which asserts the lockfile is up to date and refuses to run otherwise -- so a fresh clone of this branch could not start training at all. Regenerating adds only the three missing packages under nemo-gym's sandbox/all extras; no manifest changed. Signed-off-by: Jiaqi Zeng --- uv.lock | 43 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/uv.lock b/uv.lock index 976bc8f50b7..b9f2705218e 100644 --- a/uv.lock +++ b/uv.lock @@ -496,6 +496,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ed/4d/1392562369b1139e741b30d624f09fe7091d17dd5579fae5732f044b12bb/blobfile-3.0.0-py3-none-any.whl", hash = "sha256:48ecc3307e622804bd8fe13bf6f40e6463c4439eba7a1f9ad49fd78aa63cc658", size = 75413, upload-time = "2024-08-27T00:02:51.518Z" }, ] +[[package]] +name = "boto3" +version = "1.43.67" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "jmespath", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "s3transfer", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ad/bf/ef5de9b55523bc2141d072fbe6614627088e3e6f97b47850216e99de6a1c/boto3-1.43.67.tar.gz", hash = "sha256:75fe983b70d39cfdc274dc51f9bb02b8a0a104bdad4fa073c1af85a35e707c91", size = 112665, upload-time = "2026-08-07T19:30:20.418Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/77/0f/b41f8968452dc6bf3b83f15f5e9f76a27fd93e246906aeb2e0555f494375/boto3-1.43.67-py3-none-any.whl", hash = "sha256:082cf9df068168cb44028a1703822374c1eb7e48fa49470d7e8a76f0c977d0bc", size = 140025, upload-time = "2026-08-07T19:30:18.49Z" }, +] + +[[package]] +name = "botocore" +version = "1.43.67" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jmespath", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "python-dateutil", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "urllib3", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/53/1c/3a75deae60e36bd0ee5c27d040384756b2ea1c90bd7c8c9658335a18b4f5/botocore-1.43.67.tar.gz", hash = "sha256:6fe5cfa0c8676ba809efe505b618ec00f30d1af2d014bf316a7aa4ee86accb20", size = 15889514, upload-time = "2026-08-07T19:30:15.638Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/21/be/38af8e96f3d200c9d34eab8e9f69a4468388ed6030ea04ff012974e5af5a/botocore-1.43.67-py3-none-any.whl", hash = "sha256:48ab8e9fac26fbc2a700d57010251003e6a5f731cf74d8540fb796bc8f3fc0ef", size = 15575924, upload-time = "2026-08-07T19:30:12.116Z" }, +] + [[package]] name = "build" version = "1.5.0" @@ -3176,6 +3204,7 @@ dependencies = [ [package.optional-dependencies] all = [ + { name = "boto3", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "coverage", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "mypy", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "opensandbox", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, @@ -3200,6 +3229,7 @@ dev = [ { name = "ruff", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, ] sandbox = [ + { name = "boto3", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "opensandbox", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "tenacity", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, ] @@ -3222,6 +3252,7 @@ docs = [ requires-dist = [ { name = "aiohttp", specifier = ">=3.14.1" }, { name = "anthropic", specifier = "<=0.109.2" }, + { name = "boto3", marker = "extra == 'sandbox'", specifier = ">=1.34" }, { name = "coverage", extras = ["toml"], marker = "extra == 'dev'" }, { name = "datasets" }, { name = "devtools" }, @@ -5350,6 +5381,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8b/45/8e5fd559bea0d2f57c4e12bf197a2fade2fac465aa518284f157dfbca92b/ruff-0.9.9-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:ab90a7944c5a1296f3ecb08d1cbf8c2da34c7e68114b1271a431a3ad30cb660e", size = 11327490, upload-time = "2025-02-28T10:16:27.654Z" }, ] +[[package]] +name = "s3transfer" +version = "0.19.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/76/43/35e4d8aa320bffe8287fe8f65f578fa2d2db0a64212f0e710dce58267854/s3transfer-0.19.2.tar.gz", hash = "sha256:ba0309fd86be3c27dbf78cdd813c13c5e1df16e5874b99d2535ebbdfb9892993", size = 165592, upload-time = "2026-07-22T19:30:44.432Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/e7/5c595c75e9f41a44f30e526eda465ea0b4eec93470e074e4a111b253f13a/s3transfer-0.19.2-py3-none-any.whl", hash = "sha256:d8168eccca828cbb2cd573675333f3bddd254313a9c42494b84c76b539e8ba25", size = 90216, upload-time = "2026-07-22T19:30:43.251Z" }, +] + [[package]] name = "safetensors" version = "0.8.0" From 9e6748eac9057ac0e2e8ddcecd038f4a809a9324 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 13 Aug 2026 12:50:35 -0700 Subject: [PATCH 56/69] style: apply ruff format to ppo.py and critic_pretrain.py These two files were not formatter-clean at HEAD: earlier commits hand-wrapped calls that `ruff format` collapses. Running the project's pinned ruff (0.9.9) over HEAD's versions reproduces exactly this diff, so it is mechanical output, not hand editing. Split out from the critic work that follows so those diffs contain only behavioural change. Verified semantics-free: the parsed AST of each file is identical before and after. Signed-off-by: Jiaqi Zeng --- nemo_rl/algorithms/critic_pretrain.py | 64 +++++++++++++++------------ nemo_rl/algorithms/ppo.py | 42 ++++++++---------- 2 files changed, 54 insertions(+), 52 deletions(-) diff --git a/nemo_rl/algorithms/critic_pretrain.py b/nemo_rl/algorithms/critic_pretrain.py index 016c9fe297e..c44e85c2b97 100644 --- a/nemo_rl/algorithms/critic_pretrain.py +++ b/nemo_rl/algorithms/critic_pretrain.py @@ -73,9 +73,7 @@ def list_group_files(shards_dir: str | Path) -> list[Path]: return [f for f in files if parse_group_index(f.name) is not None] -def split_heldout( - files: list[Path], heldout_mod: int -) -> tuple[list[Path], list[Path]]: +def split_heldout(files: list[Path], heldout_mod: int) -> tuple[list[Path], list[Path]]: """Deterministic train/held-out split by dataset index. ``dataset_idx % heldout_mod == 0`` goes to held-out; ``heldout_mod <= 0`` @@ -162,11 +160,15 @@ def verify_shard_meta( metas = sorted(Path(shards_dir).glob("shard_*/meta.json")) if not metas: - print(f"⚠️ No shard meta.json found under {shards_dir}; skipping provenance check.") + print( + f"⚠️ No shard meta.json found under {shards_dir}; skipping provenance check." + ) return expected = { "model_name": master_config.policy["model_name"], - "chat_template_sha256": _sha256(getattr(tokenizer, "chat_template", None) or ""), + "chat_template_sha256": _sha256( + getattr(tokenizer, "chat_template", None) or "" + ), "max_total_sequence_length": master_config.policy["max_total_sequence_length"], } for meta_path in metas: @@ -371,9 +373,7 @@ def _forward_values_and_returns( train_data["token_mask"], ) else: - train_data["values"] = value_model.get_values(train_data)["values"].squeeze( - -1 - ) + train_data["values"] = value_model.get_values(train_data)["values"].squeeze(-1) value_model.finish_inference() initial_prompt_message_logs = extract_initial_prompt_messages( @@ -434,7 +434,11 @@ def _heldout_metrics( groups, tokenizer, master_config ) _, turn_spans = _forward_values_and_returns( - value_model, adv_estimator, train_data, repeated_batch, tokenizer, + value_model, + adv_estimator, + train_data, + repeated_batch, + tokenizer, master_config, ) values, returns = train_data["values"], train_data["returns"] @@ -447,9 +451,7 @@ def _heldout_metrics( v, r = values[mask].float(), returns[mask].float() var_r = r.var(unbiased=False) metrics["critic/explained_var"] = ( - (1.0 - (r - v).var(unbiased=False) / var_r).item() - if var_r > 1e-8 - else 0.0 + (1.0 - (r - v).var(unbiased=False) / var_r).item() if var_r > 1e-8 else 0.0 ) metrics["critic/mse"] = ((r - v) ** 2).mean().item() metrics.update(_positional_value_metrics(values, returns, scored_mask)) @@ -494,7 +496,11 @@ def _dump_heldout_values( [g], tokenizer, master_config ) _, turn_spans = _forward_values_and_returns( - value_model, adv_estimator, train_data, repeated_batch, tokenizer, + value_model, + adv_estimator, + train_data, + repeated_batch, + tokenizer, master_config, ) mask = train_data["token_mask"].bool() @@ -509,7 +515,7 @@ def _dump_heldout_values( succ = [i for i in range(len(rew)) if rew[i] > 0.5] fail = [i for i in range(len(rew)) if rew[i] <= 0.5] half = max(1, dump_token_samples // 2) - render_samples = (succ[:half] + fail[: dump_token_samples - len(succ[:half])]) + render_samples = succ[:half] + fail[: dump_token_samples - len(succ[:half])] render_samples = sorted(render_samples[:dump_token_samples]) render_set = set(render_samples) samples_msgs = [] @@ -520,9 +526,7 @@ def _dump_heldout_values( for m, s in zip(ml, spans): s["text"] = tokenizer.decode(m["token_ids"]) if want_toks and m["role"] == "assistant": - s["toks"] = [ - tokenizer.decode([int(t)]) for t in m["token_ids"] - ] + s["toks"] = [tokenizer.decode([int(t)]) for t in m["token_ids"]] samples_msgs.append(spans) anchor_mask = turn_spans.anchor_mask if turn_spans is not None else None payload = { @@ -683,14 +687,14 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: "frozen; they are ignored this run (frozen-dataset semantics)." ) else: - base_train, heldout_files = split_heldout( - all_files, cp_config["heldout_mod"] - ) + base_train, heldout_files = split_heldout(all_files, cp_config["heldout_mod"]) train_files = build_epoch_stream(base_train, num_epochs, cp_config["seed"]) missing = [ p for p in set(train_files) | set(heldout_files) if not os.path.exists(p) ] - assert not missing, f"{len(missing)} frozen group files are missing, e.g. {missing[:3]}" + assert not missing, ( + f"{len(missing)} frozen group files are missing, e.g. {missing[:3]}" + ) groups_per_step = cp_config["groups_per_step"] # drop-last within the whole multi-epoch stream @@ -781,7 +785,9 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: if eval_only: dump_dir = _Path( cp_config.get("dump_dir") - or os.path.join(master_config.checkpointing["checkpoint_dir"], "value_dumps") + or os.path.join( + master_config.checkpointing["checkpoint_dir"], "value_dumps" + ) ) print( f"🔍 Eval-only: scoring {len(heldout_files)} held-out groups -> {dump_dir}" @@ -843,16 +849,18 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: print("▶ Computing values...") critic_batch, turn_spans = _forward_values_and_returns( - value_model, adv_estimator, train_data, repeated_batch, tokenizer, + value_model, + adv_estimator, + train_data, + repeated_batch, + tokenizer, master_config, ) print("▶ Training critic...") value_model.prepare_for_training() value_train_batch = critic_batch if critic_batch is not None else train_data - value_results = value_model.train( - value_train_batch, value_loss_fn - ) + value_results = value_model.train(value_train_batch, value_loss_fn) value_model.finish_training() # ---- Metrics ---- @@ -909,9 +917,7 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: save_state["consumed_samples"] = save_state.get("consumed_samples", 0) + int( train_data["input_ids"].shape[0] ) - if checkpointing_enabled and ( - is_last_step or step % save_period == 0 - ): + if checkpointing_enabled and (is_last_step or step % save_period == 0): print(f"💾 Saving checkpoint for step {step}...") checkpoint_path = checkpointer.init_tmp_checkpoint( step, save_state, master_config diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 295806135a0..1e8ce65dd0e 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -982,9 +982,7 @@ def initialize_generation_with_policy( def _spinup_nemo_gym(base_urls, model_name): """Spin up the NeMo-Gym actor against the given vLLM server URLs.""" t0 = time.perf_counter() - nemo_gym_py_exec = get_actor_python_env( - "nemo_rl.environments.nemo_gym.NemoGym" - ) + nemo_gym_py_exec = get_actor_python_env("nemo_rl.environments.nemo_gym.NemoGym") if nemo_gym_py_exec.startswith("uv"): nemo_gym_py_exec = create_local_venv_on_each_node( nemo_gym_py_exec, "nemo_rl.environments.nemo_gym.NemoGym" @@ -1206,9 +1204,7 @@ def init_vllm_then_policy_value(): ip, port, world_size, train_world_size=train_world_size ) ray.get(futures_train + futures_inference) - worker_init_timing_metrics["collective_init_time_s"] = ( - time.perf_counter() - t0 - ) + worker_init_timing_metrics["collective_init_time_s"] = time.perf_counter() - t0 # Reload policy weights to GPU before refit (they may have been offloaded # during setup to free GPU for value model initialization). @@ -1548,9 +1544,7 @@ def _compute_critic_metrics(value_results: dict[str, Any]) -> dict[str, Any]: return critic_metrics -def _calibration_ece( - v: torch.Tensor, r: torch.Tensor, n_conf_bins: int = 10 -) -> float: +def _calibration_ece(v: torch.Tensor, r: torch.Tensor, n_conf_bins: int = 10) -> float: """Expected Calibration Error of V as an estimate of P(success). With outcome returns in [0, 1], a well-calibrated critic satisfies @@ -1597,7 +1591,9 @@ def _positional_value_metrics( # 0-based index of each response token within its sample's response, scaled to [0,1) resp_len = mask.sum(dim=1, keepdim=True).clamp(min=1) rel = (torch.cumsum(mask.long(), dim=1) - 1).float() / resp_len.float() - names = ["early", "mid", "late"] if n_bins == 3 else [f"bin{i}" for i in range(n_bins)] + names = ( + ["early", "mid", "late"] if n_bins == 3 else [f"bin{i}" for i in range(n_bins)] + ) out: dict[str, float] = {} for i in range(n_bins): lo, hi = i / n_bins, (i + 1) / n_bins @@ -1736,9 +1732,7 @@ def _pack_float(field: str) -> torch.Tensor: if "returns" in train_data: payload["returns"] = _pack_float("returns") if turn_spans is not None: - payload["is_anchor"] = ( - turn_spans.anchor_mask.detach().bool().cpu()[token_mask] - ) + payload["is_anchor"] = turn_spans.anchor_mask.detach().bool().cpu()[token_mask] if "reference_policy_logprobs" in train_data: payload["reference_policy_logprobs"] = _pack_float("reference_policy_logprobs") @@ -2117,8 +2111,9 @@ def ppo_train( # values are mapped back into the original layout. Policy/GAE/value # workers are untouched. See nemo_rl/algorithms/privileged_critic.py. privileged_critic_cfg = master_config.value.get("privileged_critic") - if privileged_critic_cfg is not None and not privileged_critic_cfg.get( - "enabled" + if ( + privileged_critic_cfg is not None + and not privileged_critic_cfg.get("enabled") ): privileged_critic_cfg = None critic_batch = None @@ -2133,9 +2128,9 @@ def ppo_train( "make_sequence_length_divisible_by" ], ) - vals_aug = value_model.get_values(critic_batch)["values"].squeeze( - -1 - ) + vals_aug = value_model.get_values(critic_batch)[ + "values" + ].squeeze(-1) # keep aug-layout old values around for the (optional) value clip critic_batch["values"] = vals_aug train_data["values"] = remap_by_response_mask( @@ -3352,8 +3347,9 @@ def _sample_max_age(s: int) -> int: # trajectories carry message_log + extra_env_info, so the augmented # batch is built at train time exactly as in sync. privileged_critic_cfg = master_config.value.get("privileged_critic") - if privileged_critic_cfg is not None and not privileged_critic_cfg.get( - "enabled" + if ( + privileged_critic_cfg is not None + and not privileged_critic_cfg.get("enabled") ): privileged_critic_cfg = None critic_batch = None @@ -3368,9 +3364,9 @@ def _sample_max_age(s: int) -> int: "make_sequence_length_divisible_by" ], ) - vals_aug = value_model.get_values(critic_batch)["values"].squeeze( - -1 - ) + vals_aug = value_model.get_values(critic_batch)[ + "values" + ].squeeze(-1) critic_batch["values"] = vals_aug train_data["values"] = remap_by_response_mask( vals_aug, From c747b6bb039930586e35f95e575199aef00aefa4 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 13 Aug 2026 12:51:31 -0700 Subject: [PATCH 57/69] fix(megatron): don't set the deprecated MoE SM-count knob on new mcore Newer mcore (NCCL 2.30.4 image) rejects more than one deprecated SM-count knob: ValueError: Conflicting deprecated SM-count knobs {'moe_deepep_num_sms': 20, 'moe_hybridep_num_sms': 32}; set a single moe_flex_dispatcher_num_sms instead. mcore supplies moe_deepep_num_sms itself, so NeMo-RL setting the hybridep knob is what tips it over. Feature-detect the replacement attribute and route the configured value there, falling back to the deprecated knob on older builds. This keeps moe_hybridep_num_sms: 32 effective in ppo_nano_v3_5_swe_cmh.yaml and the grpo/ultra configs rather than discarding it. The implementation is byte-identical to the fix already on main (ec46ab3fb), which this branch predates, so the hunk will resolve to nothing when the branch meets main. Signed-off-by: Jiaqi Zeng --- nemo_rl/models/megatron/setup.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index eea649c2355..4367ac4cada 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -702,8 +702,20 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None: model_cfg.moe_flex_dispatcher_backend = config["megatron_cfg"][ "moe_flex_dispatcher_backend" ] + # moe_hybridep_num_sms is deprecated in newer mcore (NCCL 2.30.4 image), + # which hard-errors when more than one deprecated SM-count knob is set: + # ValueError: Conflicting deprecated SM-count knobs + # {'moe_deepep_num_sms': 20, 'moe_hybridep_num_sms': 32}; + # set a single moe_flex_dispatcher_num_sms instead. + # mcore supplies moe_deepep_num_sms itself, so setting the hybridep one here + # is what triggers it. Route the value to the replacement knob when this + # mcore has it, and fall back to the deprecated one on older builds. if "moe_hybridep_num_sms" in config["megatron_cfg"]: - model_cfg.moe_hybridep_num_sms = config["megatron_cfg"]["moe_hybridep_num_sms"] + num_sms = config["megatron_cfg"]["moe_hybridep_num_sms"] + if hasattr(TransformerConfig, "moe_flex_dispatcher_num_sms"): + model_cfg.moe_flex_dispatcher_num_sms = num_sms + else: + model_cfg.moe_hybridep_num_sms = num_sms # HybridEP environment variables # These are required by DeepEP's hybrid-ep branch for NVLink domain configuration. From ed0068ab66c1db0550a14a29ffc44b4e081d7961 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 13 Aug 2026 12:58:11 -0700 Subject: [PATCH 58/69] feat(critic): decomposed group baseline + residual critic The Bayes-optimal value splits into a between-task and a within-task part, V*(s) = B(X) + C(s) with E[C | X] = 0. An absolute critic must fit both, and on this SWE workload it overwhelmingly fits the first and fits it badly: 88-93% of its output variance is between-task, yet EV ~0.157 -- below a two-value "all-fail vs rest" lookup (0.330) and far below the free leave-one-out group baseline (0.553). Substituting the critic for that baseline RAISES advantage variance 1.67-2.09x. So hand B to the rollout group and let the critic learn only C. ResidualBaselineEstimator wraps either value-based estimator (gae or turn_gae), keeps the batch in residual space end to end, and reconstructs V~ = B_LOO + C only for the inner GAE call -- which is what keeps the PPO value clip self-consistent. gamma == 1 is required and enforced: the baseline only cancels from nonterminal TD errors there. num_generations_per_prompt >= 2 is also enforced, since a single sibling makes B_LOO == R and would silently drop the reward from the policy gradient entirely. adv_estimator.residual_baseline has no default, so a sweep cannot appear to have run when it did not. It is required for gae/turn_gae and rejected for raw_reward. Diagnostics are the point of the change as much as the mechanism: the wrapper is applied even when residualization is OFF so an absolute run reports critic/ev_res on the same axis as a residual run (there 1 - ev_res is exactly the measured advantage-variance ratio). MseValueLossFn gains the per-space return statistics this needs, and critic/ev_res_mixed_group isolates EV where the target is actually nonzero -- homogeneous groups are ~58% of this pool and contribute Y = 0, so any prediction there is a pure penalty to the whole-batch number. value_loss_fn.homogeneous_group_weight can down-weight them; 1.0 keeps today's behaviour. The default path is unchanged: with residual_baseline: false the targets stay absolute and only the extra diagnostics appear. Signed-off-by: Jiaqi Zeng --- examples/configs/ppo_math_1B.yaml | 12 + examples/configs/ppo_nano_v3_5_swe_cmh.yaml | 24 ++ .../ppo_ultra_256n4g_bf16_pivotonly.yaml | 15 + nemo_rl/algorithms/advantage_estimator.py | 340 +++++++++++++++- nemo_rl/algorithms/critic_pretrain.py | 193 ++++++++- nemo_rl/algorithms/loss/loss_functions.py | 43 ++ nemo_rl/algorithms/ppo.py | 331 +++++++++++++++- .../algorithms/test_advantage_estimator.py | 369 +++++++++++++++++- tests/unit/algorithms/test_loss_functions.py | 116 ++++++ tests/unit/algorithms/test_ppo.py | 60 ++- .../ppo_math_1B_megatron.yaml | 12 + 11 files changed, 1496 insertions(+), 19 deletions(-) diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index dbf48e44f56..c96ab20f84a 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -47,6 +47,15 @@ ppo: # Length-adaptive λ_policy = 1 - 1/(α·l). 0 = disabled. length_adaptive_alpha: 0.0 # VAPO: 0.05 + # --- decomposed group baseline + residual critic --- + # research/ppo/residual_critic_report.md. true => the rollout group supplies + # the task baseline B(X) as a leave-one-out mean and the critic is trained + # only on the within-task residual C(s) = R - B_LOO; critic values/returns + # are then in RESIDUAL space and the go/no-go metric is critic/ev_res. + # Requires gae_gamma: 1. Orthogonal to token-vs-turn granularity. + # false still logs critic/ev_res + residual/* for the absolute critic. + residual_baseline: false + reward_scaling: enabled: true source_min: 0.0 @@ -104,6 +113,9 @@ loss_fn: value_loss_fn: scale: 0.4 cliprange: 0.2 + # Weight of homogeneous (zero within-group reward variance) groups in the + # value loss; only meaningful with ppo.adv_estimator.residual_baseline. + homogeneous_group_weight: 1.0 checkpointing: enabled: true diff --git a/examples/configs/ppo_nano_v3_5_swe_cmh.yaml b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml index 43d7ee9c706..a3ab7196836 100644 --- a/examples/configs/ppo_nano_v3_5_swe_cmh.yaml +++ b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml @@ -117,6 +117,21 @@ ppo: # Length-adaptive lambda_policy = 1 - 1/(alpha*l). 0 = disabled. length_adaptive_alpha: null + # --- decomposed group baseline + residual critic --- + # research/ppo/residual_critic_report.md. true => the rollout group supplies + # the task baseline B(X) as a leave-one-out mean and the critic is trained + # only on the within-task residual C(s) = R - B_LOO; critic values/returns + # are then in RESIDUAL space and the go/no-go metric is critic/ev_res. + # Requires gae_gamma: 1. Orthogonal to token-vs-turn granularity. + # + # Motivation, measured on this exact workload: the absolute critic puts + # 88-93% of its output variance between tasks yet reaches EV 0.157 — below + # even a two-value "all-fail vs rest" lookup (0.330) and far below the free + # leave-one-out group baseline (0.553). Substituting it for that baseline + # RAISES advantage variance 1.67-2.09x. false still logs critic/ev_res and + # residual/* so an absolute run is comparable on the same axis. + residual_baseline: false + # --- turn-level GAE (used only when name: "turn_gae") --- # All three lambdas/gamma are required in that mode and are never silently # defaulted, so a sweep cannot appear to have run when it did not. @@ -200,6 +215,15 @@ loss_fn: value_loss_fn: scale: 1.0 cliprange: 0.5 + # Weight of homogeneous (all-fail / all-pass) groups in the value loss. Only + # meaningful with ppo.adv_estimator.residual_baseline: those groups have Y = 0 + # for every sibling, so they contribute EXACTLY zero target variance while + # still costing ~58% of critic FLOPs on this pool (54.1% all-fail, 2.4% + # all-pass; only 43.5% of groups are mixed). Keep at 1.0 for the clean + # ablation — they are not dead weight, they are the shrinkage that enforces + # E[C | X] = 0 and suppresses between-task leakage. Lower it only to buy more + # mixed-group exposure per unit wall-clock. + homogeneous_group_weight: 1.0 # ============================================================================= # Policy diff --git a/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml b/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml index 41f15529cb1..adbc2050aab 100644 --- a/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml +++ b/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml @@ -117,6 +117,15 @@ ppo: # Length-adaptive lambda_policy = 1 - 1/(alpha*l). 0 = disabled. length_adaptive_alpha: null + # --- decomposed group baseline + residual critic --- + # research/ppo/residual_critic_report.md. true => the rollout group supplies + # the task baseline B(X) as a leave-one-out mean and the critic is trained + # only on the within-task residual C(s) = R - B_LOO; critic values/returns + # are then in RESIDUAL space and the go/no-go metric is critic/ev_res. + # Requires gae_gamma: 1. Orthogonal to token-vs-turn granularity. + # false still logs critic/ev_res + residual/* for the absolute critic. + residual_baseline: false + reward_shaping: enabled: false overlong_buffer_length: 128 @@ -184,6 +193,12 @@ loss_fn: value_loss_fn: scale: 1.0 cliprange: 0.5 + # Weight of homogeneous (zero within-group reward variance) groups in the + # value loss. Only meaningful with ppo.adv_estimator.residual_baseline: those + # groups have Y = 0 for every sibling and so contribute exactly zero target + # variance. Keep at 1.0 for the clean ablation -- they are the shrinkage that + # enforces E[C | X] = 0 and suppresses between-task leakage. + homogeneous_group_weight: 1.0 # ============================================================================= # Policy diff --git a/nemo_rl/algorithms/advantage_estimator.py b/nemo_rl/algorithms/advantage_estimator.py index f07571cba7a..2966375a894 100644 --- a/nemo_rl/algorithms/advantage_estimator.py +++ b/nemo_rl/algorithms/advantage_estimator.py @@ -21,6 +21,8 @@ - RawRewardAdvantageEstimator: Raw reward as advantage with optional batch normalization (no baseline, no value model) - GeneralizedAdvantageEstimator: Generalized Advantage Estimation (GAE) with temporal bootstrapping - TurnLevelGeneralizedAdvantageEstimator: GAE over agent TURNS instead of tokens +- ResidualBaselineEstimator: wraps a value-based estimator so the group supplies the + task baseline B(X) and the critic learns only the within-task residual C(s) - OPDAdvantageEstimator: Multi-Teacher On-Policy Distillation (MOPD) token-level distillation advantages Reference papers: - ProRLv2: https://developer.nvidia.com/blog/scaling-llm-reinforcement-learning-with-prolonged-training-using-prorl-v2/ @@ -29,6 +31,8 @@ - MOPD: https://arxiv.org/abs/2601.02780 """ +from typing import Any, Optional + import torch from nemo_rl.algorithms.loss import ClippedPGLossConfig @@ -552,7 +556,11 @@ class TurnLevelGeneralizedAdvantageEstimator: def __init__(self, estimator_config: dict, loss_config: ClippedPGLossConfig): # No silent defaults: a λ that quietly falls back to 1.0 would make a # sweep look like it ran when it did not. - for key in ("turn_gae_gamma", "turn_gae_lambda_value", "turn_gae_lambda_policy"): + for key in ( + "turn_gae_gamma", + "turn_gae_lambda_value", + "turn_gae_lambda_policy", + ): if estimator_config.get(key) is None: raise ValueError( f"adv_estimator.{key} must be set explicitly when " @@ -666,6 +674,336 @@ def compute_advantage( return advantages, returns +class ResidualBaselineEstimator: + """Decomposed group baseline + residual critic (research/ppo/residual_critic_report.md). + + The Bayes-optimal value splits exactly into a between-task and a within-task + part, ``V*(s) = B(X) + C(s)`` with ``E[C | X] = 0``. An absolute critic has + to fit both. On this SWE workload it overwhelmingly fits the first and fits + it badly: 88-93% of its output variance is between-task, yet its EV (~0.157) + is below even a two-value "all-fail vs rest" lookup (0.330) and far below the + free leave-one-out group baseline (~0.553). Substituting the critic for that + baseline *raises* advantage variance 1.67-2.09x. + + So hand ``B`` to the rollout group and let the critic learn only ``C``. This + wrapper keeps the critic in RESIDUAL space end to end: + + values in the batch = C(s) + returns in the batch = the residual lambda-return (target for C) + + and reconstructs the absolute value ``V~ = B_LOO + C`` only for the duration + of the inner GAE call. Keeping the batch residual is what makes the PPO + value clip (``old_values`` vs ``returns``) self-consistent, and it is why + this is a wrapper rather than a post-hoc transform of the returns: the report + (S26.10) says "reconstruct V~ and feed the existing GAE", but ``returns`` + flows straight into :class:`MseValueLossFn`, and at lambda=1 GAE returns are + ``R`` -- following that literally would train ``C`` against ``R``. + + With gamma = 1 the group baseline cancels from every nonterminal TD error + (``delta_t = C_{t+1} - C_t``), and at lambda = 1 the advantage telescopes to + ``A_t = R - B_LOO - C_t``: the critic becomes a state-dependent control + variate on top of the known-good group advantage, and ``C = 0`` recovers it + exactly. gamma < 1 breaks the cancellation, so it is rejected outright. + + Wrapping (rather than subclassing) means both ``gae`` and ``turn_gae`` are + covered without duplicating the leave-one-out logic. + + Args: + inner: the value-based estimator to wrap (``gae`` or ``turn_gae``). + residual_target: when False, ``B_LOO`` is still computed and exported for + metrics but the targets are left in absolute space. That is what + lets an absolute-critic run log ``critic/ev_res`` on the same axis as + a residual run -- there ``1 - ev_res`` is exactly the report's + advantage-variance ratio. + """ + + def __init__(self, inner: Any, residual_target: bool): + self.inner = inner + self.residual_target = residual_target + + # GeneralizedAdvantageEstimator names it gae_gamma; the turn-level one + # names it gamma. Neither has a default worth guessing at. + gamma = getattr(inner, "gae_gamma", getattr(inner, "gamma", None)) + if gamma is None: + raise ValueError( + f"{type(inner).__name__} exposes no discount factor, so the " + "residual baseline cannot verify the gamma = 1 condition it " + "depends on." + ) + self.gamma = float(gamma) + if residual_target and abs(self.gamma - 1.0) > 1e-8: + raise ValueError( + f"adv_estimator.residual_baseline requires gamma == 1, got {self.gamma}. " + "The task baseline only cancels from nonterminal TD errors at " + "gamma = 1; at gamma < 1 the residual value would silently pick " + "up a (gamma - 1) * B term (report S11.5)." + ) + + self.last_metrics: dict[str, float] = {} + # Per-sample offsets that move `returns` into each space. Exactly one is + # zero. Consumed by MseValueLossFn so it can report BOTH explained + # variances without knowing which mode it is in. + self.last_returns_to_abs: torch.Tensor | None = None + self.last_returns_to_res: torch.Tensor | None = None + # ``[B]`` float, 1.0 where the rollout's group is all-fail or all-pass. + self.last_group_homogeneous: torch.Tensor | None = None + # ``[B]`` long, sibling-group index. Exposed so within-group diagnostics + # partition trajectories exactly the way the baseline did. + self.last_group_ids: torch.Tensor | None = None + + # NOTE: this ``last_*`` side-channel deliberately mirrors the existing + # ``last_metrics`` contract that ppo.py and critic_pretrain.py already + # read via getattr on the estimator; keeping one convention beats adding + # a second way to hand per-step tensors back to the caller. + + def compute_advantage( + self, + prompt_ids: torch.Tensor, + rewards: torch.Tensor, + mask: torch.Tensor, + values: torch.Tensor, + turn_spans: Optional[Any] = None, + sample_mask: Optional[torch.Tensor] = None, + **kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run the inner estimator against ``V~ = B_LOO + C`` and return residual targets. + + Returns: + ``(advantages, returns)``. ``advantages`` are byte-for-byte what the + inner estimator produces for a critic predicting ``V~`` -- only the + returns are moved back into residual space. + """ + baseline = self._leave_one_out_baseline(prompt_ids, rewards) + b_col = baseline.unsqueeze(-1).to(values.dtype) + + inner_values = values + b_col if self.residual_target else values + adv_kwargs = dict( + prompt_ids=prompt_ids, + rewards=rewards, + mask=mask, + values=inner_values, + sample_mask=sample_mask, + **kwargs, + ) + if turn_spans is not None: + adv_kwargs["turn_spans"] = turn_spans + advantages, returns = self.inner.compute_advantage(**adv_kwargs) + + # Turn-level returns live at ONE anchor per turn and are structurally + # zero elsewhere (scatter_turns_to_anchors), matched by token_mask = + # anchor_mask in the critic's batch. Subtracting the baseline unmasked + # would write -B into every non-anchor position. + return_mask = ( + turn_spans.anchor_mask.to(returns.dtype) + if turn_spans is not None + else mask.to(returns.dtype) + ) + if self.residual_target: + returns = returns - b_col * return_mask + self.last_returns_to_abs = baseline + self.last_returns_to_res = torch.zeros_like(baseline) + else: + self.last_returns_to_abs = torch.zeros_like(baseline) + self.last_returns_to_res = -baseline + + self.last_metrics = dict(getattr(self.inner, "last_metrics", None) or {}) + self.last_metrics.update( + self._group_metrics( + prompt_ids, rewards, baseline, returns, return_mask, sample_mask + ) + ) + return advantages, returns + + def _leave_one_out_baseline( + self, prompt_ids: torch.Tensor, rewards: torch.Tensor + ) -> torch.Tensor: + """``B_LOO[i,j] = (sum_k R[i,k] - R[i,j]) / (G_i - 1)``, in fp32. + + ``valid_mask`` is deliberately all-ones, matching + :class:`GRPOAdvantageEstimator` exactly: a rollout dropped from the + critic loss by ``sample_mask`` still feeds its siblings' baselines. That + keeps the PPO actor A/B against the working DAPO/GRPO control on + identical group semantics, and keeps critic pretraining (stage B) aligned + with what PPO (stage C) will feed the same checkpoint. The divergence is + made visible by ``residual/frac_traj_sample_masked`` rather than hidden. + """ + rewards_f32 = rewards.float() + baseline, _ = calculate_baseline_and_std_per_prompt( + prompt_ids, + rewards_f32, + torch.ones_like(rewards_f32), + leave_one_out_baseline=True, + ) + return baseline + + @staticmethod + def _group_ids(prompt_ids: torch.Tensor) -> torch.Tensor: + """``[B]`` group index, using the same identity rule as the LOO baseline.""" + _, inverse = torch.unique(prompt_ids, dim=0, return_inverse=True) + return inverse.reshape(-1) + + def _group_metrics( + self, + prompt_ids: torch.Tensor, + rewards: torch.Tensor, + baseline: torch.Tensor, + returns: torch.Tensor, + return_mask: torch.Tensor, + sample_mask: Optional[torch.Tensor], + ) -> dict[str, float]: + """Group-composition diagnostics for the residual target. + + ``frac_groups_mixed`` is the load-bearing one: homogeneous groups have + ``Y = 0`` for every sibling and therefore contribute EXACTLY zero target + variance, so this fraction -- not the dataset size -- bounds what the + residual critic can learn. On the pi0 SWE pool it is 0.435. + + "Homogeneous" is defined as ZERO WITHIN-GROUP REWARD VARIANCE, not as a + group sum of 0 or G. Rewards arriving here are post-scaling, post-shaping + and post-penalty: ``reward_scaling`` maps ``[0,1] -> [-1,1]`` in the math + configs, and judge rewards can be fractional. A sum-based test would then + call an all-fail group "mixed" and report ``frac_groups_mixed = 1.0`` on a + pool that is mostly homogeneous. Zero within-group variance is exactly the + ``Y = 0`` condition being claimed, under any reward scale. + """ + with torch.no_grad(): + group_ids = self._group_ids(prompt_ids).to(rewards.device) + self.last_group_ids = group_ids + rewards_f32 = rewards.float() + n_groups = int(group_ids.max().item()) + 1 if group_ids.numel() else 0 + if n_groups == 0: + return {} + + ones = torch.ones_like(rewards_f32) + counts = torch.zeros(n_groups, device=rewards.device).index_add_( + 0, group_ids, ones + ) + sums = torch.zeros(n_groups, device=rewards.device).index_add_( + 0, group_ids, rewards_f32 + ) + means = sums / counts.clamp(min=1) + sq_dev = torch.zeros(n_groups, device=rewards.device).index_add_( + 0, group_ids, (rewards_f32 - means[group_ids]) ** 2 + ) + homogeneous = (sq_dev <= 1e-12).float() + mixed = 1.0 - homogeneous + # all-fail / all-pass are only well defined for a binary reward; they + # are reported as "homogeneous at the group min / max reward" so they + # stay meaningful (and still sum to `homogeneous`) under any scaling. + r_min, r_max = rewards_f32.min(), rewards_f32.max() + all_fail = homogeneous * (means <= r_min + 1e-12).float() + all_pass = homogeneous * (means >= r_max - 1e-12).float() + self.last_group_homogeneous = homogeneous[group_ids] + # <2 valid siblings: calculate_baseline_and_std_per_prompt falls back + # to baseline = reward, i.e. a silent Y = 0. Surface it. + singletons = int((counts < 2).sum().item()) + + m = return_mask.bool() + target_var = ( + returns[m].float().var(unbiased=False).item() + if int(m.sum()) > 1 + else 0.0 + ) + metrics = { + "residual/b_loo_mean": baseline.mean().item(), + "residual/b_loo_std": baseline.std(unbiased=False).item(), + "residual/target_var": target_var, + "residual/frac_groups_mixed": mixed.mean().item(), + "residual/frac_groups_all_fail": all_fail.mean().item(), + "residual/frac_groups_all_pass": all_pass.mean().item(), + "residual/n_singleton_groups": float(singletons), + "residual/group_size_min": counts.min().item(), + "residual/group_size_max": counts.max().item(), + } + if sample_mask is not None: + metrics["residual/frac_traj_sample_masked"] = ( + 1.0 - sample_mask.float().mean().item() + ) + if singletons: + print( + f" ⚠️ residual baseline: {singletons} group(s) have <2 rollouts; " + "their targets collapse to 0 (baseline = own reward)." + ) + return metrics + + +def attach_value_baseline_keys(batch: Any, adv_estimator: Any) -> None: + """Copy the estimator's per-sample return-space offsets onto a critic batch. + + No-op unless a :class:`ResidualBaselineEstimator` produced them. These let + :class:`MseValueLossFn` report explained variance in BOTH absolute and + residual space from one pass, without the loss needing to know which space + ``returns`` is in. + """ + to_abs = getattr(adv_estimator, "last_returns_to_abs", None) + to_res = getattr(adv_estimator, "last_returns_to_res", None) + if to_abs is None or to_res is None: + return + n = batch["returns"].shape[0] + if to_abs.shape[0] != n: + raise ValueError( + f"Return-space offsets have batch size {to_abs.shape[0]} but the " + f"critic batch has {n}. These are per-sample and would misalign " + "silently, reporting explained variance against the wrong baseline." + ) + batch["returns_to_abs"] = to_abs.to(batch["returns"].dtype) + batch["returns_to_res"] = to_res.to(batch["returns"].dtype) + + +def homogeneous_group_sample_mask( + sample_mask: torch.Tensor, adv_estimator: Any, weight: float +) -> torch.Tensor | None: + """``sample_mask`` rescaled so homogeneous groups carry ``weight``. + + Returns None when there is nothing to do (weight 1.0, or no residual + estimator), so callers can skip building a separate critic batch entirely. + + Under a residual target, all-fail and all-pass groups have ``Y = 0`` for every + sibling and so contribute exactly zero target variance -- on the pi0 SWE pool + that is 56.5% of groups and ~58% of critic FLOPs. They are still worth + keeping at some weight: the shrinkage-toward-zero they impose is the + mechanism enforcing ``E[C | X] = 0``, i.e. the regulariser against the + between-task leakage this whole change exists to remove. + + Rescaling ``sample_mask`` is a correctly renormalised weighted mean with no + effective-LR confound, because ``global_valid_toks`` is itself computed as + ``sum(token_mask * sample_mask)`` (nemo_rl/models/megatron/data.py) -- so the + denominator moves with the numerator. + + The caller MUST apply this to a separate critic batch: in token-level mode + the actor shares ``train_data['sample_mask']``. + """ + if weight == 1.0: + return None + if weight < 0.0: + raise ValueError( + f"value_loss_fn.homogeneous_group_weight must be >= 0, got {weight}." + ) + homogeneous = getattr(adv_estimator, "last_group_homogeneous", None) + if homogeneous is None: + raise ValueError( + "value_loss_fn.homogeneous_group_weight != 1.0 requires the residual " + "baseline estimator (it is what identifies homogeneous groups), but " + f"{type(adv_estimator).__name__} exposes no group composition." + ) + # The wrapper is installed even for an absolute-critic run (it computes + # B_LOO for metrics), so group composition is available there too -- but + # downweighting must NOT apply. Under an absolute target those groups have + # Y = R != 0 and do carry target variance, so silently reweighting them + # would change the objective of the control arm in an A/B that sets this + # knob in both arms to hold wall-clock fixed. + if not getattr(adv_estimator, "residual_target", False): + raise ValueError( + "value_loss_fn.homogeneous_group_weight != 1.0 is only meaningful " + "with ppo.adv_estimator.residual_baseline=true. Under an absolute " + "critic target, homogeneous groups still carry target variance " + "(Y = R != 0), so downweighting them would silently change the " + "objective rather than skip empty targets." + ) + homogeneous = homogeneous.to(sample_mask.device, sample_mask.dtype) + return sample_mask * (1.0 - homogeneous * (1.0 - weight)) + + class OPDAdvantageEstimator: """Multi-Teacher On-Policy Distillation (MOPD) advantage estimator (arXiv:2601.02780). diff --git a/nemo_rl/algorithms/critic_pretrain.py b/nemo_rl/algorithms/critic_pretrain.py index c44e85c2b97..d7f6e6a871c 100644 --- a/nemo_rl/algorithms/critic_pretrain.py +++ b/nemo_rl/algorithms/critic_pretrain.py @@ -147,6 +147,86 @@ def terminal_value_reward_auc( return float(auc) +def _rank_auc(scores: torch.Tensor, labels: torch.Tensor) -> float: + """Mann-Whitney AUC with tie correction; nan on a single-class input.""" + n_pos, n_neg = int(labels.sum()), int((1 - labels).sum()) + if n_pos == 0 or n_neg == 0: + return float("nan") + order = torch.argsort(scores) + ranks = torch.empty_like(scores) + ranks[order] = torch.arange(1, scores.numel() + 1, dtype=scores.dtype) + for val in torch.unique(scores): + tie = scores == val + if int(tie.sum()) > 1: + ranks[tie] = ranks[tie].mean() + return float( + (ranks[labels.bool()].sum().item() - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg) + ) + + +def within_group_auc( + values: torch.Tensor, + rewards: torch.Tensor, + group_ids: torch.Tensor, + token_mask: torch.Tensor, + n_buckets: int = 4, + positive_threshold: float = 0.5, +) -> dict[str, float]: + """Does the critic rank a group's WINNING siblings above its losing ones? + + ``terminal_value_reward_auc`` pools every trajectory in the batch, so it is + dominated by between-task variation — a critic that only knows "this issue is + hopeless" scores well on it while carrying no within-task information at all. + That is exactly the failure mode the residual target is meant to remove, so + the go/no-go diagnostic has to hold the task fixed. + + Restricted to MIXED-outcome groups (homogeneous ones are undefined: no + positive or no negative sibling). On the pi0 SWE pool only 43.5% of groups + qualify, and that fraction — not the dataset size — bounds what any critic + trained on terminal reward can learn about the within-task component. + + Scores are per-trajectory means of ``values`` over each progress bucket, so + this is calibration-free: it survives the scale/offset errors that depress + explained variance. + + Returns per-bucket mean AUC plus ``n_mixed_groups``. Buckets run earliest to + latest by relative position within each trajectory's own response. + """ + mask = token_mask.bool() + rel = (torch.cumsum(mask.long(), dim=1) - 1).float() / mask.sum( + dim=1, keepdim=True + ).clamp(min=1).float() + labels = (rewards.float() >= positive_threshold).float() + + out: dict[str, float] = {} + n_mixed = 0 + for b in range(n_buckets): + lo, hi = b / n_buckets, (b + 1) / n_buckets + upper = (rel < hi) if b < n_buckets - 1 else (rel <= 1.0) + bmask = mask & (rel >= lo) & upper + counts = bmask.sum(dim=1) + # Per-trajectory mean value inside this progress bucket. + scores = (values * bmask).sum(dim=1) / counts.clamp(min=1) + aucs = [] + for gid in torch.unique(group_ids): + sel = (group_ids == gid) & (counts > 0) + if int(sel.sum()) < 2: + continue + y = labels[sel] + if int(y.sum()) == 0 or int((1 - y).sum()) == 0: + continue # homogeneous group: within-group AUC undefined + auc = _rank_auc(scores[sel].float(), y) + if auc == auc: # not nan + aucs.append(auc) + if b == 0: + n_mixed = len(aucs) + out[f"critic/within_group_auc_q{b + 1}"] = ( + sum(aucs) / len(aucs) if aucs else float("nan") + ) + out["critic/n_mixed_groups"] = float(n_mixed) + return out + + def verify_shard_meta( shards_dir: str | Path, master_config: Any, tokenizer: Any ) -> None: @@ -324,6 +404,7 @@ def _forward_values_and_returns( repeated_batch: BatchedDataDict, tokenizer: Any, master_config: Any, + ) -> tuple[Optional[BatchedDataDict], Optional[Any]]: """Populate train_data['values'/'returns'] in place. @@ -345,9 +426,11 @@ def _forward_values_and_returns( remap_by_response_mask, ) + privileged_critic_cfg = master_config.value.get("privileged_critic") if privileged_critic_cfg is not None and not privileged_critic_cfg.get("enabled"): privileged_critic_cfg = None + critic_batch = None # Turn structure (None on the token-level path). Stage B must build this the @@ -365,6 +448,7 @@ def _forward_values_and_returns( "make_sequence_length_divisible_by" ], ) + vals_aug = value_model.get_values(critic_batch)["values"].squeeze(-1) critic_batch["values"] = vals_aug train_data["values"] = remap_by_response_mask( @@ -427,12 +511,17 @@ def _heldout_metrics( ~270 structural zeros — and the per-turn metrics from the estimator are merged in. """ - from nemo_rl.algorithms.ppo import _positional_value_metrics + from nemo_rl.algorithms.ppo import ( + _mixed_group_mask, + _mixed_group_value_metrics, + _positional_value_metrics, + ) groups = [load_group(p) for p in heldout_files] train_data, repeated_batch = build_value_train_data( groups, tokenizer, master_config ) + _, turn_spans = _forward_values_and_returns( value_model, adv_estimator, @@ -440,6 +529,7 @@ def _heldout_metrics( repeated_batch, tokenizer, master_config, + ) values, returns = train_data["values"], train_data["returns"] scored_mask = ( @@ -447,20 +537,83 @@ def _heldout_metrics( ) mask = scored_mask.bool() metrics: dict[str, float] = {} + + # Per-sample offsets into each return space (both zero without a residual + # estimator, i.e. exactly today's numbers). + zeros = torch.zeros(returns.shape[0], device=returns.device) + raw_to_abs = getattr(adv_estimator, "last_returns_to_abs", None) + raw_to_res = getattr(adv_estimator, "last_returns_to_res", None) + to_abs = zeros if raw_to_abs is None else raw_to_abs.to(returns.device) + to_res = zeros if raw_to_res is None else raw_to_res.to(returns.device) + if int(mask.sum()) >= 2: v, r = values[mask].float(), returns[mask].float() - var_r = r.var(unbiased=False) - metrics["critic/explained_var"] = ( - (1.0 - (r - v).var(unbiased=False) / var_r).item() if var_r > 1e-8 else 0.0 - ) + # Both explained variances, on the same convention _compute_critic_metrics + # uses in PPO: critic/explained_var is ALWAYS absolute-space and + # critic/ev_res ALWAYS residual-space, whichever space `returns` is in. + # The prediction error is shared (R - (B+C) == (R-B) - C); only the + # denominator changes. Held-out ev_res is the go/no-go number. + err_var = (r - v).var(unbiased=False) + for key, offset in (("explained_var", to_abs), ("ev_res", to_res)): + target = (returns + offset.unsqueeze(-1).to(returns.dtype))[mask].float() + var_t = target.var(unbiased=False) + metrics[f"critic/{key}"] = ( + (1.0 - err_var / var_t).item() if var_t > 1e-8 else 0.0 + ) metrics["critic/mse"] = ((r - v) ** 2).mean().item() - metrics.update(_positional_value_metrics(values, returns, scored_mask)) + metrics.update( + _positional_value_metrics( + values, + returns, + scored_mask, + returns_to_abs=raw_to_abs, + returns_to_res=raw_to_res, + ) + ) + # Residual EV restricted to mixed-outcome groups. critic/ev_res stays the + # whole-batch go/no-go number; this says whether a near-zero ev_res means + # "no within-task signal" or "signal, taxed by the ~58% homogeneous groups + # where Y = 0 and any prediction is a pure penalty". + metrics.update( + _mixed_group_value_metrics( + values, + returns, + scored_mask, + _mixed_group_mask(adv_estimator), + returns_to_res=raw_to_res, + ) + ) # Scored on `scored_mask`: in turn mode the last RESPONSE token carries an # untrained value, while the last anchor is the supervised V(s_K). + # + # Scored on ABSOLUTE values (V~ = C + B_LOO), like every other metric here. + # This AUC pools all trajectories, so it is largely a between-task ranking; + # in residual space the values are C with E[C | X] = 0, which strips exactly + # that component out and would read as a large regression versus the + # absolute arm when nothing regressed. + abs_values = values + to_abs.unsqueeze(-1).to(values.dtype) metrics["critic/terminal_auc"] = terminal_value_reward_auc( - values, train_data["rewards"], scored_mask + abs_values, train_data["rewards"], scored_mask ) + # Sibling ranking with the task held fixed — the calibration-free go/no-go + # complement to explained variance, and the only AUC that is not confounded + # by between-task difficulty. + # + # Deliberately scored on RAW values, unlike terminal_auc above. B_LOO is + # leave-one-out, so it is NOT constant within a group: adding it would fold + # each sibling's own reward into that sibling's score and leak the label, + # inflating this AUC. Raw values are already the right quantity in both arms + # (C in residual mode, V in absolute mode), since the task-level component is + # common to the group and cancels from a within-group ranking either way. + group_ids = getattr(adv_estimator, "last_group_ids", None) + if group_ids is not None: + metrics.update( + within_group_auc( + values, train_data["rewards"], group_ids.cpu(), scored_mask + ) + ) metrics.update(getattr(adv_estimator, "last_metrics", {}) or {}) + metrics["reward"] = train_data["rewards"].float().mean().item() metrics["num_heldout_samples"] = float(train_data["input_ids"].shape[0]) return metrics @@ -571,7 +724,10 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: from nemo_rl.algorithms.ppo import ( _compute_critic_metrics, _create_advantage_estimator, + _mixed_group_mask, + _mixed_group_value_metrics, _positional_value_metrics, + _prepare_value_train_batch, _resolve_resume_optimizer_path, ) from nemo_rl.distributed.virtual_cluster import RayVirtualCluster @@ -848,6 +1004,7 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: ) print("▶ Computing values...") + critic_batch, turn_spans = _forward_values_and_returns( value_model, adv_estimator, @@ -855,11 +1012,19 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: repeated_batch, tokenizer, master_config, + ) print("▶ Training critic...") value_model.prepare_for_training() value_train_batch = critic_batch if critic_batch is not None else train_data + # Same residual bookkeeping the PPO loops apply: without it the value + # loss sees no return-space offsets and critic/ev_res silently + # duplicates critic/explained_var, and homogeneous_group_weight would be + # a no-op that the launcher nonetheless advertises. + value_train_batch = _prepare_value_train_batch( + value_train_batch, adv_estimator, master_config + ) value_results = value_model.train(value_train_batch, value_loss_fn) value_model.finish_training() @@ -877,9 +1042,23 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: turn_spans.anchor_mask if turn_spans is not None else train_data["token_mask"], + returns_to_abs=getattr(adv_estimator, "last_returns_to_abs", None), + returns_to_res=getattr(adv_estimator, "last_returns_to_res", None), + ) + ) + metrics.update( + _mixed_group_value_metrics( + train_data["values"], + train_data["returns"], + turn_spans.anchor_mask + if turn_spans is not None + else train_data["token_mask"], + _mixed_group_mask(adv_estimator), + returns_to_res=getattr(adv_estimator, "last_returns_to_res", None), ) ) metrics.update(getattr(adv_estimator, "last_metrics", {}) or {}) + metrics["reward"] = train_data["rewards"].float().mean().item() metrics["num_samples"] = float(train_data["input_ids"].shape[0]) metrics["total_step_time"] = time.perf_counter() - step_start diff --git a/nemo_rl/algorithms/loss/loss_functions.py b/nemo_rl/algorithms/loss/loss_functions.py index 5b8b0bf9cd0..6487f98c100 100755 --- a/nemo_rl/algorithms/loss/loss_functions.py +++ b/nemo_rl/algorithms/loss/loss_functions.py @@ -1209,6 +1209,14 @@ class MseValueLossConfig(BaseModel, extra="forbid"): scale: float = 1.0 # Clipping range for value predictions (PPO-style). Set to None to disable clipping. cliprange: Optional[float] = None + # Relative weight of homogeneous (all-fail / all-pass) groups in the value loss. + # Only meaningful with adv_estimator.residual_baseline: those groups have + # Y = 0 for every sibling and so contribute EXACTLY zero target variance, + # while still costing ~58% of critic FLOPs on the SWE pool. 1.0 (default) + # keeps every sample -- they are not dead weight, they are the shrinkage + # that enforces E[C | X] = 0 and suppresses between-task leakage. Lower it + # only to buy more mixed-group exposure per unit wall-clock. + homogeneous_group_weight: float = 1.0 class MseValueLossFn(LossFunction): @@ -1322,6 +1330,37 @@ def __call__( global_normalization_factor=global_valid_toks, ).item() + # Explained variance in BOTH return spaces. The numerator is shared: + # the prediction error is the same quantity either way, since + # R - (B_LOO + C) = (R - B_LOO) - C = Y - C, + # so only the denominator (Var of the absolute vs residual return) + # differs. The estimator hands us the two per-sample offsets that + # convert `returns` into each space -- exactly one is zero -- which + # keeps this loss agnostic to which mode the run is in. Absent (an + # ordinary GAE run with no residual estimator), both are zero and + # both variances collapse to Var(returns), i.e. today's behaviour. + zeros = torch.zeros( + returns.shape[0], device=returns.device, dtype=returns.dtype + ) + to_abs = data.get("returns_to_abs") + to_res = data.get("returns_to_res") + to_abs = zeros if to_abs is None else to_abs.to(returns.dtype) + to_res = zeros if to_res is None else to_res.to(returns.dtype) + abs_returns = returns + to_abs.unsqueeze(-1) + res_returns = returns + to_res.unsqueeze(-1) + abs_returns_mean = masked_mean( + abs_returns, mask, global_normalization_factor=global_valid_toks + ).item() + abs_returns_sq_mean = masked_mean( + abs_returns**2, mask, global_normalization_factor=global_valid_toks + ).item() + res_returns_mean = masked_mean( + res_returns, mask, global_normalization_factor=global_valid_toks + ).item() + res_returns_sq_mean = masked_mean( + res_returns**2, mask, global_normalization_factor=global_valid_toks + ).item() + metrics = { "loss": float(loss.item()), "vf_clipfrac": vf_clipfrac, @@ -1331,6 +1370,10 @@ def __call__( "values_max": values_max, "returns_sq_mean": returns_sq_mean, "residual_sq_mean": residual_sq_mean, + "abs_returns_mean": abs_returns_mean, + "abs_returns_sq_mean": abs_returns_sq_mean, + "res_returns_mean": res_returns_mean, + "res_returns_sq_mean": res_returns_sq_mean, "num_valid_samples": int(values.shape[0]), } diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 1e8ce65dd0e..05f65e2a35c 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -31,7 +31,10 @@ from nemo_rl.algorithms.advantage_estimator import ( GeneralizedAdvantageEstimator, RawRewardAdvantageEstimator, + ResidualBaselineEstimator, TurnLevelGeneralizedAdvantageEstimator, + attach_value_baseline_keys, + homogeneous_group_sample_mask, ) from nemo_rl.algorithms.grpo import ( RewardPenaltyConfig, @@ -202,6 +205,14 @@ class AdvEstimatorConfig(TypedDict): turn_gae_gamma: NotRequired[Optional[float]] turn_gae_lambda_value: NotRequired[Optional[float]] turn_gae_lambda_policy: NotRequired[Optional[float]] + # Decomposed group baseline + residual critic + # (research/ppo/residual_critic_report.md). The rollout group supplies the + # task baseline B(X) via a leave-one-out mean and the critic is trained only + # on the within-task residual C(s) = V*(s) - B(X). Requires gamma = 1. + # Independent of token-vs-turn granularity: it wraps either estimator. + # Required for the value-model estimators ("gae"/"turn_gae"), which read it + # directly; "raw_reward" trains no critic and may omit it. + residual_baseline: NotRequired[bool] class PPOConfig(TypedDict): @@ -1443,6 +1454,65 @@ def _create_advantage_estimator(master_config: MasterConfig): f"PPO only supports 'gae', 'turn_gae' or 'raw_reward'." ) + if adv_estimator_name == "raw_reward": + # Read only as a presence check: raw_reward configs have no critic, so + # they are not expected to carry the key at all. + if adv_estimator_config.get("residual_baseline"): + raise ValueError( + "adv_estimator.residual_baseline requires a value model, but " + "adv_estimator.name='raw_reward' trains no critic." + ) + return adv_estimator + + # Explicit message rather than a bare KeyError, matching how the turn_gae + # lambdas reject a config that silently omits them. + if "residual_baseline" not in adv_estimator_config: + raise ValueError( + "adv_estimator.residual_baseline must be set explicitly when " + f"adv_estimator.name={adv_estimator_name!r} (no default is assumed). " + "Set it to false to keep the current absolute-critic behaviour. " + "See research/ppo/residual_critic_report.md." + ) + residual_baseline = bool(adv_estimator_config["residual_baseline"]) + + # G = 1 makes the leave-one-out baseline degenerate: with a single sibling, + # calculate_baseline_and_std_per_prompt falls back to baseline = own reward, + # so B_LOO == R exactly. The advantage then telescopes to A_t = -C_t and the + # target to Y = 0 -- the reward drops out of the policy gradient entirely and + # the critic is trained to predict zero. That is silently wrong rather than + # merely degraded, so refuse it at setup. + if residual_baseline: + gens_per_prompt = ppo_config["num_generations_per_prompt"] + if gens_per_prompt < 2: + raise ValueError( + "adv_estimator.residual_baseline requires " + f"ppo.num_generations_per_prompt >= 2, got {gens_per_prompt}. " + "With one rollout per prompt the leave-one-out baseline " + "degenerates to the rollout's own reward (B_LOO == R), so every " + "advantage becomes -C(s) and every critic target becomes 0: the " + "reward would be dropped from the policy gradient without any " + "error." + ) + + # Always wrap the value-based estimators, even when residualization is OFF: + # the wrapper then computes B_LOO for metrics only, which is what lets an + # absolute-critic run report critic/ev_res on the same axis as a residual + # run (there, 1 - ev_res is exactly the advantage-variance ratio the report + # measured at 1.67-2.09x). It never touches the targets in that mode. + adv_estimator = ResidualBaselineEstimator(adv_estimator, residual_baseline) + if residual_baseline: + print( + " ✓ RESIDUAL baseline ON: the rollout group supplies the task " + "baseline B_LOO and the critic is trained only on the within-task " + "residual C(s) = R - B_LOO. Critic values/returns are in RESIDUAL " + "space; go/no-go metric is critic/ev_res." + ) + else: + print( + " ✓ Residual baseline OFF (absolute critic targets); B_LOO is still " + "computed for critic/ev_res + residual/* diagnostics." + ) + return adv_estimator @@ -1501,6 +1571,36 @@ def _value_metric_mask( return turn_spans.anchor_mask +def _prepare_value_train_batch( + value_train_batch: BatchedDataDict[Any], + adv_estimator: Any, + master_config: MasterConfig, +) -> BatchedDataDict[Any]: + """Attach residual bookkeeping to the batch the critic actually trains on. + + Two things, both no-ops for a default run: + + 1. ``returns_to_abs`` / ``returns_to_res`` so :class:`MseValueLossFn` can + report explained variance in both return spaces. + 2. ``value_loss_fn.homogeneous_group_weight``, applied by rescaling + ``sample_mask``. This COPIES rather than mutates: in token-level mode the + critic trains on ``train_data`` itself, whose ``sample_mask`` the actor + also reads. + + Called after the privileged-critic path re-assigns ``critic_batch['sample_mask']``, + so the weighting cannot be silently clobbered. + """ + weight = master_config.value_loss_fn.homogeneous_group_weight + scaled = homogeneous_group_sample_mask( + value_train_batch["sample_mask"], adv_estimator, weight + ) + if scaled is not None: + value_train_batch = BatchedDataDict({**value_train_batch}) + value_train_batch["sample_mask"] = scaled + attach_value_baseline_keys(value_train_batch, adv_estimator) + return value_train_batch + + def _compute_critic_metrics(value_results: dict[str, Any]) -> dict[str, Any]: """Build the ``critic/*`` metrics dict from a value model's train results. @@ -1540,7 +1640,40 @@ def _compute_critic_metrics(value_results: dict[str, Any]) -> dict[str, Any]: res_sq = critic_metrics.get("critic/residual_sq_mean", 0) var_returns = r_sq - r_mean**2 var_residual = res_sq - (r_mean - v_mean) ** 2 - critic_metrics["critic/explained_var"] = 1.0 - var_residual / max(var_returns, 1e-8) + + # The prediction error is the SAME quantity in both return spaces -- + # R - (B_LOO + C) = (R - B_LOO) - C -- so `var_residual` above serves as the + # numerator for both explained variances and only the denominator changes: + # + # critic/explained_var = 1 - Var(R - V~) / Var(R) absolute space + # critic/ev_res = 1 - Var(Y - C) / Var(Y) residual space + # + # Both are reported in BOTH arms so an absolute and a residual run are + # directly comparable on one axis. critic/explained_var keeps its historical + # meaning exactly; for an absolute critic 1 - critic/ev_res is the + # advantage-variance ratio against a group-only advantage (report S14.5 + # measured 1.67-2.09x, i.e. ev_res well below zero). Runs predating these + # stats fall back to the single-space computation. + abs_mean = critic_metrics.get("critic/abs_returns_mean") + abs_sq = critic_metrics.get("critic/abs_returns_sq_mean") + res_r_mean = critic_metrics.get("critic/res_returns_mean") + res_r_sq = critic_metrics.get("critic/res_returns_sq_mean") + if abs_mean is None or res_r_mean is None: + var_abs_returns, var_res_returns = var_returns, var_returns + else: + var_abs_returns = abs_sq - abs_mean**2 + var_res_returns = res_r_sq - res_r_mean**2 + + critic_metrics["critic/explained_var"] = 1.0 - var_residual / max( + var_abs_returns, 1e-8 + ) + # Unlike Var(R), Var(Y) legitimately hits zero when every group in a step is + # homogeneous (54% of groups are all-fail on the SWE pool), and 1e-8 would + # then emit ~-1e8 and destroy the panel's autoscale. Report 0.0 -- "the + # critic explains none of a target that has nothing to explain". + critic_metrics["critic/ev_res"] = ( + 1.0 - var_residual / var_res_returns if var_res_returns > 1e-8 else 0.0 + ) return critic_metrics @@ -1571,9 +1704,12 @@ def _positional_value_metrics( returns: torch.Tensor, token_mask: torch.Tensor, n_bins: int = 3, + returns_to_abs: Optional[torch.Tensor] = None, + returns_to_res: Optional[torch.Tensor] = None, ) -> dict[str, float]: - """Critic-quality diagnostics bucketed by RELATIVE position within each response - (early / mid / late thirds): explained variance, calibration (ECE), and signed + """Critic-quality diagnostics bucketed by relative position within each response. + + Early / mid / late thirds: explained variance, calibration (ECE), and signed bias of V(s_t) vs the GAE return. EV is expected to rise early->late (the outcome is nearly determined near the @@ -1584,10 +1720,31 @@ def _positional_value_metrics( calibration axis: V should literally be P(success | prefix), and miscalibration distorts advantage magnitudes even when EV/ranking looks fine. Cheap: driver-side tensor ops on tensors already in ``train_data``. + + ``returns_to_abs`` / ``returns_to_res`` are the per-sample ``[B]`` offsets + from :class:`ResidualBaselineEstimator` that shift the whole space (values + AND returns together), so each bucket reports explained variance in both: + ``critic/ev_*`` stays absolute-space and unchanged, ``critic/ev_res_*`` is + the residual-space counterpart. ECE is scored in absolute space in both + arms, since it bins on ``[0,1]`` and only means anything for a predictor of + ``P(success)`` -- a residual value lives in ``[-1,1]`` around zero. + + Because the numerator ``returns - values`` is offset-invariant, ``abs_err`` + and ``bias`` are identical in both spaces and are reported once. """ mask = token_mask.bool() if int(mask.sum()) < 2: return {} + zeros = torch.zeros(returns.shape[0], device=returns.device, dtype=returns.dtype) + to_abs = zeros if returns_to_abs is None else returns_to_abs.to(returns.dtype) + to_res = zeros if returns_to_res is None else returns_to_res.to(returns.dtype) + abs_values, abs_returns = ( + values + to_abs.unsqueeze(-1), + returns + to_abs.unsqueeze(-1), + ) + # Only the residual RETURNS are needed: the prediction error is + # offset-invariant, so the residual values never enter any statistic here. + res_returns = returns + to_res.unsqueeze(-1) # 0-based index of each response token within its sample's response, scaled to [0,1) resp_len = mask.sum(dim=1, keepdim=True).clamp(min=1) rel = (torch.cumsum(mask.long(), dim=1) - 1).float() / resp_len.float() @@ -1604,19 +1761,126 @@ def _positional_value_metrics( continue v = values[bmask].float() r = returns[bmask].float() - var_r = r.var(unbiased=False) - ev = (1.0 - (r - v).var(unbiased=False) / var_r).item() if var_r > 1e-8 else 0.0 - out[f"critic/ev_{names[i]}"] = ev + # Offset-invariant: (r - v) is the same in both spaces, so one numerator + # serves both explained variances. + err_var = (r - v).var(unbiased=False) + av, ar = abs_values[bmask].float(), abs_returns[bmask].float() + rr = res_returns[bmask].float() + var_abs, var_res = ar.var(unbiased=False), rr.var(unbiased=False) + out[f"critic/ev_{names[i]}"] = ( + (1.0 - err_var / var_abs).item() if var_abs > 1e-8 else 0.0 + ) + out[f"critic/ev_res_{names[i]}"] = ( + (1.0 - err_var / var_res).item() if var_res > 1e-8 else 0.0 + ) out[f"critic/abs_err_{names[i]}"] = (r - v).abs().mean().item() + # Raw head output, NOT reconstructed: in residual mode this is C(s), and + # whether it actually spans a signed range is the thing worth watching. out[f"critic/mean_v_{names[i]}"] = v.mean().item() out[f"critic/n_tokens_{names[i]}"] = float(n) # Calibration: ECE (magnitude of miscalibration across confidence bins) and # signed bias (direction: >0 = overconfident/optimistic, <0 = pessimistic). - out[f"critic/ece_{names[i]}"] = _calibration_ece(v, r) + # Scored in absolute space -- ECE bins on [0,1] and is only meaningful for + # a predictor of P(success). + out[f"critic/ece_{names[i]}"] = _calibration_ece(av, ar) out[f"critic/bias_{names[i]}"] = (v - r).mean().item() return out +def _mixed_group_value_metrics( + values: torch.Tensor, + returns: torch.Tensor, + token_mask: torch.Tensor, + mixed_mask: Optional[torch.Tensor], + returns_to_res: Optional[torch.Tensor] = None, + n_bins: int = 3, +) -> dict[str, float]: + """Residual explained variance restricted to MIXED-outcome groups. + + ``critic/ev_res`` is computed over the whole batch, where homogeneous + (all-fail / all-pass) groups have ``Y = 0`` for every sibling. Those groups + therefore contribute ZERO covariance but a positive ``Var(C)`` to + + EV_res = [2 Cov(Y, C) - Var(C)] / Var(Y), + + i.e. any nonzero prediction there is a pure penalty. On this SWE pool they + are ~58% of groups, so a critic with real within-task signal on the mixed + groups can still show ``ev_res`` near zero once that tax is paid. + + These keys isolate the other side of that split: how much residual variance + the critic explains where the target is actually nonzero. Reading them + together separates two very different states of the world: + + ev_res ~ 0, ev_res_mixed_group > 0 signal exists; homogeneous groups + are taxing it (consider gating the + critic to mixed groups in PPO) + ev_res ~ 0, ev_res_mixed_group ~ 0 no within-task signal to find + + This is diagnostic only and deliberately does NOT change ``critic/ev_res``, + which remains the whole-batch go/no-go number: PPO applies the critic to + every group, so the untruncated version is what the actor actually sees. + + Args: + values: ``[B, S]`` raw critic output (``C`` in residual mode, ``V`` in + absolute mode). + returns: ``[B, S]`` critic regression target, in whichever space the run + uses. + token_mask: ``[B, S]`` positions the critic is supervised at. + mixed_mask: ``[B]``, nonzero for rollouts whose sibling group contains + both a success and a failure. ``None`` -> no keys are emitted. + returns_to_res: ``[B]`` offset carrying ``returns`` into residual space + (zero in a residual run, ``-B_LOO`` in an absolute one). + n_bins: trajectory-progress buckets, matching _positional_value_metrics. + """ + if mixed_mask is None: + return {} + mask = token_mask.bool() & mixed_mask.bool().unsqueeze(-1) + if int(mask.sum()) < 2: + return {} + zeros = torch.zeros(returns.shape[0], device=returns.device, dtype=returns.dtype) + to_res = zeros if returns_to_res is None else returns_to_res.to(returns.dtype) + res_returns = returns + to_res.unsqueeze(-1) + + def _ev(m: torch.Tensor) -> Optional[float]: + if int(m.sum()) < 2: + return None + # The prediction error is offset-invariant, so the raw difference is + # already the residual-space error; only the denominator needs shifting. + err_var = (returns[m] - values[m]).float().var(unbiased=False) + var_t = res_returns[m].float().var(unbiased=False) + return (1.0 - err_var / var_t).item() if var_t > 1e-8 else 0.0 + + out: dict[str, float] = {} + overall = _ev(mask) + if overall is not None: + out["critic/ev_res_mixed_group"] = overall + out["critic/n_mixed_group_tokens"] = float(int(mask.sum())) + + # Bucket by relative position within each response, using the FULL response + # mask for the denominator so bucket boundaries match _positional_value_metrics + # exactly (a trajectory's progress does not depend on its group's composition). + full = token_mask.bool() + rel = (torch.cumsum(full.long(), dim=1) - 1).float() / full.sum( + dim=1, keepdim=True + ).clamp(min=1).float() + names = ( + ["early", "mid", "late"] if n_bins == 3 else [f"bin{i}" for i in range(n_bins)] + ) + for i in range(n_bins): + lo, hi = i / n_bins, (i + 1) / n_bins + upper = (rel < hi) if i < n_bins - 1 else (rel <= 1.0) + ev = _ev(mask & (rel >= lo) & upper) + if ev is not None: + out[f"critic/ev_res_mixed_group_{names[i]}"] = ev + return out + + +def _mixed_group_mask(adv_estimator: Any) -> Optional[torch.Tensor]: + """``[B]`` 1.0 for rollouts in a mixed-outcome group, or None if unavailable.""" + homogeneous = getattr(adv_estimator, "last_group_homogeneous", None) + return None if homogeneous is None else 1.0 - homogeneous + + def _should_log_ppo_rollout_dump(master_config: MasterConfig, step: int) -> bool: """Whether to write the packed per-token rollout dump for this (1-based) step.""" if not master_config.ppo.get("log_rollout_dump"): @@ -2116,6 +2380,7 @@ def ppo_train( and not privileged_critic_cfg.get("enabled") ): privileged_critic_cfg = None + critic_batch = None value_model.prepare_for_inference() @@ -2128,6 +2393,7 @@ def ppo_train( "make_sequence_length_divisible_by" ], ) + vals_aug = value_model.get_values(critic_batch)[ "values" ].squeeze(-1) @@ -2255,6 +2521,10 @@ def ppo_train( train_data["advantages"] = advantages if returns is not None: train_data["returns"] = returns + # Return-space offsets for the pre-update (fresh) critic + # diagnostics below; the critic's own batch gets them in + # _prepare_value_train_batch. + attach_value_baseline_keys(train_data, adv_estimator) # Turn-level: the critic trains on one anchor per turn, # so it needs its own batch (same sequences, anchor mask). if turn_spans is not None: @@ -2319,6 +2589,12 @@ def ppo_train( value_train_batch = critic_batch else: value_train_batch = train_data + # Residual bookkeeping last, so the sample_mask + # re-assignment above cannot clobber the homogeneous + # group weighting. + value_train_batch = _prepare_value_train_batch( + value_train_batch, adv_estimator, master_config + ) value_results = value_model.train( value_train_batch, value_loss_fn, @@ -2457,12 +2733,27 @@ def ppo_train( train_data["values"], train_data["returns"], _value_metric_mask(train_data, turn_spans), + returns_to_abs=train_data.get("returns_to_abs"), + returns_to_res=train_data.get("returns_to_res"), + ) + ) + # Residual EV where the target is actually nonzero: on a + # homogeneous group Y = 0, so any prediction there is a + # pure penalty to the whole-batch critic/ev_res. + metrics.update( + _mixed_group_value_metrics( + train_data["values"], + train_data["returns"], + _value_metric_mask(train_data, turn_spans), + _mixed_group_mask(adv_estimator), + returns_to_res=train_data.get("returns_to_res"), ) ) # Turn-level counterparts (per-turn EV/bias, last-turn AUC, # turn counts). The token-level versions above are dominated # by the longest turns; these weight every decision equally. metrics.update(getattr(adv_estimator, "last_metrics", {}) or {}) + metrics.update( { "reward": rewards.numpy(), @@ -3352,6 +3643,7 @@ def _sample_max_age(s: int) -> int: and not privileged_critic_cfg.get("enabled") ): privileged_critic_cfg = None + critic_batch = None value_model.prepare_for_inference() @@ -3364,6 +3656,7 @@ def _sample_max_age(s: int) -> int: "make_sequence_length_divisible_by" ], ) + vals_aug = value_model.get_values(critic_batch)[ "values" ].squeeze(-1) @@ -3458,6 +3751,10 @@ def _sample_max_age(s: int) -> int: train_data["advantages"] = advantages if returns is not None: train_data["returns"] = returns + # Return-space offsets for the pre-update (fresh) critic + # diagnostics below; the critic's own batch gets them in + # _prepare_value_train_batch. + attach_value_baseline_keys(train_data, adv_estimator) # Turn-level: the critic trains on one anchor per turn, # so it needs its own batch (same sequences, anchor mask). if turn_spans is not None: @@ -3501,6 +3798,12 @@ def _sample_max_age(s: int) -> int: value_train_batch = critic_batch else: value_train_batch = train_data + # Residual bookkeeping last, so the sample_mask + # re-assignment above cannot clobber the homogeneous + # group weighting. + value_train_batch = _prepare_value_train_batch( + value_train_batch, adv_estimator, master_config + ) value_results = value_model.train( value_train_batch, value_loss_fn, @@ -3675,6 +3978,20 @@ def _sample_max_age(s: int) -> int: train_data["values"], train_data["returns"], _value_metric_mask(train_data, turn_spans), + returns_to_abs=train_data.get("returns_to_abs"), + returns_to_res=train_data.get("returns_to_res"), + ) + ) + # Residual EV where the target is actually nonzero: on a + # homogeneous group Y = 0, so any prediction there is a + # pure penalty to the whole-batch critic/ev_res. + metrics.update( + _mixed_group_value_metrics( + train_data["values"], + train_data["returns"], + _value_metric_mask(train_data, turn_spans), + _mixed_group_mask(adv_estimator), + returns_to_res=train_data.get("returns_to_res"), ) ) # Turn-level counterparts (per-turn EV/bias, last-turn AUC, diff --git a/tests/unit/algorithms/test_advantage_estimator.py b/tests/unit/algorithms/test_advantage_estimator.py index e7cf031bbdd..98bd1581f74 100644 --- a/tests/unit/algorithms/test_advantage_estimator.py +++ b/tests/unit/algorithms/test_advantage_estimator.py @@ -12,9 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import torch -from nemo_rl.algorithms.advantage_estimator import OPDAdvantageEstimator +from nemo_rl.algorithms.advantage_estimator import ( + GeneralizedAdvantageEstimator, + OPDAdvantageEstimator, + ResidualBaselineEstimator, + TurnLevelGeneralizedAdvantageEstimator, + homogeneous_group_sample_mask, +) def _make_estimator(): @@ -107,3 +114,363 @@ def test_opd_metrics_returned(): ) assert abs(estimator.last_metrics["on_policy_distillation/adv_mean"] - 1.0) < 1e-5 assert abs(estimator.last_metrics["on_policy_distillation/adv_std"]) < 1e-5 + + +# =============================================================================== +# Residual baseline (decomposed group baseline + residual critic) +# =============================================================================== + + +class _LossCfg: + """Minimal stand-in for ClippedPGLossConfig (KL-in-reward off).""" + + use_kl_in_reward = False + reference_policy_kl_penalty = 0.0 + reference_policy_kl_type = "low_var_kl" + + +def _gae_cfg(gamma=1, lam=1.0, normalize=False): + return { + "gae_lambda": lam, + "gae_gamma": gamma, + "normalize_advantages": normalize, + "gae_lambda_value": None, + "gae_lambda_policy": None, + "length_adaptive_alpha": 0.0, + } + + +def _make_gae(**kwargs): + return GeneralizedAdvantageEstimator(_gae_cfg(**kwargs), _LossCfg()) + + +def _one_group(rewards, seq_len=6): + """One task group of ``len(rewards)`` siblings; every token is a valid response token.""" + b = len(rewards) + prompt_ids = torch.full((b, 3), 7) # identical prompt => a single group + return ( + prompt_ids, + torch.tensor(rewards, dtype=torch.float32), + torch.ones(b, seq_len), + ) + + +def test_residual_loo_arithmetic(): + """R=[1,0,0,0] => Y=[1,-1/3,-1/3,-1/3], summing to zero across siblings.""" + est = ResidualBaselineEstimator(_make_gae(), residual_target=True) + prompt_ids, rewards, mask = _one_group([1.0, 0.0, 0.0, 0.0]) + values = torch.zeros_like(mask) + + _, returns = est.compute_advantage(prompt_ids, rewards, mask, values) + + expected = torch.tensor([1.0, -1 / 3, -1 / 3, -1 / 3]) + torch.testing.assert_close(returns[:, 0], expected) + # Constant along the trajectory (sparse terminal reward, gamma = lambda = 1). + torch.testing.assert_close(returns, expected.unsqueeze(-1).expand_as(returns)) + assert abs(returns[:, 0].sum().item()) < 1e-6 + + +@pytest.mark.parametrize("reward", [0.0, 1.0]) +def test_residual_homogeneous_group_targets_are_zero(reward): + """All-fail and all-pass groups both give Y = 0 for every sibling. + + This is why the mixed-group fraction, not the dataset size, bounds what a + residual critic can learn: homogeneous groups contribute exactly zero target + variance. + """ + est = ResidualBaselineEstimator(_make_gae(), residual_target=True) + prompt_ids, rewards, mask = _one_group([reward] * 4) + values = torch.zeros_like(mask) + + _, returns = est.compute_advantage(prompt_ids, rewards, mask, values) + + torch.testing.assert_close(returns, torch.zeros_like(returns)) + assert est.last_metrics["residual/frac_groups_mixed"] == 0.0 + + +def test_residual_matches_absolute_gae_shifted_by_baseline(): + """The correctness contract: residualization is a pure change of variables. + + Advantages must be IDENTICAL to running plain GAE on a critic that predicts + ``V~ = C + B_LOO``, and the returns must be exactly that run's returns minus + ``B_LOO``. If this holds, nothing about the advantage estimate changed -- + only which component the critic is asked to represent. + """ + prompt_ids, rewards, mask = _one_group([1.0, 0.0, 1.0, 0.0]) + torch.manual_seed(0) + c_values = torch.randn(4, 6) * 0.1 + + residual = ResidualBaselineEstimator(_make_gae(), residual_target=True) + res_adv, res_returns = residual.compute_advantage( + prompt_ids, rewards, mask, c_values + ) + + baseline = residual._leave_one_out_baseline(prompt_ids, rewards) + abs_adv, abs_returns = _make_gae().compute_advantage( + prompt_ids, rewards, mask, c_values + baseline.unsqueeze(-1) + ) + + torch.testing.assert_close(res_adv, abs_adv) + torch.testing.assert_close(res_returns, abs_returns - baseline.unsqueeze(-1)) + + +def test_residual_lambda_one_telescopes_to_r_minus_b_minus_c(): + """At gamma = lambda = 1 the advantage is exactly ``A_t = R - B_LOO - C_t``.""" + prompt_ids, rewards, mask = _one_group([1.0, 0.0, 0.0, 1.0]) + torch.manual_seed(1) + c_values = torch.randn(4, 6) * 0.2 + + est = ResidualBaselineEstimator(_make_gae(), residual_target=True) + adv, _ = est.compute_advantage(prompt_ids, rewards, mask, c_values) + + baseline = est._leave_one_out_baseline(prompt_ids, rewards) + expected = (rewards - baseline).unsqueeze(-1) - c_values + torch.testing.assert_close(adv, expected) + + +def test_residual_off_is_bitwise_identical_to_plain_gae(): + """residual_baseline=false must not perturb training at all (metrics only).""" + prompt_ids, rewards, mask = _one_group([1.0, 0.0, 1.0, 0.0]) + torch.manual_seed(2) + values = torch.randn(4, 6) * 0.3 + + wrapped = ResidualBaselineEstimator(_make_gae(), residual_target=False) + w_adv, w_returns = wrapped.compute_advantage(prompt_ids, rewards, mask, values) + p_adv, p_returns = _make_gae().compute_advantage(prompt_ids, rewards, mask, values) + + assert torch.equal(w_adv, p_adv) + assert torch.equal(w_returns, p_returns) + # ...but the return-space offsets are still exported, so critic/ev_res can be + # logged for an absolute critic on the same axis as a residual one. + torch.testing.assert_close( + wrapped.last_returns_to_res, + -wrapped._leave_one_out_baseline(prompt_ids, rewards), + ) + torch.testing.assert_close(wrapped.last_returns_to_abs, torch.zeros_like(rewards)) + + +def test_residual_rejects_gamma_below_one(): + """gamma < 1 breaks the exact cancellation of B from nonterminal TD errors.""" + with pytest.raises(ValueError, match="gamma == 1"): + ResidualBaselineEstimator(_make_gae(gamma=0.99), residual_target=True) + # ...but an absolute run is untouched by that condition. + ResidualBaselineEstimator(_make_gae(gamma=0.99), residual_target=False) + + +def test_residual_group_composition_metrics(): + """frac_groups_mixed counts GROUPS, not trajectories.""" + est = ResidualBaselineEstimator(_make_gae(), residual_target=True) + # 3 groups of 2: mixed, all-fail, all-pass. + prompt_ids = torch.tensor([[1, 1], [1, 1], [2, 2], [2, 2], [3, 3], [3, 3]]) + rewards = torch.tensor([1.0, 0.0, 0.0, 0.0, 1.0, 1.0]) + mask = torch.ones(6, 4) + + est.compute_advantage(prompt_ids, rewards, mask, torch.zeros_like(mask)) + + m = est.last_metrics + assert m["residual/frac_groups_mixed"] == pytest.approx(1 / 3) + assert m["residual/frac_groups_all_fail"] == pytest.approx(1 / 3) + assert m["residual/frac_groups_all_pass"] == pytest.approx(1 / 3) + assert m["residual/n_singleton_groups"] == 0.0 + assert m["residual/group_size_min"] == 2.0 + + +def test_homogeneous_group_sample_mask_downweights_only_homogeneous(): + est = ResidualBaselineEstimator(_make_gae(), residual_target=True) + prompt_ids = torch.tensor([[1, 1], [1, 1], [2, 2], [2, 2]]) + rewards = torch.tensor([1.0, 0.0, 0.0, 0.0]) # mixed, then all-fail + mask = torch.ones(4, 4) + est.compute_advantage(prompt_ids, rewards, mask, torch.zeros_like(mask)) + + sample_mask = torch.ones(4) + assert homogeneous_group_sample_mask(sample_mask, est, 1.0) is None + scaled = homogeneous_group_sample_mask(sample_mask, est, 0.25) + torch.testing.assert_close(scaled, torch.tensor([1.0, 1.0, 0.25, 0.25])) + + +def _turn_spans(b=4, seq_len=6, anchors=(0, 3)): + from nemo_rl.algorithms.turn_level import TurnSpans + + anchor_mask = torch.zeros(b, seq_len, dtype=torch.long) + anchor_mask[:, list(anchors)] = 1 + turn_index = torch.zeros(b, seq_len, dtype=torch.int32) + turn_index[:, anchors[1] :] = 1 + return TurnSpans( + anchor_mask=anchor_mask, + turn_index=turn_index, + anchor_pos=torch.tensor([list(anchors)] * b), + turn_valid=torch.ones(b, len(anchors), dtype=torch.bool), + num_turns=torch.full((b,), len(anchors)), + turn_ntokens=torch.full((b, len(anchors)), seq_len // len(anchors)), + ) + + +def test_residual_turn_mode_does_not_leak_baseline_off_anchor(): + """Turn returns are anchor-layout; subtracting B unmasked would write -B everywhere. + + The critic's batch pairs these returns with ``token_mask = anchor_mask``, so + a leaked ``-B_LOO`` at the ~270 non-anchor positions of a real rollout would + be silently regressed against. + """ + turn_cfg = { + "turn_gae_gamma": 1.0, + "turn_gae_lambda_value": 1.0, + "turn_gae_lambda_policy": 1.0, + "normalize_advantages": False, + } + inner = TurnLevelGeneralizedAdvantageEstimator(turn_cfg, _LossCfg()) + est = ResidualBaselineEstimator(inner, residual_target=True) + + prompt_ids, rewards, mask = _one_group([1.0, 0.0, 0.0, 0.0]) + spans = _turn_spans() + _, returns = est.compute_advantage( + prompt_ids, rewards, mask, torch.zeros_like(mask), turn_spans=spans + ) + + off_anchor = returns[spans.anchor_mask == 0] + assert torch.equal(off_anchor, torch.zeros_like(off_anchor)) + # Every anchor of a given rollout carries that rollout's residual return. + expected = torch.tensor([1.0, -1 / 3, -1 / 3, -1 / 3]) + torch.testing.assert_close(returns[:, 0], expected) + torch.testing.assert_close(returns[:, 3], expected) + + +def test_residual_homogeneity_survives_reward_scaling(): + """Homogeneity is zero within-group reward VARIANCE, not "sum is 0 or G". + + Rewards reaching the estimator are post-scaling / shaping / penalty -- + ppo_math_1B maps [0,1] -> [-1,1] -- so a sum-based test would call an + all-fail group (sum = -G) "mixed" and report frac_groups_mixed = 1.0 on a + pool that is mostly homogeneous, while homogeneous_group_weight silently + weighted nothing. + """ + est = ResidualBaselineEstimator(_make_gae(), residual_target=True) + # DAPO-scaled rewards in [-1, 1]: all-fail, all-pass, mixed. + prompt_ids = torch.tensor([[1, 1], [1, 1], [2, 2], [2, 2], [3, 3], [3, 3]]) + rewards = torch.tensor([-1.0, -1.0, 1.0, 1.0, 1.0, -1.0]) + mask = torch.ones(6, 4) + + _, returns = est.compute_advantage( + prompt_ids, rewards, mask, torch.zeros_like(mask) + ) + + m = est.last_metrics + assert m["residual/frac_groups_mixed"] == pytest.approx(1 / 3) + assert m["residual/frac_groups_all_fail"] == pytest.approx(1 / 3) + assert m["residual/frac_groups_all_pass"] == pytest.approx(1 / 3) + # The Y = 0 claim must hold for the homogeneous groups at ANY reward scale. + torch.testing.assert_close(returns[:4], torch.zeros_like(returns[:4])) + assert returns[4, 0].item() > 0 and returns[5, 0].item() < 0 + torch.testing.assert_close( + est.last_group_homogeneous, torch.tensor([1.0, 1.0, 1.0, 1.0, 0.0, 0.0]) + ) + + +def test_residual_homogeneity_with_fractional_rewards(): + """Partial-credit judge rewards: equal-but-nonbinary siblings are homogeneous.""" + est = ResidualBaselineEstimator(_make_gae(), residual_target=True) + prompt_ids = torch.tensor([[1, 1], [1, 1], [2, 2], [2, 2]]) + rewards = torch.tensor([0.4, 0.4, 0.7, 0.2]) + mask = torch.ones(4, 4) + + _, returns = est.compute_advantage( + prompt_ids, rewards, mask, torch.zeros_like(mask) + ) + + assert est.last_metrics["residual/frac_groups_mixed"] == pytest.approx(0.5) + torch.testing.assert_close(returns[:2], torch.zeros_like(returns[:2])) + + +def test_raw_reward_with_residual_baseline_raises(): + """raw_reward trains no critic, so residualization is meaningless there.""" + from nemo_rl.algorithms.ppo import _create_advantage_estimator + + class _Cfg: + pass + + cfg = _Cfg() + cfg.ppo = { + "adv_estimator": { + "name": "raw_reward", + "normalize_advantages": False, + "residual_baseline": True, + } + } + cfg.loss_fn = _LossCfg() + with pytest.raises(ValueError, match="requires a value model"): + _create_advantage_estimator(cfg) + + +def test_homogeneous_group_sample_mask_rejects_missing_group_info(): + """Fail loud rather than silently ignoring the knob on a non-residual run.""" + with pytest.raises(ValueError, match="requires the residual baseline"): + homogeneous_group_sample_mask(torch.ones(2), object(), 0.5) + with pytest.raises(ValueError, match="must be >= 0"): + homogeneous_group_sample_mask(torch.ones(2), object(), -1.0) + + +def test_attach_value_baseline_keys_rejects_size_mismatch(): + """Per-sample offsets that do not line up would misalign EV silently.""" + from nemo_rl.algorithms.advantage_estimator import attach_value_baseline_keys + + est = ResidualBaselineEstimator(_make_gae(), residual_target=True) + prompt_ids, rewards, mask = _one_group([1.0, 0.0, 0.0, 0.0]) + est.compute_advantage(prompt_ids, rewards, mask, torch.zeros_like(mask)) + + with pytest.raises(ValueError, match="misalign"): + attach_value_baseline_keys({"returns": torch.zeros(3, 6)}, est) + + +def test_ev_res_mixed_group_isolates_the_homogeneous_tax(): + """Homogeneous groups have Y=0, so any prediction there only costs ev_res. + + Constructed so the critic is PERFECT on the mixed group and merely noisy on + the homogeneous one: whole-batch ev_res must be dragged below the + mixed-group-only number, which is what makes the pair diagnostic. + """ + from nemo_rl.algorithms.ppo import ( + _mixed_group_mask, + _mixed_group_value_metrics, + ) + + est = ResidualBaselineEstimator(_make_gae(), residual_target=True) + # group 1 mixed, group 2 all-fail. + prompt_ids = torch.tensor([[1, 1], [1, 1], [2, 2], [2, 2]]) + rewards = torch.tensor([1.0, 0.0, 0.0, 0.0]) + mask = torch.ones(4, 6) + _, returns = est.compute_advantage( + prompt_ids, rewards, mask, torch.zeros_like(mask) + ) + + # Perfect on the mixed pair, wrong on the homogeneous pair (target is 0). + values = returns.clone() + values[2:] = 0.4 + + mixed = _mixed_group_mask(est) + torch.testing.assert_close(mixed, torch.tensor([1.0, 1.0, 0.0, 0.0])) + + out = _mixed_group_value_metrics( + values, returns, mask, mixed, returns_to_res=est.last_returns_to_res + ) + assert out["critic/ev_res_mixed_group"] == pytest.approx(1.0, abs=1e-5) + assert out["critic/n_mixed_group_tokens"] == 12.0 + for b in ("early", "mid", "late"): + assert out[f"critic/ev_res_mixed_group_{b}"] == pytest.approx(1.0, abs=1e-5) + + # Whole-batch ev_res over the same tensors is dragged down by the + # homogeneous group, which is exactly the effect this metric separates out. + err = (returns - values).var(unbiased=False) + whole = 1.0 - err / returns.var(unbiased=False) + assert whole < out["critic/ev_res_mixed_group"] + + +def test_ev_res_mixed_group_absent_without_group_info(): + """No residual estimator -> no keys, rather than a misleading zero.""" + from nemo_rl.algorithms.ppo import _mixed_group_value_metrics + + assert ( + _mixed_group_value_metrics( + torch.zeros(2, 4), torch.zeros(2, 4), torch.ones(2, 4), None + ) + == {} + ) diff --git a/tests/unit/algorithms/test_loss_functions.py b/tests/unit/algorithms/test_loss_functions.py index 376d4798601..057b80ba199 100644 --- a/tests/unit/algorithms/test_loss_functions.py +++ b/tests/unit/algorithms/test_loss_functions.py @@ -2834,3 +2834,119 @@ def test_split_rescale_matches_sync_normalization(): assert raw_totals["num_valid_samples"] == pytest.approx( sync_totals["num_valid_samples"] ) + + +# =============================================================================== +# Residual critic: dual-space explained variance +# =============================================================================== +def _value_loss_metrics(returns, values, to_abs=None, to_res=None): + """Run MseValueLossFn once and return its per-microbatch metrics.""" + from nemo_rl.algorithms.loss.loss_functions import ( + MseValueLossConfig, + MseValueLossFn, + ) + from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + b, s = returns.shape + data = BatchedDataDict( + { + "returns": returns, + "token_mask": torch.ones(b, s), + "sample_mask": torch.ones(b), + "values": values, + } + ) + if to_abs is not None: + data["returns_to_abs"] = to_abs + data["returns_to_res"] = to_res + _, metrics = MseValueLossFn(MseValueLossConfig())( + values.unsqueeze(-1), + data, + global_valid_seqs=torch.tensor(float(b)), + global_valid_toks=torch.tensor(float(b * s)), + ) + return metrics + + +def _explained_variances(metrics): + """Drive ppo._compute_critic_metrics off a single microbatch's statistics.""" + from nemo_rl.algorithms.ppo import _compute_critic_metrics + + out = _compute_critic_metrics( + { + "grad_norm": torch.tensor(0.0), + "loss": torch.tensor(0.0), + "all_mb_metrics": { + k: [v] + for k, v in metrics.items() + if k + in { + "returns_mean", + "values_mean", + "returns_sq_mean", + "residual_sq_mean", + "abs_returns_mean", + "abs_returns_sq_mean", + "res_returns_mean", + "res_returns_sq_mean", + } + }, + } + ) + return out["critic/explained_var"], out["critic/ev_res"] + + +def test_value_loss_dual_ev_without_offsets_is_todays_behaviour(): + """No residual estimator => both explained variances collapse to the old one.""" + torch.manual_seed(0) + returns, values = torch.rand(4, 5), torch.rand(4, 5) + ev_abs, ev_res = _explained_variances(_value_loss_metrics(returns, values)) + + direct = 1.0 - (returns - values).var(unbiased=False) / returns.var(unbiased=False) + assert ev_abs == pytest.approx(direct.item(), abs=1e-5) + assert ev_res == pytest.approx(direct.item(), abs=1e-5) + + +def test_value_loss_dual_ev_uses_the_right_denominator_per_space(): + """critic/explained_var stays absolute-space; critic/ev_res is residual-space. + + The numerator is shared -- ``R - (B+C) == (R-B) - C`` -- so the two metrics + differ only by ``Var(R)`` vs ``Var(Y)``. This is what lets an absolute-critic + run and a residual-critic run be compared on one axis. + """ + torch.manual_seed(1) + b_loo = torch.tensor([0.25, 0.5, 0.75, 0.0]) + # Residual-space batch: returns = Y, values = C. + y, c = torch.rand(4, 5) - 0.5, torch.rand(4, 5) - 0.5 + metrics = _value_loss_metrics(y, c, to_abs=b_loo, to_res=torch.zeros_like(b_loo)) + ev_abs, ev_res = _explained_variances(metrics) + + err = y - c + abs_returns = y + b_loo.unsqueeze(-1) + assert ev_abs == pytest.approx( + (1.0 - err.var(unbiased=False) / abs_returns.var(unbiased=False)).item(), + abs=1e-5, + ) + assert ev_res == pytest.approx( + (1.0 - err.var(unbiased=False) / y.var(unbiased=False)).item(), abs=1e-5 + ) + assert ev_abs != pytest.approx(ev_res, abs=1e-3) + + +def test_value_loss_absolute_arm_ev_res_matches_advantage_variance_ratio(): + """For an absolute critic, ``1 - critic/ev_res`` is Var(R-V) / Var(R-B_LOO). + + That is exactly the ratio the residual-critic report measures at 1.67-2.09x, + so logging it turns the report's headline diagnostic into a live metric. + """ + torch.manual_seed(2) + b_loo = torch.tensor([0.25, 0.5, 0.75, 0.0]) + r, v = torch.rand(4, 5), torch.rand(4, 5) + # Absolute-space batch: returns = R, values = V, so the residual offset is -B. + metrics = _value_loss_metrics(r, v, to_abs=torch.zeros_like(b_loo), to_res=-b_loo) + _, ev_res = _explained_variances(metrics) + + ratio = ( + (r - v).var(unbiased=False) / (r - b_loo.unsqueeze(-1)).var(unbiased=False) + ).item() + assert (1.0 - ev_res) == pytest.approx(ratio, abs=1e-5) diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 4e593bbd471..643363efb44 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -51,6 +51,7 @@ def _make_gae_config( length_adaptive_alpha: float = 0.0, gae_lambda_value: float | None = None, gae_lambda_policy: float | None = None, + residual_baseline: bool = False, **overrides, ) -> dict: """Build an estimator_config dict with all GAE-required keys populated. @@ -66,6 +67,7 @@ def _make_gae_config( "length_adaptive_alpha": length_adaptive_alpha, "gae_lambda_value": gae_lambda_value, "gae_lambda_policy": gae_lambda_policy, + "residual_baseline": residual_baseline, **overrides, } @@ -644,7 +646,55 @@ def test_create_advantage_estimator_gae(): ) estimator = _create_advantage_estimator(master_config) - assert isinstance(estimator, GeneralizedAdvantageEstimator) + # Value-based estimators are always wrapped in ResidualBaselineEstimator so + # critic/ev_res and the residual/* group diagnostics are available in both + # arms; with residual_baseline=false the wrapper is a no-op on the targets. + from nemo_rl.algorithms.advantage_estimator import ResidualBaselineEstimator + + assert isinstance(estimator, ResidualBaselineEstimator) + assert estimator.residual_target is False + assert isinstance(estimator.inner, GeneralizedAdvantageEstimator) + + +def test_create_advantage_estimator_requires_explicit_residual_baseline(): + """A PPO config that omits the key must fail loudly, not pick a default.""" + from types import SimpleNamespace + + from nemo_rl.algorithms.ppo import _create_advantage_estimator + + cfg = _make_gae_config() + del cfg["residual_baseline"] + master_config = SimpleNamespace( + ppo={"num_generations_per_prompt": 16, "adv_estimator": {"name": "gae", **cfg}}, + loss_fn=_make_loss_config(kl_penalty=0.0), + ) + with pytest.raises(ValueError, match="residual_baseline must be set explicitly"): + _create_advantage_estimator(master_config) + + +def test_residual_baseline_rejects_single_generation_per_prompt(): + """G=1 makes B_LOO == R, which silently drops the reward from the gradient. + + With one rollout per prompt the leave-one-out baseline degenerates to the + rollout's own reward, so A_t = -C(s_t) and the critic target is identically + zero. Nothing errors downstream, so this has to be caught at setup. + """ + from types import SimpleNamespace + + from nemo_rl.algorithms.ppo import _create_advantage_estimator + + master_config = SimpleNamespace( + ppo={ + "num_generations_per_prompt": 1, + "adv_estimator": { + "name": "gae", + **_make_gae_config(residual_baseline=True), + }, + }, + loss_fn=_make_loss_config(kl_penalty=0.0), + ) + with pytest.raises(ValueError, match="num_generations_per_prompt >= 2"): + _create_advantage_estimator(master_config) def test_create_advantage_estimator_raw_reward(): @@ -1093,9 +1143,13 @@ def test_async_warmup_collector_lead_age_boundary_at_W(): from nemo_rl.algorithms.ppo import _async_warmup_collector_lead_age W, train_age, warmup_age = 20, 1, 8 - assert _async_warmup_collector_lead_age(W - 1, W, train_age, warmup_age) == warmup_age + assert ( + _async_warmup_collector_lead_age(W - 1, W, train_age, warmup_age) == warmup_age + ) assert _async_warmup_collector_lead_age(W, W, train_age, warmup_age) == warmup_age - assert _async_warmup_collector_lead_age(W + 1, W, train_age, warmup_age) == train_age + assert ( + _async_warmup_collector_lead_age(W + 1, W, train_age, warmup_age) == train_age + ) def test_async_warmup_sample_max_age_boundary_at_W_plus_train_age(): diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 66c2284e040..f2ec59805da 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -43,6 +43,15 @@ ppo: # Length-adaptive λ_policy = 1 - 1/(α·l). 0 = disabled. length_adaptive_alpha: 0.0 # VAPO: 0.05 + # --- decomposed group baseline + residual critic --- + # research/ppo/residual_critic_report.md. true => the rollout group supplies + # the task baseline B(X) as a leave-one-out mean and the critic is trained + # only on the within-task residual C(s) = R - B_LOO; critic values/returns + # are then in RESIDUAL space and the go/no-go metric is critic/ev_res. + # Requires gae_gamma: 1. Orthogonal to token-vs-turn granularity. + # false still logs critic/ev_res + residual/* for the absolute critic. + residual_baseline: false + reward_scaling: enabled: true source_min: 0.0 @@ -94,6 +103,9 @@ loss_fn: value_loss_fn: scale: 0.4 cliprange: 0.2 + # Weight of homogeneous (zero within-group reward variance) groups in the + # value loss; only meaningful with ppo.adv_estimator.residual_baseline. + homogeneous_group_weight: 1.0 checkpointing: enabled: true From f239239121d264b47bdabbc68647c66bd9dc6ea7 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 13 Aug 2026 12:59:47 -0700 Subject: [PATCH 59/69] feat(critic): SWE reference block for the privileged critic Extends the existing privileged-critic machinery with a SWE-specific reference block: the critic sees the accepted fix and the grading tests, the policy never does. Audited over all 9262 instances -- golden patch and test patch are present for 100% of them, and none appear anywhere in the agent's context (checked against a real 408k-char rollout), so this is genuine privileged information rather than a cheaper route to something already visible. The block is prefixed BEFORE the first assistant token. That is mandatory, not a detail: the value head is causal and ~150 assistant turns are spread across the context, so appending would be a silent no-op. It is byte-identical for all 16 siblings of a group, so it cannot introduce a within-task confound. Orthogonal to both other critic axes -- it composes with residual or absolute targets, and with token-level or turn_gae. For turn_gae the anchors are remapped into the augmented layout (build_turn_value_batch_augmented), which is why this variant lifts the turn_gae-plus-privileged-critic restriction that still applies to the math privileged critic. value.max_total_sequence_length and the packing/dynamic-batching token budgets are raised automatically by the reference block's cap. The budget interpolations are resolved before setup() runs, so without raising them too the packer fails with "Sequence length N exceeds bin capacity" hours into a run rather than at startup. max_total_tokens is a single TOTAL budget spent in priority order (fail_to_pass, golden_patch, test_patch, pass_to_pass); anything cut is marked inline. Measured over 400 instances the untruncated block is median 5467 / p90 17683 / p99 77911 tokens, so the 32768 default truncates ~4.8%. Watch privilege/frac_truncated. Off by default (value.swe_privileged_critic.enabled: false). Signed-off-by: Jiaqi Zeng --- examples/configs/ppo_nano_v3_5_swe_cmh.yaml | 28 + nemo_rl/algorithms/critic_pretrain.py | 82 ++- nemo_rl/algorithms/ppo.py | 122 ++++- nemo_rl/algorithms/swe_privileged_critic.py | 481 ++++++++++++++++++ pyrefly.toml | 1 + .../algorithms/test_swe_privileged_critic.py | 443 ++++++++++++++++ 6 files changed, 1137 insertions(+), 20 deletions(-) create mode 100644 nemo_rl/algorithms/swe_privileged_critic.py create mode 100644 tests/unit/algorithms/test_swe_privileged_critic.py diff --git a/examples/configs/ppo_nano_v3_5_swe_cmh.yaml b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml index a3ab7196836..393f6ab8082 100644 --- a/examples/configs/ppo_nano_v3_5_swe_cmh.yaml +++ b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml @@ -485,6 +485,34 @@ value: max_total_sequence_length: ${policy.max_total_sequence_length} precision: "bfloat16" + # --- privileged (asymmetric) critic for SWE --- + # research/ppo/privileged_residual_critic_plan.md. The critic sees the accepted + # fix and the grading tests; the policy never does. Audited over all 9262 + # instances: golden patch + test patch are present for 100% of them, and none + # of them appear anywhere in the agent's context (verified against a real + # 408k-char rollout), so this is genuine information, not just a cheaper + # computation of something already visible. + # + # The block is prefixed BEFORE the first assistant token — mandatory, since the + # value head is causal and ~150 assistant turns are spread across the context; + # appending would be a silent no-op. It is byte-identical for all 16 siblings + # of a group, so it cannot introduce a within-task confound. + # + # Orthogonal to both other critic axes: composes with residual OR absolute + # targets, and with token-level OR turn_gae (anchors are remapped into the + # augmented layout). Enabling it raises value.max_total_sequence_length by the + # sum of the caps below automatically. + swe_privileged_critic: + enabled: false + # Single TOTAL token budget for the reference block, so the value model's + # sequence budget is exactly policy_len + this + slack. Measured over 400 + # instances the untruncated block is median 5467 / p90 17683 / p99 77911 + # tokens, so 32768 truncates ~4.8%. Budget is spent in priority order + # (fail_to_pass, golden_patch, test_patch, pass_to_pass) and anything cut is + # marked inline with "... [truncated]" plus a block-level note. Watch + # privilege/frac_truncated to see what the budget is discarding. + max_total_tokens: 32768 + reward_model_cfg: enabled: false # only used for the DTensor V2 value worker reward_model_type: "regression" diff --git a/nemo_rl/algorithms/critic_pretrain.py b/nemo_rl/algorithms/critic_pretrain.py index d7f6e6a871c..e3802bad91c 100644 --- a/nemo_rl/algorithms/critic_pretrain.py +++ b/nemo_rl/algorithms/critic_pretrain.py @@ -404,7 +404,7 @@ def _forward_values_and_returns( repeated_batch: BatchedDataDict, tokenizer: Any, master_config: Any, - + metrics_out: Optional[dict[str, float]] = None, ) -> tuple[Optional[BatchedDataDict], Optional[Any]]: """Populate train_data['values'/'returns'] in place. @@ -425,12 +425,18 @@ def _forward_values_and_returns( build_privileged_value_inputs, remap_by_response_mask, ) - + from nemo_rl.algorithms.swe_privileged_critic import ( + build_swe_privileged_value_inputs, + build_turn_value_batch_augmented, + ) + from nemo_rl.algorithms.swe_privileged_critic import ( + resolve_config as swe_privileged_resolve_config, + ) privileged_critic_cfg = master_config.value.get("privileged_critic") if privileged_critic_cfg is not None and not privileged_critic_cfg.get("enabled"): privileged_critic_cfg = None - + swe_privileged_cfg = swe_privileged_resolve_config(master_config) critic_batch = None # Turn structure (None on the token-level path). Stage B must build this the @@ -439,7 +445,17 @@ def _forward_values_and_returns( turn_spans = build_turn_spans_for_batch(master_config, repeated_batch, train_data) value_model.prepare_for_inference() - if privileged_critic_cfg is not None: + if swe_privileged_cfg is not None: + critic_batch = build_swe_privileged_value_inputs( + repeated_batch, + tokenizer, + swe_privileged_cfg, + make_seq_len_divisible_by=master_config.policy[ + "make_sequence_length_divisible_by" + ], + metrics_out=metrics_out, + ) + elif privileged_critic_cfg is not None: critic_batch = build_privileged_value_inputs( repeated_batch, tokenizer, @@ -448,7 +464,7 @@ def _forward_values_and_returns( "make_sequence_length_divisible_by" ], ) - + if critic_batch is not None: vals_aug = value_model.get_values(critic_batch)["values"].squeeze(-1) critic_batch["values"] = vals_aug train_data["values"] = remap_by_response_mask( @@ -482,7 +498,14 @@ def _forward_values_and_returns( advantages, returns = adv_estimator.compute_advantage(**adv_kwargs) del advantages # critic pretraining has no actor; only returns are used train_data["returns"] = returns - if turn_spans is not None: + if turn_spans is not None and critic_batch is not None: + # Privileged AND turn-level: anchors must be remapped into the augmented + # layout, else the critic trains on policy-layout sequences while its + # values came from privileged ones. + critic_batch = build_turn_value_batch_augmented( + critic_batch, train_data, turn_spans + ) + elif turn_spans is not None: # One supervised position per turn, equally weighted (swapping token_mask # for the anchor mask is what makes MseValueLossFn a per-turn mean). critic_batch = build_turn_value_batch(train_data, turn_spans) @@ -521,7 +544,7 @@ def _heldout_metrics( train_data, repeated_batch = build_value_train_data( groups, tokenizer, master_config ) - + priv_metrics: dict[str, float] = {} _, turn_spans = _forward_values_and_returns( value_model, adv_estimator, @@ -529,7 +552,7 @@ def _heldout_metrics( repeated_batch, tokenizer, master_config, - + metrics_out=priv_metrics, ) values, returns = train_data["values"], train_data["returns"] scored_mask = ( @@ -613,7 +636,7 @@ def _heldout_metrics( ) ) metrics.update(getattr(adv_estimator, "last_metrics", {}) or {}) - + metrics.update(priv_metrics) metrics["reward"] = train_data["rewards"].float().mean().item() metrics["num_heldout_samples"] = float(train_data["input_ids"].shape[0]) return metrics @@ -756,6 +779,41 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: # Privileged critic scores [prompt + answer + response]: raise the value # model's sequence/packing budgets exactly as ppo.setup() does, so # answer-augmented near-max-length samples fit the packing bins. + _swe_privileged = value_config.get("swe_privileged_critic") + if _swe_privileged is not None and _swe_privileged.get("enabled"): + from nemo_rl.algorithms.swe_privileged_critic import privilege_budget_tokens + + _needed = master_config.policy[ + "max_total_sequence_length" + ] + privilege_budget_tokens(_swe_privileged) + if value_config["max_total_sequence_length"] < _needed: + print( + " ↑ SWE privileged critic: raising value.max_total_sequence_length " + f"{value_config['max_total_sequence_length']} -> {_needed}", + flush=True, + ) + value_config["max_total_sequence_length"] = _needed + # The packing/dynamic-batching token budgets are OmegaConf interpolations + # of max_total_sequence_length, but the runner resolves the config + # (OmegaConf.to_container(resolve=True)) BEFORE setup() runs, so raising + # the length above does not propagate to them. Without this the packer + # raises "Sequence length N exceeds bin capacity" on the long tail -- + # minutes-to-hours into the run, not at startup. + _mbs = int(value_config.get("train_micro_batch_size", 1) or 1) + for _bcfg_key in ("sequence_packing", "dynamic_batching"): + _bcfg = value_config.get(_bcfg_key) or {} + if not _bcfg.get("enabled"): + continue + for _tok_key in ("train_mb_tokens", "logprob_mb_tokens"): + _want = _needed * _mbs + if _bcfg.get(_tok_key) is not None and _bcfg[_tok_key] < _want: + print( + f" ↑ SWE privileged critic: raising value.{_bcfg_key}.{_tok_key} " + f"{_bcfg[_tok_key]} -> {_want}", + flush=True, + ) + _bcfg[_tok_key] = _want + _privileged_critic = value_config.get("privileged_critic") if _privileged_critic is not None and _privileged_critic.get("enabled"): _needed = ( @@ -1004,7 +1062,7 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: ) print("▶ Computing values...") - + priv_metrics: dict[str, float] = {} critic_batch, turn_spans = _forward_values_and_returns( value_model, adv_estimator, @@ -1012,7 +1070,7 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: repeated_batch, tokenizer, master_config, - + metrics_out=priv_metrics, ) print("▶ Training critic...") @@ -1058,7 +1116,7 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: ) ) metrics.update(getattr(adv_estimator, "last_metrics", {}) or {}) - + metrics.update(priv_metrics) metrics["reward"] = train_data["rewards"].float().mean().item() metrics["num_samples"] = float(train_data["input_ids"].shape[0]) metrics["total_step_time"] = time.perf_counter() - step_start diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 05f65e2a35c..0370988c01d 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -390,6 +390,47 @@ def setup( # Privileged (answer-conditioned) critic — supported in BOTH the sync (ppo_train) # and async (async_ppo_train) loops; both build the answer-augmented value batch # at train time from message_log + extra_env_info. + _swe_privileged = value_config.get("swe_privileged_critic") + if _swe_privileged is not None and _swe_privileged.get("enabled"): + from nemo_rl.algorithms.swe_privileged_critic import ( + privilege_budget_tokens as swe_privilege_budget_tokens, + ) + + # The critic scores [reference block + verbatim rollout], longer than the + # policy's [rollout] by at most the sum of the per-field caps. Every field + # is capped, so this is a true worst case, not an estimate. Only ever + # RAISES; a config that already has the headroom is untouched. + _needed = policy_config[ + "max_total_sequence_length" + ] + swe_privilege_budget_tokens(_swe_privileged) + if value_config["max_total_sequence_length"] < _needed: + print( + " ↑ SWE privileged critic: raising value.max_total_sequence_length " + f"{value_config['max_total_sequence_length']} -> {_needed}", + flush=True, + ) + value_config["max_total_sequence_length"] = _needed + # The packing/dynamic-batching token budgets are OmegaConf interpolations + # of max_total_sequence_length, but the runner resolves the config + # (OmegaConf.to_container(resolve=True)) BEFORE setup() runs, so raising + # the length above does not propagate to them. Without this the packer + # raises "Sequence length N exceeds bin capacity" on the long tail -- + # minutes-to-hours into the run, not at startup. + _mbs = int(value_config.get("train_micro_batch_size", 1) or 1) + for _bcfg_key in ("sequence_packing", "dynamic_batching"): + _bcfg = value_config.get(_bcfg_key) or {} + if not _bcfg.get("enabled"): + continue + for _tok_key in ("train_mb_tokens", "logprob_mb_tokens"): + _want = _needed * _mbs + if _bcfg.get(_tok_key) is not None and _bcfg[_tok_key] < _want: + print( + f" ↑ SWE privileged critic: raising value.{_bcfg_key}.{_tok_key} " + f"{_bcfg[_tok_key]} -> {_want}", + flush=True, + ) + _bcfg[_tok_key] = _want + _privileged_critic = value_config.get("privileged_critic") if _privileged_critic is not None and _privileged_critic.get("enabled"): # The critic scores [prompt + answer + response], which is longer than the @@ -1528,7 +1569,9 @@ def _raise_if_turn_gae_unsupported(master_config: MasterConfig) -> None: "is not supported: the privileged critic scores an answer-augmented " "sequence whose token layout differs from the policy's, so turn " "anchors computed on the policy layout would land on the wrong " - "positions. Disable one of the two." + "positions. Disable one of the two, or use " + "value.swe_privileged_critic, which remaps the turn anchors into " + "the augmented layout and therefore DOES support turn_gae." ) @@ -2380,11 +2423,34 @@ def ppo_train( and not privileged_critic_cfg.get("enabled") ): privileged_critic_cfg = None + from nemo_rl.algorithms.swe_privileged_critic import ( + build_swe_privileged_value_inputs, + build_turn_value_batch_augmented, + ) + from nemo_rl.algorithms.swe_privileged_critic import ( + resolve_config as swe_privileged_resolve_config, + ) + swe_privileged_cfg = swe_privileged_resolve_config(master_config) + privilege_metrics: dict[str, float] = {} critic_batch = None value_model.prepare_for_inference() - if privileged_critic_cfg is not None: + if swe_privileged_cfg is not None: + # SWE variant: reference block prefixed to the VERBATIM + # multi-turn rollout. Same augmented-layout contract as + # the math one below, so everything downstream (remap, + # residual offsets, turn anchors) is shared. + critic_batch = build_swe_privileged_value_inputs( + repeated_batch, + tokenizer, + swe_privileged_cfg, + make_seq_len_divisible_by=master_config.policy[ + "make_sequence_length_divisible_by" + ], + metrics_out=privilege_metrics, + ) + elif privileged_critic_cfg is not None: critic_batch = build_privileged_value_inputs( repeated_batch, tokenizer, @@ -2393,7 +2459,7 @@ def ppo_train( "make_sequence_length_divisible_by" ], ) - + if critic_batch is not None: vals_aug = value_model.get_values(critic_batch)[ "values" ].squeeze(-1) @@ -2527,7 +2593,15 @@ def ppo_train( attach_value_baseline_keys(train_data, adv_estimator) # Turn-level: the critic trains on one anchor per turn, # so it needs its own batch (same sequences, anchor mask). - if turn_spans is not None: + if turn_spans is not None and critic_batch is not None: + # Privileged AND turn-level: the anchor batch must be + # expressed in the AUGMENTED layout, or the critic + # would train on policy-layout sequences while its + # values came from privileged ones. + critic_batch = build_turn_value_batch_augmented( + critic_batch, train_data, turn_spans + ) + elif turn_spans is not None: from nemo_rl.algorithms.turn_level import ( build_turn_value_batch, ) @@ -2753,7 +2827,7 @@ def ppo_train( # turn counts). The token-level versions above are dominated # by the longest turns; these weight every decision equally. metrics.update(getattr(adv_estimator, "last_metrics", {}) or {}) - + metrics.update(privilege_metrics) metrics.update( { "reward": rewards.numpy(), @@ -3643,11 +3717,34 @@ def _sample_max_age(s: int) -> int: and not privileged_critic_cfg.get("enabled") ): privileged_critic_cfg = None + from nemo_rl.algorithms.swe_privileged_critic import ( + build_swe_privileged_value_inputs, + build_turn_value_batch_augmented, + ) + from nemo_rl.algorithms.swe_privileged_critic import ( + resolve_config as swe_privileged_resolve_config, + ) + swe_privileged_cfg = swe_privileged_resolve_config(master_config) + privilege_metrics: dict[str, float] = {} critic_batch = None value_model.prepare_for_inference() - if privileged_critic_cfg is not None: + if swe_privileged_cfg is not None: + # SWE variant: reference block prefixed to the VERBATIM + # multi-turn rollout. Same augmented-layout contract as + # the math one below, so everything downstream (remap, + # residual offsets, turn anchors) is shared. + critic_batch = build_swe_privileged_value_inputs( + repeated_batch, + tokenizer, + swe_privileged_cfg, + make_seq_len_divisible_by=master_config.policy[ + "make_sequence_length_divisible_by" + ], + metrics_out=privilege_metrics, + ) + elif privileged_critic_cfg is not None: critic_batch = build_privileged_value_inputs( repeated_batch, tokenizer, @@ -3656,7 +3753,7 @@ def _sample_max_age(s: int) -> int: "make_sequence_length_divisible_by" ], ) - + if critic_batch is not None: vals_aug = value_model.get_values(critic_batch)[ "values" ].squeeze(-1) @@ -3757,7 +3854,15 @@ def _sample_max_age(s: int) -> int: attach_value_baseline_keys(train_data, adv_estimator) # Turn-level: the critic trains on one anchor per turn, # so it needs its own batch (same sequences, anchor mask). - if turn_spans is not None: + if turn_spans is not None and critic_batch is not None: + # Privileged AND turn-level: the anchor batch must be + # expressed in the AUGMENTED layout, or the critic + # would train on policy-layout sequences while its + # values came from privileged ones. + critic_batch = build_turn_value_batch_augmented( + critic_batch, train_data, turn_spans + ) + elif turn_spans is not None: from nemo_rl.algorithms.turn_level import ( build_turn_value_batch, ) @@ -3998,6 +4103,7 @@ def _sample_max_age(s: int) -> int: # turn counts). The token-level versions above are dominated # by the longest turns; these weight every decision equally. metrics.update(getattr(adv_estimator, "last_metrics", {}) or {}) + metrics.update(privilege_metrics) for k, v in metrics.items(): if k in {"probs_ratio_min", "probs_ratio_clamped_min"}: diff --git a/nemo_rl/algorithms/swe_privileged_critic.py b/nemo_rl/algorithms/swe_privileged_critic.py new file mode 100644 index 00000000000..2833b024d22 --- /dev/null +++ b/nemo_rl/algorithms/swe_privileged_critic.py @@ -0,0 +1,481 @@ +# 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. +"""Privileged critic inputs for SWE agentic rollouts. + +The critic sees the accepted fix and the grading tests; the policy never does. +This is the asymmetric actor-critic setup from sim2real robotics, adapted to a +terminal-reward, ~150-turn agentic workload. + +Why this exists separately from :mod:`nemo_rl.algorithms.privileged_critic` +(which targets single-turn math RLVR and cannot be reused here): + + * it reads ``extra_env_info["ground_truth"]`` -- a key SWE rollouts do not + have -- so it would silently inject an empty block; + * it asserts single-turn and rejects interleaved multi-turn, which is every + SWE rollout (~296 messages); + * ``max_answer_tokens`` defaults to 256, sized for a math answer, not a patch. + +Only :func:`nemo_rl.algorithms.privileged_critic.remap_by_response_mask` is +shared, and it is row-order preserving, so per-sample tensors stay aligned. + +Placement is a hard correctness constraint, not a tuning knob. The value head is +causal, so ``V(s_t)`` attends only to tokens before ``t``. A SWE rollout +interleaves ~150 assistant turns across the whole context, so the ONLY placement +that reaches every supervised position is BEFORE the first assistant token. +Appending would be a silent no-op that looks like a negative result. + +Unbiasedness is preserved: the policy cannot see the reference, so +``a_t ⊥ z | s_t`` and ``E[∇log π(a_t|s_t) · V(s_t, z)] = 0``. The privileged +batch must never reach the policy worker. + +Field availability was audited over all 9262 instances of the SWE curriculum: +golden patch and test patch are present for 100% of them, but under four +different key names across four source datasets and usually nested inside the +``instance_dict`` JSON string rather than at the top level of ``metadata``. +``rubric``/``requirements``/``interface`` exist for only 6.9% (and ``interface`` +is already in the agent's prompt), so nothing here depends on them. +""" + +import json +from typing import Any, Optional + +import torch + +from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + +# The accepted fix, in priority order. swe-bench-ext and SWE-rebench use +# ``patch``, nv-internal-1 uses ``gold_patch``, SWE-Gym carries ``golden_patch`` +# alongside ``patch``. +GOLD_KEYS: tuple[str, ...] = ("golden_patch", "gold_patch", "patch") + +# Emitted in this order. Least discriminative first: the block sits ~100k tokens +# before the late values that need it, and this model is a 52-layer hybrid with +# only 6 attention layers (MEMEM*EMEMEM*...), so mamba state retains RECENT +# context best. FAIL_TO_PASS is the most compact and most discriminative field, +# so it goes last -- nearest the trajectory. +SECTION_ORDER: tuple[str, ...] = ( + "golden_patch", + "test_patch", + "pass_to_pass", + "fail_to_pass", +) + +# Single TOTAL budget for the reference block. Fixed by construction, so the +# value model's sequence budget is exactly policy_len + this + slack -- no +# dependence on a per-field cap sum staying in sync with the seqlen bump. +# +# Measured over 400 random instances (real tokenizer): the untruncated block is +# median 5467 / p90 17683 / p99 77911 / max 326605 tokens, so 32768 truncates +# ~4.8% of instances. Watch privilege/frac_truncated: it is the exact measure of +# how much privileged information the budget is discarding. +DEFAULT_MAX_TOTAL_TOKENS = 32768 + +# Per-field ceilings, applied WITHIN the total budget so one pathological field +# cannot consume it (one corpus instance carries a 323k-token golden patch, and +# pass_to_pass reaches 125k). +# Deliberately sum to MORE than the total budget, so the single total is the +# binding constraint and these only stop one pathological field from eating it. +DEFAULT_CAPS: dict[str, int] = { + "golden_patch": 24576, + "test_patch": 16384, + "pass_to_pass": 4096, + "fail_to_pass": 4096, +} + +# Budget is allocated in THIS order, which is deliberately not the emission +# order. fail_to_pass is tiny (median 56 tokens) and states the acceptance +# criterion; golden_patch says which files should change, the signal that is +# 93% unknown at t=0; test_patch is expensive and aimed at the late region the +# blind critic already handles; pass_to_pass is regression noise with a brutal +# tail (p90 9798). Emission order stays least->most discriminative for recency. +ALLOCATION_PRIORITY: tuple[str, ...] = ( + "fail_to_pass", + "golden_patch", + "test_patch", + "pass_to_pass", +) + +TRUNCATION_MARKER = "\n... [truncated]" +# A field that had content but got no budget is marked rather than silently +# omitted: the critic trains on a fixed schema, so "this exists but was cut" and +# "this instance has none" must not look identical. +OMITTED_MARKER = "... [omitted: token budget]" + + +CONFIG_KEY = "swe_privileged_critic" + + +def resolve_config(master_config: Any) -> Optional[dict[str, Any]]: + """The ``value.swe_privileged_critic`` block, or None when disabled/absent. + + Mirrors how ``value.privileged_critic`` is resolved so the two features read + identically at every call site. Absence is the default: an existing config + that has never heard of this key behaves exactly as before. + """ + value_config = getattr(master_config, "value", None) or {} + cfg = value_config.get(CONFIG_KEY) + if cfg is None or not cfg.get("enabled"): + return None + other = value_config.get("privileged_critic") + if other is not None and other.get("enabled"): + raise ValueError( + "value.swe_privileged_critic and value.privileged_critic are both " + "enabled. They build different augmented critic layouts; enable " + "exactly one. (privileged_critic targets single-turn math RLVR and " + "cannot handle SWE rollouts anyway.)" + ) + return cfg + + +def privilege_budget_tokens(cfg: dict[str, Any]) -> int: + """Upper bound on the tokens the reference block can add to a sequence. + + Used at setup to raise the VALUE model's sequence budget, so the augmented + sequences fit its packing bins. Bounded by construction: every field is + capped, so this is a true worst case rather than an estimate. + """ + total = int((cfg or {}).get("max_total_tokens") or DEFAULT_MAX_TOTAL_TOKENS) + return total + 256 # + markup / chat-template slack + + +def build_turn_value_batch_augmented( + critic_batch: BatchedDataDict, + train_data: BatchedDataDict, + turn_spans: Any, +) -> BatchedDataDict: + """Turn-level anchor batch expressed in the AUGMENTED (privileged) layout. + + ``build_turn_value_batch`` builds the anchor batch from ``train_data``, i.e. + the POLICY layout. With a privileged critic the input ids differ, so using it + would train the critic on non-privileged sequences while its values were + computed from privileged ones. This is the reason the older + ``privileged_critic`` is hard-rejected in combination with ``turn_gae``. + + Both the anchor mask and the anchor-layout returns are carried across with + :func:`remap_by_response_mask`, which is valid because anchors are a subset + of the response tokens and the response tokens are preserved verbatim. + + Order matters: ``token_mask`` is the response mask that keys every remap, so + it is replaced by the anchor mask only after the returns have been moved. + """ + from nemo_rl.algorithms.privileged_critic import remap_by_response_mask + + resp_aug = critic_batch["token_mask"] + resp_pol = train_data["token_mask"] + + critic_batch["returns"] = remap_by_response_mask( + train_data["returns"], resp_pol, resp_aug + ) + anchor_aug = remap_by_response_mask( + turn_spans.anchor_mask.to(torch.float32), resp_pol, resp_aug + ) + critic_batch["token_mask"] = anchor_aug.to(resp_aug.dtype) + critic_batch["sample_mask"] = train_data["sample_mask"] + return critic_batch + + +def _as_lines(value: Any) -> str: + """Normalise a FAIL_TO_PASS / PASS_TO_PASS field to one entry per line. + + nv-internal-1 stores these as whitespace-separated strings, the other three + datasets as JSON lists, and some entries are themselves JSON-encoded lists. + One matchable item per line is what makes the critic's job at token ``t`` + ("how many reference items has the agent hit so far?") a per-item lookup. + """ + if value is None: + return "" + if isinstance(value, str): + s = value.strip() + if s.startswith("["): + try: + value = json.loads(s) + except json.JSONDecodeError: + return s + else: + return s + if isinstance(value, (list, tuple)): + return "\n".join(str(x) for x in value) + return str(value) + + +def resolve_privilege_fields(env_info: dict[str, Any]) -> dict[str, str]: + """Pull the privileged fields out of one rollout's ``extra_env_info``. + + The shards already carry the full instance metadata per rollout under + ``extra_env_info[i]["responses_create_params"]["metadata"]``, so nothing has + to be re-joined against the source JSONL at critic-build time. + + Top-level ``metadata`` wins over ``instance_dict`` when both carry a key. + """ + md = (env_info or {}).get("responses_create_params", {}).get("metadata", {}) + try: + idict = json.loads(md.get("instance_dict") or "{}") + except (json.JSONDecodeError, TypeError): + idict = {} + if not isinstance(idict, dict): + idict = {} + src: dict[str, Any] = {**idict, **{k: v for k, v in md.items() if v}} + + gold = "" + for key in GOLD_KEYS: + v = src.get(key) + if isinstance(v, str) and v.strip(): + gold = v + break + test_patch = src.get("test_patch") + return { + "instance_id": str(src.get("instance_id") or md.get("instance_id") or ""), + "golden_patch": gold, + "test_patch": test_patch if isinstance(test_patch, str) else "", + "fail_to_pass": _as_lines(src.get("FAIL_TO_PASS")), + "pass_to_pass": _as_lines(src.get("PASS_TO_PASS")), + } + + +def _cap_tokens(text: str, max_tokens: int, tokenizer: Any) -> str: + """Truncate ``text`` to ``max_tokens``, deterministically. + + Truncation must be a pure function of the instance and never of the rollout, + or sibling rollouts in a group would receive different reference blocks and + the privilege signal would vary WITHIN a task -- manufacturing exactly the + within-group length confound that already cripples the blind critic + (Spearman(value, length) = -0.82). + """ + if not text: + return "" + # Cheap char pre-cut so a 4MB patch is never fully tokenized. 20 chars/token + # is far above any observed ratio (diffs measure 4-9), so this cannot cut + # anything the token cap would have kept. + ids = tokenizer.encode(text[: max_tokens * 20], add_special_tokens=False) + if len(ids) <= max_tokens: + return text[: max_tokens * 20] + # Reserve room for the marker so the RESULT respects max_tokens; otherwise + # every truncated field overshoots the total budget by the marker length. + marker_len = len(tokenizer.encode(TRUNCATION_MARKER, add_special_tokens=False)) + keep = max(max_tokens - marker_len, 0) + return tokenizer.decode(ids[:keep]) + TRUNCATION_MARKER + + +def _count_tokens(text: str, tokenizer: Any, hint: int) -> int: + """Token count, with a char pre-cut so a multi-MB field is never fully encoded.""" + if not text: + return 0 + return len(tokenizer.encode(text[: max(hint, 1) * 20], add_special_tokens=False)) + + +def build_reference_block( + fields: dict[str, str], + tokenizer: Any, + caps: Optional[dict[str, int]] = None, + max_total_tokens: int = DEFAULT_MAX_TOTAL_TOKENS, +) -> tuple[str, dict[str, Any]]: + """Assemble the reference block for one instance, within a FIXED token budget. + + Fixed section order and fixed markup on every instance: the critic trains on + this format for thousands of steps, so a learnable, byte-stable schema + matters far more than prose. Deliberately carries no instructions or + roleplay -- a scalar value head does not follow them. + + Fields are emitted VERBATIM (v1). Diff compression -- stripping index lines, + hunk headers and context -- is a deliberate follow-up, kept out of the first + experiment so it cannot confound the privileged-vs-blind comparison. + + Budget is spent in ALLOCATION_PRIORITY order, NOT emission order: what the + budget cannot cover is dropped from the least useful field first. A truncated + instance therefore still carries fail_to_pass and as much golden_patch as + fits, and loses pass_to_pass -- rather than losing the acceptance criterion + because it happened to be emitted last. + + Returns ``(block_text, stats)``; stats feed the ``privilege/*`` metrics so + the information being discarded is measured rather than assumed. + """ + caps = {**DEFAULT_CAPS, **(caps or {})} + remaining = int(max_total_tokens) + kept: dict[str, str] = {} + stats: dict[str, Any] = { + "truncated_fields": [], + "wanted_tokens": 0, + "kept_tokens": 0, + } + + for name in ALLOCATION_PRIORITY: + raw = fields.get(name, "") or "" + if not raw: + continue + want = _count_tokens(raw, tokenizer, caps[name]) + stats["wanted_tokens"] += want + budget = min(caps[name], max(remaining, 0)) + body = _cap_tokens(raw, budget, tokenizer) if budget > 0 else "" + got = _count_tokens(body, tokenizer, budget) if body else 0 + if got < want: + stats["truncated_fields"].append(name) + kept[name] = body if body else OMITTED_MARKER + stats["kept_tokens"] += got + remaining -= got + + parts = [""] + for name in SECTION_ORDER: # emission order stays least->most discriminative + body = kept.get(name, "") + if body: + parts.append(f"<{name}>\n{body}\n") + if stats["truncated_fields"]: + # Block-level note so the critic can tell a complete reference from a + # partial one, instead of inferring it from a missing section. + parts.append( + "[truncated: " + ", ".join(sorted(stats["truncated_fields"])) + "]" + ) + parts.append("") + stats["truncated"] = bool(stats["truncated_fields"]) + stats["dropped_tokens"] = max(stats["wanted_tokens"] - stats["kept_tokens"], 0) + return "\n".join(parts), stats + + +def build_swe_privileged_value_inputs( + repeated_batch: BatchedDataDict, + tokenizer: Any, + pcfg: dict[str, Any], + make_seq_len_divisible_by: int = 1, + metrics_out: Optional[dict[str, float]] = None, +) -> BatchedDataDict: + """Critic input batch: reference block prefixed to each verbatim rollout. + + Row-aligned with ``train_data``. ``token_mask`` marks exactly the same + response tokens, so :func:`remap_by_response_mask` can carry values back to + the policy layout (it asserts equal per-row response counts, which is the + construction check that the rollout was preserved verbatim). + + Unlike the math implementation this does NOT split prompt from response -- + SWE rollouts interleave ~150 assistant turns with tool output, and every one + of those turns is a supervised position. The whole message log is kept + untouched and only a prefix is added. + """ + caps = {**DEFAULT_CAPS, **(pcfg.get("caps") or {})} + max_total = int(pcfg.get("max_total_tokens") or DEFAULT_MAX_TOTAL_TOKENS) + message_logs = repeated_batch["message_log"] + env_infos = repeated_batch.get("extra_env_info", None) + if env_infos is None: + raise ValueError( + "SWE privileged critic: extra_env_info is absent from the batch, so " + "there is no privileged information to inject. The run would " + "silently train a blind critic while labelled privileged." + ) + + # Cache by instance: a step has ~32 unique tasks but ~512 rollouts, and all + # 16 siblings of a group MUST receive a byte-identical block (that is what + # keeps the privilege constant within a group, so it cannot introduce a + # within-task confound). + block_cache: dict[str, torch.Tensor] = {} + block_stats: dict[str, dict[str, Any]] = {} + missing: list[int] = [] + critic_message_logs: list[list[dict[str, Any]]] = [] + + for i, (msgs, info) in enumerate(zip(message_logs, env_infos)): + fields = resolve_privilege_fields(info) + if not fields["golden_patch"]: + missing.append(i) + key = fields["instance_id"] or f"__row{i}" + if key not in block_cache: + block, bstats = build_reference_block(fields, tokenizer, caps, max_total) + block_stats[key] = bstats + rendered = tokenizer.apply_chat_template( + [{"role": "system", "content": block}], + tokenize=False, + add_generation_prompt=False, + add_special_tokens=False, + ) + block_cache[key] = tokenizer( + rendered, return_tensors="pt", add_special_tokens=False + )["input_ids"][0].to(dtype=torch.long) + prefix = block_cache[key] + + critic_msgs: list[dict[str, Any]] = [ + { + "role": "system", + "content": "", # token_ids provided; content is unused by the flattener + "token_ids": prefix, + "token_loss_mask": torch.zeros_like(prefix), + } + ] + for m in msgs: + tid = torch.as_tensor(m["token_ids"], dtype=torch.long).flatten() + # Preserve the caller's mask when it set one (both PPO loops and + # critic pretraining unmask all assistant messages before this + # point); fall back to the same role rule if not. + if "token_loss_mask" in m: + mask = torch.as_tensor(m["token_loss_mask"], dtype=tid.dtype).flatten() + else: + mask = ( + torch.ones_like(tid) + if m["role"] == "assistant" + else torch.zeros_like(tid) + ) + critic_msgs.append( + { + "role": m["role"], + "content": "", + "token_ids": tid, + "token_loss_mask": mask, + } + ) + critic_message_logs.append(critic_msgs) + + # Golden patch is present for 100% of the audited corpus, so any miss is a + # data-path bug, not a straggler. Fail loudly rather than degrade to a blind + # critic under a privileged label. + if missing: + raise ValueError( + f"SWE privileged critic: no golden patch resolved for {len(missing)}/" + f"{len(message_logs)} rollouts (rows {missing[:8]}...). Expected one of " + f"{GOLD_KEYS} in extra_env_info metadata or its instance_dict. " + "Audited coverage of the SWE curriculum is 100%, so this indicates the " + "metadata did not survive the data path." + ) + + if metrics_out is not None and block_stats: + # Reported per INSTANCE (not per rollout): the block is byte-identical + # across a group's 16 siblings, so rollout-weighting would just restate + # the group size. frac_truncated is the honest measure of how much + # privileged information the fixed budget is throwing away. + s = list(block_stats.values()) + n = len(s) + metrics_out["privilege/frac_truncated"] = sum(x["truncated"] for x in s) / n + metrics_out["privilege/block_tokens_mean"] = ( + sum(x["kept_tokens"] for x in s) / n + ) + metrics_out["privilege/block_tokens_max"] = max(x["kept_tokens"] for x in s) + metrics_out["privilege/dropped_tokens_mean"] = ( + sum(x["dropped_tokens"] for x in s) / n + ) + metrics_out["privilege/wanted_tokens_mean"] = ( + sum(x["wanted_tokens"] for x in s) / n + ) + metrics_out["privilege/n_instances"] = float(n) + for _f in ALLOCATION_PRIORITY: + metrics_out[f"privilege/frac_truncated_{_f}"] = ( + sum(_f in x["truncated_fields"] for x in s) / n + ) + + flat, input_lengths = batched_message_log_to_flat_message( + critic_message_logs, + pad_value_dict={"token_ids": tokenizer.pad_token_id}, + make_sequence_length_divisible_by=make_seq_len_divisible_by, + ) + return BatchedDataDict( + { + "input_ids": flat["token_ids"], + "input_lengths": input_lengths, + "token_mask": flat["token_loss_mask"], + } + ) diff --git a/pyrefly.toml b/pyrefly.toml index 5842310100c..98c0f4a637f 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -51,6 +51,7 @@ project-includes = [ "nemo_rl/algorithms/loss/utils.py", "nemo_rl/algorithms/opd.py", "nemo_rl/algorithms/reward_functions.py", + "nemo_rl/algorithms/swe_privileged_critic.py", "nemo_rl/algorithms/utils.py", "nemo_rl/algorithms/x_token/__init__.py", "nemo_rl/algorithms/x_token/utils.py", diff --git a/tests/unit/algorithms/test_swe_privileged_critic.py b/tests/unit/algorithms/test_swe_privileged_critic.py new file mode 100644 index 00000000000..1d3af0d943d --- /dev/null +++ b/tests/unit/algorithms/test_swe_privileged_critic.py @@ -0,0 +1,443 @@ +# 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. + +import json + +import pytest +import torch + +from nemo_rl.algorithms.swe_privileged_critic import ( + DEFAULT_CAPS, + SECTION_ORDER, + _as_lines, + _cap_tokens, + build_reference_block, + build_swe_privileged_value_inputs, + resolve_privilege_fields, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +class _FakeTokenizer: + """Whitespace tokenizer with a ChatML-ish template; enough for structure tests.""" + + pad_token_id = 0 + + def encode(self, text, add_special_tokens=False): + return list(range(len(text.split()))) + + def decode(self, ids): + return " ".join("w" for _ in ids) + + def apply_chat_template(self, msgs, **kwargs): + return "".join(f"{m['role']}\n{m['content']}\n" for m in msgs) + + def __call__(self, text, return_tensors=None, add_special_tokens=False): + n = max(len(text.split()), 1) + return {"input_ids": torch.arange(n).unsqueeze(0)} + + +def _env_info(**meta): + return {"responses_create_params": {"metadata": meta}} + + +# ---------------------------------------------------------------- field resolution +@pytest.mark.parametrize("gold_key", ["golden_patch", "gold_patch", "patch"]) +def test_resolves_every_gold_key_variant(gold_key): + """Four source datasets name the accepted fix three different ways.""" + info = _env_info( + instance_id="x", instance_dict=json.dumps({gold_key: "diff --git a/f b/f"}) + ) + assert resolve_privilege_fields(info)["golden_patch"] == "diff --git a/f b/f" + + +def test_resolves_from_instance_dict_and_top_level_wins(): + """Most instances carry the fields only inside the instance_dict JSON string.""" + info = _env_info( + instance_id="x", + test_patch="TOP", + instance_dict=json.dumps({"patch": "G", "test_patch": "NESTED"}), + ) + f = resolve_privilege_fields(info) + assert f["golden_patch"] == "G" + assert f["test_patch"] == "TOP" + + +def test_resolve_tolerates_unparseable_instance_dict(): + info = _env_info(instance_id="x", patch="G", instance_dict="{not json") + assert resolve_privilege_fields(info)["golden_patch"] == "G" + + +@pytest.mark.parametrize( + "raw,expected", + [ + (["a::b", "c::d"], "a::b\nc::d"), # swe-bench-ext / SWE-Gym / rebench + ('["a::b", "c::d"]', "a::b\nc::d"), # JSON-encoded list + ("t1 t2", "t1 t2"), # nv-internal-1 stores a plain string + (None, ""), + ], +) +def test_fail_to_pass_normalisation(raw, expected): + """One matchable item per line, whichever way the dataset stored it.""" + assert _as_lines(raw) == expected + + +# ---------------------------------------------------------------- block assembly +def test_block_section_order_is_fixed(): + """The critic trains on this schema for thousands of steps; order must be stable.""" + tok = _FakeTokenizer() + fields = {k: f"body-{k}" for k in SECTION_ORDER} + block, _ = build_reference_block(fields, tok) + positions = [block.index(f"<{name}>") for name in SECTION_ORDER] + assert positions == sorted(positions) + # fail_to_pass last => nearest the trajectory (mamba retains recent best) + assert block.rindex("") > block.rindex("") + + +def test_absent_sections_are_omitted_not_emptied(): + tok = _FakeTokenizer() + block, _ = build_reference_block({"golden_patch": "g"}, tok) + assert "" in block and "" not in block + + +def test_caps_bound_output_and_are_deterministic(): + """Truncation must depend only on the instance, never on the rollout.""" + tok = _FakeTokenizer() + huge = " ".join(str(i) for i in range(50_000)) + a = _cap_tokens(huge, 100, tok) + b = _cap_tokens(huge, 100, tok) + assert a == b + assert len(tok.encode(a)) <= 100 + 4 # +marker + + +def test_default_caps_cover_every_section(): + assert set(DEFAULT_CAPS) == set(SECTION_ORDER) + + +# ---------------------------------------------------------------- batch construction +def _batch(n_rows=4, instance="inst-1", turns=3): + logs, infos = [], [] + for r in range(n_rows): + msgs = [ + { + "role": "user", + "token_ids": torch.arange(5 + r), + "token_loss_mask": torch.zeros(5 + r, dtype=torch.long), + } + ] + for t in range(turns): + k = 4 + t + r + msgs.append( + { + "role": "assistant", + "token_ids": torch.arange(k), + "token_loss_mask": torch.ones(k, dtype=torch.long), + } + ) + msgs.append( + { + "role": "user", + "token_ids": torch.arange(3), + "token_loss_mask": torch.zeros(3, dtype=torch.long), + } + ) + logs.append(msgs) + infos.append( + _env_info( + instance_id=instance, + instance_dict=json.dumps( + {"patch": "GOLD", "test_patch": "TESTS", "FAIL_TO_PASS": ["t::a"]} + ), + ) + ) + return BatchedDataDict({"message_log": logs, "extra_env_info": infos}) + + +def test_response_tokens_preserved_verbatim(): + """remap_by_response_mask asserts equal per-row counts; this is that contract.""" + from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message + + tok = _FakeTokenizer() + rb = _batch() + cb = build_swe_privileged_value_inputs(rb, tok, {}, 1) + flat, _ = batched_message_log_to_flat_message( + rb["message_log"], pad_value_dict={"token_ids": 0} + ) + assert torch.equal(cb["token_mask"].sum(-1), flat["token_loss_mask"].sum(-1)) + + +def test_multi_turn_is_supported(): + """The math implementation asserts single-turn; SWE rollouts are ~296 messages.""" + rb = _batch(turns=8) + cb = build_swe_privileged_value_inputs(rb, _FakeTokenizer(), {}, 1) + assert cb["token_mask"].sum() > 0 + + +def test_block_identical_across_siblings(): + """A block that varied within a group would recreate the length confound. + + Compares the reference prefix only. Slicing up to the first response token + would also drag in the rollout's own prompt, which is identical for real + siblings but deliberately varies row-to-row in this fixture. + """ + tok = _FakeTokenizer() + rb = _batch(n_rows=6) + fields = resolve_privilege_fields(rb["extra_env_info"][0]) + rendered = tok.apply_chat_template( + [{"role": "system", "content": build_reference_block(fields, tok)[0]}], + tokenize=False, + add_generation_prompt=False, + add_special_tokens=False, + ) + n_prefix = tok(rendered)["input_ids"].shape[1] + cb = build_swe_privileged_value_inputs(rb, tok, {}, 1) + prefixes = {tuple(cb["input_ids"][i][:n_prefix].tolist()) for i in range(6)} + assert len(prefixes) == 1 + assert n_prefix > 0 + + +def test_privilege_precedes_every_supervised_token(): + """Causality: values before the block could not see it.""" + tok = _FakeTokenizer() + cb = build_swe_privileged_value_inputs(_batch(), tok, {}, 1) + for i in range(cb["input_ids"].shape[0]): + assert int(cb["token_mask"][i].nonzero()[0, 0]) > 0 + + +def test_missing_gold_patch_raises(): + """100% coverage was audited, so a miss is a data-path bug, not a straggler.""" + rb = _batch() + rb["extra_env_info"][2] = _env_info(instance_id="broken", instance_dict="{}") + with pytest.raises(ValueError, match="no golden patch resolved"): + build_swe_privileged_value_inputs(rb, _FakeTokenizer(), {}, 1) + + +def test_missing_extra_env_info_raises(): + rb = BatchedDataDict({"message_log": _batch()["message_log"]}) + with pytest.raises(ValueError, match="extra_env_info is absent"): + build_swe_privileged_value_inputs(rb, _FakeTokenizer(), {}, 1) + + +# ------------------------------------------------- composition: residual x turn-level +def _turn_spans_for(mask): + """Anchor at the first token of each contiguous response run.""" + from nemo_rl.algorithms.turn_level import TurnSpans + + b, s = mask.shape + anchor = torch.zeros_like(mask) + for i in range(b): + prev = 0 + for j in range(s): + cur = int(mask[i, j]) + if cur and not prev: + anchor[i, j] = 1 + prev = cur + k = int(anchor.sum(-1).max()) + pos = torch.zeros(b, k, dtype=torch.long) + valid = torch.zeros(b, k, dtype=torch.bool) + for i in range(b): + idx = anchor[i].nonzero().flatten() + pos[i, : len(idx)] = idx + valid[i, : len(idx)] = True + return TurnSpans( + anchor_mask=anchor, + turn_index=torch.zeros_like(mask, dtype=torch.int32), + anchor_pos=pos, + turn_valid=valid, + num_turns=valid.sum(-1), + turn_ntokens=torch.ones(b, k, dtype=torch.long), + ) + + +def test_turn_level_anchors_survive_into_the_augmented_layout(): + """Privileged + turn_gae: anchors must land on augmented response positions. + + build_turn_value_batch builds from train_data (POLICY layout); using it with + a privileged critic would train on non-privileged sequences. This is the + combination the older privileged_critic hard-rejects. + """ + from nemo_rl.algorithms.swe_privileged_critic import ( + build_turn_value_batch_augmented, + ) + from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message + + tok = _FakeTokenizer() + rb = _batch(n_rows=3, turns=4) + cb = build_swe_privileged_value_inputs(rb, tok, {}, 1) + flat, _ = batched_message_log_to_flat_message( + rb["message_log"], pad_value_dict={"token_ids": 0} + ) + pol_mask = flat["token_loss_mask"] + spans = _turn_spans_for(pol_mask) + + resp_aug = cb["token_mask"].clone() + train_data = { + "token_mask": pol_mask, + "sample_mask": torch.ones(3), + # anchor-layout returns, as turn GAE emits them + "returns": spans.anchor_mask.float() * 7.0, + } + out = build_turn_value_batch_augmented(cb, train_data, spans) + + # same number of anchors per row, and all of them inside the aug response mask + assert torch.equal(out["token_mask"].sum(-1), spans.anchor_mask.sum(-1)) + assert bool((out["token_mask"].bool() & ~resp_aug.bool()).sum() == 0) + # the anchor returns survived the layout change + assert torch.allclose( + out["returns"][out["token_mask"].bool()], + torch.full((int(out["token_mask"].sum()),), 7.0), + ) + + +@pytest.mark.parametrize("residual", [False, True]) +def test_privileged_composes_with_both_critic_targets(residual): + """Privilege is an INPUT-layout change; residual is a TARGET change. + + They are orthogonal: the per-sample return-space offsets are [B] and + row-order is preserved by the augmentation, so they stay aligned. + """ + from nemo_rl.algorithms.advantage_estimator import ( + GeneralizedAdvantageEstimator, + ResidualBaselineEstimator, + ) + from nemo_rl.algorithms.privileged_critic import remap_by_response_mask + from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message + + tok = _FakeTokenizer() + rb = _batch(n_rows=4, turns=2) + cb = build_swe_privileged_value_inputs(rb, tok, {}, 1) + flat, _ = batched_message_log_to_flat_message( + rb["message_log"], pad_value_dict={"token_ids": 0} + ) + pol_mask = flat["token_loss_mask"].float() + + # values live in the AUGMENTED layout, then come back to the policy layout + vals_aug = torch.randn_like(cb["token_mask"], dtype=torch.float32) + vals_pol = remap_by_response_mask(vals_aug, cb["token_mask"], pol_mask) + + class _L: + use_kl_in_reward = False + reference_policy_kl_penalty = 0.0 + reference_policy_kl_type = "low_var_kl" + + inner = GeneralizedAdvantageEstimator( + { + "gae_lambda": 1.0, + "gae_gamma": 1, + "normalize_advantages": False, + "gae_lambda_value": None, + "gae_lambda_policy": None, + "length_adaptive_alpha": 0.0, + }, + _L(), + ) + est = ResidualBaselineEstimator(inner, residual_target=residual) + adv, returns = est.compute_advantage( + prompt_ids=torch.full((4, 2), 3), + rewards=torch.tensor([1.0, 0.0, 1.0, 0.0]), + mask=pol_mask, + values=vals_pol, + ) + # returns go forward into the augmented layout for the value loss + ret_aug = remap_by_response_mask(returns, pol_mask, cb["token_mask"]) + assert ret_aug.shape == cb["token_mask"].shape + assert torch.allclose(ret_aug[cb["token_mask"].bool()], returns[pol_mask.bool()]) + # per-sample offsets are row-aligned with the augmented batch + assert est.last_returns_to_abs.shape[0] == cb["input_ids"].shape[0] + assert adv.shape == pol_mask.shape + + +# ------------------------------------------------- fixed budget + truncation marks +def test_total_budget_binds_and_is_marked_inline(): + """Over-budget instances are truncated, and the block SAYS SO. + + A silently short block is indistinguishable from an instance that simply has + less reference material, which would make the critic's schema ambiguous. + """ + tok = _FakeTokenizer() + big = " ".join(f"w{i}" for i in range(5000)) + fields = {k: big for k in SECTION_ORDER} + block, stats = build_reference_block(fields, tok, max_total_tokens=200) + + assert stats["truncated"] is True + assert stats["kept_tokens"] <= 200 + assert stats["dropped_tokens"] > 0 + assert "[truncated" in block # block-level note + assert "... [truncated]" in block # per-field cut marker + for f in stats["truncated_fields"]: + assert f in block + + +def test_budget_is_spent_in_priority_order_not_emission_order(): + """fail_to_pass must survive a tight budget; pass_to_pass is sacrificed first.""" + tok = _FakeTokenizer() + big = " ".join(f"w{i}" for i in range(5000)) + fields = {k: big for k in SECTION_ORDER} + _, stats = build_reference_block(fields, tok, max_total_tokens=300) + # fail_to_pass is allocated first, so it is not among the starved fields + assert "pass_to_pass" in stats["truncated_fields"] + assert ( + stats["truncated_fields"].index("pass_to_pass") + == len(stats["truncated_fields"]) - 1 + or "pass_to_pass" in stats["truncated_fields"] + ) + + +def test_starved_field_is_marked_not_dropped(): + """A field with content but no budget gets a placeholder, keeping the schema stable.""" + from nemo_rl.algorithms.swe_privileged_critic import OMITTED_MARKER + + tok = _FakeTokenizer() + big = " ".join(f"w{i}" for i in range(5000)) + block, _ = build_reference_block( + {"fail_to_pass": big, "pass_to_pass": big}, tok, max_total_tokens=50 + ) + assert OMITTED_MARKER in block + assert "" in block # section present, content marked as omitted + + +def test_untruncated_block_has_no_markers(): + tok = _FakeTokenizer() + block, stats = build_reference_block( + {k: "short body" for k in SECTION_ORDER}, tok, max_total_tokens=10_000 + ) + assert stats["truncated"] is False + assert "[truncated" not in block + + +def test_truncation_metrics_are_emitted_per_instance(): + """privilege/* is reported per INSTANCE, not per rollout (block is shared).""" + tok = _FakeTokenizer() + rb = _batch(n_rows=6) + m = {} + build_swe_privileged_value_inputs( + rb, tok, {"max_total_tokens": 10_000}, 1, metrics_out=m + ) + assert m["privilege/n_instances"] == 1.0 # 6 rollouts, one instance + assert m["privilege/frac_truncated"] == 0.0 + assert m["privilege/block_tokens_mean"] > 0 + for f in SECTION_ORDER: + assert f"privilege/frac_truncated_{f}" in m + + big = " ".join(f"w{i}" for i in range(5000)) + for info in rb["extra_env_info"]: + info["responses_create_params"]["metadata"]["instance_dict"] = json.dumps( + {"patch": big, "test_patch": big, "FAIL_TO_PASS": [big]} + ) + m2 = {} + build_swe_privileged_value_inputs( + rb, tok, {"max_total_tokens": 200}, 1, metrics_out=m2 + ) + assert m2["privilege/frac_truncated"] == 1.0 + assert m2["privilege/dropped_tokens_mean"] > 0 From 42452efb3f251c5b099a0d9ac8afb75c7cf6b322 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 13 Aug 2026 13:08:07 -0700 Subject: [PATCH 60/69] fix(critic): report pre-update EV in both return spaces, on the anchor mask e7ac37bc4 made critic/explained_var pre-update by overwriting it driver-side from the rollout-time values. Merged with the residual critic that is wrong in two ways, neither of which git could flag -- the two changes are in different functions and merged cleanly. 1. SPACE. ResidualBaselineEstimator keeps the batch in whichever space the run uses: with residual_baseline=true train_data holds values = C and returns = Y, so pooling it as-is yields 1 - Var(Y - C)/Var(Y), i.e. the RESIDUAL EV, stored under critic/explained_var. The control arm stores the absolute EV under the same key. The headline metric of the A/B would mean two different things in the two arms, and critic/ev_res -- which is derived from the value loss, hence a different critic snapshot -- would not be comparable to it either. _pooled_explained_var now returns both spaces from one numerator (the prediction error is offset-invariant) using the estimator's returns_to_abs / returns_to_res offsets, and both critic/explained_var and critic/ev_res are set from it. Both are pre-update, both on a fixed space. Without a residual estimator both offsets are absent and both collapse to today's number. 2. MASK. It pooled over train_data["token_mask"], the full response mask. In turn_gae mode returns are structurally zero off-anchor (scatter_turns_to_anchors) while the critic emits a dense value, so the EV would be dominated by V^2 at positions the critic is never supervised at. Use _value_metric_mask, which already exists for exactly this and is what the positional diagnostics use. critic/explained_var_post_update is untouched: it is still the loss-derived number from the post-update forward, which is the point of that key. Signed-off-by: Jiaqi Zeng --- nemo_rl/algorithms/ppo.py | 67 +++++++++++++++++++++++++++++++-------- 1 file changed, 54 insertions(+), 13 deletions(-) diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index b3f7e3c7448..b08183ec28b 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -1789,24 +1789,47 @@ def _pooled_explained_var( returns: torch.Tensor, token_mask: torch.Tensor, sample_mask: torch.Tensor, -) -> float: - """Pooled EV = 1 - Var(returns - values) / Var(returns) over valid tokens. + returns_to_abs: Optional[torch.Tensor] = None, + returns_to_res: Optional[torch.Tensor] = None, +) -> tuple[float, float]: + """Pooled pre-update EV in BOTH return spaces, as ``(absolute, residual)``. Computed driver-side from the rollout-time values -- the exact tensors GAE consumed -- so it describes the critic BEFORE any update this step, regardless of critic_ppo_epochs or gbs-vs-rollout microbatching. Masked like the value loss (token_mask * sample_mask) so it is directly comparable to critic/explained_var_post_update. + + Both spaces are returned because :class:`ResidualBaselineEstimator` leaves + the batch in whichever space the run uses (``values = C``, ``returns = Y`` + under residual_baseline; absolute otherwise). Pooling the batch as-is would + therefore put *absolute* EV in one arm of an A/B and *residual* EV in the + other under a single key. ``returns_to_abs`` / ``returns_to_res`` are that + estimator's per-sample ``[B]`` offsets -- exactly one is zero -- and both are + None without a residual estimator, which collapses this to the historical + number in both slots. + + The prediction error is offset-invariant (``R - (B_LOO + C) == (R - B_LOO) - C``), + so one numerator serves both and only the denominator changes. """ mask = (token_mask * sample_mask.unsqueeze(-1)).bool() if int(mask.sum()) < 2: - return 0.0 - r = returns[mask].float() - v = values[mask].float() - var_returns = r.var(unbiased=False) - if var_returns <= 1e-8: - return 0.0 - return (1.0 - (r - v).var(unbiased=False) / var_returns).item() + return 0.0, 0.0 + err_var = (returns[mask].float() - values[mask].float()).var(unbiased=False) + + zeros = torch.zeros(returns.shape[0], device=returns.device, dtype=returns.dtype) + to_abs = zeros if returns_to_abs is None else returns_to_abs.to(returns.dtype) + to_res = zeros if returns_to_res is None else returns_to_res.to(returns.dtype) + + evs = [] + for offset in (to_abs, to_res): + target = (returns + offset.unsqueeze(-1))[mask].float() + var_t = target.var(unbiased=False) + # Var(Y) legitimately hits zero when every group in the step is + # homogeneous (~54% all-fail on the SWE pool); report 0.0 rather than + # dividing by an epsilon and emitting ~-1e8. + evs.append((1.0 - err_var / var_t).item() if var_t > 1e-8 else 0.0) + return evs[0], evs[1] def _calibration_ece(v: torch.Tensor, r: torch.Tensor, n_conf_bins: int = 10) -> float: @@ -2966,12 +2989,21 @@ def ppo_train( # critic/loss and critic/grad_norm still describe the last # training pass. if "values" in train_data and "returns" in train_data: - metrics["critic/explained_var"] = _pooled_explained_var( + # Anchor mask in turn mode: turn-level returns are + # structurally zero off-anchor (scatter_turns_to_anchors) + # while values are dense, so pooling over the full + # response mask would score ~270 structural zeros per + # sample against a live critic output. + ev_abs, ev_res = _pooled_explained_var( train_data["values"], train_data["returns"], - train_data["token_mask"], + _value_metric_mask(train_data, turn_spans), train_data["sample_mask"], + train_data.get("returns_to_abs"), + train_data.get("returns_to_res"), ) + metrics["critic/explained_var"] = ev_abs + metrics["critic/ev_res"] = ev_res # The other end of the bracket: the same batch re-scored AFTER # every update this step (ppo.log_post_update_critic_metrics). if post_value_results is not None: @@ -4313,12 +4345,21 @@ def _sample_max_age(s: int) -> int: # Overwrite critic/explained_var with the pre-update EV (the # values GAE consumed); see the sync loop for why. if "values" in train_data and "returns" in train_data: - metrics["critic/explained_var"] = _pooled_explained_var( + # Anchor mask in turn mode: turn-level returns are + # structurally zero off-anchor (scatter_turns_to_anchors) + # while values are dense, so pooling over the full + # response mask would score ~270 structural zeros per + # sample against a live critic output. + ev_abs, ev_res = _pooled_explained_var( train_data["values"], train_data["returns"], - train_data["token_mask"], + _value_metric_mask(train_data, turn_spans), train_data["sample_mask"], + train_data.get("returns_to_abs"), + train_data.get("returns_to_res"), ) + metrics["critic/explained_var"] = ev_abs + metrics["critic/ev_res"] = ev_res # The other end of the bracket: re-scored AFTER every update # this step (ppo.log_post_update_critic_metrics). if post_value_results is not None: From dab69b14724cd8ef30937f200ddc0defe983745f Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 13 Aug 2026 13:09:12 -0700 Subject: [PATCH 61/69] config: add the critic keys to the hsg SWE config ppo_nano_v3_5_swe_hsg.yaml arrived on this branch (d8ec18929) in parallel with the residual-critic work, so it never got the keys that work added to the other four configs. adv_estimator.residual_baseline is the one that matters: _create_advantage_estimator requires it explicitly for gae/turn_gae with no default, so every hsg run would have died at setup -- after cluster spin-up and gym venv creation, i.e. minutes of allocation per attempt. The other two have safe fallbacks (homogeneous_group_weight defaults on the BaseModel, swe_privileged_critic is read with .get()) but are added for parity with the cmh config, so the two clusters can run the same arms. Values are identical to ppo_nano_v3_5_swe_cmh.yaml; nothing else in the file changed (verified by comparing the parsed YAML before and after). Signed-off-by: Jiaqi Zeng --- examples/configs/ppo_nano_v3_5_swe_hsg.yaml | 52 +++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/examples/configs/ppo_nano_v3_5_swe_hsg.yaml b/examples/configs/ppo_nano_v3_5_swe_hsg.yaml index aff7e41f8dd..14adda36682 100644 --- a/examples/configs/ppo_nano_v3_5_swe_hsg.yaml +++ b/examples/configs/ppo_nano_v3_5_swe_hsg.yaml @@ -121,6 +121,21 @@ ppo: # Length-adaptive lambda_policy = 1 - 1/(alpha*l). 0 = disabled. length_adaptive_alpha: null + # --- decomposed group baseline + residual critic --- + # research/ppo/residual_critic_report.md. true => the rollout group supplies + # the task baseline B(X) as a leave-one-out mean and the critic is trained + # only on the within-task residual C(s) = R - B_LOO; critic values/returns + # are then in RESIDUAL space and the go/no-go metric is critic/ev_res. + # Requires gae_gamma: 1. Orthogonal to token-vs-turn granularity. + # + # Motivation, measured on this exact workload: the absolute critic puts + # 88-93% of its output variance between tasks yet reaches EV 0.157 — below + # even a two-value "all-fail vs rest" lookup (0.330) and far below the free + # leave-one-out group baseline (0.553). Substituting it for that baseline + # RAISES advantage variance 1.67-2.09x. false still logs critic/ev_res and + # residual/* so an absolute run is comparable on the same axis. + residual_baseline: false + reward_shaping: enabled: false overlong_buffer_length: 128 @@ -188,6 +203,15 @@ loss_fn: value_loss_fn: scale: 1.0 cliprange: 0.5 + # Weight of homogeneous (all-fail / all-pass) groups in the value loss. Only + # meaningful with ppo.adv_estimator.residual_baseline: those groups have Y = 0 + # for every sibling, so they contribute EXACTLY zero target variance while + # still costing ~58% of critic FLOPs on this pool (54.1% all-fail, 2.4% + # all-pass; only 43.5% of groups are mixed). Keep at 1.0 for the clean + # ablation — they are not dead weight, they are the shrinkage that enforces + # E[C | X] = 0 and suppresses between-task leakage. Lower it only to buy more + # mixed-group exposure per unit wall-clock. + homogeneous_group_weight: 1.0 # ============================================================================= # Policy @@ -453,6 +477,34 @@ value: max_total_sequence_length: ${policy.max_total_sequence_length} precision: "bfloat16" + # --- privileged (asymmetric) critic for SWE --- + # research/ppo/privileged_residual_critic_plan.md. The critic sees the accepted + # fix and the grading tests; the policy never does. Audited over all 9262 + # instances: golden patch + test patch are present for 100% of them, and none + # of them appear anywhere in the agent's context (verified against a real + # 408k-char rollout), so this is genuine information, not just a cheaper + # computation of something already visible. + # + # The block is prefixed BEFORE the first assistant token — mandatory, since the + # value head is causal and ~150 assistant turns are spread across the context; + # appending would be a silent no-op. It is byte-identical for all 16 siblings + # of a group, so it cannot introduce a within-task confound. + # + # Orthogonal to both other critic axes: composes with residual OR absolute + # targets, and with token-level OR turn_gae (anchors are remapped into the + # augmented layout). Enabling it raises value.max_total_sequence_length by the + # sum of the caps below automatically. + swe_privileged_critic: + enabled: false + # Single TOTAL token budget for the reference block, so the value model's + # sequence budget is exactly policy_len + this + slack. Measured over 400 + # instances the untruncated block is median 5467 / p90 17683 / p99 77911 + # tokens, so 32768 truncates ~4.8%. Budget is spent in priority order + # (fail_to_pass, golden_patch, test_patch, pass_to_pass) and anything cut is + # marked inline with "... [truncated]" plus a block-level note. Watch + # privilege/frac_truncated to see what the budget is discarding. + max_total_tokens: 32768 + reward_model_cfg: enabled: false # only used for the DTensor V2 value worker reward_model_type: "regression" From ed119cbf8d3a4d55a091e25696527ae7134407bd Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 13 Aug 2026 13:23:24 -0700 Subject: [PATCH 62/69] test: cover the two-space return of _pooled_explained_var 42452efb3 changed _pooled_explained_var to return (ev_abs, ev_res) but left test_pooled_explained_var_pre_update -- which arrived with e7ac37bc4 -- asserting against a scalar. Comparing a tuple to pytest.approx(...) does not raise; approx swallows the TypeError and returns False, so all four assertions simply failed. Update them to assert both spaces, and add test_pooled_explained_var_residual_offsets for the returns_to_abs/returns_to_res arms, which had no coverage at all. It pins the properties the merge fix relies on: a per-sample offset that is constant across samples moves neither EV (variance is shift-invariant); an offset applied only to the absolute space raises ev_abs and leaves ev_res untouched; and offsets never enter the shared numerator, so a perfect predictor stays at 1.0 in both spaces. Signed-off-by: Jiaqi Zeng --- tests/unit/algorithms/test_ppo.py | 62 ++++++++++++++++++++++++++++--- 1 file changed, 57 insertions(+), 5 deletions(-) diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 0c58e9b4962..e66369b1dda 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -1207,7 +1207,11 @@ def test_resolve_critic_ppo_epochs_rejects_zero_ppo_epochs(): def test_pooled_explained_var_pre_update(): """Pre-update EV: perfect prediction -> 1, mean prediction -> 0, and both - masks are respected (values at masked positions must not leak in).""" + masks are respected (values at masked positions must not leak in). + + Returns ``(ev_abs, ev_res)``. With no residual offsets the two spaces are + the same number, so every case here asserts both. + """ from nemo_rl.algorithms.ppo import _pooled_explained_var token_mask = torch.tensor([[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]]) @@ -1219,19 +1223,24 @@ def test_pooled_explained_var_pre_update(): values[0, 3] = 7.0 values[1, 2] = 7.0 assert _pooled_explained_var(values, returns, token_mask, sample_mask) == ( - pytest.approx(1.0) + pytest.approx(1.0), + pytest.approx(1.0), ) # Constant prediction at the valid-token mean explains nothing. mean_pred = torch.full_like(returns, 0.6) # mean of [1,1,1,0,0] assert _pooled_explained_var(mean_pred, returns, token_mask, sample_mask) == ( - pytest.approx(0.0, abs=1e-6) + pytest.approx(0.0, abs=1e-6), + pytest.approx(0.0, abs=1e-6), ) # sample_mask knocks out the second sequence; the survivor has constant # returns (zero variance), which reports 0.0 rather than dividing by ~0. solo = torch.tensor([1.0, 0.0]) - assert _pooled_explained_var(returns.clone(), returns, token_mask, solo) == 0.0 + assert _pooled_explained_var(returns.clone(), returns, token_mask, solo) == ( + 0.0, + 0.0, + ) # Matches the sufficient-statistics form used by the loss-side metrics. mask = (token_mask * sample_mask.unsqueeze(-1)).bool() @@ -1243,8 +1252,51 @@ def test_pooled_explained_var_pre_update(): r_sq, res_sq = (r**2).mean(), ((r - v) ** 2).mean() ev_suff = 1.0 - (res_sq - (r_mean - v_mean) ** 2) / (r_sq - r_mean**2) assert _pooled_explained_var(noisy, returns, token_mask, sample_mask) == ( - pytest.approx(ev_suff.item(), abs=1e-6) + pytest.approx(ev_suff.item(), abs=1e-6), + pytest.approx(ev_suff.item(), abs=1e-6), + ) + + +def test_pooled_explained_var_residual_offsets(): + """The two return spaces share one numerator and differ only in denominator. + + This is what keeps critic/explained_var absolute-space in BOTH arms of the + residual A/B: without it the same key would carry Var(Y)-normalised EV under + residual_baseline and Var(R)-normalised EV without it. + """ + from nemo_rl.algorithms.ppo import _pooled_explained_var + + token_mask = torch.tensor([[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]]) + sample_mask = torch.tensor([1.0, 1.0]) + returns = torch.tensor([[1.0, 1.0, 1.0, 9.0], [0.0, 0.0, 9.0, 9.0]]) + noisy = returns + 0.25 * torch.tensor( + [[1.0, -1.0, 0.5, 0.0], [-0.5, 1.0, 0.0, 0.0]] + ) + base_abs, base_res = _pooled_explained_var(noisy, returns, token_mask, sample_mask) + + # A per-sample offset that is CONSTANT across samples shifts both targets + # without changing either variance, so neither EV moves. + const = torch.tensor([2.0, 2.0]) + ev_abs, ev_res = _pooled_explained_var( + noisy, returns, token_mask, sample_mask, const, const + ) + assert ev_abs == pytest.approx(base_abs, abs=1e-6) + assert ev_res == pytest.approx(base_res, abs=1e-6) + + # A spreading offset applied to the ABSOLUTE space only: Var(R) grows while + # the shared error does not, so ev_abs rises and ev_res is left untouched. + spread = torch.tensor([2.0, -2.0]) + ev_abs, ev_res = _pooled_explained_var( + noisy, returns, token_mask, sample_mask, spread, None ) + assert ev_abs > base_abs + assert ev_res == pytest.approx(base_res, abs=1e-6) + + # Offsets never touch the numerator: a perfect predictor stays at 1.0 in + # both spaces however the targets are shifted. + assert _pooled_explained_var( + returns.clone(), returns, token_mask, sample_mask, spread, const + ) == (pytest.approx(1.0), pytest.approx(1.0)) # NeMo-Gym IS supported for async PPO: the AsyncTrajectoryCollector runs the gym From 78a12dc4313922f6e69987bfebafe88f9ed67a82 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 13 Aug 2026 13:23:35 -0700 Subject: [PATCH 63/69] fix(critic): read the required train_micro_batch_size without a fallback value.train_micro_batch_size is a required field of ValueConfig (nemo_rl/models/value/config.py:47, not NotRequired), so value_config.get("train_micro_batch_size", 1) is a call-site default of the kind the config conventions forbid: the value is supposed to come from the YAML and nowhere else. It is not harmless here. The fallback feeds the SWE privileged critic's packing-budget bump, so if the key were ever absent or null the bins would be sized to _needed * 1 instead of _needed * mbs -- under-sized exactly where that block exists to prevent it, and surfacing as "Sequence length N exceeds bin capacity" hours into a run rather than at startup. Fail loudly at config load instead. Signed-off-by: Jiaqi Zeng --- nemo_rl/algorithms/critic_pretrain.py | 5 ++++- nemo_rl/algorithms/ppo.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/nemo_rl/algorithms/critic_pretrain.py b/nemo_rl/algorithms/critic_pretrain.py index e3802bad91c..d4b113ca07b 100644 --- a/nemo_rl/algorithms/critic_pretrain.py +++ b/nemo_rl/algorithms/critic_pretrain.py @@ -799,7 +799,10 @@ def critic_pretrain(master_config: Any, tokenizer: Any) -> None: # the length above does not propagate to them. Without this the packer # raises "Sequence length N exceeds bin capacity" on the long tail -- # minutes-to-hours into the run, not at startup. - _mbs = int(value_config.get("train_micro_batch_size", 1) or 1) + # Required field of ValueConfig; a call-site fallback here would + # silently under-size the packing bins to _needed * 1 and resurface as + # "Sequence length N exceeds bin capacity" deep into a run. + _mbs = int(value_config["train_micro_batch_size"]) for _bcfg_key in ("sequence_packing", "dynamic_batching"): _bcfg = value_config.get(_bcfg_key) or {} if not _bcfg.get("enabled"): diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index b08183ec28b..8e2c83f17b2 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -452,7 +452,10 @@ def setup( # the length above does not propagate to them. Without this the packer # raises "Sequence length N exceeds bin capacity" on the long tail -- # minutes-to-hours into the run, not at startup. - _mbs = int(value_config.get("train_micro_batch_size", 1) or 1) + # Required field of ValueConfig; a call-site fallback here would + # silently under-size the packing bins to _needed * 1 and resurface as + # "Sequence length N exceeds bin capacity" deep into a run. + _mbs = int(value_config["train_micro_batch_size"]) for _bcfg_key in ("sequence_packing", "dynamic_batching"): _bcfg = value_config.get(_bcfg_key) or {} if not _bcfg.get("enabled"): From edfce826dff7928c0bd5dd57f95ee56d0515eb10 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 13 Aug 2026 14:04:56 -0700 Subject: [PATCH 64/69] test: match the turn_gae-aware unsupported-estimator message _create_advantage_estimator has raised "PPO only supports 'gae', 'turn_gae' or 'raw_reward'." since 2f5ff5081 added the turn-level estimator, but test_create_advantage_estimator_rejects_unsupported_name still matched the older "only supports 'gae' or 'raw_reward'". The test has been red on this branch since that commit -- it fails identically on backup/ppo-dev-pre-merge, so neither the residual/privileged work nor the merge introduced it. Worth fixing rather than leaving: pyproject's addopts carries -x, so this single stale assertion aborted the whole tests/unit/algorithms run after 144 of 177 items, hiding everything ordered after it. Signed-off-by: Jiaqi Zeng --- tests/unit/algorithms/test_ppo.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index e66369b1dda..f961b7df8fa 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -726,7 +726,9 @@ def test_create_advantage_estimator_rejects_unsupported_name(): loss_fn={"reference_policy_kl_penalty": 0.0}, ) - with pytest.raises(ValueError, match="only supports 'gae' or 'raw_reward'"): + with pytest.raises( + ValueError, match="only supports 'gae', 'turn_gae' or 'raw_reward'" + ): _create_advantage_estimator(master_config) From fb2f2a115461c6456ebd725ff785a0c85edd03a7 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Thu, 13 Aug 2026 14:10:12 -0700 Subject: [PATCH 65/69] feat(ppo): report adv_raw/* on the turn-level path too TurnLevelGeneralizedAdvantageEstimator assigned last_metrics wholesale from turn_level_metrics(), so the pre-whitening advantage stats that GeneralizedAdvantageEstimator reports never appeared under turn_gae. The two features landed in parallel -- adv_raw/* on the remote side, turn_gae here -- so the gap only became visible once they were merged. It is not only a missing diagnostic. b5eb09e57 reconstructs unwhitened advantages in the rollout dump as `advantages * adv_raw_std + adv_raw_mean`, and `if adv_raw_metrics:` sees only turn/* keys, so the dump silently degraded to "absent" for exactly the runs the turn estimator exists for. advantage/turn_std_prenorm is not a substitute: it is computed over the [B, K] turn tensor, while the whitening is a token-level masked statistic, so it is not the affine constant the reconstruction needs. Hoist the helper to a module-level raw_advantage_metrics() -- it only ever used self.normalize_advantages -- and call it from both estimators on the TOKEN-level advantages under the same mask the whitening uses. In turn mode the per-turn advantage has already been broadcast to every token of its turn (scatter_turns_to_tokens), so the two arms are directly comparable and whiten_gain describes the rescale actually applied. Merge into last_metrics instead of assigning, so turn_level_metrics can no longer drop it; the two key families are disjoint (adv_raw/* vs advantage/*, critic/*, turn/*). Signed-off-by: Jiaqi Zeng --- nemo_rl/algorithms/advantage_estimator.py | 95 ++++++++++++------- .../algorithms/test_advantage_estimator.py | 68 +++++++++++++ 2 files changed, 127 insertions(+), 36 deletions(-) diff --git a/nemo_rl/algorithms/advantage_estimator.py b/nemo_rl/algorithms/advantage_estimator.py index 143c0e702b7..afb90993011 100644 --- a/nemo_rl/algorithms/advantage_estimator.py +++ b/nemo_rl/algorithms/advantage_estimator.py @@ -267,6 +267,50 @@ def compute_advantage(self, prompt_ids, rewards, mask, **kwargs): return adv, None +def raw_advantage_metrics( + advantages: torch.Tensor, + mask: torch.Tensor, + normalize_advantages: bool, +) -> dict[str, float]: + """Stats on the PRE-whitening advantages, over valid tokens. + + ``normalize_advantages`` rescales advantages to unit std every step, so the + post-whitening spread is 1.0 by construction and carries no information. + The pre-whitening spread does: with lambda=1 the advantage is ``R - V(s)``, + so ``adv_raw/std`` is the critic's residual scale and should SHRINK as the + critic improves. A flat ``adv_raw/std`` alongside a flat + ``critic/explained_var`` is independent confirmation that the critic is not + learning. ``whiten_gain`` is the amplification whitening applies (1/std); a + large and growing value means the loss is being scaled up to compensate for + a shrinking signal. + + Shared by the token-level and turn-level estimators so the two arms report + the same diagnostic on the same axis. Both call it on the TOKEN-level + advantages under the same ``mask`` the whitening uses -- in turn mode the + per-turn advantage has already been broadcast to every token of its turn + (``scatter_turns_to_tokens``), so the two are directly comparable, and + ``whiten_gain`` describes the rescaling that is actually applied. + + Note: under :class:`ResidualBaselineEstimator` the inner estimator is fed + ``V~ = B_LOO + C``, so this measures the residual scale of the combined + group baseline plus critic, not of the critic alone. + """ + m = mask.bool() + if int(m.sum()) < 2: + return {} + a = advantages[m].float() + std = a.std(unbiased=False) + metrics = { + "adv_raw/mean": a.mean().item(), + "adv_raw/std": std.item(), + "adv_raw/abs_mean": a.abs().mean().item(), + "adv_raw/max_abs": a.abs().max().item(), + } + if normalize_advantages: + metrics["adv_raw/whiten_gain"] = (1.0 / (std + 1e-8)).item() + return metrics + + class GeneralizedAdvantageEstimator: """Generalized Advantage Estimation (GAE) with temporal bootstrapping. @@ -313,38 +357,6 @@ def __init__(self, estimator_config: dict, loss_config: ClippedPGLossConfig): # Populated by compute_advantage; surfaced by ppo.py into rollout metrics. self.last_metrics: dict[str, float] = {} - def _compute_raw_advantage_metrics( - self, - advantages: torch.Tensor, - mask: torch.Tensor, - ) -> dict[str, float]: - """Stats on the PRE-whitening advantages, over valid tokens. - - ``normalize_advantages`` rescales advantages to unit std every step, so - the post-whitening spread is 1.0 by construction and carries no - information. The pre-whitening spread does: with lambda=1 the advantage - is ``R - V(s_t)``, so ``adv_raw/std`` is the critic's residual scale and - should SHRINK as the critic improves. A flat ``adv_raw/std`` alongside a - flat ``critic/explained_var`` is independent confirmation that the critic - is not learning. ``whiten_gain`` is the amplification whitening applies - (1/std); a large and growing value means the loss is being scaled up to - compensate for a shrinking signal. - """ - m = mask.bool() - if int(m.sum()) < 2: - return {} - a = advantages[m].float() - std = a.std(unbiased=False) - metrics = { - "adv_raw/mean": a.mean().item(), - "adv_raw/std": std.item(), - "adv_raw/abs_mean": a.abs().mean().item(), - "adv_raw/max_abs": a.abs().max().item(), - } - if self.normalize_advantages: - metrics["adv_raw/whiten_gain"] = (1.0 / (std + 1e-8)).item() - return metrics - def _reward_whiten( self, rewards: torch.Tensor, @@ -494,7 +506,9 @@ def compute_advantage( # Capture the pre-whitening advantage scale (critic-quality diagnostic) # BEFORE normalize_advantages pins the std to 1.0. - self.last_metrics = self._compute_raw_advantage_metrics(advantages, mask) + self.last_metrics = raw_advantage_metrics( + advantages, mask, self.normalize_advantages + ) # Whiten advantages (optional) and zero out masked positions (always) if self.normalize_advantages: @@ -701,15 +715,24 @@ def compute_advantage( # the critic's mask does not cover are silently discarded. returns = scatter_turns_to_anchors(turn_returns, turn_spans, seq_len) + # Pre-whitening advantage scale, on the same token-level mask the + # whitening below uses -- captured BEFORE normalize_advantages pins the + # std to 1.0, exactly as the token-level estimator does. + raw_metrics = raw_advantage_metrics(advantages, mask, self.normalize_advantages) + if self.normalize_advantages: adv_mean = masked_mean(advantages, mask) adv_var = masked_var(advantages, mask, adv_mean) advantages = (advantages - adv_mean) * torch.rsqrt(adv_var + 1e-8) advantages = torch.masked_fill(advantages, ~(mask.bool()), 0) - self.last_metrics = turn_level_metrics( - turn_values, turn_advantages, turn_spans, sample_mask - ) + # Merge rather than assign: turn_level_metrics() would otherwise drop + # adv_raw/*, leaving the rollout dump unable to reconstruct unwhitened + # advantages on turn-level runs. + self.last_metrics = { + **raw_metrics, + **turn_level_metrics(turn_values, turn_advantages, turn_spans, sample_mask), + } return advantages, returns diff --git a/tests/unit/algorithms/test_advantage_estimator.py b/tests/unit/algorithms/test_advantage_estimator.py index 98bd1581f74..07616c08a9f 100644 --- a/tests/unit/algorithms/test_advantage_estimator.py +++ b/tests/unit/algorithms/test_advantage_estimator.py @@ -474,3 +474,71 @@ def test_ev_res_mixed_group_absent_without_group_info(): ) == {} ) + + +def _turn_cfg(normalize): + return { + "turn_gae_gamma": 1.0, + "turn_gae_lambda_value": 1.0, + "turn_gae_lambda_policy": 1.0, + "normalize_advantages": normalize, + } + + +def test_turn_level_reports_raw_advantage_metrics(): + """adv_raw/* must survive on the turn-level path, not just token-level. + + turn_level_metrics() used to REPLACE last_metrics wholesale, so adv_raw/* + never reached the rollout dump under turn_gae -- and the dump reconstructs + unwhitened advantages as ``advantages * adv_raw_std + adv_raw_mean``, so it + was silently degraded to "absent" for exactly the runs the turn estimator + exists for. advantage/turn_std_prenorm is not a substitute: it is computed + over the ``[B, K]`` turn tensor, whereas the whitening applied here is a + token-level masked statistic. + """ + prompt_ids, rewards, mask = _one_group([1.0, 0.0, 0.0, 0.0]) + spans = _turn_spans() + values = torch.zeros_like(mask) + + est = TurnLevelGeneralizedAdvantageEstimator(_turn_cfg(False), _LossCfg()) + raw_adv, _ = est.compute_advantage( + prompt_ids, rewards, mask, values, turn_spans=spans + ) + unnormalized = est.last_metrics + + # Both families are present; neither clobbers the other. + assert {"adv_raw/mean", "adv_raw/std", "adv_raw/abs_mean", "adv_raw/max_abs"} <= ( + unnormalized.keys() + ) + assert "turn/total_turns" in unnormalized + # normalize_advantages off => no whitening is applied, so no gain is reported. + assert "adv_raw/whiten_gain" not in unnormalized + + # With whitening off the returned advantages ARE the raw ones. + m = mask.bool() + torch.testing.assert_close( + unnormalized["adv_raw/std"], raw_adv[m].float().std(unbiased=False).item() + ) + torch.testing.assert_close( + unnormalized["adv_raw/mean"], raw_adv[m].float().mean().item() + ) + + # Same inputs with whitening ON: adv_raw/* is captured BEFORE the rescale, so + # it is unchanged, while the returned advantages are now unit-std. + est_norm = TurnLevelGeneralizedAdvantageEstimator(_turn_cfg(True), _LossCfg()) + white_adv, _ = est_norm.compute_advantage( + prompt_ids, rewards, mask, values, turn_spans=spans + ) + normalized = est_norm.last_metrics + for key in ("adv_raw/mean", "adv_raw/std", "adv_raw/abs_mean", "adv_raw/max_abs"): + assert normalized[key] == pytest.approx(unnormalized[key]) + assert normalized["adv_raw/whiten_gain"] == pytest.approx( + 1.0 / (unnormalized["adv_raw/std"] + 1e-8) + ) + # The whitening divides by sqrt(masked_var(...)), which is UNBIASED, so the + # post-whitening unbiased std is 1.0 while the biased one sits at + # sqrt((n-1)/n). Assert the convention the estimator actually applies. + assert white_adv[m].float().std(unbiased=True).item() == pytest.approx( + 1.0, abs=1e-4 + ) + assert unnormalized["adv_raw/std"] != pytest.approx(1.0) From 7d85a64d85e5bf48db4b8fc9c025dd78efd0aa59 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Fri, 28 Aug 2026 11:11:38 -0700 Subject: [PATCH 66/69] fix privilege critic for r2egym Signed-off-by: Jiaqi Zeng --- nemo_rl/algorithms/swe_privileged_critic.py | 185 ++++++++++++++++-- .../algorithms/test_swe_privileged_critic.py | 148 ++++++++++++++ 2 files changed, 322 insertions(+), 11 deletions(-) diff --git a/nemo_rl/algorithms/swe_privileged_critic.py b/nemo_rl/algorithms/swe_privileged_critic.py index 2833b024d22..3ee31ab2c5c 100644 --- a/nemo_rl/algorithms/swe_privileged_critic.py +++ b/nemo_rl/algorithms/swe_privileged_critic.py @@ -39,10 +39,22 @@ ``a_t ⊥ z | s_t`` and ``E[∇log π(a_t|s_t) · V(s_t, z)] = 0``. The privileged batch must never reach the policy worker. -Field availability was audited over all 9262 instances of the SWE curriculum: -golden patch and test patch are present for 100% of them, but under four -different key names across four source datasets and usually nested inside the -``instance_dict`` JSON string rather than at the top level of ``metadata``. +Field availability: the curriculum draws on FIVE source datasets, and the fix is +recoverable for 100% of them -- but not uniformly. + + * swe-bench-ext, SWE-rebench-V2, nv-internal-1, SWE-Gym carry a patch STRING, + under three different key names, usually nested inside the ``instance_dict`` + JSON string rather than at the top level of ``metadata``. + * R2E-Gym carries no patch string at all. It stores the commit structurally + under ``parsed_commit_content``, which :func:`_resolve_r2e_gym` reassembles + into a unified diff so the critic sees one schema everywhere. + +An earlier audit covered only the first four and concluded "100%", which is why +the first privileged launch died with "no golden patch resolved for 48/512 +rollouts". Re-measured directly over the 7394 collected rollout groups: +SWE-rebench-V2 4769, swe-bench-ext 1613, R2E-Gym 456, nv-internal-1 377, +SWE-Gym 179 -- gold now resolves for all 7394. + ``rubric``/``requirements``/``interface`` exist for only 6.9% (and ``interface`` is already in the agent's prompt), so nothing here depends on them. """ @@ -60,6 +72,32 @@ # alongside ``patch``. GOLD_KEYS: tuple[str, ...] = ("golden_patch", "gold_patch", "patch") +# R2E-Gym is the fifth dataset in this curriculum and the one exception to +# "every instance carries a patch string": it stores NO unified diff at all. +# 456 of the 7394 collected groups (6.2%) are R2E-Gym, and the original audit +# missed them, which is what made the first privileged launch die with +# "no golden patch resolved for 48/512 rollouts". +# +# The reference fix IS present, just structured rather than textual, under +# ``parsed_commit_content`` -- a JSON blob of per-file hunks that we reassemble +# into a real unified diff below. Two other sources were considered and +# rejected: +# * the ``prompt`` field embeds a ```diff block, but measured over all 456 +# instances it covers only PART of the non-test files in 38% of them (it is +# the issue-writer's prompt, not the patch of record) -- so it silently +# under-reports the fix; +# * ``old_file_content``/``new_file_content`` are whole files (median ~190 KB), +# which would blow the token budget on unchanged lines. +# Reassembling the hunks gives 100% coverage AND keeps the critic on ONE schema +# across all five datasets, which is the thing that actually has to be learnable. +R2E_COMMIT_KEY = "parsed_commit_content" +# R2E-Gym leaves FAIL_TO_PASS/PASS_TO_PASS empty (0/456) and states its +# acceptance criterion as ``expected_output_json``: test name -> expected status +# (PASSED / ERROR / FAILED). That is the same role FAIL_TO_PASS plays for the +# other four datasets, so it is emitted in that slot. +R2E_EXPECTED_KEY = "expected_output_json" +_DIFF_LINE_PREFIX = {"context": " ", "deleted": "-", "added": "+"} + # Emitted in this order. Least discriminative first: the block sits ~100k tokens # before the late values that need it, and this model is a 52-layer hybrid with # only 6 attention layers (MEMEM*EMEMEM*...), so mamba state retains RECENT @@ -210,6 +248,120 @@ def _as_lines(value: Any) -> str: return str(value) +def _is_test_path(path: str) -> bool: + """Path heuristic for the fix/tests split -- the FALLBACK only. + + ``relevant_files`` (see :func:`_resolve_r2e_gym`) states the split + authoritatively and covers 100% of this corpus, so this only runs on a + record that lacks it. Deliberately conservative: an earlier version also + treated ``test*.py`` as a test file and thereby swallowed pandas' + ``pandas/util/testing.py`` -- a source module -- leaving that instance with + an empty golden patch. + """ + low = path.lower() + base = low.rsplit("/", 1)[-1] + return ( + base.startswith("test_") + or base.endswith("_test.py") + or base.endswith("_test.go") + or "/tests/" in f"/{low}" + or "/test/" in f"/{low}" + ) + + +def _unified_diff_from_file_diff(fd: dict[str, Any]) -> str: + """Rebuild one file's unified diff from R2E-Gym's structured hunks. + + Emits exactly the ``diff --git`` / ``index`` / ``---`` / ``+++`` / ``@@`` + shape the other four datasets supply verbatim, so the critic sees a single + patch format everywhere. ``modified_entities`` (whole function bodies, which + dwarf the hunks) is deliberately dropped. + """ + path = ((fd.get("header") or {}).get("file") or {}).get("path") or "" + if not path: + return "" + minus = (fd.get("minus_file") or {}).get("path") or f"a/{path}" + plus = (fd.get("plus_file") or {}).get("path") or f"b/{path}" + out = [f"diff --git a/{path} b/{path}"] + idx = fd.get("index_line") or {} + if idx.get("old_commit_hash") and idx.get("new_commit_hash"): + mode = f" {idx['mode']}" if idx.get("mode") else "" + out.append(f"index {idx['old_commit_hash']}..{idx['new_commit_hash']}{mode}") + if fd.get("is_binary_file"): + out.append(fd.get("binary_line") or f"Binary files {minus} and {plus} differ") + return "\n".join(out) + out += [f"--- {minus}", f"+++ {plus}"] + for hunk in fd.get("hunks") or []: + d = hunk.get("descriptor") or {} + o, n = d.get("old_range") or {}, d.get("new_range") or {} + section = d.get("section") or "" + out.append( + f"@@ -{o.get('start', 0)},{o.get('length', 0)} " + f"+{n.get('start', 0)},{n.get('length', 0)} @@" + + (f" {section}" if section else "") + ) + for line in ((hunk.get("line_group") or {}).get("all_lines") or []): + prefix = _DIFF_LINE_PREFIX.get(line.get("type"), " ") + out.append(prefix + (line.get("content") or "")) + return "\n".join(out) + + +def _resolve_r2e_gym(src: dict[str, Any]) -> dict[str, str]: + """Privileged fields for an R2E-Gym instance (no patch string in metadata). + + Returns ``{}`` when this is not an R2E-Gym-shaped record, so the caller can + keep failing loudly on a genuinely broken data path rather than papering + over it with empty strings. + """ + try: + commit = json.loads(src.get(R2E_COMMIT_KEY) or "{}") + except (json.JSONDecodeError, TypeError): + return {} + file_diffs = commit.get("file_diffs") if isinstance(commit, dict) else None + if not file_diffs: + return {} + + # Gold/test split from BOTH available signals, because neither alone is + # right. ``relevant_files`` is R2E-Gym's "primary" file(s) and is narrower + # than the fix -- on one Pillow instance it names ImagePalette.py while the + # commit also fixes Image.py, and R2E-Gym's OWN diff rendering includes + # both. The path heuristic is broader but misfires on source modules that + # merely look test-shaped (pandas/util/testing.py). Union of the two: a file + # is part of the fix unless it looks like a test AND R2E-Gym did not call it + # relevant. Cross-checked against R2E-Gym's own rendering over all 456 + # instances (see the docstring's note on _unified_diff_from_file_diff). + relevant = src.get("relevant_files") + relevant = set(relevant) if isinstance(relevant, list) else set() + + gold_parts, test_parts = [], [] + for fd in file_diffs: + if not isinstance(fd, dict): + continue + text = _unified_diff_from_file_diff(fd) + if not text: + continue + path = ((fd.get("header") or {}).get("file") or {}).get("path") or "" + is_test = _is_test_path(path) and path not in relevant + (test_parts if is_test else gold_parts).append(text) + + # expected_output_json is the acceptance criterion; "name: STATUS" per line + # matches the one-item-per-line shape _as_lines() gives the other datasets. + expected = "" + try: + eo = json.loads(src.get(R2E_EXPECTED_KEY) or "{}") + if isinstance(eo, dict): + expected = "\n".join(f"{k}: {v}" for k, v in eo.items()) + except (json.JSONDecodeError, TypeError): + expected = "" + + return { + "golden_patch": "\n".join(gold_parts), + "test_patch": "\n".join(test_parts), + "fail_to_pass": expected, + "pass_to_pass": "", + } + + def resolve_privilege_fields(env_info: dict[str, Any]) -> dict[str, str]: """Pull the privileged fields out of one rollout's ``extra_env_info``. @@ -234,9 +386,19 @@ def resolve_privilege_fields(env_info: dict[str, Any]) -> dict[str, str]: if isinstance(v, str) and v.strip(): gold = v break + instance_id = str(src.get("instance_id") or md.get("instance_id") or "") + + # R2E-Gym: no patch string anywhere, but the commit is present structurally. + # Only consulted when the textual keys came up empty, so the other four + # datasets take exactly the path they always did. + if not gold: + r2e = _resolve_r2e_gym(src) + if r2e.get("golden_patch"): + return {"instance_id": instance_id, **r2e} + test_patch = src.get("test_patch") return { - "instance_id": str(src.get("instance_id") or md.get("instance_id") or ""), + "instance_id": instance_id, "golden_patch": gold, "test_patch": test_patch if isinstance(test_patch, str) else "", "fail_to_pass": _as_lines(src.get("FAIL_TO_PASS")), @@ -431,16 +593,17 @@ def build_swe_privileged_value_inputs( ) critic_message_logs.append(critic_msgs) - # Golden patch is present for 100% of the audited corpus, so any miss is a - # data-path bug, not a straggler. Fail loudly rather than degrade to a blind - # critic under a privileged label. + # The fix resolves for 100% of the five source datasets, so any miss is a + # data-path or new-dataset bug, not a straggler. Fail loudly rather than + # degrade to a blind critic under a privileged label. if missing: raise ValueError( f"SWE privileged critic: no golden patch resolved for {len(missing)}/" f"{len(message_logs)} rollouts (rows {missing[:8]}...). Expected one of " - f"{GOLD_KEYS} in extra_env_info metadata or its instance_dict. " - "Audited coverage of the SWE curriculum is 100%, so this indicates the " - "metadata did not survive the data path." + f"{GOLD_KEYS} in extra_env_info metadata / its instance_dict, or an " + f"R2E-Gym-style {R2E_COMMIT_KEY!r}. Either the metadata did not " + "survive the data path, or the campaign mixes in a SIXTH dataset with " + "yet another schema — check dataset_name on the failing rows." ) if metrics_out is not None and block_stats: diff --git a/tests/unit/algorithms/test_swe_privileged_critic.py b/tests/unit/algorithms/test_swe_privileged_critic.py index 1d3af0d943d..9bbd68ae043 100644 --- a/tests/unit/algorithms/test_swe_privileged_critic.py +++ b/tests/unit/algorithms/test_swe_privileged_critic.py @@ -79,6 +79,154 @@ def test_resolve_tolerates_unparseable_instance_dict(): assert resolve_privilege_fields(info)["golden_patch"] == "G" +# ------------------------------------------------------- R2E-Gym (no patch string) +def _r2e_file_diff(path, added, deleted=(), section=""): + lines = [{"content": c, "type": "deleted"} for c in deleted] + lines += [{"content": c, "type": "added"} for c in added] + return { + "header": {"file": {"path": path}}, + "index_line": {"old_commit_hash": "aaa", "new_commit_hash": "bbb", "mode": "100644"}, + "minus_file": {"path": f"a/{path}"}, + "plus_file": {"path": f"b/{path}"}, + "is_binary_file": False, + "hunks": [ + { + "descriptor": { + "old_range": {"start": 1, "length": 3}, + "new_range": {"start": 1, "length": 4}, + "section": section, + }, + "line_group": {"all_lines": [{"content": "ctx", "type": "context"}] + lines}, + # whole function bodies; must NOT reach the block + "modified_entities": [{"content": "def f():\n " + "x" * 5000}], + } + ], + } + + +def _r2e_env_info(file_diffs, relevant_files=None, expected_output=None): + """R2E-Gym stores no patch string: the commit is structured, under + parsed_commit_content, and the acceptance criterion is expected_output_json.""" + return _env_info( + instance_id="r2e-1", + dataset_name="R2E-Gym/R2E-Gym-Subset", + instance_dict=json.dumps( + { + "instance_id": "r2e-1", + "FAIL_TO_PASS": [], + "PASS_TO_PASS": [], + "relevant_files": relevant_files or [], + "parsed_commit_content": json.dumps({"file_diffs": file_diffs}), + "expected_output_json": json.dumps(expected_output or {}), + } + ), + ) + + +def test_r2e_gym_reconstructs_a_unified_diff(): + """R2E-Gym is the 6.2% of the corpus with no patch string anywhere; the fix + has to be rebuilt from structured hunks or the whole group is unusable.""" + info = _r2e_env_info( + [_r2e_file_diff("src/app.py", added=["new line"], deleted=["old line"], section="def f():")], + relevant_files=["src/app.py"], + ) + gold = resolve_privilege_fields(info)["golden_patch"] + assert gold.startswith("diff --git a/src/app.py b/src/app.py") + assert "index aaa..bbb 100644" in gold + assert "--- a/src/app.py" in gold and "+++ b/src/app.py" in gold + assert "@@ -1,3 +1,4 @@ def f():" in gold + assert "+new line" in gold and "-old line" in gold and " ctx" in gold + # modified_entities are whole function bodies and would dwarf the hunks + assert "xxxxx" not in gold + + +def test_r2e_gym_splits_fix_from_grading_tests(): + info = _r2e_env_info( + [ + _r2e_file_diff("src/app.py", added=["fix"]), + _r2e_file_diff("tests/test_app.py", added=["assert True"]), + ], + relevant_files=["src/app.py"], + ) + f = resolve_privilege_fields(info) + assert "src/app.py" in f["golden_patch"] and "tests/test_app.py" not in f["golden_patch"] + assert "tests/test_app.py" in f["test_patch"] and "src/app.py" not in f["test_patch"] + + +def test_r2e_gym_relevant_files_does_not_narrow_the_fix(): + """relevant_files names the PRIMARY file only. A second non-test file that it + omits is still part of the accepted fix (R2E-Gym's own diff rendering + includes it), so it must not be misfiled as a grading test.""" + info = _r2e_env_info( + [ + _r2e_file_diff("src/app.py", added=["fix a"]), + _r2e_file_diff("src/helper.py", added=["fix b"]), + ], + relevant_files=["src/app.py"], + ) + gold = resolve_privilege_fields(info)["golden_patch"] + assert "src/app.py" in gold and "src/helper.py" in gold + assert resolve_privilege_fields(info)["test_patch"] == "" + + +def test_r2e_gym_test_shaped_source_module_stays_in_the_fix(): + """pandas/util/testing.py is a source module. A pure path heuristic files it + as a test and leaves that instance with an EMPTY golden patch -- which is + exactly the failure the whole R2E-Gym branch exists to prevent.""" + info = _r2e_env_info( + [ + _r2e_file_diff("pandas/util/testing.py", added=["fix"]), + _r2e_file_diff("pandas/tests/test_testing.py", added=["assert True"]), + ], + relevant_files=["pandas/util/testing.py"], + ) + f = resolve_privilege_fields(info) + assert "pandas/util/testing.py" in f["golden_patch"] + assert "pandas/tests/test_testing.py" in f["test_patch"] + + +def test_r2e_gym_expected_output_becomes_the_acceptance_criterion(): + """FAIL_TO_PASS/PASS_TO_PASS are empty for every R2E-Gym instance; the + expected test statuses play that role and belong in that budget slot.""" + info = _r2e_env_info( + [_r2e_file_diff("src/app.py", added=["fix"])], + relevant_files=["src/app.py"], + expected_output={"test_a": "PASSED", "test_b": "ERROR"}, + ) + f = resolve_privilege_fields(info) + assert f["fail_to_pass"] == "test_a: PASSED\ntest_b: ERROR" + assert f["pass_to_pass"] == "" + + +def test_r2e_gym_branch_does_not_touch_the_other_datasets(): + """A record that HAS a patch string must take exactly the path it always did, + even if it also happens to carry parsed_commit_content.""" + info = _env_info( + instance_id="x", + instance_dict=json.dumps( + { + "patch": "REAL", + "test_patch": "REAL_TESTS", + "FAIL_TO_PASS": ["a::b"], + "parsed_commit_content": json.dumps( + {"file_diffs": [_r2e_file_diff("src/app.py", added=["ignored"])]} + ), + } + ), + ) + f = resolve_privilege_fields(info) + assert f["golden_patch"] == "REAL" + assert f["test_patch"] == "REAL_TESTS" + assert f["fail_to_pass"] == "a::b" + + +def test_unresolvable_record_still_raises(): + """The loud failure must survive: a genuinely broken data path (no patch, no + parsed commit) must not be papered over with an empty block.""" + info = _env_info(instance_id="x", instance_dict=json.dumps({"repo": "r"})) + assert resolve_privilege_fields(info)["golden_patch"] == "" + + @pytest.mark.parametrize( "raw,expected", [ From f338853ffc59483a6b0986a437da06c66e0e7d45 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Sat, 29 Aug 2026 22:11:13 -0700 Subject: [PATCH 67/69] feat(precision): fp32 LM head on both engines (opt-in), ported from super-v3.5 Port of the MiniMax-M1-style fp32 LM head fix (arXiv:2506.13585) from the super-v3.5-posttraining branch. All inert unless the flags are set: - megatron/setup.py: apply_fp32_lm_head() wraps output_layer.forward so the logits GEMM runs in fp32 (megatron_cfg.fp32_lm_head=true|tf32; tf32 uses tensor cores, numerically identical here at near-bf16 speed). Also passes through megatron_cfg.fp32_residual_connection (WARNING: crashes TE LayerNormLinear on this stack - dtype assert in linear_qkv; leave unset). - policy worker: applies the wrapper after model setup. - vllm/patches.py: source-patches vLLM nemotron_h.py compute_logits to keep an fp32 copy of the LM head, gated by NRL_VLLM_FP32_LM_HEAD=1 (set via policy.generation.vllm_cfg.env_vars). Raises if the flag is set but the patch anchor is missing: a one-sided fp32 head is worse than all-bf16. Forwards the env var into vLLM's ray TP workers via ADDITIONAL_ENV_VARS. - vllm_backend.py: invalidates (per _load_weights batch) and eagerly refreshes (post-refit) the cached fp32 head so it tracks refitted weights; engines start on dummy weights, so a stale cache is catastrophic. Validated in-loop on the nano GA opencode PPO run: both engines active, refit drift 0.0487 (dummy->real) then 0 (cache tracking), train/ token_mult_prob_error 1.044 -> 1.037 (steps 1-5 avg vs pre-fix run). Note the gain is model-dependent: nano's logits are small (top ~2.8, bf16 ulp ~0.008) so head rounding is a minor term here; on super-v3.5 (logits 15-30, ulp 0.125-0.25) the same fix cuts the error 1.035 -> 1.015. Co-Authored-By: Claude Fable 5 --- .../ppo_nano_v3_5_swe_opencode_cmh.yaml | 799 ++++++++++++++++++ nemo_rl/models/generation/vllm/patches.py | 101 +++ .../models/generation/vllm/vllm_backend.py | 104 +++ nemo_rl/models/megatron/setup.py | 64 ++ .../policy/workers/megatron_policy_worker.py | 5 + 5 files changed, 1073 insertions(+) create mode 100644 examples/configs/ppo_nano_v3_5_swe_opencode_cmh.yaml diff --git a/examples/configs/ppo_nano_v3_5_swe_opencode_cmh.yaml b/examples/configs/ppo_nano_v3_5_swe_opencode_cmh.yaml new file mode 100644 index 00000000000..ced4836bf8e --- /dev/null +++ b/examples/configs/ppo_nano_v3_5_swe_opencode_cmh.yaml @@ -0,0 +1,799 @@ +# ============================================================================= +# PPO Ultra E2E — SWE Config for GB200 NVL72 (PPO/VAPO port of the GRPO config) +# ============================================================================= +# Duplicated from grpo_nano_v3_5_swe_hsg.yaml and adapted for PPO with a GAE +# critic (run with examples/nemo_gym/run_ppo_nemo_gym.py). Differences vs the +# GRPO parent: +# - `grpo:` section replaced by `ppo:` (GAE adv_estimator, ppo_epochs, +# critic warmup via policy_training_start_step, async_ppo instead of +# async_grpo). max_val_samples / val_batch_size MUST stay null — the +# NeMo-Gym PPO runner rejects preset values and sets both to +# len(val_dataset). +# - new `value:` block — Megatron-backend value model initialized from the +# same nano checkpoint (fresh scalar value head), mirroring the policy's +# parallelism/MoE settings. NOTE: the megatron value worker has only been +# validated on GPT-style dense models; the nano hybrid (mamba+MoE) critic +# is experimental. +# - new `value_loss_fn:` block (clipped MSE, VAPO-style scale=1.0/clip=0.5 +# from the validated 7B math PPO run). +# - loss_fn.force_on_policy_ratio false: PPO trains on the real prev/current +# logprob ratio (GRPO nano forced ratio=1 and leaned on TIS only). +# - grpo's invalid_tool_call_advantage / malformed_thinking_advantage are +# DROPPED: message-level advantage-overwrite penalties are rejected on the +# PPO path (they would break GAE advantage/return consistency). The +# reward-zeroing `reward_penalties` block is the PPO-compatible path; left +# all-off (default) for parity with the GRPO baseline, which also ran with +# them off (its top-level penalize_* keys are inert in this checkout). +# +# Static algorithm, loss, and environment settings live here. Per-run +# overrides (model path, data paths, parallelism, batch sizes) are applied by +# the launch script (scripts/swe/ppo/nano_on_main.sh) or Hydra CLI. +# ============================================================================= + +# ============================================================================= +# Cluster — overridden by launch script for parallelism changes +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 128 + segment_size: 16 + +# ============================================================================= +# Checkpointing +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/ppo_ultra_e2e" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 3 + ft_keep_latest_k: 1 + ft_save_period: 1 + save_optimizer: true + checkpoint_must_save_by: "00:03:30:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# PPO Algorithm (replaces the parent's grpo: section) +# ============================================================================= +ppo: + num_prompts_per_step: 16 + num_generations_per_prompt: 32 + max_rollout_turns: 1 + max_num_epochs: 4 + max_num_steps: 1000000 + # One optimizer pass per rollout batch (PPO-epochs>1 reuses the batch). + # This is the ACTOR's epoch count. + ppo_epochs: 1 + # Critic (value) passes over each rollout batch. null = same as ppo_epochs + # (coupled, the historical behavior); must be >= ppo_epochs. The surplus runs + # as extra critic-only passes, fitting the critic harder without + # over-training the actor. + critic_ppo_epochs: null + # Re-score the batch with a forward-only critic pass after the final update, + # adding critic/explained_var_post_update + critic/loss_post_update. + # critic/explained_var is always the PRE-update EV (from the rollout-time + # values GAE consumed); critic/loss is from the last training pass. Costs one + # extra critic forward per step. + log_post_update_critic_metrics: false + # Critic-only warmup: policy training is skipped for the first N steps while + # the (fresh-head) value model trains on rollouts from the frozen policy. + # Overridden by the launch script. + policy_training_start_step: 50 + val_period: 10000 + val_at_start: false + val_at_end: false + overlong_filtering: false + # NeMo-Gym principle: the validation set is used verbatim (no hidden pre/post + # processing). run_ppo_nemo_gym.py rejects a preset max_val_samples and sets + # both to len(val_dataset), so they must start null. + max_val_samples: null + val_batch_size: null + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + # loss_fn.reference_policy_kl_penalty is 0 (SWE baseline), so skip the + # reference-policy logprob forward. The launch script flips this to false + # automatically when a nonzero KL penalty is requested. + skip_reference_policy_logprobs_calculation: true + seq_logprob_error_threshold: 2 + log_rollout_dump: true + rollout_dump_period: 5 # dump every N PPO steps (when log_rollout_dump is true) + + # GAE with VAPO decoupled lambdas (defaults match the validated 7B math PPO + # run: lambda_value=1.0, length-adaptive lambda_policy with alpha=0.05). + # + # name: "gae" -> token-level GAE (the historical path, unchanged). + # name: "turn_gae" -> TURN-level GAE: one assistant message is one action, + # V(s_k) is read at that turn's first token (the value head is right-shifted, + # so that position has seen the whole preceding observation and none of the + # action), the advantage is constant across the turn's tokens, and the critic + # is supervised at one anchor per turn. Only the turn_gae_* keys apply. + # + # Why it exists: at the token level lambda has effective horizon 1/(1-lambda) + # TOKENS, so on a ~45k-token SWE rollout anything below 1 severs the terminal + # reward — which is why length_adaptive_alpha pushes lambda to 1 - 1.5e-5 and + # GAE degenerates to the pure baseline A_t = R - V(s_t) with no temporal + # credit assignment at all. Over ~92 turns, lambda=0.97 is a 33-turn horizon. + # See research/ppo/turn_level_critic_plan.md. + adv_estimator: + name: "gae" + gae_lambda: 1.0 + gae_gamma: 1 + normalize_advantages: true + gae_lambda_value: 1.0 + gae_lambda_policy: 1 + # Length-adaptive lambda_policy = 1 - 1/(alpha*l). 0 = disabled. + length_adaptive_alpha: null + + # --- decomposed group baseline + residual critic --- + # research/ppo/residual_critic_report.md. true => the rollout group supplies + # the task baseline B(X) as a leave-one-out mean and the critic is trained + # only on the within-task residual C(s) = R - B_LOO; critic values/returns + # are then in RESIDUAL space and the go/no-go metric is critic/ev_res. + # Requires gae_gamma: 1. Orthogonal to token-vs-turn granularity. + # + # Motivation, measured on this exact workload: the absolute critic puts + # 88-93% of its output variance between tasks yet reaches EV 0.157 — below + # even a two-value "all-fail vs rest" lookup (0.330) and far below the free + # leave-one-out group baseline (0.553). Substituting it for that baseline + # RAISES advantage variance 1.67-2.09x. false still logs critic/ev_res and + # residual/* so an absolute run is comparable on the same axis. + residual_baseline: false + + # --- turn-level GAE (used only when name: "turn_gae") --- + # All three lambdas/gamma are required in that mode and are never silently + # defaulted, so a sweep cannot appear to have run when it did not. + turn_gae_gamma: 1.0 + # lambda_value=1.0 keeps the critic's target Monte-Carlo (G_k == R): the same + # target it regresses on today, just at ~92 anchors instead of ~45k tokens. + # Lower it only deliberately — bootstrapped targets feed the critic's own + # (still weak) estimates back into its own supervision. + turn_gae_lambda_value: 1.0 + # lambda_policy is THE knob. 1.0 reproduces today's A_k = R - V(s_k); + # 0.99 ~ 100-turn horizon, 0.97 ~ 33 turns, 0.95 ~ 20 turns. + turn_gae_lambda_policy: 1.0 + # The critic is supervised at ONE anchor per turn, every turn weighted + # equally (under token weighting the longest decile of turns owns ~43% of + # the critic loss). + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + # NOTE: grpo's penalize_invalid_tool_call / invalid_tool_call_advantage / + # penalize_malformed_thinking / malformed_thinking_advantage are intentionally + # absent: setup() raises NotImplementedError if the *_advantage keys are set + # on the PPO path (advantage overwrites break GAE consistency). + + async_ppo: + enabled: true + # lag-1: at most one policy version off-policy (recommended/validated). + # During critic warmup the frozen actor makes any-age trajectories + # on-policy for free (warmup_max_trajectory_age_steps; null => same value). + max_trajectory_age_steps: 1 + warmup_max_trajectory_age_steps: 8 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + # Drop the incomplete (survivorship-biased) frontier target on resume and + # regenerate it fresh — unbiased batches at the cost of a one-target + # generation bubble per resume. + drop_incomplete_targets_on_restore: false + # Throttle collector/replay-buffer per-rollout heartbeat prints. + log_every: 500 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + # Defaults from the validated 7B math PPO run (symmetric 0.2 clip + dual-clip + # c=3), not the GRPO-DAPO asymmetric 0.2/0.28. Overridable from the script. + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: 3 + use_on_policy_kl_approximation: true + # Async PPO requires the off-policy importance-sampling correction. + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + # PPO trains on the real prev/current-logprob ratio; forcing ratio=1 (as the + # GRPO nano config does) would disable PPO clipping entirely. + force_on_policy_ratio: false + use_kl_in_reward: false + +# ============================================================================= +# Value Loss (MSE with PPO-style clipping) +# ============================================================================= +value_loss_fn: + scale: 1.0 + cliprange: 0.5 + # Weight of homogeneous (all-fail / all-pass) groups in the value loss. Only + # meaningful with ppo.adv_estimator.residual_baseline: those groups have Y = 0 + # for every sibling, so they contribute EXACTLY zero target variance while + # still costing ~58% of critic FLOPs on this pool (54.1% all-fail, 2.4% + # all-pass; only 43.5% of groups are mixed). Keep at 1.0 for the clean + # ablation — they are not dead weight, they are the shrinkage that enforces + # E[C | X] = 0 and suppresses between-task leakage. Lower it only to buy more + # mixed-group exposure per unit wall-clock. + homogeneous_group_weight: 1.0 + +# ============================================================================= +# Policy +# ============================================================================= +policy: + model_name: "/lustre/fsw/portfolios/llmservice/users/wdai/megatron-lm-ultra/checkpoints/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-288k-nano-loss-032026/iter_0007000/hf" + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 512 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 131072 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 32 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + use_gloo_process_groups: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + # community_import.py reads megatron_cfg["gradient_accumulation_fusion"] directly + # (no default) on the HF->megatron conversion path; absent here it raises KeyError. + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled + mtp_loss_scaling_factor: 0.3 + mtp_use_repeated_layer: true + # null disables MTP entirely (matches the upstream mtp=0 run's behavior). This + # checkout's hybrid model gates MTP on `mtp_num_layers is not None` + # (hybrid_model.py:544): 0 still enters the MTP block and asserts > 0, and 1 runs + # MTP forward on only the last-stage ranks (suspected cause of the EXPERT_MODEL_ + # PARALLEL NCCL collective timeout). null makes that check False -> MTP is skipped + # on all ranks, no partial-rank EP collective. Requires the schema change making + # MegatronConfig.mtp_num_layers accept None (nemo_rl/models/policy/__init__.py). + mtp_num_layers: 5 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 3.0e-6 + min_lr: 3.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 3e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + # FP8 — overridden by precision recipe in launch script + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo + # were dropped from this checkout's megatron-bridge CheckpointConfig (submodule + # 554c7b9); leaving them in raises InstantiationException "Unexpected config keys". + # Removed for the current-repo reproduction. + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + # Must be a multiple of minimum_pad_factor = cp_size*2*tp_size (see + # nemo_rl/models/megatron/data.py). With CP=16 + TP=4 + sequence_parallel that is + # 128; the old value (=TP=4) made get_logprobs assert and abort every training + # step. Formula matches examples/nemo_gym/grpo_nanov3.yaml (TP*CP*2). + make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 11001 + port_range_high: 15000 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 1 + gpu_memory_utilization: 0.8 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + # reasoning_parser_plugin lives at vllm_cfg top level in this checkout + # (vllm_worker_async reads vllm_cfg["reasoning_parser_plugin"]); passing it + # inside http_server_serving_chat_kwargs makes OpenAIServingChat raise + # "unexpected keyword argument 'reasoning_parser_plugin'". Path updated to this + # repo's location (the parser moved out of nemo_rl/utils/). + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + + vllm_kwargs: + attention_backend: FLASH_ATTN + moe_backend: triton + mamba_ssm_cache_dtype: "float32" + # Route TP allreduces through NCCL instead of vLLM's custom CUDA-IPC + # spin kernel. Jobs 37606/38016 each lost a generation engine ~15 min in: + # the NCCL flight recorder showed all 4 TP ranks with their next + # collective enqueued-but-never-started, i.e. all four streams parked in + # the custom-allreduce kernel ahead of it (symmetric cross-GPU deadlock, + # different node each time). Costs a few % decode latency on TP-4. + disable_custom_all_reduce: true + compilation_config: + cudagraph_mode: PIECEWISE + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 64 + +# ============================================================================= +# Value Model (GAE critic) — Megatron backend, initialized from the policy +# checkpoint with a fresh scalar value head. Parallelism/MoE settings mirror +# the policy block; the launch script overrides them together. +# ============================================================================= +value: + model_name: ${policy.model_name} + tokenizer: + name: ${value.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: ${policy.train_global_batch_size} + train_micro_batch_size: ${policy.train_micro_batch_size} + logprob_batch_size: ${policy.logprob_batch_size} + max_total_sequence_length: ${policy.max_total_sequence_length} + precision: "bfloat16" + + # --- privileged (asymmetric) critic for SWE --- + # research/ppo/privileged_residual_critic_plan.md. The critic sees the accepted + # fix and the grading tests; the policy never does. Audited over all 9262 + # instances: golden patch + test patch are present for 100% of them, and none + # of them appear anywhere in the agent's context (verified against a real + # 408k-char rollout), so this is genuine information, not just a cheaper + # computation of something already visible. + # + # The block is prefixed BEFORE the first assistant token — mandatory, since the + # value head is causal and ~150 assistant turns are spread across the context; + # appending would be a silent no-op. It is byte-identical for all 16 siblings + # of a group, so it cannot introduce a within-task confound. + # + # Orthogonal to both other critic axes: composes with residual OR absolute + # targets, and with token-level OR turn_gae (anchors are remapped into the + # augmented layout). Enabling it raises value.max_total_sequence_length by the + # sum of the caps below automatically. + swe_privileged_critic: + enabled: false + # Single TOTAL token budget for the reference block, so the value model's + # sequence budget is exactly policy_len + this + slack. Measured over 400 + # instances the untruncated block is median 5467 / p90 17683 / p99 77911 + # tokens, so 32768 truncates ~4.8%. Budget is spent in priority order + # (fail_to_pass, golden_patch, test_patch, pass_to_pass) and anything cut is + # marked inline with "... [truncated]" plus a block-level note. Watch + # privilege/frac_truncated to see what the budget is discarding. + max_total_tokens: 32768 + + reward_model_cfg: + enabled: false # only used for the DTensor V2 value worker + reward_model_type: "regression" + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 32 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${value.precision} + sequence_parallel: true + + # MoE — mirror the policy + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + use_gloo_process_groups: false + + # Compute — mirror the policy + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — fully disabled for the critic. The value model replaces the LM + # head with a scalar head, so MTP next-token heads are dead weight; null + # skips the MTP block on all ranks (no partial-rank EP collective). + mtp_loss_scaling_factor: 0.0 + mtp_use_repeated_layer: true + mtp_num_layers: null + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + # Critic LR from the validated 7B math PPO run (fresh value head trains + # faster than the pretrained actor). Overridden by the launch script. + lr: 6.0e-6 + min_lr: 6.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${value.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${value.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${value.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 6.0e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + + env_vars: null + + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + + make_sequence_length_divisible_by: ${mul:${mul:${value.megatron_cfg.tensor_model_parallel_size}, ${value.megatron_cfg.context_parallel_size}}, 2} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + use_multiple_dataloader: false + train: + data_path: null # Set by launch script + validation: + data_path: null # Set by launch script + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + SWE Agents +# ============================================================================= +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + # Reuse pre-built gym venvs. With false, all 3 gym servers (policy_model, + # swe_agents_{train,val}) rebuild the editable nemo-gym package in parallel and + # contend on the same uv distribution-cache lock on Lustre; policy_model then + # times out (default 300s) and NemoGym spinup fails. Mirrors the qwen-30b SWE run. + skip_venv_if_present: true + num_gpu_nodes: 0 + port_range_low: 15001 + port_range_high: 20000 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - responses_api_agents/swe_agents/configs/swebench_opencode_training.yaml + swe_agents_train: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}} + swebench_agent_timeout: 2700 + swebench_tests_timeout: 900 + run_with_mixed_prompts: true + dataset_path: ${data.train.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swerebench/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/nv_internal/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/r2e_gym/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/sweb.eval.arm64.{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swebench/swe-bench.eval.arm64.{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/swebenchpro_ots/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/images_0501/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/images_0402/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/all_images/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swegym/sweb.eval.arm64.{instance_id}.sif" + swe_agents_val: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}} + swebench_agent_timeout: 3600 + dataset_path: ${data.validation.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter} + use_absolute_ip: true + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "logs" + num_val_samples_to_print: 0 + wandb_enabled: true + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "ultra-v3-swe-e2e" + name: "ppo-ultra-e2e" + tensorboard: {} + mlflow: + experiment_name: "ppo-ultra-e2e" + run_name: "ppo-ultra-e2e" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels / Token IDs / penalize_* — carried over from the GRPO config +# for parity. NOTE: these top-level keys are INERT in this checkout (nothing +# reads them; the GRPO baseline also ran with them having no effect). The +# PPO-compatible penalty mechanism is the structured `reward_penalties:` block +# (reward-zeroing, pre-GAE) — left at its all-off default to match the GRPO +# baseline's effective behavior. +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + + +# Reward-zeroing penalties applied to NeMo-Gym rollout results. +reward_penalties: + penalize_duplicated_reasoning: true + penalize_empty_final_answer: true + penalize_unwanted_tokens: true + penalize_malformed_think_tag: true + # Optional model/tokenizer-specific token IDs. Add token IDs that should + # be penalized when emitted by the model. Unwanted token IDs must be + # specified explicitly; think-tag IDs are inferred when each tag encodes + # to one token. Example: {unwanted: [2], think_open: 12, think_close: 13} + token_ids: {unwanted: [2], think_open: 12, think_close: 13} \ No newline at end of file diff --git a/nemo_rl/models/generation/vllm/patches.py b/nemo_rl/models/generation/vllm/patches.py index 1d7fbfbdbb0..f12b88345a5 100644 --- a/nemo_rl/models/generation/vllm/patches.py +++ b/nemo_rl/models/generation/vllm/patches.py @@ -92,6 +92,9 @@ def _patch_vllm_init_workers_ray( "NCCL_CUMEM_ENABLE", "NCCL_NVLS_ENABLE", "RAY_ENABLE_UV_RUN_RUNTIME_ENV", + # fp32 LM head opt-in must reach the TP ray workers that execute the + # model, not just the leader actor (see _patch_vllm_nemotron_h_fp32_lm_head). + "NRL_VLLM_FP32_LM_HEAD", *(extra_env_vars or []), ] additional_env_str = ", ".join(f'"{env_var}"' for env_var in additional_env_vars) @@ -167,6 +170,103 @@ def _patch_vllm_llama_eagle3_own_lm_head(logger) -> None: logger.info("Successfully patched llama_eagle3 lm_head ownership.") +def _patch_vllm_nemotron_h_fp32_lm_head(logger) -> None: + """Compute NemotronH logits with an fp32 LM head (MiniMax-M1-style). + + bf16 rounding of the logits GEMM output is the dominant contributor to + generation/training logprob mismatch (train/token_mult_prob_error). With + this patch the sampled-token logprobs come from fp32 logits, matching a + trainer that enables megatron_cfg.fp32_lm_head. + + This must be a source patch (not a monkeypatch): the model executes in + vLLM's EngineCore worker subprocesses, which import vllm independently of + this process. The patched code is opt-in at runtime via + NRL_VLLM_FP32_LM_HEAD=1 (set it through policy.generation.vllm_cfg.env_vars + so worker processes inherit it). Costs one fp32 copy of the vocab-sharded + head weight per rank. + + IMPORTANT: a one-sided fp32 head (trainer fp32, vLLM bf16) is WORSE than + all-bf16 — two bf16 heads round to the same grid and their errors partially + cancel. So if the feature is requested but the patch cannot be applied, + fail loudly instead of degrading silently. + """ + fp32_head_requested = os.environ.get("NRL_VLLM_FP32_LM_HEAD", "0") == "1" + + try: + file_to_patch = _get_vllm_file("model_executor/models/nemotron_h.py") + except RuntimeError: + if fp32_head_requested: + raise RuntimeError( + "NRL_VLLM_FP32_LM_HEAD=1 is set but nemotron_h.py could not be " + "located in the vLLM installation. Refusing to run with a " + "one-sided fp32 LM head (it is worse than all-bf16)." + ) + logger.warning("Could not locate nemotron_h.py for the fp32 LM head patch.") + return + + old_snippet = """ logits = self.logits_processor(self.lm_head, hidden_states) + return logits""" + new_snippet = """ import os as _os + + if _os.environ.get("NRL_VLLM_FP32_LM_HEAD", "0") == "1": + # NeMo-RL patch: fp32 LM head (MiniMax-M1-style). bf16 rounding of + # the logits is the dominant gen/train logprob mismatch source. + _fp32_head = getattr(self, "_nrl_lm_head_fp32", None) + if _fp32_head is None and not torch.cuda.is_current_stream_capturing(): + # Skipped under graph capture: an allocation there lives in the + # graph's memory pool and is not valid for later eager replays. + # Capture output is discarded anyway, so bf16 is fine for it. + import copy as _copy + + _fp32_head = _copy.deepcopy(self.lm_head).float() + # object.__setattr__ bypasses nn.Module.__setattr__: registering + # this as a submodule would add a vocab-sized parameter to + # named_parameters(), which the refit weight mapping is built from. + object.__setattr__(self, "_nrl_lm_head_fp32", _fp32_head) + self._nrl_lm_head_fp32_dirty = False + print( + "[fp32_lm_head] built fp32 head in forward shape=%s" + % (tuple(_fp32_head.weight.shape),), + flush=True, + ) + elif _fp32_head is not None and getattr( + self, "_nrl_lm_head_fp32_dirty", False + ): + # Refreshed in place: replacing the module would leave any + # captured CUDA graph pointing at the old storage. + _fp32_head.weight.data.copy_(self.lm_head.weight) + if getattr(_fp32_head, "bias", None) is not None: + _fp32_head.bias.data.copy_(self.lm_head.bias) + self._nrl_lm_head_fp32_dirty = False + print("[fp32_lm_head] refreshed cached head in forward", flush=True) + if _fp32_head is not None: + return self.logits_processor(_fp32_head, hidden_states.float()) + logits = self.logits_processor(self.lm_head, hidden_states) + return logits""" + + with _locked_file_patch(file_to_patch) as (content, write_back): + if "NRL_VLLM_FP32_LM_HEAD" in content: + logger.info("NemotronH fp32 LM head patch already present.") + return + if content.count(old_snippet) != 1: + if fp32_head_requested: + raise RuntimeError( + "NRL_VLLM_FP32_LM_HEAD=1 is set but the fp32 LM head patch " + f"anchor was not found exactly once in {file_to_patch} " + "(vLLM version changed?). Refusing to run with a one-sided " + "fp32 LM head (it is worse than all-bf16)." + ) + logger.warning( + "NemotronH fp32 LM head patch anchor not found exactly once " + "in %s; patch not applied.", + file_to_patch, + ) + return + write_back(content.replace(old_snippet, new_snippet, 1)) + + logger.info("Applied NemotronH fp32 LM head source patch.") + + def _patch_vllm_hermes_tool_parser_thread_safety(logger) -> None: """Patch Hermes2ProToolParser.__init__ to cache tokenizer calls. @@ -295,3 +395,4 @@ def _apply_vllm_patches( _patch_vllm_llama_eagle3_own_lm_head(patch_logger) _patch_vllm_hermes_tool_parser_thread_safety(patch_logger) + _patch_vllm_nemotron_h_fp32_lm_head(patch_logger) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 4b693403c99..bb7082dea3e 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import gc +import os import re import traceback from typing import Any @@ -352,6 +353,107 @@ def _load_weights(self, weights): self.model_runner.model.load_weights(weights=policy_weights) self._load_draft_weights(draft_weights) + # Every refit transport funnels through here (IPC ZMQ and collective on + # this branch), so this is where the fp32 LM head cache is invalidated. + self._mark_fp32_lm_head_dirty() + + def _get_drafter_model(self) -> Any: + draft_owner = getattr(self.model_runner, "drafter", None) + return getattr(draft_owner, "model", None) if draft_owner else None + + def _mark_fp32_lm_head_dirty(self) -> None: + """Flag the NRL_VLLM_FP32_LM_HEAD cached head for refresh. + + The patch keeps an fp32 copy of lm_head that the just-loaded weights + supersede; compute_logits refreshes it in place on the next forward. + Set on every batch of a multi-batch refit so the copy stays dirty until + the whole update lands. + """ + if os.environ.get("NRL_VLLM_FP32_LM_HEAD", "0") != "1": + return + marked = [] + for label, model in ( + ("policy", self.model_runner.model), + ("drafter", self._get_drafter_model()), + ): + if model is None: + continue + if getattr(model, "_nrl_lm_head_fp32", None) is not None: + model._nrl_lm_head_fp32_dirty = True + marked.append(label) + else: + # No cache yet: the first compute_logits after this refit builds + # it, which is already correct (post-refit) weights. + marked.append(f"{label}:none-yet") + # Logged once per worker: _load_weights runs per refit batch. + if not getattr(self, "_nrl_fp32_dirty_logged", False): + self._nrl_fp32_dirty_logged = True + print(f"[fp32_lm_head] refit marked: {marked}", flush=True) + + def _sync_fp32_lm_head(self) -> None: + """Rebuild the fp32 LM head cache from the weights refit just loaded. + + The NRL_VLLM_FP32_LM_HEAD source patch keeps an fp32 copy of lm_head. + Engines start on dummy weights, so that copy is stale after every + refit. Building it here (eager, post-refit) rather than lazily in + compute_logits also keeps the allocation out of any CUDA graph pool. + + Covers the drafter too when speculative decoding is on: it is a + separate module with its own head and its own refit stream. + """ + if os.environ.get("NRL_VLLM_FP32_LM_HEAD", "0") != "1": + return + for label, model in ( + ("policy", self.model_runner.model), + ("drafter", self._get_drafter_model()), + ): + if model is not None: + self._sync_fp32_lm_head_for(label, model) + + def _sync_fp32_lm_head_for(self, label: str, model: Any) -> None: + lm_head = getattr(model, "lm_head", None) + if lm_head is None: + # Drafters commonly tie their head to the policy's; nothing to sync. + print( + f"[fp32_lm_head] {label} ({type(model).__name__}) has no lm_head; skipping", + flush=True, + ) + return + + cached = getattr(model, "_nrl_lm_head_fp32", None) + if cached is None: + import copy + + cached = copy.deepcopy(lm_head).float() + # Bypass nn.Module.__setattr__: registering this as a submodule + # would add a vocab-sized parameter to named_parameters(), which + # the refit weight mapping is built from. + object.__setattr__(model, "_nrl_lm_head_fp32", cached) + print( + f"[fp32_lm_head] {label}: built fp32 head cache after refit " + f"shape={tuple(cached.weight.shape)}", + flush=True, + ) + else: + probe = slice(0, 16) + drift = ( + ( + cached.weight.flatten()[probe] + - lm_head.weight.flatten()[probe].float() + ) + .abs() + .max() + .item() + ) + cached.weight.data.copy_(lm_head.weight) + if getattr(cached, "bias", None) is not None: + cached.bias.data.copy_(lm_head.bias) + print( + f"[fp32_lm_head] {label}: refreshed fp32 head cache after refit " + f"(pre-refresh drift={drift:.6g})", + flush=True, + ) + model._nrl_lm_head_fp32_dirty = False @wrap_with_nvtx_name("vllm_internal_worker_extension/update_weights_via_ipc_zmq") def update_weights_via_ipc_zmq(self) -> bool: @@ -380,6 +482,7 @@ def update_weights_via_ipc_zmq(self) -> bool: process_weights_after_loading( self.model_runner.model, self.model_config, self.device ) + self._sync_fp32_lm_head() self.zmq_socket.send(IPCProtocol.ACK.value.encode()) break @@ -468,6 +571,7 @@ def update_weights_from_collective(self) -> bool: process_weights_after_loading( self.model_runner.model, self.model_config, self.device ) + self._sync_fp32_lm_head() self._maybe_process_fp8_kv_cache() except Exception as e: diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 4367ac4cada..09414da8b42 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -356,6 +356,62 @@ def _resolve_iter_dir_from_root(path: str, not_found_msg: str) -> str: return os.path.join(path, iter_subdirs[-1]) +def apply_fp32_lm_head(model_chunks: list, use_tf32: bool = False) -> None: + """Run the LM output-layer GEMM in fp32 (MiniMax-M1-style, arXiv:2506.13585). + + bf16 rounding of the logits (magnitude ~15-30, bf16 ulp 0.125-0.25) is the + dominant contributor to generation/training logprob mismatch + (train/token_mult_prob_error). Upcasting the head input and weight to fp32 + removes that rounding. The casts are part of the autograd graph, so + training gradients flow to the bf16 weight through the fp32 cast. + + Note: has no effect on the fused linear+CE path + (megatron_cfg.use_fused_linear_logprobs), which bypasses output_layer's + standalone forward. + + With ``use_tf32`` (megatron_cfg.fp32_lm_head: "tf32"), the fp32 head GEMM + runs with TF32 tensor cores. The inputs are exact bf16 values (<= 8-bit + mantissa), so TF32's 10-bit input rounding loses nothing; accumulation and + output stay fp32. Numerically equivalent to full fp32 here, at near-bf16 + tensor-core throughput. + """ + if not isinstance(model_chunks, (list, tuple)): + model_chunks = [model_chunks] + for chunk in model_chunks: + module = chunk + while hasattr(module, "module"): + module = module.module + output_layer = getattr(module, "output_layer", None) + if output_layer is None: + continue # not the last pipeline stage + original_forward = output_layer.forward + + def _fp32_forward( + input_, + *args, + weight=None, + _orig_forward=original_forward, + _layer=output_layer, + _tf32=use_tf32, + **kwargs, + ): + w = weight if weight is not None else _layer.weight + if not _tf32: + return _orig_forward(input_.float(), *args, weight=w.float(), **kwargs) + prev = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = True + try: + return _orig_forward(input_.float(), *args, weight=w.float(), **kwargs) + finally: + torch.backends.cuda.matmul.allow_tf32 = prev + + output_layer.forward = _fp32_forward + print( + "[fp32_lm_head] output layer will compute logits in fp32" + + (" (tf32 tensor cores)" if use_tf32 else "") + ) + + def validate_model_paths(config: PolicyConfig) -> tuple[str, str, bool]: """Validate and setup model paths. @@ -561,6 +617,14 @@ def setup_model_config( if "layernorm_epsilon" in config["megatron_cfg"]: model_cfg.layernorm_epsilon = config["megatron_cfg"]["layernorm_epsilon"] + # Optional fp32 residual stream. Reduces generation/training logprob + # mismatch (train/token_mult_prob_error) by accumulating the residual in + # fp32; supported by both transformer and mamba layers. + if "fp32_residual_connection" in config["megatron_cfg"]: + model_cfg.fp32_residual_connection = config["megatron_cfg"][ + "fp32_residual_connection" + ] + # Validate chunking configuration _validate_chunking_config(config) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index ea605b00591..908b4cfd9e7 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -388,6 +388,11 @@ def __init__( self.mcore_state = model_and_optimizer_state.state self.model = model_and_optimizer_state.model + _fp32_lm_head = self.cfg["megatron_cfg"].get("fp32_lm_head", False) + if _fp32_lm_head: + from nemo_rl.models.megatron.setup import apply_fp32_lm_head + + apply_fp32_lm_head(self.model, use_tf32=(_fp32_lm_head == "tf32")) self.optimizer = model_and_optimizer_state.optimizer self.scheduler = model_and_optimizer_state.scheduler self.checkpointing_context = model_and_optimizer_state.checkpointing_context From 23229cdfa6783b9b06912f171b462e3d69b9f946 Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Wed, 9 Sep 2026 20:56:08 -0700 Subject: [PATCH 68/69] fix(vllm-http): return 400 for engine-side context-length ValueError instead of opaque 500 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The /v1/chat/completions endpoint already converts VLLMValidationError (vLLM's tokenization-time length check) into a 400 with the error message, so the Gym proxy can recognize context-length overflow and end the turn gracefully with finish_reason="length". But the same overflow also surfaces as a plain ValueError from the engine-side length check ("Input length (N) exceeds model's maximum context length (M)"), e.g. after required_prefix_token_ids replacement resizes the prompt past max_model_len. That exception type sailed past the handler and escaped as an HTTP 500 whose body is Starlette's generic "Internal Server Error" — the actual message never reaches the client. The Gym retry layer treats any 500 as transient, so every over-length agent turn burned 3 identical failed requests and the graceful handler could never match (147 retried 500s vs 123 such ValueErrors in a single 4h SWE training job). Add an except ValueError clause that converts messages containing "maximum context length" to the same 400 contract and re-raises everything else unchanged. End-to-end: 400 is not in the Gym client's retry codes and the body carries the message, so the existing out-of-context handler fires on the first attempt. Co-Authored-By: Claude Fable 5 --- .../generation/vllm/vllm_worker_async.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index e7d2b1eaa6b..8e984c4920a 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -796,6 +796,27 @@ async def create_chat_completion( }, status_code=400, ) + except ValueError as e: + # The same overflow can also surface as a plain ValueError from + # the engine-side length check (e.g. "Input length (N) exceeds + # model's maximum context length (M)" after prefix-token + # replacement resizes the prompt). Without this clause it + # escapes as an opaque HTTP 500 ("Internal Server Error" body), + # which the Gym client blindly retries 3x and can never map to + # its graceful finish_reason="length" handling. Convert to the + # same 400 contract as VLLMValidationError above. + if "maximum context length" not in str(e): + raise + return JSONResponse( + content={ + "error": { + "message": str(e), + "type": "invalid_request_error", + "code": 400, + } + }, + status_code=400, + ) if isinstance(generator, ErrorResponse): return JSONResponse( From 8396fd2c61b11321f5a8643fe9206f6d7b64d14c Mon Sep 17 00:00:00 2001 From: Jiaqi Zeng Date: Sat, 12 Sep 2026 23:06:15 -0700 Subject: [PATCH 69/69] configs: alignfix variant of swe opencode cmh yaml Identical to ppo_nano_v3_5_swe_opencode_cmh.yaml except config_paths points at Gym's swebench_opencode_training_alignfix.yaml (network time-travel block enabled + 50/50 deterministic prompt mix with the verbatim SWE-V eval scaffold). Used by super_v24_ppo_opencode_scopemix1300_alignfix.sh. Must be committed: the launch snapshot only carries tracked files (first alignfix links 3697201/3697209 failed on the missing file; snapshot was hot-patched with an identical copy). Co-Authored-By: Claude Fable 5 --- ...o_nano_v3_5_swe_opencode_cmh_alignfix.yaml | 809 ++++++++++++++++++ 1 file changed, 809 insertions(+) create mode 100644 examples/configs/ppo_nano_v3_5_swe_opencode_cmh_alignfix.yaml diff --git a/examples/configs/ppo_nano_v3_5_swe_opencode_cmh_alignfix.yaml b/examples/configs/ppo_nano_v3_5_swe_opencode_cmh_alignfix.yaml new file mode 100644 index 00000000000..908e88f4b28 --- /dev/null +++ b/examples/configs/ppo_nano_v3_5_swe_opencode_cmh_alignfix.yaml @@ -0,0 +1,809 @@ +# ============================================================================= +# PPO Ultra E2E — SWE Config for GB200 NVL72 (PPO/VAPO port of the GRPO config) +# ============================================================================= +# Duplicated from grpo_nano_v3_5_swe_hsg.yaml and adapted for PPO with a GAE +# critic (run with examples/nemo_gym/run_ppo_nemo_gym.py). Differences vs the +# GRPO parent: +# - `grpo:` section replaced by `ppo:` (GAE adv_estimator, ppo_epochs, +# critic warmup via policy_training_start_step, async_ppo instead of +# async_grpo). max_val_samples / val_batch_size MUST stay null — the +# NeMo-Gym PPO runner rejects preset values and sets both to +# len(val_dataset). +# - new `value:` block — Megatron-backend value model initialized from the +# same nano checkpoint (fresh scalar value head), mirroring the policy's +# parallelism/MoE settings. NOTE: the megatron value worker has only been +# validated on GPT-style dense models; the nano hybrid (mamba+MoE) critic +# is experimental. +# - new `value_loss_fn:` block (clipped MSE, VAPO-style scale=1.0/clip=0.5 +# from the validated 7B math PPO run). +# - loss_fn.force_on_policy_ratio false: PPO trains on the real prev/current +# logprob ratio (GRPO nano forced ratio=1 and leaned on TIS only). +# - grpo's invalid_tool_call_advantage / malformed_thinking_advantage are +# DROPPED: message-level advantage-overwrite penalties are rejected on the +# PPO path (they would break GAE advantage/return consistency). The +# reward-zeroing `reward_penalties` block is the PPO-compatible path; left +# all-off (default) for parity with the GRPO baseline, which also ran with +# them off (its top-level penalize_* keys are inert in this checkout). +# +# Static algorithm, loss, and environment settings live here. Per-run +# overrides (model path, data paths, parallelism, batch sizes) are applied by +# the launch script (scripts/swe/ppo/nano_on_main.sh) or Hydra CLI. +# ============================================================================= + +# ============================================================================= +# Cluster — overridden by launch script for parallelism changes +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 128 + segment_size: 16 + +# ============================================================================= +# Checkpointing +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/ppo_ultra_e2e" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 3 + ft_keep_latest_k: 1 + ft_save_period: 1 + save_optimizer: true + checkpoint_must_save_by: "00:03:30:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# PPO Algorithm (replaces the parent's grpo: section) +# ============================================================================= +ppo: + num_prompts_per_step: 16 + num_generations_per_prompt: 32 + max_rollout_turns: 1 + max_num_epochs: 4 + max_num_steps: 1000000 + # One optimizer pass per rollout batch (PPO-epochs>1 reuses the batch). + # This is the ACTOR's epoch count. + ppo_epochs: 1 + # Critic (value) passes over each rollout batch. null = same as ppo_epochs + # (coupled, the historical behavior); must be >= ppo_epochs. The surplus runs + # as extra critic-only passes, fitting the critic harder without + # over-training the actor. + critic_ppo_epochs: null + # Re-score the batch with a forward-only critic pass after the final update, + # adding critic/explained_var_post_update + critic/loss_post_update. + # critic/explained_var is always the PRE-update EV (from the rollout-time + # values GAE consumed); critic/loss is from the last training pass. Costs one + # extra critic forward per step. + log_post_update_critic_metrics: false + # Critic-only warmup: policy training is skipped for the first N steps while + # the (fresh-head) value model trains on rollouts from the frozen policy. + # Overridden by the launch script. + policy_training_start_step: 50 + val_period: 10000 + val_at_start: false + val_at_end: false + overlong_filtering: false + # NeMo-Gym principle: the validation set is used verbatim (no hidden pre/post + # processing). run_ppo_nemo_gym.py rejects a preset max_val_samples and sets + # both to len(val_dataset), so they must start null. + max_val_samples: null + val_batch_size: null + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + # loss_fn.reference_policy_kl_penalty is 0 (SWE baseline), so skip the + # reference-policy logprob forward. The launch script flips this to false + # automatically when a nonzero KL penalty is requested. + skip_reference_policy_logprobs_calculation: true + seq_logprob_error_threshold: 2 + log_rollout_dump: true + rollout_dump_period: 5 # dump every N PPO steps (when log_rollout_dump is true) + + # GAE with VAPO decoupled lambdas (defaults match the validated 7B math PPO + # run: lambda_value=1.0, length-adaptive lambda_policy with alpha=0.05). + # + # name: "gae" -> token-level GAE (the historical path, unchanged). + # name: "turn_gae" -> TURN-level GAE: one assistant message is one action, + # V(s_k) is read at that turn's first token (the value head is right-shifted, + # so that position has seen the whole preceding observation and none of the + # action), the advantage is constant across the turn's tokens, and the critic + # is supervised at one anchor per turn. Only the turn_gae_* keys apply. + # + # Why it exists: at the token level lambda has effective horizon 1/(1-lambda) + # TOKENS, so on a ~45k-token SWE rollout anything below 1 severs the terminal + # reward — which is why length_adaptive_alpha pushes lambda to 1 - 1.5e-5 and + # GAE degenerates to the pure baseline A_t = R - V(s_t) with no temporal + # credit assignment at all. Over ~92 turns, lambda=0.97 is a 33-turn horizon. + # See research/ppo/turn_level_critic_plan.md. + adv_estimator: + name: "gae" + gae_lambda: 1.0 + gae_gamma: 1 + normalize_advantages: true + gae_lambda_value: 1.0 + gae_lambda_policy: 1 + # Length-adaptive lambda_policy = 1 - 1/(alpha*l). 0 = disabled. + length_adaptive_alpha: null + + # --- decomposed group baseline + residual critic --- + # research/ppo/residual_critic_report.md. true => the rollout group supplies + # the task baseline B(X) as a leave-one-out mean and the critic is trained + # only on the within-task residual C(s) = R - B_LOO; critic values/returns + # are then in RESIDUAL space and the go/no-go metric is critic/ev_res. + # Requires gae_gamma: 1. Orthogonal to token-vs-turn granularity. + # + # Motivation, measured on this exact workload: the absolute critic puts + # 88-93% of its output variance between tasks yet reaches EV 0.157 — below + # even a two-value "all-fail vs rest" lookup (0.330) and far below the free + # leave-one-out group baseline (0.553). Substituting it for that baseline + # RAISES advantage variance 1.67-2.09x. false still logs critic/ev_res and + # residual/* so an absolute run is comparable on the same axis. + residual_baseline: false + + # --- turn-level GAE (used only when name: "turn_gae") --- + # All three lambdas/gamma are required in that mode and are never silently + # defaulted, so a sweep cannot appear to have run when it did not. + turn_gae_gamma: 1.0 + # lambda_value=1.0 keeps the critic's target Monte-Carlo (G_k == R): the same + # target it regresses on today, just at ~92 anchors instead of ~45k tokens. + # Lower it only deliberately — bootstrapped targets feed the critic's own + # (still weak) estimates back into its own supervision. + turn_gae_lambda_value: 1.0 + # lambda_policy is THE knob. 1.0 reproduces today's A_k = R - V(s_k); + # 0.99 ~ 100-turn horizon, 0.97 ~ 33 turns, 0.95 ~ 20 turns. + turn_gae_lambda_policy: 1.0 + # The critic is supervised at ONE anchor per turn, every turn weighted + # equally (under token weighting the longest decile of turns owns ~43% of + # the critic loss). + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + # NOTE: grpo's penalize_invalid_tool_call / invalid_tool_call_advantage / + # penalize_malformed_thinking / malformed_thinking_advantage are intentionally + # absent: setup() raises NotImplementedError if the *_advantage keys are set + # on the PPO path (advantage overwrites break GAE consistency). + + async_ppo: + enabled: true + # lag-1: at most one policy version off-policy (recommended/validated). + # During critic warmup the frozen actor makes any-age trajectories + # on-policy for free (warmup_max_trajectory_age_steps; null => same value). + max_trajectory_age_steps: 1 + warmup_max_trajectory_age_steps: 8 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + # Drop the incomplete (survivorship-biased) frontier target on resume and + # regenerate it fresh — unbiased batches at the cost of a one-target + # generation bubble per resume. + drop_incomplete_targets_on_restore: false + # Throttle collector/replay-buffer per-rollout heartbeat prints. + log_every: 500 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + # KL(pi_gen || pi_curr) penalty toward the behaviour (rollout) policy. + # Bounds cumulative drift across async steps; clamps must stay set (k3 + # blows up exponentially on bf16 logprob outliers). These clamps are + # deliberately separate from kl_{input,output}_clamp_value above, which are + # null here — that is why enabling this term is safe in this config. + # Uses curr_logprobs (nucleus-filtered when policy.generation.top_p < 1), + # matching the other pi_gen-facing terms (actor IS weights, approx_entropy). + behaviour_kl_penalty: 0.0 + behaviour_kl_type: "k3" + behaviour_kl_input_clamp_value: 20.0 + behaviour_kl_output_clamp_value: 10.0 + + # Defaults from the validated 7B math PPO run (symmetric 0.2 clip + dual-clip + # c=3), not the GRPO-DAPO asymmetric 0.2/0.28. Overridable from the script. + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: 3 + use_on_policy_kl_approximation: true + # Async PPO requires the off-policy importance-sampling correction. + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + # PPO trains on the real prev/current-logprob ratio; forcing ratio=1 (as the + # GRPO nano config does) would disable PPO clipping entirely. + force_on_policy_ratio: false + use_kl_in_reward: false + +# ============================================================================= +# Value Loss (MSE with PPO-style clipping) +# ============================================================================= +value_loss_fn: + scale: 1.0 + cliprange: 0.5 + # Weight of homogeneous (all-fail / all-pass) groups in the value loss. Only + # meaningful with ppo.adv_estimator.residual_baseline: those groups have Y = 0 + # for every sibling, so they contribute EXACTLY zero target variance while + # still costing ~58% of critic FLOPs on this pool (54.1% all-fail, 2.4% + # all-pass; only 43.5% of groups are mixed). Keep at 1.0 for the clean + # ablation — they are not dead weight, they are the shrinkage that enforces + # E[C | X] = 0 and suppresses between-task leakage. Lower it only to buy more + # mixed-group exposure per unit wall-clock. + homogeneous_group_weight: 1.0 + +# ============================================================================= +# Policy +# ============================================================================= +policy: + model_name: "/lustre/fsw/portfolios/llmservice/users/wdai/megatron-lm-ultra/checkpoints/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-288k-nano-loss-032026/iter_0007000/hf" + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 512 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 131072 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 32 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + use_gloo_process_groups: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + # community_import.py reads megatron_cfg["gradient_accumulation_fusion"] directly + # (no default) on the HF->megatron conversion path; absent here it raises KeyError. + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled + mtp_loss_scaling_factor: 0.3 + mtp_use_repeated_layer: true + # null disables MTP entirely (matches the upstream mtp=0 run's behavior). This + # checkout's hybrid model gates MTP on `mtp_num_layers is not None` + # (hybrid_model.py:544): 0 still enters the MTP block and asserts > 0, and 1 runs + # MTP forward on only the last-stage ranks (suspected cause of the EXPERT_MODEL_ + # PARALLEL NCCL collective timeout). null makes that check False -> MTP is skipped + # on all ranks, no partial-rank EP collective. Requires the schema change making + # MegatronConfig.mtp_num_layers accept None (nemo_rl/models/policy/__init__.py). + mtp_num_layers: 5 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 3.0e-6 + min_lr: 3.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1e-10 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 3e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + # FP8 — overridden by precision recipe in launch script + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo + # were dropped from this checkout's megatron-bridge CheckpointConfig (submodule + # 554c7b9); leaving them in raises InstantiationException "Unexpected config keys". + # Removed for the current-repo reproduction. + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + # Must be a multiple of minimum_pad_factor = cp_size*2*tp_size (see + # nemo_rl/models/megatron/data.py). With CP=16 + TP=4 + sequence_parallel that is + # 128; the old value (=TP=4) made get_logprobs assert and abort every training + # step. Formula matches examples/nemo_gym/grpo_nanov3.yaml (TP*CP*2). + make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 11001 + port_range_high: 15000 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 1 + gpu_memory_utilization: 0.8 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + # reasoning_parser_plugin lives at vllm_cfg top level in this checkout + # (vllm_worker_async reads vllm_cfg["reasoning_parser_plugin"]); passing it + # inside http_server_serving_chat_kwargs makes OpenAIServingChat raise + # "unexpected keyword argument 'reasoning_parser_plugin'". Path updated to this + # repo's location (the parser moved out of nemo_rl/utils/). + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + + vllm_kwargs: + attention_backend: FLASH_ATTN + moe_backend: triton + mamba_ssm_cache_dtype: "float32" + # Route TP allreduces through NCCL instead of vLLM's custom CUDA-IPC + # spin kernel. Jobs 37606/38016 each lost a generation engine ~15 min in: + # the NCCL flight recorder showed all 4 TP ranks with their next + # collective enqueued-but-never-started, i.e. all four streams parked in + # the custom-allreduce kernel ahead of it (symmetric cross-GPU deadlock, + # different node each time). Costs a few % decode latency on TP-4. + disable_custom_all_reduce: true + compilation_config: + cudagraph_mode: PIECEWISE + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 64 + +# ============================================================================= +# Value Model (GAE critic) — Megatron backend, initialized from the policy +# checkpoint with a fresh scalar value head. Parallelism/MoE settings mirror +# the policy block; the launch script overrides them together. +# ============================================================================= +value: + model_name: ${policy.model_name} + tokenizer: + name: ${value.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: ${policy.train_global_batch_size} + train_micro_batch_size: ${policy.train_micro_batch_size} + logprob_batch_size: ${policy.logprob_batch_size} + max_total_sequence_length: ${policy.max_total_sequence_length} + precision: "bfloat16" + + # --- privileged (asymmetric) critic for SWE --- + # research/ppo/privileged_residual_critic_plan.md. The critic sees the accepted + # fix and the grading tests; the policy never does. Audited over all 9262 + # instances: golden patch + test patch are present for 100% of them, and none + # of them appear anywhere in the agent's context (verified against a real + # 408k-char rollout), so this is genuine information, not just a cheaper + # computation of something already visible. + # + # The block is prefixed BEFORE the first assistant token — mandatory, since the + # value head is causal and ~150 assistant turns are spread across the context; + # appending would be a silent no-op. It is byte-identical for all 16 siblings + # of a group, so it cannot introduce a within-task confound. + # + # Orthogonal to both other critic axes: composes with residual OR absolute + # targets, and with token-level OR turn_gae (anchors are remapped into the + # augmented layout). Enabling it raises value.max_total_sequence_length by the + # sum of the caps below automatically. + swe_privileged_critic: + enabled: false + # Single TOTAL token budget for the reference block, so the value model's + # sequence budget is exactly policy_len + this + slack. Measured over 400 + # instances the untruncated block is median 5467 / p90 17683 / p99 77911 + # tokens, so 32768 truncates ~4.8%. Budget is spent in priority order + # (fail_to_pass, golden_patch, test_patch, pass_to_pass) and anything cut is + # marked inline with "... [truncated]" plus a block-level note. Watch + # privilege/frac_truncated to see what the budget is discarding. + max_total_tokens: 32768 + + reward_model_cfg: + enabled: false # only used for the DTensor V2 value worker + reward_model_type: "regression" + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 32 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${value.precision} + sequence_parallel: true + + # MoE — mirror the policy + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + use_gloo_process_groups: false + + # Compute — mirror the policy + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + gradient_accumulation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — fully disabled for the critic. The value model replaces the LM + # head with a scalar head, so MTP next-token heads are dead weight; null + # skips the MTP block on all ranks (no partial-rank EP collective). + mtp_loss_scaling_factor: 0.0 + mtp_use_repeated_layer: true + mtp_num_layers: null + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + # Critic LR from the validated 7B math PPO run (fresh value head trains + # faster than the pretrained actor). Overridden by the launch script. + lr: 6.0e-6 + min_lr: 6.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1e-15 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${value.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${value.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${value.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 6.0e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + + env_vars: null + + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + + make_sequence_length_divisible_by: ${mul:${mul:${value.megatron_cfg.tensor_model_parallel_size}, ${value.megatron_cfg.context_parallel_size}}, 2} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + use_multiple_dataloader: false + train: + data_path: null # Set by launch script + validation: + data_path: null # Set by launch script + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + SWE Agents +# ============================================================================= +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + # Reuse pre-built gym venvs. With false, all 3 gym servers (policy_model, + # swe_agents_{train,val}) rebuild the editable nemo-gym package in parallel and + # contend on the same uv distribution-cache lock on Lustre; policy_model then + # times out (default 300s) and NemoGym spinup fails. Mirrors the qwen-30b SWE run. + skip_venv_if_present: true + num_gpu_nodes: 0 + port_range_low: 15001 + port_range_high: 20000 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - responses_api_agents/swe_agents/configs/swebench_opencode_training_alignfix.yaml + swe_agents_train: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}} + swebench_agent_timeout: 3600 + swebench_tests_timeout: 900 + run_with_mixed_prompts: true + dataset_path: ${data.train.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swerebench/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/nv_internal/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/r2e_gym/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/sweb.eval.arm64.{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swebench/swe-bench.eval.arm64.{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/swebenchpro_ots/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/sweagentv2/{instance_id}.sif" + - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swegym/sweb.eval.arm64.{instance_id}.sif" + swe_agents_val: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}} + swebench_agent_timeout: 3600 + dataset_path: ${data.validation.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter} + use_absolute_ip: true + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "logs" + num_val_samples_to_print: 0 + wandb_enabled: true + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "ultra-v3-swe-e2e" + name: "ppo-ultra-e2e" + tensorboard: {} + mlflow: + experiment_name: "ppo-ultra-e2e" + run_name: "ppo-ultra-e2e" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels / Token IDs / penalize_* — carried over from the GRPO config +# for parity. NOTE: these top-level keys are INERT in this checkout (nothing +# reads them; the GRPO baseline also ran with them having no effect). The +# PPO-compatible penalty mechanism is the structured `reward_penalties:` block +# (reward-zeroing, pre-GAE) — left at its all-off default to match the GRPO +# baseline's effective behavior. +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + + +# Reward-zeroing penalties applied to NeMo-Gym rollout results. +reward_penalties: + penalize_duplicated_reasoning: true + penalize_empty_final_answer: true + penalize_unwanted_tokens: true + penalize_malformed_think_tag: true + # Optional model/tokenizer-specific token IDs. Add token IDs that should + # be penalized when emitted by the model. Unwanted token IDs must be + # specified explicitly; think-tag IDs are inferred when each tag encodes + # to one token. Example: {unwanted: [2], think_open: 12, think_close: 13} + token_ids: {unwanted: [2], think_open: 12, think_close: 13} \ No newline at end of file