From 7f013f8019cee682d3708c1fdaae026abb89c827 Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 12 Aug 2026 10:34:27 -0700 Subject: [PATCH 01/27] Port GDPO length adjustment features to GitHub main --- examples/configs/grpo_ultra_256n4g_bf16.yaml | 784 +++++++++++ examples/configs/grpo_ultra_512n4g_bf16.yaml | 743 +++++++++++ ...grpo_ultra_64n4g_length_adj_pipeclean.yaml | 645 +++++++++ .../configs/grpo_ultra_64n4g_pipeclean.yaml | 715 ++++++++++ examples/nemo_gym/grpo_nanov3.yaml | 2 +- .../length_adj_pipeclean_launch.sh | 35 + nemo_rl/algorithms/advantage_estimator.py | 246 ++++ nemo_rl/algorithms/grpo.py | 2 + nemo_rl/experience/rollouts.py | 167 ++- nemo_rl/utils/length_adjustments.md | 598 +++++++++ nemo_rl/utils/length_adjustments.py | 1178 +++++++++++++++++ scripts/build_profile_band_dataset.py | 204 +++ scripts/profile_band.yaml | 28 + tests/unit/algorithms/test_grpo.py | 49 + .../unit/experience/test_reward_penalties.py | 88 ++ 15 files changed, 5481 insertions(+), 3 deletions(-) create mode 100644 examples/configs/grpo_ultra_256n4g_bf16.yaml create mode 100644 examples/configs/grpo_ultra_512n4g_bf16.yaml create mode 100644 examples/configs/grpo_ultra_64n4g_length_adj_pipeclean.yaml create mode 100644 examples/configs/grpo_ultra_64n4g_pipeclean.yaml create mode 100755 examples/nemo_gym/nemotron-3-ultra/length_adj_pipeclean_launch.sh create mode 100644 nemo_rl/utils/length_adjustments.md create mode 100644 nemo_rl/utils/length_adjustments.py create mode 100644 scripts/build_profile_band_dataset.py create mode 100644 scripts/profile_band.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..f1695b8766b --- /dev/null +++ b/examples/configs/grpo_ultra_256n4g_bf16.yaml @@ -0,0 +1,784 @@ +# ============================================================================= +# 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: 512 + 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_period: 8 + 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: 512 + 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: -20 + advantage_clip_high: 20 + 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 + + length_bonus: + verbose: true + default: + enabled: false + length_type: "tokens" + reasoning_bonus: 0.0 + answer_bonus: 0.0 + total_bonus: 0.0 + longest_reasoning_penalty: 0.0 + longest_answer_penalty: 0.0 + longest_total_penalty: 0.0 + top_percentile: 0.5 + group_reasoning_length_penalty_coeff: 0.0 + group_answer_length_penalty_coeff: 0.0 + group_total_length_penalty_coeff: 0.0 + reasoning_zmad_threshold: 0.0 + reasoning_zmad_penalty: 0.0 + answer_zmad_threshold: 0.0 + answer_zmad_penalty: 0.0 + total_zmad_threshold: 0.0 + total_zmad_penalty: 0.0 + profiled_length_penalty: 0.0 + profiled_length_n_std: 1.0 + profiled_length_min_samples: 2 + agent_overrides: + abstention_simple_agent: + enabled: false + calendar_simple_agent: + enabled: false + code_gen_simple_agent: + enabled: true + total_zmad_threshold: 2.5 + total_zmad_penalty: 0.1 + equivalence_llm_judge_simple_agent: + enabled: true + total_zmad_threshold: 2.5 + total_zmad_penalty: 0.1 + ether0_simple_agent: + enabled: true + total_zmad_threshold: 2.0 + total_zmad_penalty: 0.1 + genrm_simple_agent: + enabled: false + genrm_simple_agent_reasoning_off: + enabled: false + instruction_following_simple_agent: + enabled: true + total_zmad_threshold: 2.0 + total_zmad_penalty: 0.1 + jailbreak_engagement_with_disclaimer: + enabled: false + jailbreak_hard_refusal_no_redirection: + enabled: false + jailbreak_hard_refusal_with_helplines: + enabled: false + jailbreak_refusal_with_explanation: + enabled: false + lc_equivalence_rule_simple_agent: + enabled: false + lc_judge_simple_agent: + enabled: true + total_zmad_threshold: 2.0 + total_zmad_penalty: 0.1 + math_formal_lean_refinement_agent: + enabled: true + total_zmad_threshold: 2.5 + total_zmad_penalty: 0.1 + math_with_judge_simple_agent: + enabled: true + total_zmad_threshold: 2.0 + total_zmad_penalty: 0.1 + mcqa_simple_agent: + enabled: true + longest_total_penalty: 0.1 + multichallenge_simple_agent: + enabled: true + total_zmad_threshold: 2.5 + total_zmad_penalty: 0.1 + ns_tools_simple_agent: + enabled: false + nvarc_inductive_simple_agent: + enabled: false + nvarc_transductive_simple_agent: + enabled: false + reasoning_gym_simple_agent: + enabled: false + search_pivot_single_step_tool_use_with_argument_comparison_agent: + enabled: false + single_step_tool_use_with_argument_comparison_agent: + enabled: true + total_zmad_threshold: 2.0 + total_zmad_penalty: 0.1 + structured_outputs_simple_agent: + enabled: true + total_zmad_threshold: 2.0 + total_zmad_penalty: 0.1 + swe_pivot_single_step_tool_use_with_argument_comparison_agent: + enabled: false + terminus_judge_string_only_simple_agent: + enabled: true + total_zmad_threshold: 2.0 + total_zmad_penalty: 0.1 + toolcall_schema_single_step_tool_use_with_argument_comparison_agent: + enabled: true + total_zmad_threshold: 2.0 + total_zmad_penalty: 0.1 + + 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: null # set by the launcher via MODEL_PATH + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 8192 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 49152 + 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" #"flex" + moe_flex_dispatcher_backend: "alltoall" #"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 + 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 + 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: 65536 + 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: 65536 + 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 + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py + + vllm_kwargs: + attention_backend: FLASH_ATTN + 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: 182 # Overridden by launch script + +# ============================================================================= +# 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 + 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 + + # Increase the num workers to help with the serialization/deserialization overhead of token IDs + policy_model: + responses_api_models: + vllm_model: + num_workers: 16 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + policy_model_reasoning_off: + responses_api_models: + vllm_model: + num_workers: 4 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + + 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 + + vllm_serve_kwargs: + attention_backend: TRITON_ATTN + tensor_parallel_size: 4 + data_parallel_size: 1 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + gpu_memory_utilization: 0.85 + max_model_len: 96000 + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32,64,128,256] + + + # 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 + + vllm_serve_kwargs: + attention_backend: TRITON_ATTN + tensor_parallel_size: 4 + data_parallel_size: 9 + 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 + enable_prefix_caching: true + enable_chunked_prefill: true + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32,64,128,256] + 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" + 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 + genrm_parse_retries: 0 + group_reasoning_length_penalty_coeff: 0.1 + group_answer_length_penalty_coeff: 0.1 + group_style_penalty_coeff: 0.0 + default_score: 3.0 + default_ranking: 3.5 + + # 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 + ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + tensor_parallel_size: 4 + data_parallel_size: 4 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + reasoning_parser: deepseek_r1 + gpu_memory_utilization: 0.85 + max_model_len: 60000 + enable_prefix_caching: true + enable_chunked_prefill: true + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + load_format: instanttensor + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + 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 + + math_formal_lean_refinement_agent: + responses_api_agents: + proof_refinement_agent: + max_correction_turns: 0 + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "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 diff --git a/examples/configs/grpo_ultra_512n4g_bf16.yaml b/examples/configs/grpo_ultra_512n4g_bf16.yaml new file mode 100644 index 00000000000..46669662daf --- /dev/null +++ b/examples/configs/grpo_ultra_512n4g_bf16.yaml @@ -0,0 +1,743 @@ +# ============================================================================= +# 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: 512 + 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_period: 8 + 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: 512 + 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: -20 + advantage_clip_high: 20 + 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 + + length_bonus: + verbose: true + default: + enabled: false + length_type: "tokens" + reasoning_bonus: 0.0 + answer_bonus: 0.0 + longest_reasoning_penalty: 0.0 + longest_answer_penalty: 0.0 + top_percentile: 0.5 + group_reasoning_length_penalty_coeff: 0.0 + group_answer_length_penalty_coeff: 0.0 + reasoning_zmad_threshold: 0.0 + reasoning_zmad_penalty: 0.0 + answer_zmad_threshold: 0.0 + answer_zmad_penalty: 0.0 + agent_overrides: + abstention_simple_agent: + enabled: false + calendar_simple_agent: + enabled: false + code_gen_simple_agent: + enabled: true + reasoning_zmad_threshold: 2.5 + reasoning_zmad_penalty: 0.1 + equivalence_llm_judge_simple_agent: + enabled: true + reasoning_zmad_threshold: 2.5 + reasoning_zmad_penalty: 0.1 + ether0_simple_agent: + enabled: true + reasoning_zmad_threshold: 2.0 + reasoning_zmad_penalty: 0.1 + genrm_simple_agent: + enabled: false + genrm_simple_agent_reasoning_off: + enabled: false + instruction_following_simple_agent: + enabled: true + reasoning_zmad_threshold: 2.0 + reasoning_zmad_penalty: 0.1 + jailbreak_engagement_with_disclaimer: + enabled: false + jailbreak_hard_refusal_no_redirection: + enabled: false + jailbreak_hard_refusal_with_helplines: + enabled: false + jailbreak_refusal_with_explanation: + enabled: false + lc_equivalence_rule_simple_agent: + enabled: false + lc_judge_simple_agent: + enabled: true + reasoning_zmad_threshold: 2.0 + reasoning_zmad_penalty: 0.1 + math_formal_lean_refinement_agent: + enabled: true + reasoning_zmad_threshold: 2.5 + reasoning_zmad_penalty: 0.1 + math_with_judge_simple_agent: + enabled: true + reasoning_zmad_threshold: 2.0 + reasoning_zmad_penalty: 0.1 + mcqa_simple_agent: + enabled: true + longest_reasoning_penalty: 0.1 + multichallenge_simple_agent: + enabled: true + reasoning_zmad_threshold: 2.5 + reasoning_zmad_penalty: 0.1 + ns_tools_simple_agent: + enabled: false + nvarc_inductive_simple_agent: + enabled: false + nvarc_transductive_simple_agent: + enabled: false + reasoning_gym_simple_agent: + enabled: false + search_pivot_single_step_tool_use_with_argument_comparison_agent: + enabled: false + single_step_tool_use_with_argument_comparison_agent: + enabled: true + reasoning_zmad_threshold: 2.0 + reasoning_zmad_penalty: 0.1 + structured_outputs_simple_agent: + enabled: true + reasoning_zmad_threshold: 2.0 + reasoning_zmad_penalty: 0.1 + swe_pivot_single_step_tool_use_with_argument_comparison_agent: + enabled: false + terminus_judge_string_only_simple_agent: + enabled: true + reasoning_zmad_threshold: 2.0 + reasoning_zmad_penalty: 0.1 + toolcall_schema_single_step_tool_use_with_argument_comparison_agent: + enabled: true + reasoning_zmad_threshold: 2.0 + reasoning_zmad_penalty: 0.1 + + 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: null # set by the launcher via MODEL_PATH + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 8192 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 49152 + 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" #"flex" + moe_flex_dispatcher_backend: "alltoall" #"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 + 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 + 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: ${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 + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py + + vllm_kwargs: + attention_backend: FLASH_ATTN + max_num_seqs: 64 + 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: 182 # Overridden by launch script + +# ============================================================================= +# 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 + 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 + + vllm_serve_kwargs: + attention_backend: TRITON_ATTN + tensor_parallel_size: 4 + data_parallel_size: 1 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + gpu_memory_utilization: 0.85 + max_model_len: 96000 + 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,64,128,256] + + + # 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 + + vllm_serve_kwargs: + attention_backend: FLASH_ATTN + tensor_parallel_size: 4 + data_parallel_size: 9 + 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,64,128,256] + 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: 32768 + temperature: 1.0 + top_p: 0.95 + comparison_strategy: "circular" + 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.1 + group_answer_length_penalty_coeff: 0.25 + group_style_penalty_coeff: 0.0 + default_score: 3.0 + default_ranking: 3.5 + + # 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: "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 diff --git a/examples/configs/grpo_ultra_64n4g_length_adj_pipeclean.yaml b/examples/configs/grpo_ultra_64n4g_length_adj_pipeclean.yaml new file mode 100644 index 00000000000..065d8a08a10 --- /dev/null +++ b/examples/configs/grpo_ultra_64n4g_length_adj_pipeclean.yaml @@ -0,0 +1,645 @@ +# ============================================================================= +# GRPO Ultra V3 — 64-node GB200 NVL72 Pipeclean Config +# ============================================================================= +# Pipeclean config for validating the end-to-end GRPO pipeline on +# 64 nodes × 4 GPUs/node. Batch sizes and sequence lengths are +# intentionally small to complete steps quickly. +# +# Defaults to 10 training steps (~2h on 64 nodes). Override at launch +# time with NRL_MAX_STEPS for shorter CI runs (e.g. NRL_MAX_STEPS=4). +# +# - gpus_per_node: 4 +# - TP: 8 +# - EP: 32 (16 experts/rank) +# - vLLM TP: 8 +# - Judge TP: 4 +# ============================================================================= + +# ============================================================================= +# Cluster — overridden by launch script +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 64 + 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_period: 4 + checkpoint_must_save_by: "00:03:30:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# GRPO Algorithm +# ============================================================================= +grpo: + num_prompts_per_step: 16 # Pipeclean: reduced from 128 + num_generations_per_prompt: 4 # Pipeclean: reduced from 16 + num_val_generations_per_prompt: 2 + max_rollout_turns: 1 + max_num_epochs: 1 + max_num_steps: 10 # Pipeclean: reduced from 1000000; override via NRL_MAX_STEPS + 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 + 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 + + length_bonus: + default: + enabled: true + length_type: "tokens" + reasoning_bonus: 0.1 + answer_bonus: 0.1 + longest_reasoning_penalty: 0.0 + longest_answer_penalty: 0.0 + top_percentile: 0.5 + group_reasoning_length_penalty_coeff: 0.05 + group_answer_length_penalty_coeff: 0.05 + agent_overrides: + # --- Agents from training dataset (curriculum_v29_no_judge) --- + calendar_simple_agent: null + code_gen_simple_agent: + group_answer_length_penalty_coeff: 0.1 + instruction_following_simple_agent: null + mcqa_simple_agent: null + reasoning_gym_simple_agent: null + single_step_tool_use_with_argument_comparison_agent: + group_answer_length_penalty_coeff: 0.0 + structured_outputs_simple_agent: null + terminal_pivot_simple_agent: + enabled: false + workplace_assistant_simple_agent: null + # --- Additional agents from ultra config environments --- + math_with_judge_simple_agent: + enabled: false + genrm_simple_agent: + enabled: false + genrm_simple_agent_reasoning_off: + enabled: false + lc_judge_simple_agent: + enabled: false + nl2bash_equivalency_simple_agent: + enabled: false + equivalence_llm_judge_simple_agent: + enabled: false + ns_tools_simple_agent: + enabled: false + math_formal_lean_simple_agent: + enabled: false + jailbreak_detection_simple_agent: + enabled: false + over_refusal_detection_simple_agent: + enabled: false + multichallenge_simple_agent: + enabled: false + inverse_if_simple_agent: + enabled: false + search_pivot_single_step_tool_use_with_argument_comparison_agent: + enabled: false + toolcall_schema_single_step_tool_use_with_argument_comparison_agent: + enabled: false + + 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: 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 +# ============================================================================= +policy: + model_name: null # set by the launcher via MODEL_PATH + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 64 # Pipeclean: num_prompts × num_generations = 16 × 4 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 32768 # Pipeclean: reduced from 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: 1 + 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: ${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_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + use_gloo_process_groups: false + + # Compute + apply_rope_fusion: true + bias_activation_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 + mtp_num_layers: 0 + 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: 0 + lr_warmup_init: 3e-7 + + 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 + + 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 + + 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: 32768 # Pipeclean: reduced from 65536 + 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.85 + max_model_len: 32768 # Must be <= 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 + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py + + vllm_kwargs: + attention_backend: FLASH_ATTN + max_num_seqs: 16 + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + # Disable FlashInfer TRTLLM AllReduce fusion. It requires + # SymmDeviceMemory (MNNVL IPC) to initialise successfully on every TP + # rank. When TP spans multiple nodes some ranks succeed while others + # fail (CUDA_ERROR_INVALID_DEVICE), causing a split-brain: half the TP + # group uses the MNNVL-based TRTLLM allreduce path while the other half + # falls back to standard NCCL. They never rendezvous → permanent hang. + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 26 # Overridden by launch script + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: 16384 # NOTE: no-op for NemoGymDataset (processor ignores max_seq_length) + 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 + Judge Models +# ============================================================================= +env: + should_use_nemo_gym: true + use_genrm_compare: true + genrm_agent_names: + - "genrm_simple_agent" + - "genrm_simple_agent_reasoning_off" + genrm_compare_server_name: "genrm_compare" + nemo_gym: + skip_venv_if_present: true + num_gpu_nodes: 6 # Overridden by launch script + 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/structured_outputs/configs/structured_outputs_json.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/terminal_pivot/configs/terminal_pivot.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/jailbreak_detection/configs/jailbreak_detection_nemotron_combined_reward_tp8.yaml + - resources_servers/over_refusal_detection/configs/over_refusal_detection_nemotron_tp8.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 + + # Safety Model: 4B, so TP=1, so fits single GPU + jailbreak_detection: + resources_servers: + jailbreak_detection: + judge_model_server: + type: responses_api_models + name: safety_judge_model + + safety_judge_model: + responses_api_models: + vllm_model: + entrypoint: app.py + base_url: http://127.0.0.1:8001/v1 + api_key: dummy_key + model: null # Set by launch script (or omit for no safety judge) + return_token_id_information: false + uses_reasoning_parser: false + spinup_server: true + router_dp_size: 4 + server_args: + attention_backend: TRITON_ATTN + tensor_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: 2 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16] + + terminal_pivot_simple_agent: + responses_api_agents: + simple_agent: + model_server: + name: policy_model + + # nl2bash / General Judge: TP=4 on GB200 192GB + nl2bash_judge_model: + responses_api_models: + vllm_model: + entrypoint: app.py + base_url: http://127.0.0.1:10000/v1 + api_key: dummy_key + model: null # Set by launch script + return_token_id_information: false + uses_reasoning_parser: false + spinup_server: true + router_dp_size: 2 + server_args: + attention_backend: FLASH_ATTN + tensor_parallel_size: 4 + data_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: 16 + 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_servers: + genrm_compare: + genrm_model_server: + type: responses_api_models + name: genrm_model + genrm_responses_create_params: + max_output_tokens: 16384 + temperature: 0.6 + top_p: 0.95 + comparison_strategy: "circular" + num_judges_per_comparison: 1 + use_principle: true + default_principle: "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 + group_answer_length_penalty_coeff: 0 + group_style_penalty_coeff: 0.1 + default_score: 3.0 + default_ranking: 3.5 + + genrm_model: + responses_api_models: + vllm_model: + entrypoint: app.py + base_url: http://127.0.0.1:8000/v1 + api_key: dummy_key + model: null # Set by launch script + uses_reasoning_parser: true + return_token_id_information: false + spinup_server: true + router_dp_size: 2 + server_args: + attention_backend: FLASH_ATTN + tensor_parallel_size: 4 + reasoning_parser: deepseek_r1 + gpu_memory_utilization: 0.85 + max_model_len: 60000 + 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,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + 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: 1024 + unit_test_timeout_secs: 10 + debug: false + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "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" + tensorboard: {} + mlflow: + experiment_name: "grpo-ultra-v3" + run_name: "grpo-ultra-v3" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: low}" + low_weight: 0.2 + low_penalty: 1 + low_ub: 3000 diff --git a/examples/configs/grpo_ultra_64n4g_pipeclean.yaml b/examples/configs/grpo_ultra_64n4g_pipeclean.yaml new file mode 100644 index 00000000000..4aab4c9b369 --- /dev/null +++ b/examples/configs/grpo_ultra_64n4g_pipeclean.yaml @@ -0,0 +1,715 @@ +# ============================================================================= +# GRPO Ultra V3 — 64-node GB200 NVL72 Pipeclean Config +# ============================================================================= +# Pipeclean config for validating the end-to-end GRPO pipeline on +# 64 nodes × 4 GPUs/node. Batch sizes and sequence lengths are +# intentionally small to complete steps quickly. +# +# Defaults to 10 training steps (~2h on 64 nodes). Override at launch +# time with NRL_MAX_STEPS for shorter CI runs (e.g. NRL_MAX_STEPS=4). +# +# - gpus_per_node: 4 +# - TP: 8 +# - EP: 32 (16 experts/rank) +# - vLLM TP: 8 +# - Judge TP: 4 +# ============================================================================= + +# ============================================================================= +# Cluster — overridden by launch script +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 64 + 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_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: 16 # Pipeclean: reduced from 128 + num_generations_per_prompt: 4 # Pipeclean: reduced from 16 + num_val_generations_per_prompt: 2 + max_rollout_turns: 1 + max_num_epochs: 1 + max_num_steps: 10 # Pipeclean: reduced from 1000000; override via NRL_MAX_STEPS + 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 + 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 + + length_bonus: + verbose: false + default: + enabled: false + length_type: "tokens" + reasoning_bonus: 0.0 + answer_bonus: 0.0 + longest_reasoning_penalty: 0.0 + longest_answer_penalty: 0.0 + top_percentile: 0.5 + group_reasoning_length_penalty_coeff: 0.0 + group_answer_length_penalty_coeff: 0.0 + reasoning_zmad_threshold: 0.0 + reasoning_zmad_penalty: 0.0 + answer_zmad_threshold: 0.0 + answer_zmad_penalty: 0.0 + agent_overrides: + # --- Agents from training dataset (curriculum_v29_no_judge) --- + calendar_simple_agent: + enabled: false + code_gen_simple_agent: + enabled: false + instruction_following_simple_agent: + enabled: false + mcqa_simple_agent: + enabled: false + reasoning_gym_simple_agent: + enabled: false + single_step_tool_use_with_argument_comparison_agent: + enabled: false + structured_outputs_simple_agent: + enabled: false + terminal_pivot_simple_agent: + enabled: false + workplace_assistant_simple_agent: + enabled: false + # --- Additional agents from ultra config environments --- + math_with_judge_simple_agent: + enabled: false + genrm_simple_agent: + enabled: false + genrm_simple_agent_reasoning_off: + enabled: false + lc_judge_simple_agent: + enabled: false + nl2bash_equivalency_simple_agent: + enabled: false + equivalence_llm_judge_simple_agent: + enabled: false + ns_tools_simple_agent: + enabled: false + math_formal_lean_simple_agent: + enabled: false + jailbreak_detection_simple_agent: + enabled: false + over_refusal_detection_simple_agent: + enabled: false + multichallenge_simple_agent: + enabled: false + inverse_if_simple_agent: + enabled: false + search_pivot_single_step_tool_use_with_argument_comparison_agent: + enabled: false + toolcall_schema_single_step_tool_use_with_argument_comparison_agent: + enabled: false + + 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: 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 +# ============================================================================= +policy: + model_name: null # set by the launcher via MODEL_PATH + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 64 # Pipeclean: num_prompts × num_generations = 16 × 4 + 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: 1 + 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: ${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 + use_gloo_process_groups: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_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: 0 + lr_warmup_init: 3e-7 + + 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 + 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: 65536 + 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.85 + max_model_len: 65536 + 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 + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py + + vllm_kwargs: + attention_backend: FLASH_ATTN + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + # Disable FlashInfer TRTLLM AllReduce fusion. It requires + # SymmDeviceMemory (MNNVL IPC) to initialise successfully on every TP + # rank. When TP spans multiple nodes some ranks succeed while others + # fail (CUDA_ERROR_INVALID_DEVICE), causing a split-brain: half the TP + # group uses the MNNVL-based TRTLLM allreduce path while the other half + # falls back to standard NCCL. They never rendezvous → permanent hang. + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 26 # Overridden by launch script + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: 16384 # NOTE: no-op for NemoGymDataset (processor ignores max_seq_length) + 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 + 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/over_refusal_detection/configs/over_refusal_detection_nemotron_tp8.yaml + - resources_servers/indirect_prompt_injection/configs/indirect_prompt_injection.yaml + + # Increase the num workers to help with the serialization/deserialization overhead of token IDs + policy_model: + responses_api_models: + vllm_model: + num_workers: 16 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + policy_model_reasoning_off: + responses_api_models: + vllm_model: + num_workers: 4 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + + 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, so TP=1, so fits single GPU + 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 (or omit for no safety judge) + return_token_id_information: false + uses_reasoning_parser: false + debug: true + show_vllm_engine_stats: true + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + attention_backend: TRITON_ATTN + tensor_parallel_size: 4 + data_parallel_size: 1 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + gpu_memory_utilization: 0.85 + max_model_len: 96000 + 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 + show_vllm_engine_stats: true + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + attention_backend: FLASH_ATTN + tensor_parallel_size: 4 + data_parallel_size: 2 + 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 + 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: 16384 + temperature: 0.6 + top_p: 0.95 + comparison_strategy: "circular" + 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 + genrm_parse_retries: 0 + group_reasoning_length_penalty_coeff: 0.1 + group_answer_length_penalty_coeff: 0.2 + group_style_penalty_coeff: 0.0 + default_score: 3.0 + default_ranking: 3.5 + + genrm_model: + responses_api_models: + genrm_model: + entrypoint: app.py + api_key: dummy_key + model: null # Set by launch script + uses_reasoning_parser: true + return_token_id_information: false + debug: true + show_vllm_engine_stats: true + ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + tensor_parallel_size: 4 + data_parallel_size: 2 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + reasoning_parser: deepseek_r1 + gpu_memory_utilization: 0.85 + max_model_len: 60000 + enable_prefix_caching: true + enable_chunked_prefill: true + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + load_format: instanttensor + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + 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: 1024 + unit_test_timeout_secs: 10 + debug: false + + math_formal_lean_refinement_agent: + responses_api_agents: + proof_refinement_agent: + max_correction_turns: 0 + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "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-pipeclean" + name: "grpo-ultra-v3-pipeclean" + tensorboard: {} + mlflow: + experiment_name: "grpo-ultra-v3-pipeclean" + run_name: "grpo-ultra-v3-pipeclean" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: low}" + low_weight: 0.2 + low_penalty: 1 + low_ub: 3000 diff --git a/examples/nemo_gym/grpo_nanov3.yaml b/examples/nemo_gym/grpo_nanov3.yaml index 6bb2f6963a6..b7fced772f1 100644 --- a/examples/nemo_gym/grpo_nanov3.yaml +++ b/examples/nemo_gym/grpo_nanov3.yaml @@ -312,7 +312,7 @@ data: env: should_use_nemo_gym: true # true: skip expensive train_data_step*.jsonl; false: write full jsonl. - should_log_nemo_gym_responses: true + should_log_nemo_gym_responses: false # false: ignore env mask_sample flags so the loss trains on every sample. should_mask_flagged_samples: true nemo_gym: diff --git a/examples/nemo_gym/nemotron-3-ultra/length_adj_pipeclean_launch.sh b/examples/nemo_gym/nemotron-3-ultra/length_adj_pipeclean_launch.sh new file mode 100755 index 00000000000..875e3ee8d68 --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/length_adj_pipeclean_launch.sh @@ -0,0 +1,35 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Convenience launcher for the length-adjusted Ultra pipeclean recipe. +# +# This intentionally follows examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh: +# callers provide cluster, container, model, data, cache, and optional judge +# model paths through environment variables. This wrapper only selects the +# length-adjusted config by default. +# +# Required by ultra_launch.sh: +# EXP_NAME +# MODEL_PATH +# TRAIN_PATH +# VAL_PATH +# CONTAINER +# SANDBOX_CONTAINER +# PERSISTENT_CACHE +# SLURM_PARTITION +# SLURM_ACCOUNT +# +# Optional: +# CONFIG_PATH Override the recipe config. +# WANDB_PROJ W&B project name. +# GENRM_MODEL GenRM model path or HF id. +# NL2BASH_JUDGE_MODEL +# SAFETY_JUDGE_MODEL +# EXTRA_MOUNTS Comma-separated host:container mount pairs. + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" + +export CONFIG_PATH="${CONFIG_PATH:-examples/configs/grpo_ultra_64n4g_length_adj_pipeclean.yaml}" +export WANDB_PROJ="${WANDB_PROJ:-nemotron-3-ultra-length-adjusted}" + +exec "${SCRIPT_DIR}/ultra_launch.sh" "$@" diff --git a/nemo_rl/algorithms/advantage_estimator.py b/nemo_rl/algorithms/advantage_estimator.py index 491de0d9414..293a9f0f485 100644 --- a/nemo_rl/algorithms/advantage_estimator.py +++ b/nemo_rl/algorithms/advantage_estimator.py @@ -28,6 +28,8 @@ - MOPD: https://arxiv.org/abs/2601.02780 """ +from typing import Any + import torch from pydantic import BaseModel @@ -50,10 +52,89 @@ class AdvEstimatorConfig(BaseModel, extra="allow"): use_leave_one_out_baseline: bool = True # GDPO specific: optional per-component weights w_n for the aggregation. reward_weights: list[float] | None = None + reward_features: list[str] | dict[str, Any] | None = None + verbose: bool = False # Reinforce++ specific minus_baseline: bool = True +def _feature_reward(feature_entry: Any) -> float: + if feature_entry is None: + return 0.0 + if isinstance(feature_entry, dict): + return float(feature_entry.get("reward", 0.0)) + return float(feature_entry) + + +def _format_gdpo_feature_for_log( + name: str, feature_entry: Any, feature_weight: float = 1.0 +) -> str: + if feature_entry is None: + return f"{name}=MISSING" + if not isinstance(feature_entry, dict): + return f"{name}={float(feature_entry):+.4f}" + + details = [] + adjustment = feature_entry.get("adjustment") + if adjustment is not None: + details.append(str(adjustment)) + multiplier = feature_entry.get("multiplier") + if multiplier is not None: + details.append(f"mult={float(multiplier):.4f}") + close_thinks = feature_entry.get("num_close_think_tags") + if close_thinks is not None: + details.append(f"close_thinks={int(close_thinks)}") + if feature_weight != 1.0: + details.append(f"weight={feature_weight:.4f}") + + suffix = f"({','.join(details)})" if details else "" + return f"{name}={_feature_reward(feature_entry):+.4f}{suffix}" + + +def _resolve_gdpo_reward_feature_selection( + estimator_config: AdvEstimatorConfig, agent_name: str | None +) -> list[str] | dict[str, Any]: + cfg = estimator_config.reward_features or ["env_reward"] + if isinstance(cfg, (list, tuple)): + return list(cfg) + if not isinstance(cfg, dict): + raise TypeError("gdpo reward_features must be a list or a dict") + + if "default" not in cfg and "agent_overrides" not in cfg: + return cfg + + default = cfg.get("default", ["env_reward"]) + overrides = cfg.get("agent_overrides", {}) + return overrides.get(agent_name, default) + + +def _resolve_gdpo_reward_features( + estimator_config: AdvEstimatorConfig, agent_name: str | None +) -> list[str]: + selected = _resolve_gdpo_reward_feature_selection(estimator_config, agent_name) + if isinstance(selected, dict): + return list(selected) + return list(selected) + + +def _resolve_gdpo_reward_feature_weights( + estimator_config: AdvEstimatorConfig, agent_name: str | None +) -> dict[str, float]: + selected = _resolve_gdpo_reward_feature_selection(estimator_config, agent_name) + if not isinstance(selected, dict): + return {name: 1.0 for name in selected} + + weights = {} + for name, feature_cfg in selected.items(): + if isinstance(feature_cfg, dict): + weights[name] = float(feature_cfg.get("weight", 1.0)) + elif isinstance(feature_cfg, (int, float)): + weights[name] = float(feature_cfg) + else: + weights[name] = 1.0 + return weights + + class GRPOAdvantageEstimator: """GRPO-style advantage estimator with leave-one-out baseline. @@ -112,6 +193,8 @@ def __init__( # Optional per-reward weights w_n for the aggregation A = sum_n w_n * A_n # (paper: https://arxiv.org/abs/2601.05242). None => equal weights (all 1.0). self.reward_weights = estimator_config.reward_weights + self.estimator_config = estimator_config + self.verbose = estimator_config.verbose def compute_advantage( self, @@ -133,6 +216,14 @@ def compute_advantage( Returns: Advantages tensor of shape [batch_size, seq_len]. """ + if "gdpo_reward_features" in repeated_batch: + return self._compute_feature_advantage( + prompt_ids=prompt_ids, + rewards=rewards, + mask=mask, + repeated_batch=repeated_batch, + ) + reward_component_keys = get_gdpo_reward_component_keys(repeated_batch) if len(reward_component_keys) < 2: raise ValueError( @@ -189,6 +280,161 @@ def compute_advantage( return advantages.expand(mask.shape) + def _compute_feature_advantage( + self, + prompt_ids, + rewards, + mask, + repeated_batch, + ): + gdpo_features = repeated_batch["gdpo_reward_features"] + agent_refs = repeated_batch.get("agent_ref", [{} for _ in gdpo_features]) + if len(gdpo_features) != prompt_ids.shape[0]: + raise ValueError( + "gdpo_reward_features must match batch size; " + f"got {len(gdpo_features)} vs {prompt_ids.shape[0]}" + ) + + resolved_features: list[list[str]] = [] + resolved_feature_weights: list[dict[str, float]] = [] + for agent_ref in agent_refs: + agent_name = agent_ref.get("name") if isinstance(agent_ref, dict) else None + resolved_features.append( + _resolve_gdpo_reward_features(self.estimator_config, agent_name) + ) + resolved_feature_weights.append( + _resolve_gdpo_reward_feature_weights( + self.estimator_config, agent_name + ) + ) + self._validate_prompt_group_features( + prompt_ids, resolved_features, resolved_feature_weights + ) + + feature_names: list[str] = [] + seen = set() + for features in resolved_features: + for name in features: + if name not in seen: + seen.add(name) + feature_names.append(name) + if not feature_names: + raise ValueError("GDPO requires at least one reward feature") + + if self.verbose: + self._print_reward_feature_summary( + prompt_ids=prompt_ids, + rewards=rewards, + gdpo_features=gdpo_features, + resolved_features=resolved_features, + resolved_feature_weights=resolved_feature_weights, + ) + + advantage_parts = [] + reward_device = prompt_ids.device + reward_dtype = rewards.dtype if rewards.is_floating_point() else torch.float32 + for name in feature_names: + vals = [] + weights = [] + for selected, feature_dict, feature_weight_dict in zip( + resolved_features, gdpo_features, resolved_feature_weights + ): + if name not in selected: + vals.append(0.0) + weights.append(0.0) + continue + + entry = feature_dict.get(name) if isinstance(feature_dict, dict) else None + vals.append(_feature_reward(entry)) + weights.append(feature_weight_dict.get(name, 1.0)) + + feature_rewards = torch.tensor( + vals, dtype=reward_dtype, device=reward_device + ) + feature_weights = torch.tensor( + weights, dtype=reward_dtype, device=reward_device + ) + baseline, std = calculate_baseline_and_std_per_prompt( + prompt_ids, + feature_rewards, + torch.ones_like(feature_rewards), + leave_one_out_baseline=self.use_leave_one_out_baseline, + ) + adv = (feature_rewards - baseline).unsqueeze(-1) + + if self.normalize_rewards: + epsilon = 1e-6 + non_zero_std_mask = std > 0 + adv[non_zero_std_mask] = adv[non_zero_std_mask] / ( + std.unsqueeze(-1)[non_zero_std_mask] + epsilon + ) + adv = adv * feature_weights.unsqueeze(-1) + advantage_parts.append(adv) + + advantages = sum(advantage_parts) + adv_std = advantages.std() + if adv_std > 0: + advantages = (advantages - advantages.mean()) / adv_std + else: + advantages = advantages - advantages.mean() + + return advantages.to(mask.device).expand(mask.shape) + + @staticmethod + def _validate_prompt_group_features( + prompt_ids, + resolved_features: list[list[str]], + resolved_feature_weights: list[dict[str, float]], + ) -> None: + _, inverse = torch.unique(prompt_ids, dim=0, return_inverse=True) + for group_idx in torch.unique(inverse).cpu().tolist(): + indices = (inverse == group_idx).nonzero(as_tuple=True)[0].cpu().tolist() + group_feature_lists = { + tuple( + (name, resolved_feature_weights[i].get(name, 1.0)) + for name in resolved_features[i] + ) + for i in indices + } + if len(group_feature_lists) > 1: + raise ValueError( + "GDPO reward_features must resolve to one feature spec per prompt group; " + f"got {sorted(group_feature_lists)}" + ) + + @staticmethod + def _print_reward_feature_summary( + prompt_ids, + rewards, + gdpo_features, + resolved_features: list[list[str]], + resolved_feature_weights: list[dict[str, float]], + ) -> None: + print(f"\n{'=' * 70}", flush=True) + print("[Advantage] GDPO reward features", flush=True) + _, inverse = torch.unique(prompt_ids, dim=0, return_inverse=True) + scalar_rewards = rewards.detach().cpu().tolist() + for group_idx in torch.unique(inverse).cpu().tolist(): + indices = (inverse == group_idx).nonzero(as_tuple=True)[0].cpu().tolist() + print(f"\n group {group_idx}", flush=True) + for local_idx, batch_idx in enumerate(indices): + feature_dict = gdpo_features[batch_idx] + if not isinstance(feature_dict, dict): + feature_dict = {} + parts = [ + f" [{local_idx}] scalar_reward={float(scalar_rewards[batch_idx]):.4f}" + ] + for name in resolved_features[batch_idx]: + parts.append( + _format_gdpo_feature_for_log( + name, + feature_dict.get(name), + resolved_feature_weights[batch_idx].get(name, 1.0), + ) + ) + print(" ".join(parts), flush=True) + print(f"{'=' * 70}\n", flush=True) + class ReinforcePlusPlusAdvantageEstimator: """Reinforce++ advantage estimator with optional baseline subtraction and KL penalty in reward. diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 9639c884293..f297bf304b7 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -3015,6 +3015,7 @@ def grpo_train( greedy=False, effort_config=_get_effort_config(master_config), reward_penalty_config=master_config.reward_penalties, + length_adjustment_config=master_config.grpo.model_dump(), thinking_tags=get_nemo_gym_thinking_tags(master_config.env), mask_env_flagged_samples=should_mask_flagged_samples( master_config.env @@ -3938,6 +3939,7 @@ def validate( greedy=False, effort_config=_get_effort_config(master_config), reward_penalty_config=master_config.reward_penalties, + length_adjustment_config=master_config.grpo.model_dump(), thinking_tags=get_nemo_gym_thinking_tags(master_config.env), mask_env_flagged_samples=should_mask_flagged_samples( master_config.env diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 4b03741f199..7e678f30d28 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -72,6 +72,7 @@ collect_multimodal_payload_metrics, print_multimodal_payload_metrics, ) +from nemo_rl.utils.length_adjustments import apply_group_length_adjustments from nemo_rl.utils.timer import Timer TokenizerType = PreTrainedTokenizerBase @@ -2219,6 +2220,130 @@ def apply_reward_penalties( return counts +def _as_token_id_list(token_ids: Any) -> list[int]: + if token_ids is None: + return [] + if isinstance(token_ids, torch.Tensor): + return [int(x) for x in token_ids.detach().cpu().flatten().tolist()] + if isinstance(token_ids, (list, tuple)): + return [int(x) for x in token_ids] + return [] + + +def _assistant_generated_token_ids(result: dict) -> list[int]: + ids: list[int] = [] + for msg in result.get("message_log", []): + if msg.get("role") == "assistant": + ids.extend(_as_token_id_list(msg.get("token_ids"))) + return ids + + +def _output_generation_text(result: dict) -> str: + chunks: list[str] = [] + output_items = result["full_result"].get("response", {}).get("output", []) + for item in output_items: + gen_str = item.get("generation_str", "") + if isinstance(gen_str, str) and gen_str: + chunks.append(gen_str) + return "".join(chunks) + + +def _count_close_think_tags( + result: dict, + reward_penalty_config: dict[str, Any] | BaseModel | None, + tokenizer: TokenizerType | None = None, +) -> dict[str, Any]: + think_close_token_id = _get_reward_penalty_token_id( + reward_penalty_config or {}, "think_close" + ) + if think_close_token_id is None: + think_close_token_id = 13 + + assistant_ids = _assistant_generated_token_ids(result) + token_count = sum(1 for token_id in assistant_ids if token_id == think_close_token_id) + + decoded_count = None + if tokenizer is not None and assistant_ids: + try: + decoded = tokenizer.decode(assistant_ids, skip_special_tokens=False) + decoded_count = decoded.count("") + except Exception: + decoded_count = None + + generation_text = _output_generation_text(result) + generation_str_count = generation_text.count("") if generation_text else 0 + + if decoded_count is not None: + close_count = decoded_count + source = "decoded_assistant_tokens" + elif generation_text: + close_count = max(token_count, generation_str_count) + source = "max_token_or_generation_str" + else: + close_count = token_count + source = "assistant_token_ids" + + return { + "count": int(close_count), + "source": source, + "token_count": int(token_count), + "generation_str_count": int(generation_str_count), + "decoded_count": decoded_count, + } + + +def _record_gdpo_think_count_features( + results: list[dict], + reward_penalty_config: dict[str, Any] | BaseModel | None, + tokenizer: TokenizerType | None = None, +) -> None: + for result in results: + close_info = _count_close_think_tags(result, reward_penalty_config, tokenizer) + close_count = close_info["count"] + features = result["full_result"].setdefault("gdpo_reward_features", {}) + features["think_count_delta"] = { + "reward": float(-abs(close_count - 1)), + "adjustment": "format", + "num_close_think_tags": close_count, + "count_source": close_info["source"], + "token_close_think_tags": close_info["token_count"], + "generation_str_close_think_tags": close_info["generation_str_count"], + "decoded_close_think_tags": close_info["decoded_count"], + } + + +def _record_gdpo_env_reward_features(results: list[dict]) -> None: + for result in results: + features = result["full_result"].setdefault("gdpo_reward_features", {}) + reward = float(result["full_result"]["reward"]) + features["env_reward"] = {"reward": reward, "adjustment": None} + features.setdefault( + "length_adjusted_reward", + {"reward": reward, "adjustment": "combined"}, + ) + + +def _calculate_gdpo_reward_feature_metrics(results: list[dict]) -> dict[str, float]: + values_by_feature: dict[str, list[float]] = defaultdict(list) + for result in results: + features = result["full_result"].get("gdpo_reward_features", {}) + if not isinstance(features, dict): + continue + for name, entry in features.items(): + if isinstance(entry, dict) and "reward" in entry: + values_by_feature[name].append(float(entry["reward"])) + + metrics: dict[str, float] = {} + for name, values in values_by_feature.items(): + if not values: + continue + metric_prefix = f"gdpo_{name}" + metrics[f"{metric_prefix}/mean"] = sum(values) / len(values) + metrics[f"{metric_prefix}/min"] = min(values) + metrics[f"{metric_prefix}/max"] = max(values) + return metrics + + def _prepare_nemo_gym_rows( rows: list[dict], generation_config: GenerationConfig, @@ -2271,6 +2396,7 @@ async def run_async_nemo_gym_rollout( greedy: bool = False, effort_config: Optional[EffortLevelsConfig] = None, reward_penalty_config: dict[str, Any] | BaseModel | None = None, + length_adjustment_config: dict[str, Any] | BaseModel | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, mask_env_flagged_samples: bool = True, returns_entire_batch: bool = False, @@ -2302,6 +2428,7 @@ async def run_async_nemo_gym_rollout( greedy: Must be ``False`` because this path does not support greedy mode. effort_config: Optional configuration for effort-based reward shaping. reward_penalty_config: Optional reward-penalty configuration. + length_adjustment_config: Optional GRPO config block for length adjustments. thinking_tags: Optional opening and closing tags used by thinking penalties. mask_env_flagged_samples: Whether to carry env-driven ``mask_sample`` flags in the rollout batch for loss masking. @@ -2476,8 +2603,9 @@ async def run_async_nemo_gym_rollout( tokenizer=tokenizer, log_full_result_tables=log_full_result_tables, effort_config=effort_config, - reward_penalty_config=reward_penalty_config, - thinking_tags=thinking_tags, + reward_penalty_config=reward_penalty_config, + length_adjustment_config=length_adjustment_config, + thinking_tags=thinking_tags, mask_env_flagged_samples=mask_env_flagged_samples, ) if accumulator.is_complete: @@ -2514,6 +2642,7 @@ def run_nemo_gym_rollout_sync( greedy: bool = False, effort_config: Optional[EffortLevelsConfig] = None, reward_penalty_config: dict[str, Any] | BaseModel | None = None, + length_adjustment_config: dict[str, Any] | BaseModel | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, sampling_params: Optional[GenerationSamplingParams] = None, mask_env_flagged_samples: bool = True, @@ -2578,6 +2707,7 @@ async def _consume_rollout() -> NemoGymRolloutResult: greedy=greedy, effort_config=effort_config, reward_penalty_config=reward_penalty_config, + length_adjustment_config=length_adjustment_config, thinking_tags=thinking_tags, mask_env_flagged_samples=mask_env_flagged_samples, returns_entire_batch=True, @@ -2604,6 +2734,7 @@ def _postprocess_single_nemo_gym_group( log_full_result_tables: bool, effort_config: Optional[EffortLevelsConfig] = None, reward_penalty_config: dict[str, Any] | BaseModel | None = None, + length_adjustment_config: dict[str, Any] | BaseModel | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, mask_env_flagged_samples: bool = True, ) -> NemoGymRolloutResult: @@ -2618,7 +2749,35 @@ def _postprocess_single_nemo_gym_group( resolved_reward_penalty_config = resolve_reward_penalty_config( reward_penalty_config, tokenizer, thinking_tags=thinking_tags ) + _record_gdpo_think_count_features( + results, resolved_reward_penalty_config, tokenizer + ) penalty_counts = apply_reward_penalties(results, resolved_reward_penalty_config) + _record_gdpo_env_reward_features(results) + + for nemo_gym_row, result in zip(nemo_gym_rows, results): + result["agent_ref"] = nemo_gym_row["agent_ref"] + result["profiled_rewards"] = nemo_gym_row.get("profiled_rewards") + result["profiled_output_lengths"] = nemo_gym_row.get( + "profiled_output_lengths" + ) + result["profiled_reasoning_lengths"] = nemo_gym_row.get( + "profiled_reasoning_lengths" + ) + result["profiled_answer_lengths"] = nemo_gym_row.get( + "profiled_answer_lengths" + ) + result["profile_band"] = nemo_gym_row.get("profile_band") + + if length_adjustment_config is not None: + grpo_config = ( + length_adjustment_config.model_dump() + if isinstance(length_adjustment_config, BaseModel) + else dict(length_adjustment_config) + ) + apply_group_length_adjustments( + results, {"grpo": grpo_config}, tokenizer=tokenizer + ) # Prepare for the rollout metrics calculation below. Not strictly necessary here, but good to have parity with `run_async_multi_turn_rollout` with timer.time(f"{timer_prefix}/prepare_for_metrics_calculation"): @@ -2748,6 +2907,7 @@ def _postprocess_single_nemo_gym_group( ) rollout_metrics.update(per_agent_metrics) + rollout_metrics.update(_calculate_gdpo_reward_feature_metrics(results)) # Necessary for downstream nemo rl logging/printing. rollout_metrics["mean_gen_tokens_per_sample"] = rollout_metrics[ @@ -2782,6 +2942,9 @@ def _postprocess_single_nemo_gym_group( # stop_strings: NotRequired[list[str]] # Optional stop strings for generation # Extra information not in the DatumSpec used by the GRPO algorithm "total_reward": torch.tensor([r["full_result"]["reward"] for r in results]), + "gdpo_reward_features": [ + r["full_result"].get("gdpo_reward_features", {}) for r in results + ], # Add truncated field to match other rollout paths (reusing hit_max_tokens logic) "truncated": torch.tensor( [m["hit_max_tokens"] for m in all_sample_metrics], dtype=torch.bool diff --git a/nemo_rl/utils/length_adjustments.md b/nemo_rl/utils/length_adjustments.md new file mode 100644 index 00000000000..35dd936a8dc --- /dev/null +++ b/nemo_rl/utils/length_adjustments.md @@ -0,0 +1,598 @@ +# Length Adjustment Algorithms + +This file documents the length-penalty and length-bonus algorithms implemented in: + +`nemo_rl/utils/length_adjustments.py` + +The code supports two usage modes: + +1. **Reward mutation mode** + Configure `grpo.length_bonus`. The length adjustment mutates `full_result["reward"]`. + +2. **GDPO feature mode** + Configure `grpo.adv_estimator.name: gdpo` and put length feature knobs under + `grpo.adv_estimator.reward_features`. The same calculations are recorded in + `full_result["gdpo_reward_features"]` and consumed by GDPO without mutating the scalar + environment reward. + +All algorithms are resolved per prompt group. Unless otherwise stated, only rollouts with +`reward > 0` participate in length comparisons and receive length-based adjustments. + +## Common Config + +```yaml +grpo: + length_bonus: + verbose: true + default: + enabled: true + length_type: tokens + top_percentile: 0.5 + reasoning_bonus: 0.0 + answer_bonus: 0.0 + total_bonus: 0.0 + longest_reasoning_penalty: 0.0 + longest_answer_penalty: 0.0 + longest_total_penalty: 0.0 + group_reasoning_length_penalty_coeff: 0.0 + group_answer_length_penalty_coeff: 0.0 + group_total_length_penalty_coeff: 0.0 + reasoning_zmad_threshold: 0.0 + reasoning_zmad_penalty: 0.0 + answer_zmad_threshold: 0.0 + answer_zmad_penalty: 0.0 + total_zmad_threshold: 0.0 + total_zmad_penalty: 0.0 + profiled_length_penalty: 0.0 + profiled_length_n_std: 1.0 + profiled_length_min_samples: 2 + profile_band_total: false + profile_band_reasoning: false + profile_band_answer: false + agent_overrides: + math_with_judge_simple_agent: + enabled: true + group_total_length_penalty_coeff: 0.1 + instruction_following_simple_agent: + enabled: false + genrm_simple_agent: + enabled: false + code_gen_simple_agent: + enabled: true + total_zmad_threshold: 2.0 + total_zmad_penalty: 0.1 +``` + +`length_type` may be: + +- `tokens`: lengths are tokenizer token counts. +- anything else: lengths fall back to character counts. + +`agent_overrides` can override any supported parameter per agent. If an agent is missing from +`agent_overrides`, the implementation falls back to `default`. To disable length adjustments for +a specific environment or agent while keeping the default enabled, set `enabled: false` for that +agent: + +```yaml +grpo: + length_bonus: + default: + enabled: true + group_total_length_penalty_coeff: 0.1 + agent_overrides: + instruction_following_simple_agent: + enabled: false + genrm_simple_agent: + enabled: false +``` + +## Flag Reference + +| Flag | Short description | +| --- | --- | +| `verbose` | Prints per-group length adjustment details during rollout processing. | +| `default` | Default length-adjustment config used for agents without an override. | +| `agent_overrides` | Per-agent config overrides keyed by agent name. | +| `enabled` | Enables length adjustment for this config block. | +| `length_type` | Selects length unit: `tokens` uses tokenizer counts; other values use character counts. | +| `top_percentile` | Fraction of positive scorers treated as top scorers for longest-penalty selection. | +| `reasoning_bonus` | Flat bonus for the shortest positive/top-scoring reasoning trace in a prompt group. | +| `answer_bonus` | Flat bonus for the shortest positive/top-scoring answer in a prompt group. | +| `total_bonus` | Flat bonus for the shortest positive/top-scoring reasoning + answer total length. | +| `longest_reasoning_penalty` | Flat penalty for the longest reasoning trace among top scorers. | +| `longest_answer_penalty` | Flat penalty for the longest answer among top scorers. | +| `longest_total_penalty` | Flat penalty for the longest reasoning + answer total length among top scorers. | +| `group_reasoning_length_penalty_coeff` | Dense group-relative coefficient for reasoning length; shorter positive rollouts get higher adjustment. | +| `group_answer_length_penalty_coeff` | Dense group-relative coefficient for answer length. | +| `group_total_length_penalty_coeff` | Dense group-relative coefficient for total reasoning + answer length. | +| `reasoning_zmad_threshold` | Modified-Z threshold for flagging long reasoning outliers. | +| `reasoning_zmad_penalty` | Flat penalty applied to reasoning lengths above the zMAD threshold. | +| `answer_zmad_threshold` | Modified-Z threshold for flagging long answer outliers. | +| `answer_zmad_penalty` | Flat penalty applied to answer lengths above the zMAD threshold. | +| `total_zmad_threshold` | Modified-Z threshold for flagging long total-length outliers. | +| `total_zmad_penalty` | Flat penalty applied to total lengths above the zMAD threshold. | +| `profiled_length_penalty` | Flat penalty for rollouts longer than a per-prompt profiled-length threshold. | +| `profiled_length_n_std` | Number of standard deviations used in `mean + n_std * std` for profiled-length thresholding. | +| `profiled_length_min_samples` | Minimum profiled length samples needed before computing the profiled threshold. | +| `profile_band_total` | Enables per-prompt `{a,b,f}` multiplier on total length for correct rollouts. | +| `profile_band_reasoning` | Enables per-prompt `{a,b,f}` multiplier on reasoning length for correct rollouts. | +| `profile_band_answer` | Enables per-prompt `{a,b,f}` multiplier on answer length for correct rollouts. | +| `group_length_penalty_profile_gate` | Gates group-relative length coefficients using a per-prompt `profile_band` threshold. | +| `group_length_penalty_profile_gate_channel` | Selects which profile-band channel to gate on: `reasoning`, `answer`, or `total`. | +| `group_length_penalty_profile_gate_field` | Selects which field from the chosen profile-band channel to use as the gate threshold, usually `a`. | +| `group_length_penalty_profile_gate_positive_only` | If true, computes the gate mean using only `reward > 0` rollouts; if false, uses all rollouts. | + +## GDPO Feature Mode + +In GDPO feature mode, the same feature names can be selected under `reward_features`. + +```yaml +grpo: + adv_estimator: + name: gdpo + reward_features: + default: + env_reward: 1.0 + length_adjusted_reward: + group_total_length_penalty_coeff: 0.1 + think_count_delta: 1.0 +``` + +Feature entries can also be weighted: + +```yaml +grpo: + adv_estimator: + name: gdpo + reward_features: + default: + env_reward: 1.0 + length_adjusted_reward: + group_total_length_penalty_coeff: 0.1 + think_count_delta: 0.5 +``` + +The rollout code records feature metrics to WandB using names like: + +```text +train/gdpo_length_adjusted_reward/mean +train/gdpo_length_adjusted_reward/min +train/gdpo_length_adjusted_reward/max +train/gdpo_think_count_delta/mean +``` + +## Per-Prompt Data Format + +Some algorithms depend on metadata stored on each training-data row. The rollout code copies +these fields from `extra_env_info` into each rollout result before applying length adjustments. + +At minimum, a row still looks like a normal NeMo-Gym training example. The length-related fields +are extra keys: + +```json +{ + "problem": "Solve ...", + "expected_answer": "42", + "agent_name": "math_with_judge_simple_agent", + "extra_env_info": { + "profiled_rewards": [1, 1, 0, 1, 0, 1, 1, 1], + "profiled_output_lengths": [18342, 17110, 32768, 19004, 28991, 16820, 17455, 18101], + "profile_band": { + "total": {"a": 18138.6667, "b": 23756.0123, "f": 0.9}, + "reasoning": {"a": 17686.3333, "b": 23111.0123, "f": 0.9}, + "answer": {"a": 452.3333, "b": 1097.3333, "f": 0.9} + } + } +} +``` + +Some data files store these fields at top level instead of inside `extra_env_info`; the important +part is that by rollout time the result has: + +```json +{ + "profiled_rewards": [1, 1, 0, 1, 0, 1, 1, 1], + "profiled_output_lengths": [18342, 17110, 32768, 19004, 28991, 16820, 17455, 18101], + "profile_band": { + "total": {"a": 18138.6667, "b": 23756.0123, "f": 0.9}, + "reasoning": {"a": 17686.3333, "b": 23111.0123, "f": 0.9}, + "answer": {"a": 452.3333, "b": 1097.3333, "f": 0.9} + } +} +``` + +Field usage: + +- `profiled_rewards`: used by `profiled_length_penalty` to identify passing profiled rollouts. +- `profiled_output_lengths`: used by `profiled_length_penalty` to compute + `mean + n_std * std`. +- `profile_band.total`: used by `profile_band_total` and by profile-gated group-relative + penalties when `group_length_penalty_profile_gate_channel: total`. +- `profile_band.reasoning`: used by `profile_band_reasoning` and by profile-gated + group-relative penalties when the gate channel is `reasoning`. +- `profile_band.answer`: used by `profile_band_answer` and by profile-gated group-relative + penalties when the gate channel is `answer`. + +The profile-band values mean: + +```text +a: full reward / no penalty up to this length +b: multiplier reaches f at this length +f: multiplier at b +``` + +For profile-gated group-relative penalties, `a`, `b`, or `f` can be selected as the gate field, +though `a` is the normal choice: + +```yaml +group_length_penalty_profile_gate: true +group_length_penalty_profile_gate_channel: total +group_length_penalty_profile_gate_field: a +``` + +## Implemented Algorithms + +### 1. Shortest Rollout Bonus + +Config keys: + +- `reasoning_bonus` +- `answer_bonus` +- `total_bonus` + +For each prompt group, the code finds the shortest positive rollout for the selected channel +and adds a flat bonus if that rollout is also a top scorer. + +Channels: + +- `reasoning_bonus`: shortest non-empty reasoning length. +- `answer_bonus`: shortest non-empty answer length. +- `total_bonus`: shortest non-empty reasoning + answer length. + +This is a sparse adjustment: usually only one rollout per group gets the bonus for each enabled +channel. + +Example: + +```yaml +grpo: + length_bonus: + default: + enabled: true + total_bonus: 0.1 +``` + +### 2. Longest Top-Scorer Penalty + +Config keys: + +- `longest_reasoning_penalty` +- `longest_answer_penalty` +- `longest_total_penalty` +- `top_percentile` + +For each prompt group, the code first selects positive rollouts in the top score percentile. +Among those, it subtracts a flat penalty from the longest rollout for the selected channel. + +Channels: + +- `longest_reasoning_penalty` +- `longest_answer_penalty` +- `longest_total_penalty` + +The implementation requires at least two eligible top-scorer rollouts to compare. + +Example: + +```yaml +grpo: + length_bonus: + default: + enabled: true + top_percentile: 0.5 + longest_total_penalty: 0.1 +``` + +### 3. Group Relative-Length Scaling + +Config keys: + +- `group_reasoning_length_penalty_coeff` +- `group_answer_length_penalty_coeff` +- `group_total_length_penalty_coeff` + +This is a dense group-relative adjustment over positive rollouts. + +For each enabled channel: + +1. Find the shortest and longest positive rollout lengths in the group. +2. Convert each length to a raw weight where shorter is larger: + + ```text + raw_weight = 1 - (length - min_length) / (max_length - min_length) + ``` + +3. Zero-center the weights by subtracting the mean raw weight. +4. Multiply by the configured coefficient. + +Shorter positive rollouts receive positive adjustment; longer positive rollouts receive negative +adjustment. If all lengths are equal, the adjustment is zero. + +Example: + +```yaml +grpo: + length_bonus: + default: + enabled: true + group_total_length_penalty_coeff: 0.1 +``` + +### 4. zMAD Long-Outlier Penalty + +Config keys: + +- `reasoning_zmad_threshold` +- `reasoning_zmad_penalty` +- `answer_zmad_threshold` +- `answer_zmad_penalty` +- `total_zmad_threshold` +- `total_zmad_penalty` + +This flags high-side length outliers among positive rollouts using the Iglewicz-Hoaglin modified +Z score: + +```text +modified_z = 0.6745 * (length - median_length) / MAD +``` + +If `modified_z > threshold`, the corresponding flat penalty is subtracted. + +Only long-side outliers are penalized. Short outliers are not penalized. + +The implementation also has a fixed MAD floor: + +```text +MAD / median >= 0.015 +``` + +If the MAD is too small, no zMAD outliers are flagged. + +Example: + +```yaml +grpo: + length_bonus: + default: + enabled: true + total_zmad_threshold: 2.5 + total_zmad_penalty: 0.1 +``` + +### 5. Profiled Length Threshold Penalty + +Config keys: + +- `profiled_length_penalty` +- `profiled_length_n_std` +- `profiled_length_min_samples` + +This uses per-prompt profiling metadata: + +- `profiled_rewards` +- `profiled_output_lengths` + +For each prompt group: + +1. Prefer profiled lengths from passing rollouts. +2. If there are fewer than `profiled_length_min_samples` passing rollouts, fall back to all + profiled lengths. +3. Compute: + + ```text + threshold = mean(profiled_lengths) + profiled_length_n_std * std(profiled_lengths) + ``` + +4. Penalize rollouts whose total generated length is greater than or equal to the threshold. + +Example: + +```yaml +grpo: + length_bonus: + default: + enabled: true + profiled_length_penalty: 0.1 + profiled_length_n_std: 1.0 + profiled_length_min_samples: 2 +``` + +### 6. Profile-Band Multiplier + +Config keys: + +- `profile_band_total` +- `profile_band_reasoning` +- `profile_band_answer` + +This uses per-row `profile_band` metadata with channel-specific `{a, b, f}` values: + +```json +{ + "profile_band": { + "total": {"a": 10000, "b": 15000, "f": 0.9}, + "reasoning": {"a": 9000, "b": 14000, "f": 0.9}, + "answer": {"a": 500, "b": 1000, "f": 0.9} + } +} +``` + +For an enabled channel, the multiplier is: + +```text +length <= a: multiplier = 1 +length > a: multiplier = max(0, 1 - (length - a) / (b - a) * (1 - f)) +``` + +So the multiplier is `f` at `b`, then the same slope continues past `b` until clamped at `0`. + +Profile-band multipliers are applied only to rollouts whose original environment reward is +positive. + +Example: + +```yaml +grpo: + length_bonus: + default: + enabled: true + profile_band_total: true +``` + +### 7. Profile-Gated Group Relative-Length Scaling + +Config keys: + +- `group_length_penalty_profile_gate` +- `group_length_penalty_profile_gate_channel` +- `group_length_penalty_profile_gate_field` +- `group_length_penalty_profile_gate_positive_only` +- plus one or more group-relative coefficients: + - `group_reasoning_length_penalty_coeff` + - `group_answer_length_penalty_coeff` + - `group_total_length_penalty_coeff` + +This is a gate on group-relative length scaling. It does not define a separate penalty by itself. + +For each prompt group: + +1. Read a threshold from `profile_band[channel][field]`, for example `profile_band["total"]["a"]`. +2. Compute the mean rollout length for the selected channel. +3. If `group_length_penalty_profile_gate_positive_only` is true, use only positive rollouts in + that mean. +4. Enable group-relative length scaling only if: + + ```text + mean_length > profile_band[channel][field] + ``` + +If the gate is closed, all group-relative coefficients are set to zero for that prompt group. + +`group_length_penalty_profile_gate_positive_only` only affects the gate decision. It does not +change the rollouts that receive the group-relative adjustment after the gate opens. In the +current implementation, group-relative length scaling itself still applies only to positive +rollouts. + +Example: + +```yaml +grpo: + length_bonus: + default: + enabled: true + group_total_length_penalty_coeff: 0.1 + group_length_penalty_profile_gate: true + group_length_penalty_profile_gate_channel: total + group_length_penalty_profile_gate_field: a + group_length_penalty_profile_gate_positive_only: true +``` + +## Recorded GDPO Feature Names + +The implementation records these feature names in `full_result["gdpo_reward_features"]`: + +- `env_reward` +- `reasoning_bonus` +- `answer_bonus` +- `total_bonus` +- `longest_reasoning_penalty` +- `longest_answer_penalty` +- `longest_total_penalty` +- `group_reasoning_length_penalty_coeff` +- `group_answer_length_penalty_coeff` +- `group_total_length_penalty_coeff` +- `reasoning_zmad_penalty` +- `answer_zmad_penalty` +- `total_zmad_penalty` +- `profiled_length_penalty` +- `profile_band_total` +- `profile_band_reasoning` +- `profile_band_answer` +- `profile_band_delta` +- `length_additive_delta` +- `length_total_delta` +- `length_adjusted_reward` + +`length_adjusted_reward` is the combined length-adjusted scalar that GDPO can use as one reward +feature. `length_additive_delta`, `profile_band_delta`, and `length_total_delta` are derived +summary features. + +## Formatting Feature: think_count_delta + +The GDPO branch also supports a formatting feature: + +```text +think_count_delta = -abs(num_close_think_tags - 1) +``` + +This is not a length-penalty algorithm, but it can be selected alongside length features in GDPO: + +```yaml +grpo: + adv_estimator: + name: gdpo + reward_features: + default: + env_reward: 1.0 + length_adjusted_reward: + group_total_length_penalty_coeff: 0.1 + think_count_delta: 1.0 +``` + +For a format-only GDPO setup, use only the environment reward and the malformed-format feature: + +```yaml +grpo: + adv_estimator: + name: gdpo + reward_features: + default: + env_reward: 1.0 + think_count_delta: 1.0 +``` + +To make the malformed-format feature weaker than the task reward, lower its weight: + +```yaml +grpo: + adv_estimator: + name: gdpo + reward_features: + default: + env_reward: 1.0 + think_count_delta: 0.5 +``` + +## Practical Notes + +- Most length algorithms act only on positive rollouts (`reward > 0`). +- `profile_band_*` multipliers also apply only to originally correct rollouts. +- Group-relative scaling can reduce average length aggressively because it gives dense per-group + pressure. +- zMAD is more selective: it only hits high-side outliers. +- GDPO feature mode is useful when you want length or formatting behavior to be represented as a + separate feature instead of mixing it into the scalar environment reward. + +## Final Recommendations + +1. For domains where longer reasoning does not strongly correlate with higher accuracy, apply + group-relative length scaling. This gives steady pressure toward shorter correct rollouts. + +2. For domains where longer reasoning does correlate with higher accuracy, leave length + unpenalized. Penalizing length in those domains can remove useful reasoning and hurt task + performance. + +3. If you have a target length in mind, for example when this model should be less verbose on a + domain than another reference model, use profile-gated group-relative length scaling. The + profile gate lets the penalty activate only when the prompt group's rollout lengths exceed a + per-prompt target. diff --git a/nemo_rl/utils/length_adjustments.py b/nemo_rl/utils/length_adjustments.py new file mode 100644 index 00000000000..289f78d647a --- /dev/null +++ b/nemo_rl/utils/length_adjustments.py @@ -0,0 +1,1178 @@ +# 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. + +"""Per-prompt-group length bonuses/penalties for rollout rewards. + +Rewards conciseness among high-quality generations by applying: +1. A flat bonus to the shortest generation among top scorers in each prompt group. +2. Optional flat penalties on the longest reasoning / longest answer among + top-percentile scorers (with at least two eligible rollouts to compare). +3. Independent zero-centered penalties for reasoning and answer length. +4. Optional Iglewicz–Hoaglin modified-Z (MAD) high-side outliers among positive + scorers. Config keys: ``reasoning_zmad_threshold``, ``reasoning_zmad_penalty``, + ``answer_zmad_threshold``, ``answer_zmad_penalty``. + If ``reasoning_zmad_threshold`` or ``answer_zmad_threshold`` is ≤ 0, that + channel is off and its penalty is ignored (no flagging). If threshold > 0 + but the matching penalty is 0, nothing is subtracted. + +Supports per-agent filtering and parameter overrides via config. +""" + +from __future__ import annotations + +import logging +import statistics +from typing import Any + +logger = logging.getLogger(__name__) + +# MAD/median floor for zMAD (fixed; matches ``flag_reasoning_length_outliers`` default). +_ZMAD_MIN_MAD_REL = 0.015 + +_PARAM_KEYS = ( + "enabled", + "reasoning_bonus", + "answer_bonus", + "total_bonus", + "longest_reasoning_penalty", + "longest_answer_penalty", + "longest_total_penalty", + "top_percentile", + "group_reasoning_length_penalty_coeff", + "group_answer_length_penalty_coeff", + "group_total_length_penalty_coeff", + "length_type", + "reasoning_zmad_threshold", + "reasoning_zmad_penalty", + "answer_zmad_threshold", + "answer_zmad_penalty", + "total_zmad_threshold", + "total_zmad_penalty", + "profiled_length_penalty", + "profiled_length_n_std", + "profiled_length_min_samples", + "profile_band_total", + "profile_band_reasoning", + "profile_band_answer", + "group_length_penalty_profile_gate", + "group_length_penalty_profile_gate_channel", + "group_length_penalty_profile_gate_field", + "group_length_penalty_profile_gate_positive_only", +) + +# Param keys that should be merged as bools rather than floats. +_BOOL_PARAM_KEYS = frozenset({ + "enabled", + "profile_band_total", + "profile_band_reasoning", + "profile_band_answer", + "group_length_penalty_profile_gate", + "group_length_penalty_profile_gate_positive_only", +}) + +_STR_PARAM_KEYS = frozenset({ + "length_type", + "group_length_penalty_profile_gate_channel", + "group_length_penalty_profile_gate_field", +}) + +_GDPO_LENGTH_FEATURE_PARAM_KEYS = frozenset({ + "reasoning_bonus", + "answer_bonus", + "total_bonus", + "longest_reasoning_penalty", + "longest_answer_penalty", + "longest_total_penalty", + "group_reasoning_length_penalty_coeff", + "group_answer_length_penalty_coeff", + "group_total_length_penalty_coeff", + "reasoning_zmad_penalty", + "answer_zmad_penalty", + "total_zmad_penalty", + "profiled_length_penalty", + "profile_band_total", + "profile_band_reasoning", + "profile_band_answer", + "group_length_penalty_profile_gate", + "group_length_penalty_profile_gate_channel", + "group_length_penalty_profile_gate_field", + "group_length_penalty_profile_gate_positive_only", +}) + + +def _resolve_gdpo_feature_log_names_for_agent( + adv_cfg: dict[str, Any], agent_name: str +) -> list[str]: + reward_features = adv_cfg.get("reward_features", ["env_reward"]) + if isinstance(reward_features, (list, tuple)): + return list(reward_features) + if not isinstance(reward_features, dict): + return ["env_reward"] + if "default" not in reward_features and "agent_overrides" not in reward_features: + return list(reward_features) + + selected = reward_features.get("agent_overrides", {}).get( + agent_name, reward_features.get("default", ["env_reward"]) + ) + if isinstance(selected, dict): + return list(selected) + if isinstance(selected, (list, tuple)): + return list(selected) + return ["env_reward"] + + +def _set_gdpo_reward_feature( + result: dict[str, Any], + name: str, + reward: float, + adjustment: str | None, + **metadata: Any, +) -> None: + features = result["full_result"].setdefault("gdpo_reward_features", {}) + entry = {"reward": float(reward), "adjustment": adjustment} + entry.update(metadata) + features[name] = entry + + +def _extract_reasoning_and_answer_text(result: dict[str, Any]) -> tuple[str, str]: + """Extract reasoning and answer text from the Response API output items.""" + fr = result.get("full_result", {}) + response_obj = fr.get("response", {}) + output_items = ( + response_obj.get("output", []) + if isinstance(response_obj, dict) + else getattr(response_obj, "output", []) + ) + + reasoning_text = "" + answer_text = "" + for item in output_items: + item_type = item.get("type", "") if isinstance(item, dict) else getattr(item, "type", "") + if item_type == "reasoning": + summaries = item.get("summary", []) if isinstance(item, dict) else getattr(item, "summary", []) + for s in summaries: + t = s.get("text", "") if isinstance(s, dict) else getattr(s, "text", "") + reasoning_text += t + elif item_type == "message": + content = item.get("content", []) if isinstance(item, dict) else getattr(item, "content", []) + if isinstance(content, list): + for c in content: + t = c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") + answer_text += t + elif isinstance(content, str): + answer_text += content + + return reasoning_text, answer_text + + +def _extract_gdpo_length_feature_params(feature_cfg: Any) -> dict[str, Any]: + if not isinstance(feature_cfg, dict): + return {} + params = {} + for key, value in feature_cfg.items(): + if key not in _GDPO_LENGTH_FEATURE_PARAM_KEYS: + continue + if key in _BOOL_PARAM_KEYS: + if isinstance(value, bool): + params[key] = value + continue + if key in _STR_PARAM_KEYS: + if isinstance(value, str): + params[key] = value + continue + if isinstance(value, bool) or not isinstance(value, (int, float)): + continue + params[key] = float(value) + return params + + +def _merge_gdpo_reward_features_into_length_cfg( + grpo_config: dict[str, Any], +) -> dict[str, Any]: + length_cfg = dict(grpo_config.get("length_bonus", {}) or {}) + has_explicit_length_cfg = bool(length_cfg) + adv_cfg = grpo_config.get("adv_estimator", {}) or {} + if adv_cfg.get("name") != "gdpo": + return length_cfg + + reward_features = adv_cfg.get("reward_features") + if not isinstance(reward_features, dict): + return length_cfg + + if "default" in reward_features or "agent_overrides" in reward_features: + default_features = reward_features.get("default", {}) + agent_features = reward_features.get("agent_overrides", {}) + else: + default_features = reward_features + agent_features = {} + + default_params = _extract_gdpo_length_feature_params(default_features) + agent_params = { + agent_name: _extract_gdpo_length_feature_params(features) + for agent_name, features in agent_features.items() + } + agent_params = {k: v for k, v in agent_params.items() if v} + if not default_params and not agent_params: + return length_cfg + + if not has_explicit_length_cfg: + length_cfg["_gdpo_feature_only"] = True + + default_cfg = dict(length_cfg.get("default", {}) or {}) + default_cfg.update(default_params) + default_cfg.setdefault("enabled", True) + length_cfg["default"] = default_cfg + + overrides = dict(length_cfg.get("agent_overrides", {}) or {}) + for agent_name, params in agent_params.items(): + agent_cfg = dict(overrides.get(agent_name, {}) or {}) + agent_cfg.update(params) + agent_cfg.setdefault("enabled", True) + overrides[agent_name] = agent_cfg + if overrides: + length_cfg["agent_overrides"] = overrides + + if adv_cfg.get("verbose", False): + length_cfg["_gdpo_feature_verbose"] = True + return length_cfg + + +def apply_group_length_adjustments( + results: list[dict[str, Any]], + master_config: dict[str, Any], + tokenizer: Any = None, +) -> None: + """Apply per-prompt-group length bonuses/penalties. + + Reads ``grpo.length_bonus`` and GDPO ``reward_features`` for configuration. + No-ops when no length-adjustment feature is enabled. + + Calculates all adjustments first and records GDPO reward features. Explicit + ``grpo.length_bonus`` configs mutate scalar rewards; GDPO-only feature + configs do not. + + Args: + results: List of per-generation result dicts. + master_config: Full training config dict. + tokenizer: Tokenizer for computing reasoning/answer token counts. + """ + grpo_config = master_config.get("grpo", {}) + length_cfg = _merge_gdpo_reward_features_into_length_cfg(grpo_config) + if not length_cfg: + return + + default_cfg = length_cfg.get("default", {}) + agents_cfg = length_cfg.get("agent_overrides") + gdpo_feature_only = bool(length_cfg.get("_gdpo_feature_only", False)) + verbose = bool(length_cfg.get("verbose", False)) and not gdpo_feature_only + gdpo_feature_verbose = bool( + length_cfg.get("_gdpo_feature_verbose", False) + or (length_cfg.get("verbose", False) and gdpo_feature_only) + ) + should_mutate_reward = not gdpo_feature_only + if not default_cfg.get("enabled", False) and not agents_cfg: + return + + num_gens = master_config["grpo"]["num_generations_per_prompt"] + defaults: dict[str, Any] = {} + for k in _PARAM_KEYS: + if k == "length_type": + defaults[k] = default_cfg.get(k, "tokens") + elif k == "group_length_penalty_profile_gate_channel": + defaults[k] = default_cfg.get(k, "total") + elif k == "group_length_penalty_profile_gate_field": + defaults[k] = default_cfg.get(k, "a") + elif k == "enabled": + defaults[k] = default_cfg.get(k, True) + elif k == "group_length_penalty_profile_gate_positive_only": + defaults[k] = default_cfg.get(k, True) + elif k in _BOOL_PARAM_KEYS: + defaults[k] = default_cfg.get(k, False) + elif k == "profiled_length_min_samples": + defaults[k] = default_cfg.get(k, 2) + elif k == "profiled_length_n_std": + defaults[k] = default_cfg.get(k, 1.0) + else: + defaults[k] = default_cfg.get(k, 0.0) + defaults.setdefault("top_percentile", 0.2) + + n = len(results) + original_rewards = [r["full_result"]["reward"] for r in results] + agent_names = [r["agent_ref"]["name"] for r in results] + + # Extract text once; lengths computed per-group based on resolved length_type + texts: list[tuple[str, str]] = [] + for r in results: + texts.append(_extract_reasoning_and_answer_text(r)) + + # Phase 1: calculate all adjustments per-group + all_adjustments = [0.0] * n + all_reasoning_adj = [0.0] * n + all_answer_adj = [0.0] * n + all_total_adj = [0.0] * n + all_reasoning_bonus = [0.0] * n + all_answer_bonus = [0.0] * n + all_total_bonus = [0.0] * n + all_reasoning_longest_pen = [0.0] * n + all_answer_longest_pen = [0.0] * n + all_total_longest_pen = [0.0] * n + all_zmad_reasoning_adj = [0.0] * n + all_zmad_answer_adj = [0.0] * n + all_zmad_total_adj = [0.0] * n + all_profiled_length_adj = [0.0] * n + reasoning_lengths = [0] * n + answer_lengths = [0] * n + total_lengths = [0] * n + groups_adjusted = 0 + group_gate_infos: dict[int, dict[str, Any]] = {} + + for g in range(0, n, num_gens): + agent_name = agent_names[g] + group_size = min(num_gens, n - g) + if any(results[g + k].get("low_effort_applied") for k in range(group_size)): + continue + params = _resolve_agent_params(agent_name, agents_cfg, defaults) + if params is None: + continue + if not params.pop("enabled", True): + continue + + group_lt = params.pop("length_type", "tokens") + use_tokens = group_lt == "tokens" + + for k in range(group_size): + idx = g + k + r_text, a_text = texts[idx] + if use_tokens and tokenizer is not None: + reasoning_lengths[idx] = len(tokenizer.encode(r_text, add_special_tokens=False)) if r_text else 0 + answer_lengths[idx] = len(tokenizer.encode(a_text, add_special_tokens=False)) if a_text else 0 + else: + reasoning_lengths[idx] = len(r_text) + answer_lengths[idx] = len(a_text) + + group_reasoning = reasoning_lengths[g : g + num_gens] + group_answer = answer_lengths[g : g + num_gens] + group_total = [r + a for r, a in zip(group_reasoning, group_answer)] + total_lengths[g : g + group_size] = group_total[:group_size] + group_rewards = original_rewards[g : g + num_gens] + gate_info = _group_length_profile_gate_info( + band=results[g].get("profile_band"), + params=params, + rewards=group_rewards[:group_size], + reasoning_lengths=group_reasoning[:group_size], + answer_lengths=group_answer[:group_size], + total_lengths=group_total[:group_size], + ) + group_gate_infos[g] = gate_info + if gate_info["enabled"] and not gate_info["open"]: + params["group_reasoning_length_penalty_coeff"] = 0.0 + params["group_answer_length_penalty_coeff"] = 0.0 + params["group_total_length_penalty_coeff"] = 0.0 + ( + _, + adjustments, + reasoning_adjs, + answer_adjs, + total_adjs, + r_bonus, + a_bonus, + t_bonus, + r_lpen, + a_lpen, + t_lpen, + zmad_r_adj, + zmad_a_adj, + zmad_t_adj, + ) = _apply_length_bonuses_and_penalties( + group_rewards, group_reasoning, group_answer, group_total, **params + ) + + for k in range(len(adjustments)): + all_adjustments[g + k] = adjustments[k] + all_reasoning_adj[g + k] = reasoning_adjs[k] + all_answer_adj[g + k] = answer_adjs[k] + all_total_adj[g + k] = total_adjs[k] + all_reasoning_bonus[g + k] = r_bonus[k] + all_answer_bonus[g + k] = a_bonus[k] + all_total_bonus[g + k] = t_bonus[k] + all_reasoning_longest_pen[g + k] = r_lpen[k] + all_answer_longest_pen[g + k] = a_lpen[k] + all_total_longest_pen[g + k] = t_lpen[k] + all_zmad_reasoning_adj[g + k] = zmad_r_adj[k] + all_zmad_answer_adj[g + k] = zmad_a_adj[k] + all_zmad_total_adj[g + k] = zmad_t_adj[k] + + if any(a != 0.0 for a in adjustments): + groups_adjusted += 1 + + # Profiled length penalty: penalize rollouts longer than mean + n_std of + # passing profiled lengths for this prompt. + plp = params.get("profiled_length_penalty", 0.0) + if plp > 0.0: + p_rewards = results[g].get("profiled_rewards") + p_lengths = results[g].get("profiled_output_lengths") + if p_rewards is not None and p_lengths is not None: + min_samples = int(params.get("profiled_length_min_samples", 2)) + passing = [l for r, l in zip(p_rewards, p_lengths) if r > 0] + if len(passing) < min_samples: + passing = list(p_lengths) + if len(passing) >= min_samples: + mean_l = statistics.mean(passing) + std_l = statistics.stdev(passing) if len(passing) >= 2 else 0.0 + n_std = float(params.get("profiled_length_n_std", 1.0)) + threshold = mean_l + n_std * std_l + for k in range(group_size): + idx = g + k + if total_lengths[idx] >= threshold: + all_profiled_length_adj[idx] = -plp + + # Phase 2: debug print (only when verbose flag is set) + if verbose: + num_groups = n // num_gens if num_gens > 0 else 0 + print(f"\n{'=' * 70}", flush=True) + print( + f"[Rollout] {n} samples, {num_groups} groups, {groups_adjusted} adjusted" + f" default longest_reasoning_penalty={defaults['longest_reasoning_penalty']}" + f" longest_answer_penalty={defaults['longest_answer_penalty']}", + flush=True, + ) + + for g in range(0, n, num_gens): + agent_name = agent_names[g] + group_size = min(num_gens, n - g) + low_effort = any(results[g + k].get("low_effort_applied") for k in range(group_size)) + params = _resolve_agent_params(agent_name, agents_cfg, defaults) + skipped = params is None + disabled = params is not None and not params.get("enabled", True) + + if low_effort: + print( + f"\n group {g // num_gens} agent={agent_name} [low_effort — skipped]", + flush=True, + ) + elif skipped: + print( + f"\n group {g // num_gens} agent={agent_name} [skipped]" + f" (default longest_reasoning_penalty={defaults['longest_reasoning_penalty']}" + f" longest_answer_penalty={defaults['longest_answer_penalty']})", + flush=True, + ) + elif disabled: + print( + f"\n group {g // num_gens} agent={agent_name} [disabled]" + f" longest_reasoning_penalty={params['longest_reasoning_penalty']}" + f" longest_answer_penalty={params['longest_answer_penalty']}", + flush=True, + ) + else: + lt = params.get("length_type", "tokens") + unit = "tok" if lt == "tokens" else "chr" + print( + f"\n group {g // num_gens} agent={agent_name}" + f" length_type={unit}" + f" reasoning_bonus={params['reasoning_bonus']} answer_bonus={params['answer_bonus']}" + f" total_bonus={params['total_bonus']}" + f" longest_reasoning_penalty={params['longest_reasoning_penalty']}" + f" longest_answer_penalty={params['longest_answer_penalty']}" + f" longest_total_penalty={params['longest_total_penalty']}" + f" top_pct={params['top_percentile']}" + f" reasoning_coeff={params['group_reasoning_length_penalty_coeff']}" + f" answer_coeff={params['group_answer_length_penalty_coeff']}" + f" total_coeff={params['group_total_length_penalty_coeff']}" + f" reasoning_zmad_threshold={params['reasoning_zmad_threshold']}" + f" reasoning_zmad_penalty={params['reasoning_zmad_penalty']}" + f" answer_zmad_threshold={params['answer_zmad_threshold']}" + f" answer_zmad_penalty={params['answer_zmad_penalty']}" + f" total_zmad_threshold={params['total_zmad_threshold']}" + f" total_zmad_penalty={params['total_zmad_penalty']}" + f" profiled_length_penalty={params['profiled_length_penalty']}" + f" profiled_length_n_std={params['profiled_length_n_std']}" + f" profiled_length_min_samples={params['profiled_length_min_samples']}", + flush=True, + ) + gate = group_gate_infos.get(g) + if gate and gate["enabled"]: + print( + f" profile_gate channel={gate['channel']} field={gate['field']}" + f" positive_only={gate['positive_only']}" + f" mean={gate['mean']}" + f" limit={gate['limit']}" + f" open={gate['open']}" + f" reason={gate['reason']}", + flush=True, + ) + for k in range(group_size): + idx = g + k + orig = original_rewards[idx] + profiled_adj = all_profiled_length_adj[idx] if all_adjustments[idx] >= 0 else 0.0 + final = orig + all_adjustments[idx] + profiled_adj + print( + f" [{k}] reward={orig:.4f}" + f" reasoning_len={reasoning_lengths[idx]}" + f" reasoning_adj={all_reasoning_adj[idx]:+.4f}" + f" reasoning_bonus={all_reasoning_bonus[idx]:+.4f}" + f" longest_reasoning_penalty_adj={all_reasoning_longest_pen[idx]:+.4f}" + f" answer_len={answer_lengths[idx]}" + f" answer_adj={all_answer_adj[idx]:+.4f}" + f" answer_bonus={all_answer_bonus[idx]:+.4f}" + f" longest_answer_penalty_adj={all_answer_longest_pen[idx]:+.4f}" + f" total_len={total_lengths[idx]}" + f" total_adj={all_total_adj[idx]:+.4f}" + f" total_bonus={all_total_bonus[idx]:+.4f}" + f" longest_total_penalty_adj={all_total_longest_pen[idx]:+.4f}" + f" zmad_r={all_zmad_reasoning_adj[idx]:+.4f}" + f" zmad_a={all_zmad_answer_adj[idx]:+.4f}" + f" zmad_t={all_zmad_total_adj[idx]:+.4f}" + f" profiled_len_adj={all_profiled_length_adj[idx]:+.4f}" + f" final_reward={final:.4f}", + flush=True, + ) + + print(f"{'=' * 70}\n", flush=True) + + # Phase 3: apply additive adjustments and record GDPO reward features + additive_base_rewards = [0.0] * n + for i, r in enumerate(results): + profiled_adj = all_profiled_length_adj[i] if all_adjustments[i] >= 0 else 0.0 + additive_delta = all_adjustments[i] + profiled_adj + additive_base_rewards[i] = original_rewards[i] + additive_delta + if "env_reward" not in r["full_result"].setdefault( + "gdpo_reward_features", {} + ): + _set_gdpo_reward_feature(r, "env_reward", original_rewards[i], None) + + for name, value in ( + ("reasoning_bonus", all_reasoning_bonus[i]), + ("answer_bonus", all_answer_bonus[i]), + ("total_bonus", all_total_bonus[i]), + ("longest_reasoning_penalty", all_reasoning_longest_pen[i]), + ("longest_answer_penalty", all_answer_longest_pen[i]), + ("longest_total_penalty", all_total_longest_pen[i]), + ("group_reasoning_length_penalty_coeff", all_reasoning_adj[i]), + ("group_answer_length_penalty_coeff", all_answer_adj[i]), + ("group_total_length_penalty_coeff", all_total_adj[i]), + ("reasoning_zmad_penalty", all_zmad_reasoning_adj[i]), + ("answer_zmad_penalty", all_zmad_answer_adj[i]), + ("total_zmad_penalty", all_zmad_total_adj[i]), + ("profiled_length_penalty", profiled_adj), + ): + _set_gdpo_reward_feature(r, name, value, "additive") + + _set_gdpo_reward_feature( + r, "profile_band_total", 0.0, "multiplicative", multiplier=1.0 + ) + _set_gdpo_reward_feature( + r, "profile_band_reasoning", 0.0, "multiplicative", multiplier=1.0 + ) + _set_gdpo_reward_feature( + r, "profile_band_answer", 0.0, "multiplicative", multiplier=1.0 + ) + _set_gdpo_reward_feature(r, "profile_band_delta", 0.0, "derived_sum") + _set_gdpo_reward_feature( + r, "length_additive_delta", additive_delta, "derived_sum" + ) + _set_gdpo_reward_feature(r, "length_total_delta", additive_delta, "derived_sum") + _set_gdpo_reward_feature( + r, "length_adjusted_reward", additive_base_rewards[i], "combined" + ) + + if should_mutate_reward: + r["full_result"]["reward"] = additive_base_rewards[i] + + # Phase 4: apply per-prompt profile_band multipliers (correct rollouts only). + _apply_profile_band_multipliers( + results=results, + original_rewards=original_rewards, + base_rewards=additive_base_rewards, + total_lengths=total_lengths, + reasoning_lengths=reasoning_lengths, + answer_lengths=answer_lengths, + agent_names=agent_names, + agents_cfg=agents_cfg, + defaults=defaults, + num_gens=num_gens, + should_mutate_reward=should_mutate_reward, + ) + + if verbose or gdpo_feature_verbose: + _print_gdpo_reward_feature_summary( + results=results, + agent_names=agent_names, + original_rewards=original_rewards, + num_gens=num_gens, + adv_cfg=grpo_config.get("adv_estimator", {}) or {}, + ) + + +def _print_gdpo_reward_feature_summary( + results: list[dict[str, Any]], + agent_names: list[str], + original_rewards: list[float], + num_gens: int, + adv_cfg: dict[str, Any], +) -> None: + print(f"\n{'=' * 70}", flush=True) + print("[Rollout] GDPO reward features", flush=True) + n = len(results) + for g in range(0, n, num_gens): + group_size = min(num_gens, n - g) + agent_name = agent_names[g] + print( + f"\n group {g // num_gens} agent={agent_name}", + flush=True, + ) + feature_names = _resolve_gdpo_feature_log_names_for_agent(adv_cfg, agent_name) + for k in range(group_size): + idx = g + k + fr = results[idx]["full_result"] + features = fr.get("gdpo_reward_features", {}) + parts = [ + f" [{k}] original_reward={float(original_rewards[idx]):.4f}", + f"final_reward={float(fr['reward']):.4f}", + ] + for name in feature_names: + entry = features.get(name) + if not isinstance(entry, dict): + continue + reward = float(entry.get("reward", 0.0)) + adjustment = entry.get("adjustment") + multiplier = entry.get("multiplier") + suffix = ( + f":{float(multiplier):.4f}" + if multiplier is not None + else "" + ) + parts.append( + f"{name}={reward:+.4f}" + f"({adjustment}{suffix})" + ) + print(" ".join(parts), flush=True) + print(f"{'=' * 70}\n", flush=True) + + +def _apply_profile_band_multipliers( + results: list[dict[str, Any]], + original_rewards: list[float], + base_rewards: list[float], + total_lengths: list[int], + reasoning_lengths: list[int], + answer_lengths: list[int], + agent_names: list[str], + agents_cfg: dict[str, Any] | None, + defaults: dict[str, Any], + num_gens: int, + should_mutate_reward: bool, +) -> None: + """Apply per-channel profile_band multipliers to correct rollouts. + + Each enabled channel contributes a multiplier in [0.0, 1.0] derived from the + per-row {a, b, f} block. Records multiplicative deltas as additive-equivalent + GDPO reward features, and mutates scalar rewards for length-bonus configs. + + Skips any group where the low-effort bypass already replaced the reward + (parity with Phase 1 of ``apply_group_length_adjustments``). + """ + n = len(results) + for g in range(0, n, num_gens): + agent_name = agent_names[g] + group_size = min(num_gens, n - g) + if any(results[g + k].get("low_effort_applied") for k in range(group_size)): + continue + params = _resolve_agent_params(agent_name, agents_cfg, defaults) + if params is None: + continue + use_total = bool(params.get("profile_band_total", False)) + use_rsn = bool(params.get("profile_band_reasoning", False)) + use_ans = bool(params.get("profile_band_answer", False)) + if not (use_total or use_rsn or use_ans): + continue + band = results[g].get("profile_band") + if not band: + continue + ch_total = band.get("total") if use_total else None + ch_rsn = band.get("reasoning") if use_rsn else None + ch_ans = band.get("answer") if use_ans else None + for k in range(group_size): + idx = g + k + # Gate on the env reward (correct rollouts only). + if original_rewards[idx] <= 0: + continue + current_reward = base_rewards[idx] + + total_m = _band_multiplier(total_lengths[idx], ch_total) + total_delta = current_reward * total_m - current_reward + current_reward += total_delta + + reasoning_m = _band_multiplier(reasoning_lengths[idx], ch_rsn) + reasoning_delta = current_reward * reasoning_m - current_reward + current_reward += reasoning_delta + + answer_m = _band_multiplier(answer_lengths[idx], ch_ans) + answer_delta = current_reward * answer_m - current_reward + current_reward += answer_delta + + profile_band_delta = current_reward - base_rewards[idx] + _set_gdpo_reward_feature( + results[idx], + "profile_band_total", + total_delta, + "multiplicative", + multiplier=total_m, + ) + _set_gdpo_reward_feature( + results[idx], + "profile_band_reasoning", + reasoning_delta, + "multiplicative", + multiplier=reasoning_m, + ) + _set_gdpo_reward_feature( + results[idx], + "profile_band_answer", + answer_delta, + "multiplicative", + multiplier=answer_m, + ) + _set_gdpo_reward_feature( + results[idx], "profile_band_delta", profile_band_delta, "derived_sum" + ) + _set_gdpo_reward_feature( + results[idx], + "length_total_delta", + current_reward - original_rewards[idx], + "derived_sum", + ) + _set_gdpo_reward_feature( + results[idx], "length_adjusted_reward", current_reward, "combined" + ) + if should_mutate_reward: + results[idx]["full_result"]["reward"] = current_reward + + +def _band_multiplier(rl: int, ch: dict[str, Any] | None) -> float: + """Per-channel profile_band reward multiplier. + + Returns 1.0 if the channel block is missing or malformed (no-op). + Otherwise: + rl <= a -> 1.0 + rl == b -> f + rl > a -> same linear slope continues past b, floored at 0.0 + """ + if not ch: + return 1.0 + a = ch.get("a") + b = ch.get("b") + f = ch.get("f") + if a is None or b is None or f is None or b <= a: + return 1.0 + if rl <= a: + return 1.0 + return max(0.0, 1.0 - (rl - a) / (b - a) * (1.0 - float(f))) + + +def _group_length_profile_gate_info( + *, + band: dict[str, Any] | None, + params: dict[str, Any], + rewards: list[float], + reasoning_lengths: list[int], + answer_lengths: list[int], + total_lengths: list[int], +) -> dict[str, Any]: + """Prompt-level gate for group-relative length penalties. + + When enabled, group-relative coefficients are applied only if the mean + rollout length exceeds a prompt-specific threshold from ``profile_band``. + """ + enabled = bool(params.get("group_length_penalty_profile_gate", False)) + channel = str(params.get("group_length_penalty_profile_gate_channel", "total")) + field = str(params.get("group_length_penalty_profile_gate_field", "a")) + positive_only = bool(params.get("group_length_penalty_profile_gate_positive_only", True)) + info = { + "enabled": enabled, + "open": True, + "channel": channel, + "field": field, + "positive_only": positive_only, + "mean": None, + "limit": None, + "reason": "disabled", + } + if not enabled: + return info + + limit = _profile_band_numeric_value(band, channel, field) + info["limit"] = limit + if limit is None: + info["open"] = False + info["reason"] = "missing_profile_limit" + return info + + length_by_channel = { + "reasoning": reasoning_lengths, + "answer": answer_lengths, + "total": total_lengths, + } + candidate_lengths = length_by_channel.get(channel) + if candidate_lengths is None: + info["open"] = False + info["reason"] = "unknown_channel" + return info + + if positive_only: + lengths = [l for r, l in zip(rewards, candidate_lengths) if r > 0] + else: + lengths = list(candidate_lengths) + if not lengths: + info["open"] = False + info["reason"] = "no_lengths" + return info + + mean_length = float(statistics.mean(lengths)) + info["mean"] = mean_length + info["open"] = mean_length > limit + info["reason"] = "mean_gt_limit" if info["open"] else "mean_le_limit" + return info + + +def _profile_band_numeric_value( + band: dict[str, Any] | None, channel: str, field: str +) -> float | None: + if not isinstance(band, dict): + return None + channel_block = band.get(channel) + if not isinstance(channel_block, dict): + return None + value = channel_block.get(field) + if isinstance(value, bool): + return None + if isinstance(value, (int, float)): + return float(value) + return None + + +def _resolve_agent_params( + agent_name: str, + agents_cfg: dict[str, Any] | None, + defaults: dict[str, Any], +) -> dict[str, Any] | None: + """Resolve length bonus parameters for a given agent.""" + if agents_cfg is None: + return dict(defaults) + + if agent_name not in agents_cfg: + print( + f"[length_adjustments] WARNING: agent '{agent_name}' not found in " + f"agent_overrides, falling back to defaults", + flush=True, + ) + return dict(defaults) + + overrides = agents_cfg[agent_name] + if overrides is None: + return dict(defaults) + + merged = dict(defaults) + for key in _PARAM_KEYS: + if key in overrides: + if key in _STR_PARAM_KEYS: + merged[key] = overrides[key] + elif key in _BOOL_PARAM_KEYS: + merged[key] = bool(overrides[key]) + else: + merged[key] = float(overrides[key]) + return merged + + +def _zmad_local_outliers( + lengths: list[int], z_thresh: float, min_mad_rel: float +) -> set[int]: + """Indices into ``lengths`` with Iglewicz–Hoaglin modified Z (MAD) > ``z_thresh``.""" + if len(lengths) < 2: + return set() + med = statistics.median(lengths) + devs = [abs(x - med) for x in lengths] + mad = statistics.median(devs) + if mad == 0: + return set() + if min_mad_rel > 0 and mad / max(med, 1e-9) < min_mad_rel: + return set() + out: set[int] = set() + for k, x in enumerate(lengths): + mz = 0.6745 * (x - med) / mad + if mz > z_thresh: + out.add(k) + return out + + +def _apply_length_bonuses_and_penalties( + rewards: list[float], + reasoning_lengths: list[int], + answer_lengths: list[int], + total_lengths: list[int], + reasoning_bonus: float, + answer_bonus: float, + total_bonus: float, + longest_reasoning_penalty: float, + longest_answer_penalty: float, + longest_total_penalty: float, + top_percentile: float, + group_reasoning_length_penalty_coeff: float, + group_answer_length_penalty_coeff: float, + group_total_length_penalty_coeff: float, + reasoning_zmad_threshold: float = 0.0, + reasoning_zmad_penalty: float = 0.0, + answer_zmad_threshold: float = 0.0, + answer_zmad_penalty: float = 0.0, + total_zmad_threshold: float = 0.0, + total_zmad_penalty: float = 0.0, + **_kwargs, +) -> tuple[ + list[float], + list[float], + list[float], + list[float], + list[float], + list[float], + list[float], + list[float], + list[float], + list[float], + list[float], + list[float], + list[float], + list[float], +]: + """Apply length-based bonuses/penalties to a single prompt group. + + Only samples with reward > 0 participate. Samples with reward <= 0 + are left untouched and excluded from weight computation. + + 1. Reasoning bonus: shortest non-empty reasoning among positive scorers; awarded only if that + sample satisfies ``reward >= top_threshold``. + 2. Answer bonus: same pattern for shortest non-empty answer. + 3. Total bonus: same pattern for shortest combined (reasoning + answer) length. + 4. Longest penalties: subtract from longest non-empty reasoning / answer / total among + top-percentile scorers; needs at least two eligible rollouts to compare. + 5. Independent zero-centered penalties for reasoning, answer, and total lengths. + 6. Optional MAD modified-Z outliers among positives for reasoning, answer, and total lengths. + Each channel runs only if its threshold is > 0; otherwise that channel is disabled and + its penalty is ignored. + """ + n = len(rewards) + zeros = [0.0] * n + if n < 2: + return ( + list(rewards), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + ) + + positive_indices = [i for i in range(n) if rewards[i] > 0] + if len(positive_indices) < 2: + return ( + list(rewards), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + ) + + adjusted = list(rewards) + adjustments = [0.0] * n + reasoning_adjs = [0.0] * n + answer_adjs = [0.0] * n + total_adjs = [0.0] * n + r_bonus_per = [0.0] * n + a_bonus_per = [0.0] * n + t_bonus_per = [0.0] * n + r_longest_pen_per = [0.0] * n + a_longest_pen_per = [0.0] * n + t_longest_pen_per = [0.0] * n + zmad_reasoning_adj = [0.0] * n + zmad_answer_adj = [0.0] * n + zmad_total_adj = [0.0] * n + + pos_reasoning = [reasoning_lengths[i] for i in positive_indices] + pos_answer = [answer_lengths[i] for i in positive_indices] + pos_total = [total_lengths[i] for i in positive_indices] + pos_rewards = [rewards[i] for i in positive_indices] + + sorted_scores = sorted(pos_rewards, reverse=True) + threshold_idx = max(0, int(len(pos_rewards) * top_percentile) - 1) + top_threshold = sorted_scores[threshold_idx] + top_scorer_indices = [i for i in positive_indices if rewards[i] >= top_threshold] + + # Reasoning bonus: shortest non-empty reasoning among top scorers + if reasoning_bonus > 0: + valid = [(pi, pos_reasoning[k]) for k, pi in enumerate(positive_indices) if pos_reasoning[k] > 0] + if valid: + shortest_pi, _ = min(valid, key=lambda x: x[1]) + if adjusted[shortest_pi] >= top_threshold: + adjusted[shortest_pi] += reasoning_bonus + adjustments[shortest_pi] += reasoning_bonus + r_bonus_per[shortest_pi] = reasoning_bonus + + # Answer bonus: shortest non-empty answer among top scorers + if answer_bonus > 0: + valid = [(pi, pos_answer[k]) for k, pi in enumerate(positive_indices) if pos_answer[k] > 0] + if valid: + shortest_pi, _ = min(valid, key=lambda x: x[1]) + if adjusted[shortest_pi] >= top_threshold: + adjusted[shortest_pi] += answer_bonus + adjustments[shortest_pi] += answer_bonus + a_bonus_per[shortest_pi] = answer_bonus + + # Total bonus: shortest combined (reasoning + answer) among top scorers + if total_bonus > 0: + valid = [(pi, pos_total[k]) for k, pi in enumerate(positive_indices) if pos_total[k] > 0] + if valid: + shortest_pi, _ = min(valid, key=lambda x: x[1]) + if adjusted[shortest_pi] >= top_threshold: + adjusted[shortest_pi] += total_bonus + adjustments[shortest_pi] += total_bonus + t_bonus_per[shortest_pi] = total_bonus + + # Longest reasoning penalty: longest among top-percentile scorers only + if longest_reasoning_penalty > 0: + valid = [(pi, reasoning_lengths[pi]) for pi in top_scorer_indices if reasoning_lengths[pi] > 0] + if len(valid) >= 2: + longest_pi, _ = max(valid, key=lambda x: x[1]) + pen = -longest_reasoning_penalty + adjusted[longest_pi] += pen + adjustments[longest_pi] += pen + r_longest_pen_per[longest_pi] = pen + + # Longest answer penalty: longest among top-percentile scorers only + if longest_answer_penalty > 0: + valid = [(pi, answer_lengths[pi]) for pi in top_scorer_indices if answer_lengths[pi] > 0] + if len(valid) >= 2: + longest_pi, _ = max(valid, key=lambda x: x[1]) + pen = -longest_answer_penalty + adjusted[longest_pi] += pen + adjustments[longest_pi] += pen + a_longest_pen_per[longest_pi] = pen + + # Longest total penalty: longest combined length among top-percentile scorers only + if longest_total_penalty > 0: + valid = [(pi, total_lengths[pi]) for pi in top_scorer_indices if total_lengths[pi] > 0] + if len(valid) >= 2: + longest_pi, _ = max(valid, key=lambda x: x[1]) + pen = -longest_total_penalty + adjusted[longest_pi] += pen + adjustments[longest_pi] += pen + t_longest_pen_per[longest_pi] = pen + + # Independent reasoning, answer, and total length penalties (zero-centered) + if group_reasoning_length_penalty_coeff > 0 or group_answer_length_penalty_coeff > 0 or group_total_length_penalty_coeff > 0: + reasoning_weights = _compute_length_weights(pos_reasoning) + answer_weights = _compute_length_weights(pos_answer) + total_weights = _compute_length_weights(pos_total) + + for k, i in enumerate(positive_indices): + r_adj = reasoning_weights[k] * group_reasoning_length_penalty_coeff + a_adj = answer_weights[k] * group_answer_length_penalty_coeff + t_adj = total_weights[k] * group_total_length_penalty_coeff + combined_adj = r_adj + a_adj + t_adj + reasoning_adjs[i] = r_adj + answer_adjs[i] = a_adj + total_adjs[i] = t_adj + if combined_adj != 0: + adjusted[i] += combined_adj + adjustments[i] += combined_adj + + zm = _ZMAD_MIN_MAD_REL + ztr = float(reasoning_zmad_threshold) + zpr = float(reasoning_zmad_penalty) + zta = float(answer_zmad_threshold) + zpa = float(answer_zmad_penalty) + ztt = float(total_zmad_threshold) + zpt = float(total_zmad_penalty) + + if len(positive_indices) >= 2: + if ztr > 0.0: + if zpr != 0.0: + for local_k in _zmad_local_outliers(pos_reasoning, ztr, zm): + gi = positive_indices[local_k] + adjusted[gi] -= zpr + adjustments[gi] -= zpr + zmad_reasoning_adj[gi] -= zpr + if zta > 0.0: + if zpa != 0.0: + for local_k in _zmad_local_outliers(pos_answer, zta, zm): + gi = positive_indices[local_k] + adjusted[gi] -= zpa + adjustments[gi] -= zpa + zmad_answer_adj[gi] -= zpa + if ztt > 0.0: + if zpt != 0.0: + for local_k in _zmad_local_outliers(pos_total, ztt, zm): + gi = positive_indices[local_k] + adjusted[gi] -= zpt + adjustments[gi] -= zpt + zmad_total_adj[gi] -= zpt + + return ( + adjusted, + adjustments, + reasoning_adjs, + answer_adjs, + total_adjs, + r_bonus_per, + a_bonus_per, + t_bonus_per, + r_longest_pen_per, + a_longest_pen_per, + t_longest_pen_per, + zmad_reasoning_adj, + zmad_answer_adj, + zmad_total_adj, + ) + + +def _compute_length_weights(lengths: list[int]) -> list[float]: + """Compute zero-centered weights where shorter = higher weight. + + Returns all zeros if all lengths are equal. + """ + max_len = max(lengths) + min_len = min(lengths) + + if max_len == min_len: + return [0.0] * len(lengths) + + span = max_len - min_len + raw_weights = [1.0 - ((length - min_len) / span) for length in lengths] + mean_weight = sum(raw_weights) / len(raw_weights) + return [w - mean_weight for w in raw_weights] diff --git a/scripts/build_profile_band_dataset.py b/scripts/build_profile_band_dataset.py new file mode 100644 index 00000000000..89af2ddb5b8 --- /dev/null +++ b/scripts/build_profile_band_dataset.py @@ -0,0 +1,204 @@ +#!/usr/bin/env python3 +"""Attach a per-prompt ``profile_band`` block to a profiled training JSONL. + +Reads a JSONL produced by ``profile_run`` (rows must carry +``profiled_rewards``, ``profiled_output_lengths``, +``profiled_reasoning_lengths``, ``profiled_answer_lengths``, and +``pass_rate``) and writes a new JSONL where each row has an additional +``profile_band`` field consumed by Mechanism 6 in +``nemo_rl/utils/length_adjustments.py``:: + + profile_band: + total: {a, b, f} # only present if data is non-degenerate + reasoning: {a, b, f} + answer: {a, b, f} + +For each channel: + - reference set = passing profiled rollouts (reward > 0). If too few + (< ``min_passing``), fall back to all profiled rollouts. + - a = mean(reference) + - b = mean(reference) + n_std * std(reference) + - f = looked up in ``f_table`` by row's ``pass_rate``. + - channel block is OMITTED if std == 0 (degenerate / cap-clamped), + if reference set has fewer than 2 samples, or if pass_rate is not + in the f_table. + +Config yaml shape:: + + n_std: 2.0 + min_passing: 2 + channels: [total, reasoning, answer] + f_table: # pass_rate -> f; omitted pass_rates skip the row entirely + - {pass_rate: 1.000, f: 0.6} + - {pass_rate: 0.875, f: 0.7} + - {pass_rate: 0.750, f: 0.8} + - {pass_rate: 0.625, f: 0.9} + +Usage:: + + python scripts/build_profile_band_dataset.py \ + --input /path/to/dapo17k_profiled_boxed_nanov3.jsonl \ + --config /path/to/profile_band.yaml \ + --output /path/to/dapo17k_profiled_band_boxed_nanov3.jsonl +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import sys +from collections import Counter +from pathlib import Path +from typing import Any + +# yaml is in stdlib via pyyaml on the cluster; fall back to a tiny parser if absent. +try: + import yaml # type: ignore +except ImportError: + yaml = None # type: ignore + +CHANNEL_TO_LENGTHS_KEY = { + "total": "profiled_output_lengths", + "reasoning": "profiled_reasoning_lengths", + "answer": "profiled_answer_lengths", +} + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--input", required=True, help="Path to profiled JSONL (input).") + p.add_argument("--config", required=True, help="Path to profile_band yaml config.") + p.add_argument("--output", required=True, help="Path to write augmented JSONL.") + p.add_argument("--quiet", action="store_true", help="Suppress per-row diagnostics summary.") + return p.parse_args() + + +def load_config(path: str) -> dict[str, Any]: + if yaml is None: + sys.exit("PyYAML not available; install pyyaml to use this script.") + with open(path) as f: + cfg = yaml.safe_load(f) + if not isinstance(cfg, dict): + sys.exit(f"config at {path} did not parse as a dict") + cfg.setdefault("n_std", 2.0) + cfg.setdefault("min_passing", 2) + cfg.setdefault("channels", ["total", "reasoning", "answer"]) + cfg.setdefault("f_table", []) + # Normalize f_table to a {rounded_pass_rate: f} dict. + f_table: dict[float, float] = {} + for entry in cfg["f_table"]: + pr = round(float(entry["pass_rate"]), 4) + f_table[pr] = float(entry["f"]) + cfg["_f_table"] = f_table + bad_channels = [c for c in cfg["channels"] if c not in CHANNEL_TO_LENGTHS_KEY] + if bad_channels: + sys.exit(f"unknown channels in config: {bad_channels}; valid: {sorted(CHANNEL_TO_LENGTHS_KEY)}") + return cfg + + +def lookup_f(pass_rate: float, f_table: dict[float, float]) -> float | None: + return f_table.get(round(float(pass_rate), 4)) + + +def channel_block( + lengths: list[int], + rewards: list[float], + n_std: float, + min_passing: int, + f_value: float, +) -> dict[str, float] | None: + """Compute {a, b, f} for one channel; return None if degenerate.""" + passing = [l for l, r in zip(lengths, rewards) if r is not None and r > 0] + ref = passing if len(passing) >= min_passing else list(lengths) + if len(ref) < 2: + return None + mean_l = statistics.mean(ref) + std_l = statistics.stdev(ref) + if std_l <= 0: + return None + a = mean_l + b = mean_l + n_std * std_l + return {"a": float(a), "b": float(b), "f": float(f_value)} + + +def build_band(row: dict[str, Any], cfg: dict[str, Any]) -> dict[str, Any] | None: + """Construct the profile_band dict for this row, or None to skip.""" + f = lookup_f(row.get("pass_rate", -1.0), cfg["_f_table"]) + if f is None: + return None + rewards = row.get("profiled_rewards") or [] + band: dict[str, Any] = {} + for ch_name in cfg["channels"]: + lengths_key = CHANNEL_TO_LENGTHS_KEY[ch_name] + lengths = row.get(lengths_key) + if lengths is None: + continue + block = channel_block( + lengths=lengths, + rewards=rewards, + n_std=cfg["n_std"], + min_passing=cfg["min_passing"], + f_value=f, + ) + if block is not None: + band[ch_name] = block + return band or None + + +def main() -> None: + args = parse_args() + cfg = load_config(args.config) + + in_path = Path(args.input) + out_path = Path(args.output) + out_path.parent.mkdir(parents=True, exist_ok=True) + + n_total = 0 + n_with_band = 0 + n_pr_skipped = 0 + n_pr_skipped_by_passrate: Counter[float] = Counter() + n_channels_emitted: Counter[str] = Counter() + + with in_path.open() as fin, out_path.open("w") as fout: + for line in fin: + line = line.rstrip("\n") + if not line: + continue + row = json.loads(line) + n_total += 1 + band = build_band(row, cfg) + if band is None: + pr = round(float(row.get("pass_rate", -1.0)), 4) + if lookup_f(pr, cfg["_f_table"]) is None: + n_pr_skipped += 1 + n_pr_skipped_by_passrate[pr] += 1 + else: + row["profile_band"] = band + n_with_band += 1 + for ch in band: + n_channels_emitted[ch] += 1 + fout.write(json.dumps(row) + "\n") + + if not args.quiet: + print(f"input : {in_path}") + print(f"output : {out_path}") + print(f"rows in: {n_total}") + print(f"rows w/ profile_band: {n_with_band} ({100*n_with_band/n_total:.1f}%)") + print(f"rows skipped (pass_rate not in f_table): {n_pr_skipped}") + if n_pr_skipped_by_passrate: + print(" by pass_rate:") + for pr in sorted(n_pr_skipped_by_passrate): + print(f" {pr:.3f}: {n_pr_skipped_by_passrate[pr]}") + print("channels emitted (per row, summed):") + for ch in cfg["channels"]: + c = n_channels_emitted.get(ch, 0) + print(f" {ch:>10}: {c} ({100*c/n_total:.1f}% of rows)") + print(f"f_table (rounded pass_rate -> f):") + for pr in sorted(cfg["_f_table"]): + print(f" {pr:.3f} -> {cfg['_f_table'][pr]}") + print(f"n_std={cfg['n_std']} min_passing={cfg['min_passing']}") + + +if __name__ == "__main__": + main() diff --git a/scripts/profile_band.yaml b/scripts/profile_band.yaml new file mode 100644 index 00000000000..0538bff5926 --- /dev/null +++ b/scripts/profile_band.yaml @@ -0,0 +1,28 @@ +# Config for build_profile_band_dataset.py +# +# a, b are derived from the per-prompt profiled length distribution: +# reference set = passing rollouts (reward > 0); fall back to all +# profiled rollouts if fewer than ``min_passing`` pass. +# a = mean(reference) +# b = mean(reference) + n_std * std(reference) +# +# f is selected per-row by exact pass_rate match in f_table. +# Pass-rates not in the table → row gets no profile_band block (mechanism +# is a no-op for that prompt). +# +# Channels with std==0 (cap-clamped) or fewer than 2 samples are +# omitted from a row's profile_band block but other channels still apply. + +n_std: 2.0 +min_passing: 2 + +channels: + - total + - reasoning + - answer + +f_table: + - {pass_rate: 1.000, f: 0.6} + - {pass_rate: 0.875, f: 0.7} + - {pass_rate: 0.750, f: 0.8} + - {pass_rate: 0.625, f: 0.9} diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 44debcd1273..b26f8502681 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -4719,6 +4719,55 @@ def test_grpo_advantage_estimator_small_nonzero_std(): assert result[0, 0] * result[1, 0] < 0 +def test_gdpo_advantage_estimator_feature_weight_changes_advantage(): + """Feature weights are applied after per-feature normalization.""" + + def compute_advantage(think_weight): + estimator_config = AdvEstimatorConfig( + use_leave_one_out_baseline=False, + normalize_rewards=True, + reward_features={ + "default": { + "length_adjusted_reward": {}, + "think_count_delta": {"weight": think_weight}, + } + }, + ) + estimator = GDPOAdvantageEstimator(estimator_config, ClippedPGLossConfig()) + prompt_ids = torch.tensor([[0], [0], [0], [0]]) + rewards = torch.zeros(4) + mask = torch.ones(4, 1) + repeated_batch = { + "agent_ref": [{"name": "math_with_judge_simple_agent"} for _ in range(4)], + "gdpo_reward_features": [ + { + "length_adjusted_reward": {"reward": 1.0}, + "think_count_delta": {"reward": 0.0}, + }, + { + "length_adjusted_reward": {"reward": 1.0}, + "think_count_delta": {"reward": -1.0}, + }, + { + "length_adjusted_reward": {"reward": 0.0}, + "think_count_delta": {"reward": 0.0}, + }, + { + "length_adjusted_reward": {"reward": 0.0}, + "think_count_delta": {"reward": -1.0}, + }, + ], + } + return estimator.compute_advantage( + prompt_ids, rewards, mask, repeated_batch=repeated_batch + )[:, 0] + + without_think_feature = compute_advantage(0.0) + with_think_feature = compute_advantage(1.0) + + assert not torch.allclose(without_think_feature, with_think_feature) + + # ============================================================================ # Tests for ReinforcePlusPlusAdvantageEstimator class # ============================================================================ diff --git a/tests/unit/experience/test_reward_penalties.py b/tests/unit/experience/test_reward_penalties.py index ea5f4921519..753cd024fda 100644 --- a/tests/unit/experience/test_reward_penalties.py +++ b/tests/unit/experience/test_reward_penalties.py @@ -19,8 +19,10 @@ from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.rollouts import ( + _calculate_gdpo_reward_feature_metrics, _extract_mask_sample_flags, _postprocess_single_nemo_gym_group, + _record_gdpo_think_count_features, apply_reward_penalties, resolve_reward_penalty_config, should_mask_flagged_samples, @@ -827,6 +829,91 @@ def test_no_generation_str_skipped(self): assert result["full_result"]["reward"] == 1.0 +# ===================================================================== +# GDPO feature: think_count_delta +# ===================================================================== + + +class TestGdpoThinkCountFeature: + CFG = {"token_ids": {"think_close": 13}} + + def test_one_close_think_has_zero_delta(self): + result = _make_result( + message_log=[ + _msg("user", [100, 12]), + _msg("assistant", [300, 13, 400]), + ] + ) + _record_gdpo_think_count_features([result], self.CFG) + entry = result["full_result"]["gdpo_reward_features"]["think_count_delta"] + assert entry["reward"] == 0.0 + assert entry["num_close_think_tags"] == 1 + + def test_missing_close_think_has_negative_delta(self): + result = _make_result( + message_log=[ + _msg("user", [100, 12]), + _msg("assistant", [300, 400]), + ] + ) + _record_gdpo_think_count_features([result], self.CFG) + entry = result["full_result"]["gdpo_reward_features"]["think_count_delta"] + assert entry["reward"] == -1.0 + assert entry["num_close_think_tags"] == 0 + + def test_double_close_think_has_negative_delta(self): + result = _make_result( + message_log=[ + _msg("user", [100, 12]), + _msg("assistant", [300, 13, 400, 13]), + ] + ) + _record_gdpo_think_count_features([result], self.CFG) + entry = result["full_result"]["gdpo_reward_features"]["think_count_delta"] + assert entry["reward"] == -1.0 + assert entry["num_close_think_tags"] == 2 + + def test_generation_str_fallback_counts_close_thinks(self): + result = _make_result( + output_items=[ + _message_item("answer", generation_str="a b ") + ], + message_log=[ + _msg("user", [100, 12]), + _msg("assistant", [300, 400]), + ], + ) + _record_gdpo_think_count_features([result], self.CFG) + entry = result["full_result"]["gdpo_reward_features"]["think_count_delta"] + assert entry["reward"] == -1.0 + assert entry["num_close_think_tags"] == 2 + + def test_generic_gdpo_feature_metrics(self): + r1 = _make_result() + r1["full_result"]["gdpo_reward_features"] = { + "length_adjusted_reward": {"reward": 1.0}, + "think_count_delta": {"reward": 0.0}, + } + r2 = _make_result() + r2["full_result"]["gdpo_reward_features"] = { + "length_adjusted_reward": {"reward": -0.5}, + "think_count_delta": {"reward": -2.0}, + } + r3 = _make_result() + r3["full_result"]["gdpo_reward_features"] = { + "length_adjusted_reward": {"reward": 0.5}, + } + + metrics = _calculate_gdpo_reward_feature_metrics([r1, r2, r3]) + + assert metrics["gdpo_length_adjusted_reward/min"] == -0.5 + assert metrics["gdpo_length_adjusted_reward/max"] == 1.0 + assert metrics["gdpo_length_adjusted_reward/mean"] == (1.0 - 0.5 + 0.5) / 3 + assert metrics["gdpo_think_count_delta/min"] == -2.0 + assert metrics["gdpo_think_count_delta/max"] == 0.0 + assert metrics["gdpo_think_count_delta/mean"] == -1.0 + + # ===================================================================== # Cross-cutting: multiple penalties, config gating, batch behavior # ===================================================================== @@ -902,6 +989,7 @@ def test_batch_of_results_mixed(self): TestPenalizeEmptyFinalAnswer, TestPenalizeUnwantedTokens, TestPenalizeMultiEndThink, + TestGdpoThinkCountFeature, TestCrossCutting, ] From 2407f738d69addc3df43b8ff1c549adcd97cdc06 Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 12 Aug 2026 13:09:04 -0700 Subject: [PATCH 02/27] fix: require explicit think_close token id for think-count feature Token ids are tokenizer-specific, so defaulting the think_close token id to 13 could count unrelated tokens (e.g. newlines) on other tokenizers and silently inject large spurious think_count_delta penalties. Token counting now only runs when the id is configured; when no counting method is available the feature stays neutral instead of guessing. Also aligns kwargs at the rollout call site and documents the Phase-3 profiled-penalty gate in length_adjustments.py. Co-Authored-By: Claude Fable 5 --- nemo_rl/experience/rollouts.py | 34 ++++++++++++++----- nemo_rl/utils/length_adjustments.py | 4 +++ .../unit/experience/test_reward_penalties.py | 32 +++++++++++++++++ 3 files changed, 61 insertions(+), 9 deletions(-) diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 7e678f30d28..3b9e4f45cf7 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -2253,14 +2253,20 @@ def _count_close_think_tags( reward_penalty_config: dict[str, Any] | BaseModel | None, tokenizer: TokenizerType | None = None, ) -> dict[str, Any]: + # Token-based counting requires an explicit think_close token id from the + # reward-penalty config. Token ids are tokenizer-specific, so guessing a + # default here would count unrelated tokens (e.g. newlines) on other + # tokenizers and silently inject large spurious penalties. think_close_token_id = _get_reward_penalty_token_id( reward_penalty_config or {}, "think_close" ) - if think_close_token_id is None: - think_close_token_id = 13 assistant_ids = _assistant_generated_token_ids(result) - token_count = sum(1 for token_id in assistant_ids if token_id == think_close_token_id) + token_count = ( + sum(1 for token_id in assistant_ids if token_id == think_close_token_id) + if think_close_token_id is not None + else None + ) decoded_count = None if tokenizer is not None and assistant_ids: @@ -2277,16 +2283,26 @@ def _count_close_think_tags( close_count = decoded_count source = "decoded_assistant_tokens" elif generation_text: - close_count = max(token_count, generation_str_count) + close_count = ( + max(token_count, generation_str_count) + if token_count is not None + else generation_str_count + ) source = "max_token_or_generation_str" - else: + elif token_count is not None: close_count = token_count source = "assistant_token_ids" + else: + # No way to count: no tokenizer decode, no generation text, and no + # configured think_close token id. Report one close tag so the + # think_count_delta feature stays neutral rather than penalizing. + close_count = 1 + source = "unavailable" return { "count": int(close_count), "source": source, - "token_count": int(token_count), + "token_count": int(token_count) if token_count is not None else None, "generation_str_count": int(generation_str_count), "decoded_count": decoded_count, } @@ -2603,9 +2619,9 @@ async def run_async_nemo_gym_rollout( tokenizer=tokenizer, log_full_result_tables=log_full_result_tables, effort_config=effort_config, - reward_penalty_config=reward_penalty_config, - length_adjustment_config=length_adjustment_config, - thinking_tags=thinking_tags, + reward_penalty_config=reward_penalty_config, + length_adjustment_config=length_adjustment_config, + thinking_tags=thinking_tags, mask_env_flagged_samples=mask_env_flagged_samples, ) if accumulator.is_complete: diff --git a/nemo_rl/utils/length_adjustments.py b/nemo_rl/utils/length_adjustments.py index 289f78d647a..ff68335e07f 100644 --- a/nemo_rl/utils/length_adjustments.py +++ b/nemo_rl/utils/length_adjustments.py @@ -544,6 +544,10 @@ def apply_group_length_adjustments( # Phase 3: apply additive adjustments and record GDPO reward features additive_base_rewards = [0.0] * n for i, r in enumerate(results): + # The profiled-length penalty stacks only on rollouts whose group + # adjustments are non-negative: a rollout already penalized by the + # group-relative channels should not be double-penalized for the same + # excess length. profiled_adj = all_profiled_length_adj[i] if all_adjustments[i] >= 0 else 0.0 additive_delta = all_adjustments[i] + profiled_adj additive_base_rewards[i] = original_rewards[i] + additive_delta diff --git a/tests/unit/experience/test_reward_penalties.py b/tests/unit/experience/test_reward_penalties.py index 753cd024fda..70de89e2def 100644 --- a/tests/unit/experience/test_reward_penalties.py +++ b/tests/unit/experience/test_reward_penalties.py @@ -888,6 +888,38 @@ def test_generation_str_fallback_counts_close_thinks(self): assert entry["reward"] == -1.0 assert entry["num_close_think_tags"] == 2 + def test_unconfigured_token_id_is_neutral_without_text(self): + # No think_close token id configured, no tokenizer, no generation + # text: counting is impossible, so the feature must stay neutral + # (count=1 -> delta=0) instead of guessing a token id. + result = _make_result( + message_log=[ + _msg("user", [100, 12]), + _msg("assistant", [300, 13, 400, 13]), + ] + ) + _record_gdpo_think_count_features([result], {}) + entry = result["full_result"]["gdpo_reward_features"]["think_count_delta"] + assert entry["reward"] == 0.0 + assert entry["count_source"] == "unavailable" + + def test_unconfigured_token_id_uses_generation_str(self): + # Without a configured token id the generation string count is used + # alone (token ids of other tokenizers must not leak into the count). + result = _make_result( + output_items=[ + _message_item("answer", generation_str="a b") + ], + message_log=[ + _msg("user", [100, 12]), + _msg("assistant", [300, 13, 400, 13]), + ], + ) + _record_gdpo_think_count_features([result], {}) + entry = result["full_result"]["gdpo_reward_features"]["think_count_delta"] + assert entry["reward"] == 0.0 + assert entry["num_close_think_tags"] == 1 + def test_generic_gdpo_feature_metrics(self): r1 = _make_result() r1["full_result"]["gdpo_reward_features"] = { From a55eb228d0f24d1a35a3b3757aaa5ee19d9f2e7b Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 12 Aug 2026 13:22:16 -0700 Subject: [PATCH 03/27] chore: add copyright header to build_profile_band_dataset.py Co-Authored-By: Claude Fable 5 --- scripts/build_profile_band_dataset.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/scripts/build_profile_band_dataset.py b/scripts/build_profile_band_dataset.py index 89af2ddb5b8..d6352bf38cf 100644 --- a/scripts/build_profile_band_dataset.py +++ b/scripts/build_profile_band_dataset.py @@ -1,4 +1,17 @@ #!/usr/bin/env python3 +# 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. """Attach a per-prompt ``profile_band`` block to a profiled training JSONL. Reads a JSONL produced by ``profile_run`` (rows must carry From b8ca5c331c968b226868fff3e29a929e431cfe2f Mon Sep 17 00:00:00 2001 From: adithyare Date: Mon, 24 Aug 2026 12:32:26 -0700 Subject: [PATCH 04/27] fix: propagate length_adjustment_config to async NeMo-Gym rollouts AsyncTrajectoryCollector called run_async_nemo_gym_rollout without length_adjustment_config, so async_grpo_train silently skipped all configured length bonuses/penalties (both grpo.length_bonus mode and GDPO reward_features mode). Pass master_config.grpo like the sync grpo_train and validate paths already do. Co-Authored-By: Claude Fable 5 --- nemo_rl/algorithms/async_utils/trajectory_collector.py | 1 + 1 file changed, 1 insertion(+) diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index 2136a593557..4bf824c1777 100644 --- a/nemo_rl/algorithms/async_utils/trajectory_collector.py +++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py @@ -1291,6 +1291,7 @@ async def _iter_rollout_groups( max_rollout_turns=None, greedy=False, reward_penalty_config=self.master_config.reward_penalties, + length_adjustment_config=self.master_config.grpo.model_dump(), thinking_tags=get_nemo_gym_thinking_tags(self.master_config.env), mask_env_flagged_samples=should_mask_flagged_samples( self.master_config.env From 25df5bfcc553ed00d82f602845ff1d3664c660b5 Mon Sep 17 00:00:00 2001 From: adithyare Date: Mon, 24 Aug 2026 13:18:57 -0700 Subject: [PATCH 05/27] feat: profile_band multiplier clamps at f past b instead of decaying to 0 Previously the linear slope from (a, 1.0) through (b, f) continued past b until the multiplier hit 0, fully wiping rewards for long-but-correct rollouts. Now the multiplier floors at f for all lengths >= b, so the maximum penalty is bounded by the configured f. Co-Authored-By: Claude Fable 5 --- nemo_rl/utils/length_adjustments.md | 8 +++++--- nemo_rl/utils/length_adjustments.py | 8 +++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/nemo_rl/utils/length_adjustments.md b/nemo_rl/utils/length_adjustments.md index 35dd936a8dc..a7c322c9f0c 100644 --- a/nemo_rl/utils/length_adjustments.md +++ b/nemo_rl/utils/length_adjustments.md @@ -430,11 +430,13 @@ This uses per-row `profile_band` metadata with channel-specific `{a, b, f}` valu For an enabled channel, the multiplier is: ```text -length <= a: multiplier = 1 -length > a: multiplier = max(0, 1 - (length - a) / (b - a) * (1 - f)) +length <= a: multiplier = 1 +a < length < b: multiplier = 1 - (length - a) / (b - a) * (1 - f) +length >= b: multiplier = f ``` -So the multiplier is `f` at `b`, then the same slope continues past `b` until clamped at `0`. +So the multiplier interpolates linearly from `1` at `a` down to `f` at `b`, then stays at `f` for +all lengths past `b`. Profile-band multipliers are applied only to rollouts whose original environment reward is positive. diff --git a/nemo_rl/utils/length_adjustments.py b/nemo_rl/utils/length_adjustments.py index ff68335e07f..0aac9b6e2cd 100644 --- a/nemo_rl/utils/length_adjustments.py +++ b/nemo_rl/utils/length_adjustments.py @@ -770,8 +770,8 @@ def _band_multiplier(rl: int, ch: dict[str, Any] | None) -> float: Returns 1.0 if the channel block is missing or malformed (no-op). Otherwise: rl <= a -> 1.0 - rl == b -> f - rl > a -> same linear slope continues past b, floored at 0.0 + a < rl < b -> linear interpolation from 1.0 down to f + rl >= b -> f """ if not ch: return 1.0 @@ -782,7 +782,9 @@ def _band_multiplier(rl: int, ch: dict[str, Any] | None) -> float: return 1.0 if rl <= a: return 1.0 - return max(0.0, 1.0 - (rl - a) / (b - a) * (1.0 - float(f))) + if rl >= b: + return float(f) + return 1.0 - (rl - a) / (b - a) * (1.0 - float(f)) def _group_length_profile_gate_info( From 4d73c5bc06d3d5af9b4d2413030ac86a6057a55c Mon Sep 17 00:00:00 2001 From: adithyare Date: Mon, 24 Aug 2026 13:57:21 -0700 Subject: [PATCH 06/27] feat: global profile_band defaults for datasets without per-prompt bands New config block length_bonus.profile_band with enabled + defaults holding per-channel {a, b, f}. Channels listed under defaults are implicitly enabled; per-row profile_band metadata still wins per-channel when present. The global band also feeds the profile-gated group-relative penalty. Malformed channel blocks (missing keys or b <= a) are ignored with a warning. Previously profile_band silently no-oped when the dataset lacked per-prompt metadata. Co-Authored-By: Claude Fable 5 --- nemo_rl/utils/length_adjustments.md | 48 ++++++++++++++++++++++ nemo_rl/utils/length_adjustments.py | 62 ++++++++++++++++++++++++++--- 2 files changed, 105 insertions(+), 5 deletions(-) diff --git a/nemo_rl/utils/length_adjustments.md b/nemo_rl/utils/length_adjustments.md index a7c322c9f0c..2b082ea2535 100644 --- a/nemo_rl/utils/length_adjustments.md +++ b/nemo_rl/utils/length_adjustments.md @@ -451,6 +451,54 @@ grpo: profile_band_total: true ``` +#### Global Defaults (dataset without per-prompt bands) + +When the dataset has no per-prompt `profile_band` metadata, global `{a, b, f}` values can be +set directly in the config under `length_bonus.profile_band`. Only the channels listed under +`defaults` are activated: + +```yaml +grpo: + length_bonus: + profile_band: + enabled: true + defaults: + total: {a: 10000, b: 20000, f: 0.5} +``` + +```yaml +grpo: + length_bonus: + profile_band: + enabled: true + defaults: + reasoning: {a: 9000, b: 14000, f: 0.9} +``` + +```yaml +grpo: + length_bonus: + profile_band: + enabled: true + defaults: + answer: {a: 500, b: 1000, f: 0.9} +``` + +The first config applies the multiplier on total length only, the second on reasoning length +only, and the last on answer length only. Multiple channels may be listed together. + +Semantics: + +- Channels under `defaults` are implicitly enabled — no need to also set + `profile_band_total/reasoning/answer: true` under `length_bonus.default`. Per-agent + `agent_overrides` can still disable a channel (e.g. `profile_band_total: false`). +- Per-prompt `profile_band` metadata, when present on a row, takes precedence over the global + defaults on a per-channel basis (a row that only provides `total` still falls back to the + global `reasoning`/`answer` blocks if those are configured). +- The global band also feeds profile-gated group-relative penalties + (`group_length_penalty_profile_gate`) when rows lack metadata. +- A malformed channel block (missing `a`/`b`/`f`, or `b <= a`) is ignored with a warning. + ### 7. Profile-Gated Group Relative-Length Scaling Config keys: diff --git a/nemo_rl/utils/length_adjustments.py b/nemo_rl/utils/length_adjustments.py index 0aac9b6e2cd..1f818ab6301 100644 --- a/nemo_rl/utils/length_adjustments.py +++ b/nemo_rl/utils/length_adjustments.py @@ -274,6 +274,7 @@ def apply_group_length_adjustments( default_cfg = length_cfg.get("default", {}) agents_cfg = length_cfg.get("agent_overrides") + global_band = _resolve_global_profile_band(length_cfg.get("profile_band")) gdpo_feature_only = bool(length_cfg.get("_gdpo_feature_only", False)) verbose = bool(length_cfg.get("verbose", False)) and not gdpo_feature_only gdpo_feature_verbose = bool( @@ -281,7 +282,7 @@ def apply_group_length_adjustments( or (length_cfg.get("verbose", False) and gdpo_feature_only) ) should_mutate_reward = not gdpo_feature_only - if not default_cfg.get("enabled", False) and not agents_cfg: + if not default_cfg.get("enabled", False) and not agents_cfg and not global_band: return num_gens = master_config["grpo"]["num_generations_per_prompt"] @@ -306,6 +307,10 @@ def apply_group_length_adjustments( else: defaults[k] = default_cfg.get(k, 0.0) defaults.setdefault("top_percentile", 0.2) + # Channels listed under length_bonus.profile_band.defaults are implicitly + # enabled; per-agent overrides can still disable them. + for _ch in global_band: + defaults[f"profile_band_{_ch}"] = True n = len(results) original_rewards = [r["full_result"]["reward"] for r in results] @@ -367,7 +372,7 @@ def apply_group_length_adjustments( total_lengths[g : g + group_size] = group_total[:group_size] group_rewards = original_rewards[g : g + num_gens] gate_info = _group_length_profile_gate_info( - band=results[g].get("profile_band"), + band=_merged_profile_band(results[g].get("profile_band"), global_band), params=params, rewards=group_rewards[:group_size], reasoning_lengths=group_reasoning[:group_size], @@ -607,6 +612,7 @@ def apply_group_length_adjustments( defaults=defaults, num_gens=num_gens, should_mutate_reward=should_mutate_reward, + global_band=global_band, ) if verbose or gdpo_feature_verbose: @@ -665,6 +671,48 @@ def _print_gdpo_reward_feature_summary( print(f"{'=' * 70}\n", flush=True) +def _resolve_global_profile_band(pb_cfg: Any) -> dict[str, dict[str, Any]]: + """Parse ``length_bonus.profile_band`` into per-channel {a, b, f} blocks. + + Returns only channels ("total", "reasoning", "answer") present under + ``defaults`` with a complete, well-formed block. Empty dict when the + section is absent or disabled. + """ + if not isinstance(pb_cfg, dict) or not pb_cfg.get("enabled", False): + return {} + pb_defaults = pb_cfg.get("defaults") + if not isinstance(pb_defaults, dict): + return {} + band: dict[str, dict[str, Any]] = {} + for ch in ("total", "reasoning", "answer"): + ch_cfg = pb_defaults.get(ch) + if not isinstance(ch_cfg, dict): + continue + a, b, f = ch_cfg.get("a"), ch_cfg.get("b"), ch_cfg.get("f") + if a is None or b is None or f is None or b <= a: + logger.warning( + f"length_bonus.profile_band.defaults.{ch} is malformed " + f"(a={a}, b={b}, f={f}); ignoring this channel" + ) + continue + band[ch] = {"a": a, "b": b, "f": f} + return band + + +def _merged_profile_band( + row_band: dict[str, Any] | None, + global_band: dict[str, dict[str, Any]], +) -> dict[str, Any] | None: + """Merge per-row profile_band over global defaults (row channel wins).""" + if not global_band: + return row_band + if not row_band: + return dict(global_band) + merged: dict[str, Any] = dict(global_band) + merged.update(row_band) + return merged + + def _apply_profile_band_multipliers( results: list[dict[str, Any]], original_rewards: list[float], @@ -677,17 +725,21 @@ def _apply_profile_band_multipliers( defaults: dict[str, Any], num_gens: int, should_mutate_reward: bool, + global_band: dict[str, dict[str, Any]] | None = None, ) -> None: """Apply per-channel profile_band multipliers to correct rollouts. Each enabled channel contributes a multiplier in [0.0, 1.0] derived from the - per-row {a, b, f} block. Records multiplicative deltas as additive-equivalent - GDPO reward features, and mutates scalar rewards for length-bonus configs. + per-row {a, b, f} block, falling back to ``length_bonus.profile_band.defaults`` + for channels the row does not provide. Records multiplicative deltas as + additive-equivalent GDPO reward features, and mutates scalar rewards for + length-bonus configs. Skips any group where the low-effort bypass already replaced the reward (parity with Phase 1 of ``apply_group_length_adjustments``). """ n = len(results) + global_band = global_band or {} for g in range(0, n, num_gens): agent_name = agent_names[g] group_size = min(num_gens, n - g) @@ -701,7 +753,7 @@ def _apply_profile_band_multipliers( use_ans = bool(params.get("profile_band_answer", False)) if not (use_total or use_rsn or use_ans): continue - band = results[g].get("profile_band") + band = _merged_profile_band(results[g].get("profile_band"), global_band) if not band: continue ch_total = band.get("total") if use_total else None From c0f605fe1d14e5645f8bc2e7c0fbca866e515bb4 Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 11:05:26 -0700 Subject: [PATCH 07/27] test: unit tests for profile_band and group relative-length scaling 22 tests covering the two production length-penalty algorithms: - _band_multiplier shape: 1.0 up to a, linear to f at b, clamped at f past b, no-op on missing/malformed channels - profile_band via per-row metadata: total/reasoning channels, positive- reward gating, missing-band and disabled-channel no-ops - profile_band via global config defaults: implicit channel enabling, works without a length_bonus.default block, per-row precedence, disabled block and malformed (b <= a) channel no-ops - group relative-length scaling: symmetric zero-centered adjustments, mean preservation, equal-length and zero-coefficient no-ops, positive-only participation, reasoning-channel isolation, agent_overrides disable Co-Authored-By: Claude Fable 5 --- tests/unit/utils/test_length_adjustments.py | 320 ++++++++++++++++++++ 1 file changed, 320 insertions(+) create mode 100644 tests/unit/utils/test_length_adjustments.py diff --git a/tests/unit/utils/test_length_adjustments.py b/tests/unit/utils/test_length_adjustments.py new file mode 100644 index 00000000000..6c9c286ba0a --- /dev/null +++ b/tests/unit/utils/test_length_adjustments.py @@ -0,0 +1,320 @@ +# 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 profile_band multiplier and group relative-length +scaling algorithms in nemo_rl/utils/length_adjustments.py.""" + +import pytest + +from nemo_rl.utils.length_adjustments import ( + _band_multiplier, + apply_group_length_adjustments, +) + +AGENT = "math_with_judge_simple_agent" + + +def make_result(reasoning: str, answer: str, reward: float, band=None): + """Build a minimal rollout result dict in the shape the module consumes.""" + result = { + "full_result": { + "reward": reward, + "response": { + "output": [ + {"type": "reasoning", "summary": [{"text": reasoning}]}, + {"type": "message", "content": [{"text": answer}]}, + ] + }, + }, + "agent_ref": {"name": AGENT}, + } + if band is not None: + result["profile_band"] = band + return result + + +def make_config(default=None, profile_band=None, num_gens=2): + length_bonus = {} + if default is not None: + length_bonus["default"] = {"length_type": "chars", **default} + if profile_band is not None: + length_bonus["profile_band"] = profile_band + return { + "grpo": { + "num_generations_per_prompt": num_gens, + "length_bonus": length_bonus, + } + } + + +def rewards_of(results): + return [r["full_result"]["reward"] for r in results] + + +class TestBandMultiplier: + """Direct tests of the {a, b, f} multiplier shape.""" + + CH = {"a": 10, "b": 20, "f": 0.5} + + def test_at_or_below_a_is_one(self): + assert _band_multiplier(5, self.CH) == 1.0 + assert _band_multiplier(10, self.CH) == 1.0 + + def test_linear_interpolation_between_a_and_b(self): + assert _band_multiplier(15, self.CH) == pytest.approx(0.75) + + def test_exactly_b_is_f(self): + assert _band_multiplier(20, self.CH) == pytest.approx(0.5) + + def test_clamps_at_f_past_b(self): + # Past b the multiplier stays at f; it must NOT keep decaying to 0. + assert _band_multiplier(25, self.CH) == pytest.approx(0.5) + assert _band_multiplier(30, self.CH) == pytest.approx(0.5) + assert _band_multiplier(10_000, self.CH) == pytest.approx(0.5) + + def test_missing_or_malformed_channel_is_noop(self): + assert _band_multiplier(100, None) == 1.0 + assert _band_multiplier(100, {}) == 1.0 + assert _band_multiplier(100, {"a": 10, "b": 20}) == 1.0 # missing f + assert _band_multiplier(100, {"a": 20, "b": 10, "f": 0.5}) == 1.0 # b <= a + + +class TestProfileBandPerRow: + """profile_band multipliers driven by per-row dataset metadata.""" + + def test_total_channel_scales_correct_rollouts(self): + band = {"total": {"a": 10, "b": 20, "f": 0.5}} + results = [ + make_result("12345", "12345", 1.0, band=band), # total 10 -> x1.0 + make_result("1234567890", "1234567890", 1.0, band=band), # 20 -> x0.5 + ] + cfg = make_config(default={"enabled": True, "profile_band_total": True}) + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 0.5]) + + def test_zero_reward_rollouts_untouched(self): + band = {"total": {"a": 10, "b": 20, "f": 0.5}} + results = [ + make_result("1234567890", "1234567890", 0.0, band=band), + make_result("1234567890", "1234567890", 1.0, band=band), + ] + cfg = make_config(default={"enabled": True, "profile_band_total": True}) + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([0.0, 0.5]) + + def test_reasoning_channel_ignores_answer_length(self): + band = {"reasoning": {"a": 10, "b": 20, "f": 0.5}} + long_answer = "x" * 100 # must not affect the reasoning channel + results = [ + make_result("12345", long_answer, 1.0, band=band), # reasoning 5 -> x1.0 + make_result("123456789012345", long_answer, 1.0, band=band), # 15 -> x0.75 + ] + cfg = make_config(default={"enabled": True, "profile_band_reasoning": True}) + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 0.75]) + + def test_missing_row_band_is_noop(self): + results = [ + make_result("12345", "12345", 1.0), + make_result("1234567890123456789012345", "12345", 1.0), + ] + cfg = make_config(default={"enabled": True, "profile_band_total": True}) + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_channel_not_enabled_in_config_is_noop(self): + band = {"total": {"a": 10, "b": 20, "f": 0.5}} + results = [ + make_result("1234567890", "1234567890", 1.0, band=band), + make_result("12345", "12345", 1.0, band=band), + ] + cfg = make_config(default={"enabled": True}) # no profile_band_* flag + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + +class TestProfileBandGlobalDefaults: + """profile_band driven by config-level length_bonus.profile_band defaults.""" + + def test_global_total_only(self): + cfg = make_config( + default={"enabled": True}, + profile_band={ + "enabled": True, + "defaults": {"total": {"a": 10, "b": 20, "f": 0.5}}, + }, + ) + results = [ + make_result("12345", "12345", 1.0), # total 10 -> x1.0 + make_result("1234567890123456789012345", "12345", 1.0), # 30 -> x0.5 + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 0.5]) + + def test_global_works_without_default_block(self): + # Channels under defaults are implicitly enabled; no length_bonus.default + # `enabled` or profile_band_* booleans required. + cfg = make_config( + default={}, # only length_type + profile_band={ + "enabled": True, + "defaults": {"reasoning": {"a": 10, "b": 20, "f": 0.5}}, + }, + ) + results = [ + make_result("123456789012345", "xx", 1.0), # reasoning 15 -> x0.75 + make_result("12345", "xx", 1.0), # 5 -> x1.0 + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([0.75, 1.0]) + + def test_row_band_wins_over_global(self): + cfg = make_config( + default={"enabled": True}, + profile_band={ + "enabled": True, + "defaults": {"total": {"a": 10, "b": 20, "f": 0.5}}, + }, + ) + generous = {"total": {"a": 100, "b": 200, "f": 0.5}} + results = [ + make_result("12345", "12345", 1.0, band=generous), + make_result("1234567890123456789012345", "12345", 1.0, band=generous), + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_disabled_block_is_noop(self): + cfg = make_config( + default={"enabled": True}, + profile_band={ + "enabled": False, + "defaults": {"total": {"a": 10, "b": 20, "f": 0.5}}, + }, + ) + results = [ + make_result("1234567890123456789012345", "12345", 1.0), + make_result("12345", "12345", 1.0), + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_malformed_global_channel_ignored(self): + cfg = make_config( + default={"enabled": True}, + profile_band={ + "enabled": True, + "defaults": {"total": {"a": 20, "b": 10, "f": 0.5}}, # b <= a + }, + ) + results = [ + make_result("1234567890123456789012345", "12345", 1.0), + make_result("12345", "12345", 1.0), + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + +class TestGroupRelativeLengthScaling: + """Dense zero-centered group relative-length penalty.""" + + def test_two_rollouts_symmetric_adjustment(self): + # lengths 10 and 30: raw weights 1 and 0, centered +0.5/-0.5, coeff 0.1. + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.1} + ) + results = [ + make_result("12345", "12345", 1.0), # total 10 -> +0.05 + make_result("1234567890123456789012345", "12345", 1.0), # 30 -> -0.05 + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.05, 0.95]) + + def test_three_rollouts_zero_centered(self): + # lengths 10/20/30 -> raw weights 1/0.5/0 -> centered +0.5/0/-0.5. + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.1}, + num_gens=3, + ) + results = [ + make_result("12345", "12345", 1.0), # 10 + make_result("1234567890", "1234567890", 1.0), # 20 + make_result("123456789012345", "123456789012345", 1.0), # 30 + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.05, 1.0, 0.95]) + # Zero-centered: the group's mean reward is unchanged by the adjustment. + assert sum(rewards_of(results)) == pytest.approx(3.0) + + def test_equal_lengths_no_adjustment(self): + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.1} + ) + results = [ + make_result("12345", "12345", 1.0), + make_result("12345", "12345", 1.0), + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_only_positive_rollouts_participate(self): + # The zero-reward rollout is neither adjusted nor part of min/max, so + # the two positives (10 and 30) still get the symmetric +/-0.05. + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.1}, + num_gens=3, + ) + results = [ + make_result("12345", "12345", 1.0), # 10 -> +0.05 + make_result("1" * 1000, "1" * 1000, 0.0), # untouched, excluded + make_result("1234567890123456789012345", "12345", 1.0), # 30 -> -0.05 + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.05, 0.0, 0.95]) + + def test_reasoning_channel_uses_reasoning_length_only(self): + # Same total lengths, different reasoning/answer split: only the + # reasoning coefficient is on, so the shorter-reasoning rollout wins. + cfg = make_config( + default={"enabled": True, "group_reasoning_length_penalty_coeff": 0.1} + ) + results = [ + make_result("12345", "123456789012345", 1.0), # reasoning 5 -> +0.05 + make_result("123456789012345", "12345", 1.0), # reasoning 15 -> -0.05 + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.05, 0.95]) + + def test_zero_coefficient_is_noop(self): + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.0} + ) + results = [ + make_result("12345", "12345", 1.0), + make_result("1234567890123456789012345", "12345", 1.0), + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_agent_override_disables_for_agent(self): + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.1} + ) + cfg["grpo"]["length_bonus"]["agent_overrides"] = {AGENT: {"enabled": False}} + results = [ + make_result("12345", "12345", 1.0), + make_result("1234567890123456789012345", "12345", 1.0), + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) From f8724c5bb449c14066e1e148be10733e289610f8 Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 11:09:33 -0700 Subject: [PATCH 08/27] docs: document profile-gate defaults, fail-closed behavior, and global-band threshold source Co-Authored-By: Claude Fable 5 --- nemo_rl/utils/length_adjustments.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/nemo_rl/utils/length_adjustments.md b/nemo_rl/utils/length_adjustments.md index 2b082ea2535..ed2e1b55759 100644 --- a/nemo_rl/utils/length_adjustments.md +++ b/nemo_rl/utils/length_adjustments.md @@ -514,9 +514,16 @@ Config keys: This is a gate on group-relative length scaling. It does not define a separate penalty by itself. +Defaults: `group_length_penalty_profile_gate_channel: total`, +`group_length_penalty_profile_gate_field: a`, +`group_length_penalty_profile_gate_positive_only: true`. + For each prompt group: 1. Read a threshold from `profile_band[channel][field]`, for example `profile_band["total"]["a"]`. + The band is the row's `profile_band` metadata merged over the global + `length_bonus.profile_band.defaults` (row channels win), so the gate also works on datasets + without per-prompt bands when global defaults are configured. 2. Compute the mean rollout length for the selected channel. 3. If `group_length_penalty_profile_gate_positive_only` is true, use only positive rollouts in that mean. @@ -528,6 +535,12 @@ For each prompt group: If the gate is closed, all group-relative coefficients are set to zero for that prompt group. +The gate fails closed: when the threshold cannot be resolved — no `profile_band` available for +the channel (neither per-row nor global), an unknown channel name, or no eligible rollouts to +average (e.g. `positive_only: true` and the whole group scored 0) — group-relative scaling is +disabled for that prompt group. The verbose per-group logs record the reason +(`missing_profile_limit`, `unknown_channel`, `no_lengths`, `mean_le_limit`, `mean_gt_limit`). + `group_length_penalty_profile_gate_positive_only` only affects the gate decision. It does not change the rollouts that receive the group-relative adjustment after the gate opens. In the current implementation, group-relative length scaling itself still applies only to positive From 50244bcf5753665c58f639b9e72d1718716e8459 Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 11:48:44 -0700 Subject: [PATCH 09/27] chore: drop example/recipe configs from the lenpen PR Remove the ultra recipe configs, the pipeclean launcher, and the example profile_band.yaml. The PR carries only the length-adjustment feature code, tests, and docs; recipe configs ship separately. Co-Authored-By: Claude Fable 5 --- examples/configs/grpo_ultra_256n4g_bf16.yaml | 784 ------------------ examples/configs/grpo_ultra_512n4g_bf16.yaml | 743 ----------------- ...grpo_ultra_64n4g_length_adj_pipeclean.yaml | 645 -------------- .../configs/grpo_ultra_64n4g_pipeclean.yaml | 715 ---------------- .../length_adj_pipeclean_launch.sh | 35 - scripts/profile_band.yaml | 28 - 6 files changed, 2950 deletions(-) delete mode 100644 examples/configs/grpo_ultra_256n4g_bf16.yaml delete mode 100644 examples/configs/grpo_ultra_512n4g_bf16.yaml delete mode 100644 examples/configs/grpo_ultra_64n4g_length_adj_pipeclean.yaml delete mode 100644 examples/configs/grpo_ultra_64n4g_pipeclean.yaml delete mode 100755 examples/nemo_gym/nemotron-3-ultra/length_adj_pipeclean_launch.sh delete mode 100644 scripts/profile_band.yaml diff --git a/examples/configs/grpo_ultra_256n4g_bf16.yaml b/examples/configs/grpo_ultra_256n4g_bf16.yaml deleted file mode 100644 index f1695b8766b..00000000000 --- a/examples/configs/grpo_ultra_256n4g_bf16.yaml +++ /dev/null @@ -1,784 +0,0 @@ -# ============================================================================= -# 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: 512 - 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_period: 8 - 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: 512 - 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: -20 - advantage_clip_high: 20 - 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 - - length_bonus: - verbose: true - default: - enabled: false - length_type: "tokens" - reasoning_bonus: 0.0 - answer_bonus: 0.0 - total_bonus: 0.0 - longest_reasoning_penalty: 0.0 - longest_answer_penalty: 0.0 - longest_total_penalty: 0.0 - top_percentile: 0.5 - group_reasoning_length_penalty_coeff: 0.0 - group_answer_length_penalty_coeff: 0.0 - group_total_length_penalty_coeff: 0.0 - reasoning_zmad_threshold: 0.0 - reasoning_zmad_penalty: 0.0 - answer_zmad_threshold: 0.0 - answer_zmad_penalty: 0.0 - total_zmad_threshold: 0.0 - total_zmad_penalty: 0.0 - profiled_length_penalty: 0.0 - profiled_length_n_std: 1.0 - profiled_length_min_samples: 2 - agent_overrides: - abstention_simple_agent: - enabled: false - calendar_simple_agent: - enabled: false - code_gen_simple_agent: - enabled: true - total_zmad_threshold: 2.5 - total_zmad_penalty: 0.1 - equivalence_llm_judge_simple_agent: - enabled: true - total_zmad_threshold: 2.5 - total_zmad_penalty: 0.1 - ether0_simple_agent: - enabled: true - total_zmad_threshold: 2.0 - total_zmad_penalty: 0.1 - genrm_simple_agent: - enabled: false - genrm_simple_agent_reasoning_off: - enabled: false - instruction_following_simple_agent: - enabled: true - total_zmad_threshold: 2.0 - total_zmad_penalty: 0.1 - jailbreak_engagement_with_disclaimer: - enabled: false - jailbreak_hard_refusal_no_redirection: - enabled: false - jailbreak_hard_refusal_with_helplines: - enabled: false - jailbreak_refusal_with_explanation: - enabled: false - lc_equivalence_rule_simple_agent: - enabled: false - lc_judge_simple_agent: - enabled: true - total_zmad_threshold: 2.0 - total_zmad_penalty: 0.1 - math_formal_lean_refinement_agent: - enabled: true - total_zmad_threshold: 2.5 - total_zmad_penalty: 0.1 - math_with_judge_simple_agent: - enabled: true - total_zmad_threshold: 2.0 - total_zmad_penalty: 0.1 - mcqa_simple_agent: - enabled: true - longest_total_penalty: 0.1 - multichallenge_simple_agent: - enabled: true - total_zmad_threshold: 2.5 - total_zmad_penalty: 0.1 - ns_tools_simple_agent: - enabled: false - nvarc_inductive_simple_agent: - enabled: false - nvarc_transductive_simple_agent: - enabled: false - reasoning_gym_simple_agent: - enabled: false - search_pivot_single_step_tool_use_with_argument_comparison_agent: - enabled: false - single_step_tool_use_with_argument_comparison_agent: - enabled: true - total_zmad_threshold: 2.0 - total_zmad_penalty: 0.1 - structured_outputs_simple_agent: - enabled: true - total_zmad_threshold: 2.0 - total_zmad_penalty: 0.1 - swe_pivot_single_step_tool_use_with_argument_comparison_agent: - enabled: false - terminus_judge_string_only_simple_agent: - enabled: true - total_zmad_threshold: 2.0 - total_zmad_penalty: 0.1 - toolcall_schema_single_step_tool_use_with_argument_comparison_agent: - enabled: true - total_zmad_threshold: 2.0 - total_zmad_penalty: 0.1 - - 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: null # set by the launcher via MODEL_PATH - tokenizer: - name: ${policy.model_name} - chat_template_kwargs: null - hf_config_overrides: {} - - train_global_batch_size: 8192 - train_micro_batch_size: 1 - generation_batch_size: 64 - logprob_batch_size: 1 - max_total_sequence_length: 49152 - 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" #"flex" - moe_flex_dispatcher_backend: "alltoall" #"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 - 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 - 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: 65536 - 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: 65536 - 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 - http_server_serving_chat_kwargs: - enable_auto_tools: true - tool_parser: qwen3_coder - reasoning_parser: nano_v3 - reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py - - vllm_kwargs: - attention_backend: FLASH_ATTN - 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: 182 # Overridden by launch script - -# ============================================================================= -# 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 + 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 - - # Increase the num workers to help with the serialization/deserialization overhead of token IDs - policy_model: - responses_api_models: - vllm_model: - num_workers: 16 - num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} - policy_model_reasoning_off: - responses_api_models: - vllm_model: - num_workers: 4 - num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} - - 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 - - vllm_serve_kwargs: - attention_backend: TRITON_ATTN - tensor_parallel_size: 4 - data_parallel_size: 1 - data_parallel_size_local: 1 - pipeline_parallel_size: 1 - gpu_memory_utilization: 0.85 - max_model_len: 96000 - model_loader_extra_config: - enable_multithread_load: true - num_threads: 112 - compilation_config: - cudagraph_capture_sizes: [1,2,4,8,16,32,64,128,256] - - - # 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 - - vllm_serve_kwargs: - attention_backend: TRITON_ATTN - tensor_parallel_size: 4 - data_parallel_size: 9 - 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 - enable_prefix_caching: true - enable_chunked_prefill: true - model_loader_extra_config: - enable_multithread_load: true - num_threads: 112 - compilation_config: - cudagraph_capture_sizes: [1,2,4,8,16,32,64,128,256] - 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" - 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 - genrm_parse_retries: 0 - group_reasoning_length_penalty_coeff: 0.1 - group_answer_length_penalty_coeff: 0.1 - group_style_penalty_coeff: 0.0 - default_score: 3.0 - default_ranking: 3.5 - - # 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 - ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python - vllm_serve_env_vars: - VLLM_RAY_DP_PACK_STRATEGY: strict - - vllm_serve_kwargs: - tensor_parallel_size: 4 - data_parallel_size: 4 - data_parallel_size_local: 1 - pipeline_parallel_size: 1 - reasoning_parser: deepseek_r1 - gpu_memory_utilization: 0.85 - max_model_len: 60000 - enable_prefix_caching: true - enable_chunked_prefill: true - model_loader_extra_config: - enable_multithread_load: true - num_threads: 112 - load_format: instanttensor - compilation_config: - cudagraph_capture_sizes: [1,2,4,8,16,32] - server_env: - NCCL_MNNVL_ENABLE: "0" - - 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 - - math_formal_lean_refinement_agent: - responses_api_agents: - proof_refinement_agent: - max_correction_turns: 0 - -# ============================================================================= -# Logger -# ============================================================================= -logger: - log_dir: "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 diff --git a/examples/configs/grpo_ultra_512n4g_bf16.yaml b/examples/configs/grpo_ultra_512n4g_bf16.yaml deleted file mode 100644 index 46669662daf..00000000000 --- a/examples/configs/grpo_ultra_512n4g_bf16.yaml +++ /dev/null @@ -1,743 +0,0 @@ -# ============================================================================= -# 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: 512 - 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_period: 8 - 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: 512 - 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: -20 - advantage_clip_high: 20 - 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 - - length_bonus: - verbose: true - default: - enabled: false - length_type: "tokens" - reasoning_bonus: 0.0 - answer_bonus: 0.0 - longest_reasoning_penalty: 0.0 - longest_answer_penalty: 0.0 - top_percentile: 0.5 - group_reasoning_length_penalty_coeff: 0.0 - group_answer_length_penalty_coeff: 0.0 - reasoning_zmad_threshold: 0.0 - reasoning_zmad_penalty: 0.0 - answer_zmad_threshold: 0.0 - answer_zmad_penalty: 0.0 - agent_overrides: - abstention_simple_agent: - enabled: false - calendar_simple_agent: - enabled: false - code_gen_simple_agent: - enabled: true - reasoning_zmad_threshold: 2.5 - reasoning_zmad_penalty: 0.1 - equivalence_llm_judge_simple_agent: - enabled: true - reasoning_zmad_threshold: 2.5 - reasoning_zmad_penalty: 0.1 - ether0_simple_agent: - enabled: true - reasoning_zmad_threshold: 2.0 - reasoning_zmad_penalty: 0.1 - genrm_simple_agent: - enabled: false - genrm_simple_agent_reasoning_off: - enabled: false - instruction_following_simple_agent: - enabled: true - reasoning_zmad_threshold: 2.0 - reasoning_zmad_penalty: 0.1 - jailbreak_engagement_with_disclaimer: - enabled: false - jailbreak_hard_refusal_no_redirection: - enabled: false - jailbreak_hard_refusal_with_helplines: - enabled: false - jailbreak_refusal_with_explanation: - enabled: false - lc_equivalence_rule_simple_agent: - enabled: false - lc_judge_simple_agent: - enabled: true - reasoning_zmad_threshold: 2.0 - reasoning_zmad_penalty: 0.1 - math_formal_lean_refinement_agent: - enabled: true - reasoning_zmad_threshold: 2.5 - reasoning_zmad_penalty: 0.1 - math_with_judge_simple_agent: - enabled: true - reasoning_zmad_threshold: 2.0 - reasoning_zmad_penalty: 0.1 - mcqa_simple_agent: - enabled: true - longest_reasoning_penalty: 0.1 - multichallenge_simple_agent: - enabled: true - reasoning_zmad_threshold: 2.5 - reasoning_zmad_penalty: 0.1 - ns_tools_simple_agent: - enabled: false - nvarc_inductive_simple_agent: - enabled: false - nvarc_transductive_simple_agent: - enabled: false - reasoning_gym_simple_agent: - enabled: false - search_pivot_single_step_tool_use_with_argument_comparison_agent: - enabled: false - single_step_tool_use_with_argument_comparison_agent: - enabled: true - reasoning_zmad_threshold: 2.0 - reasoning_zmad_penalty: 0.1 - structured_outputs_simple_agent: - enabled: true - reasoning_zmad_threshold: 2.0 - reasoning_zmad_penalty: 0.1 - swe_pivot_single_step_tool_use_with_argument_comparison_agent: - enabled: false - terminus_judge_string_only_simple_agent: - enabled: true - reasoning_zmad_threshold: 2.0 - reasoning_zmad_penalty: 0.1 - toolcall_schema_single_step_tool_use_with_argument_comparison_agent: - enabled: true - reasoning_zmad_threshold: 2.0 - reasoning_zmad_penalty: 0.1 - - 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: null # set by the launcher via MODEL_PATH - tokenizer: - name: ${policy.model_name} - chat_template_kwargs: null - hf_config_overrides: {} - - train_global_batch_size: 8192 - train_micro_batch_size: 1 - generation_batch_size: 64 - logprob_batch_size: 1 - max_total_sequence_length: 49152 - 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" #"flex" - moe_flex_dispatcher_backend: "alltoall" #"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 - 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 - 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: ${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 - http_server_serving_chat_kwargs: - enable_auto_tools: true - tool_parser: qwen3_coder - reasoning_parser: nano_v3 - reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py - - vllm_kwargs: - attention_backend: FLASH_ATTN - max_num_seqs: 64 - 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: 182 # Overridden by launch script - -# ============================================================================= -# 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 + 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 - - vllm_serve_kwargs: - attention_backend: TRITON_ATTN - tensor_parallel_size: 4 - data_parallel_size: 1 - data_parallel_size_local: 1 - pipeline_parallel_size: 1 - gpu_memory_utilization: 0.85 - max_model_len: 96000 - 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,64,128,256] - - - # 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 - - vllm_serve_kwargs: - attention_backend: FLASH_ATTN - tensor_parallel_size: 4 - data_parallel_size: 9 - 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,64,128,256] - 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: 32768 - temperature: 1.0 - top_p: 0.95 - comparison_strategy: "circular" - 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.1 - group_answer_length_penalty_coeff: 0.25 - group_style_penalty_coeff: 0.0 - default_score: 3.0 - default_ranking: 3.5 - - # 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: "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 diff --git a/examples/configs/grpo_ultra_64n4g_length_adj_pipeclean.yaml b/examples/configs/grpo_ultra_64n4g_length_adj_pipeclean.yaml deleted file mode 100644 index 065d8a08a10..00000000000 --- a/examples/configs/grpo_ultra_64n4g_length_adj_pipeclean.yaml +++ /dev/null @@ -1,645 +0,0 @@ -# ============================================================================= -# GRPO Ultra V3 — 64-node GB200 NVL72 Pipeclean Config -# ============================================================================= -# Pipeclean config for validating the end-to-end GRPO pipeline on -# 64 nodes × 4 GPUs/node. Batch sizes and sequence lengths are -# intentionally small to complete steps quickly. -# -# Defaults to 10 training steps (~2h on 64 nodes). Override at launch -# time with NRL_MAX_STEPS for shorter CI runs (e.g. NRL_MAX_STEPS=4). -# -# - gpus_per_node: 4 -# - TP: 8 -# - EP: 32 (16 experts/rank) -# - vLLM TP: 8 -# - Judge TP: 4 -# ============================================================================= - -# ============================================================================= -# Cluster — overridden by launch script -# ============================================================================= -cluster: - gpus_per_node: 4 - num_nodes: 64 - 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_period: 4 - checkpoint_must_save_by: "00:03:30:00" - model_save_format: "safetensors" - save_consolidated: false - -# ============================================================================= -# GRPO Algorithm -# ============================================================================= -grpo: - num_prompts_per_step: 16 # Pipeclean: reduced from 128 - num_generations_per_prompt: 4 # Pipeclean: reduced from 16 - num_val_generations_per_prompt: 2 - max_rollout_turns: 1 - max_num_epochs: 1 - max_num_steps: 10 # Pipeclean: reduced from 1000000; override via NRL_MAX_STEPS - 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 - 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 - - length_bonus: - default: - enabled: true - length_type: "tokens" - reasoning_bonus: 0.1 - answer_bonus: 0.1 - longest_reasoning_penalty: 0.0 - longest_answer_penalty: 0.0 - top_percentile: 0.5 - group_reasoning_length_penalty_coeff: 0.05 - group_answer_length_penalty_coeff: 0.05 - agent_overrides: - # --- Agents from training dataset (curriculum_v29_no_judge) --- - calendar_simple_agent: null - code_gen_simple_agent: - group_answer_length_penalty_coeff: 0.1 - instruction_following_simple_agent: null - mcqa_simple_agent: null - reasoning_gym_simple_agent: null - single_step_tool_use_with_argument_comparison_agent: - group_answer_length_penalty_coeff: 0.0 - structured_outputs_simple_agent: null - terminal_pivot_simple_agent: - enabled: false - workplace_assistant_simple_agent: null - # --- Additional agents from ultra config environments --- - math_with_judge_simple_agent: - enabled: false - genrm_simple_agent: - enabled: false - genrm_simple_agent_reasoning_off: - enabled: false - lc_judge_simple_agent: - enabled: false - nl2bash_equivalency_simple_agent: - enabled: false - equivalence_llm_judge_simple_agent: - enabled: false - ns_tools_simple_agent: - enabled: false - math_formal_lean_simple_agent: - enabled: false - jailbreak_detection_simple_agent: - enabled: false - over_refusal_detection_simple_agent: - enabled: false - multichallenge_simple_agent: - enabled: false - inverse_if_simple_agent: - enabled: false - search_pivot_single_step_tool_use_with_argument_comparison_agent: - enabled: false - toolcall_schema_single_step_tool_use_with_argument_comparison_agent: - enabled: false - - 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: 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 -# ============================================================================= -policy: - model_name: null # set by the launcher via MODEL_PATH - tokenizer: - name: ${policy.model_name} - chat_template_kwargs: null - hf_config_overrides: {} - - train_global_batch_size: 64 # Pipeclean: num_prompts × num_generations = 16 × 4 - train_micro_batch_size: 1 - generation_batch_size: 64 - logprob_batch_size: 1 - max_total_sequence_length: 32768 # Pipeclean: reduced from 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: 1 - 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: ${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_aux_loss_coeff: 0.0 - moe_shared_expert_overlap: false - use_gloo_process_groups: false - - # Compute - apply_rope_fusion: true - bias_activation_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 - mtp_num_layers: 0 - 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: 0 - lr_warmup_init: 3e-7 - - 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 - - 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 - - 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: 32768 # Pipeclean: reduced from 65536 - 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.85 - max_model_len: 32768 # Must be <= 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 - http_server_serving_chat_kwargs: - enable_auto_tools: true - tool_parser: qwen3_coder - reasoning_parser: nano_v3 - reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py - - vllm_kwargs: - attention_backend: FLASH_ATTN - max_num_seqs: 16 - mamba_ssm_cache_dtype: "float32" - compilation_config: - cudagraph_capture_sizes: [1,2,4,8,16,32,64] - pass_config: - # Disable FlashInfer TRTLLM AllReduce fusion. It requires - # SymmDeviceMemory (MNNVL IPC) to initialise successfully on every TP - # rank. When TP spans multiple nodes some ranks succeed while others - # fail (CUDA_ERROR_INVALID_DEVICE), causing a split-brain: half the TP - # group uses the MNNVL-based TRTLLM allreduce path while the other half - # falls back to standard NCCL. They never rendezvous → permanent hang. - fuse_allreduce_rms: false - - colocated: - enabled: false - resources: - gpus_per_node: 4 - num_nodes: 26 # Overridden by launch script - -# ============================================================================= -# Data -# ============================================================================= -data: - max_input_seq_length: 16384 # NOTE: no-op for NemoGymDataset (processor ignores max_seq_length) - 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 + Judge Models -# ============================================================================= -env: - should_use_nemo_gym: true - use_genrm_compare: true - genrm_agent_names: - - "genrm_simple_agent" - - "genrm_simple_agent_reasoning_off" - genrm_compare_server_name: "genrm_compare" - nemo_gym: - skip_venv_if_present: true - num_gpu_nodes: 6 # Overridden by launch script - 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/structured_outputs/configs/structured_outputs_json.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/terminal_pivot/configs/terminal_pivot.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/jailbreak_detection/configs/jailbreak_detection_nemotron_combined_reward_tp8.yaml - - resources_servers/over_refusal_detection/configs/over_refusal_detection_nemotron_tp8.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 - - # Safety Model: 4B, so TP=1, so fits single GPU - jailbreak_detection: - resources_servers: - jailbreak_detection: - judge_model_server: - type: responses_api_models - name: safety_judge_model - - safety_judge_model: - responses_api_models: - vllm_model: - entrypoint: app.py - base_url: http://127.0.0.1:8001/v1 - api_key: dummy_key - model: null # Set by launch script (or omit for no safety judge) - return_token_id_information: false - uses_reasoning_parser: false - spinup_server: true - router_dp_size: 4 - server_args: - attention_backend: TRITON_ATTN - tensor_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: 2 - compilation_config: - cudagraph_capture_sizes: [1,2,4,8,16] - - terminal_pivot_simple_agent: - responses_api_agents: - simple_agent: - model_server: - name: policy_model - - # nl2bash / General Judge: TP=4 on GB200 192GB - nl2bash_judge_model: - responses_api_models: - vllm_model: - entrypoint: app.py - base_url: http://127.0.0.1:10000/v1 - api_key: dummy_key - model: null # Set by launch script - return_token_id_information: false - uses_reasoning_parser: false - spinup_server: true - router_dp_size: 2 - server_args: - attention_backend: FLASH_ATTN - tensor_parallel_size: 4 - data_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: 16 - 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_servers: - genrm_compare: - genrm_model_server: - type: responses_api_models - name: genrm_model - genrm_responses_create_params: - max_output_tokens: 16384 - temperature: 0.6 - top_p: 0.95 - comparison_strategy: "circular" - num_judges_per_comparison: 1 - use_principle: true - default_principle: "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 - group_answer_length_penalty_coeff: 0 - group_style_penalty_coeff: 0.1 - default_score: 3.0 - default_ranking: 3.5 - - genrm_model: - responses_api_models: - vllm_model: - entrypoint: app.py - base_url: http://127.0.0.1:8000/v1 - api_key: dummy_key - model: null # Set by launch script - uses_reasoning_parser: true - return_token_id_information: false - spinup_server: true - router_dp_size: 2 - server_args: - attention_backend: FLASH_ATTN - tensor_parallel_size: 4 - reasoning_parser: deepseek_r1 - gpu_memory_utilization: 0.85 - max_model_len: 60000 - 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,32] - server_env: - NCCL_MNNVL_ENABLE: "0" - - 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: 1024 - unit_test_timeout_secs: 10 - debug: false - -# ============================================================================= -# Logger -# ============================================================================= -logger: - log_dir: "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" - tensorboard: {} - mlflow: - experiment_name: "grpo-ultra-v3" - run_name: "grpo-ultra-v3" - gpu_monitoring: - collection_interval: 10 - flush_interval: 10 - -# ============================================================================= -# Effort Levels -# ============================================================================= -effort_levels: - low_string: "{reasoning effort: low}" - low_weight: 0.2 - low_penalty: 1 - low_ub: 3000 diff --git a/examples/configs/grpo_ultra_64n4g_pipeclean.yaml b/examples/configs/grpo_ultra_64n4g_pipeclean.yaml deleted file mode 100644 index 4aab4c9b369..00000000000 --- a/examples/configs/grpo_ultra_64n4g_pipeclean.yaml +++ /dev/null @@ -1,715 +0,0 @@ -# ============================================================================= -# GRPO Ultra V3 — 64-node GB200 NVL72 Pipeclean Config -# ============================================================================= -# Pipeclean config for validating the end-to-end GRPO pipeline on -# 64 nodes × 4 GPUs/node. Batch sizes and sequence lengths are -# intentionally small to complete steps quickly. -# -# Defaults to 10 training steps (~2h on 64 nodes). Override at launch -# time with NRL_MAX_STEPS for shorter CI runs (e.g. NRL_MAX_STEPS=4). -# -# - gpus_per_node: 4 -# - TP: 8 -# - EP: 32 (16 experts/rank) -# - vLLM TP: 8 -# - Judge TP: 4 -# ============================================================================= - -# ============================================================================= -# Cluster — overridden by launch script -# ============================================================================= -cluster: - gpus_per_node: 4 - num_nodes: 64 - 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_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: 16 # Pipeclean: reduced from 128 - num_generations_per_prompt: 4 # Pipeclean: reduced from 16 - num_val_generations_per_prompt: 2 - max_rollout_turns: 1 - max_num_epochs: 1 - max_num_steps: 10 # Pipeclean: reduced from 1000000; override via NRL_MAX_STEPS - 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 - 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 - - length_bonus: - verbose: false - default: - enabled: false - length_type: "tokens" - reasoning_bonus: 0.0 - answer_bonus: 0.0 - longest_reasoning_penalty: 0.0 - longest_answer_penalty: 0.0 - top_percentile: 0.5 - group_reasoning_length_penalty_coeff: 0.0 - group_answer_length_penalty_coeff: 0.0 - reasoning_zmad_threshold: 0.0 - reasoning_zmad_penalty: 0.0 - answer_zmad_threshold: 0.0 - answer_zmad_penalty: 0.0 - agent_overrides: - # --- Agents from training dataset (curriculum_v29_no_judge) --- - calendar_simple_agent: - enabled: false - code_gen_simple_agent: - enabled: false - instruction_following_simple_agent: - enabled: false - mcqa_simple_agent: - enabled: false - reasoning_gym_simple_agent: - enabled: false - single_step_tool_use_with_argument_comparison_agent: - enabled: false - structured_outputs_simple_agent: - enabled: false - terminal_pivot_simple_agent: - enabled: false - workplace_assistant_simple_agent: - enabled: false - # --- Additional agents from ultra config environments --- - math_with_judge_simple_agent: - enabled: false - genrm_simple_agent: - enabled: false - genrm_simple_agent_reasoning_off: - enabled: false - lc_judge_simple_agent: - enabled: false - nl2bash_equivalency_simple_agent: - enabled: false - equivalence_llm_judge_simple_agent: - enabled: false - ns_tools_simple_agent: - enabled: false - math_formal_lean_simple_agent: - enabled: false - jailbreak_detection_simple_agent: - enabled: false - over_refusal_detection_simple_agent: - enabled: false - multichallenge_simple_agent: - enabled: false - inverse_if_simple_agent: - enabled: false - search_pivot_single_step_tool_use_with_argument_comparison_agent: - enabled: false - toolcall_schema_single_step_tool_use_with_argument_comparison_agent: - enabled: false - - 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: 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 -# ============================================================================= -policy: - model_name: null # set by the launcher via MODEL_PATH - tokenizer: - name: ${policy.model_name} - chat_template_kwargs: null - hf_config_overrides: {} - - train_global_batch_size: 64 # Pipeclean: num_prompts × num_generations = 16 × 4 - 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: 1 - 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: ${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 - use_gloo_process_groups: false - - # Compute - apply_rope_fusion: true - use_fused_weighted_squared_relu: true - bias_activation_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: 0 - lr_warmup_init: 3e-7 - - 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 - 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: 65536 - 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.85 - max_model_len: 65536 - 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 - http_server_serving_chat_kwargs: - enable_auto_tools: true - tool_parser: qwen3_coder - reasoning_parser: nano_v3 - reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py - - vllm_kwargs: - attention_backend: FLASH_ATTN - mamba_ssm_cache_dtype: "float32" - compilation_config: - cudagraph_capture_sizes: [1,2,4,8,16,32,64] - pass_config: - # Disable FlashInfer TRTLLM AllReduce fusion. It requires - # SymmDeviceMemory (MNNVL IPC) to initialise successfully on every TP - # rank. When TP spans multiple nodes some ranks succeed while others - # fail (CUDA_ERROR_INVALID_DEVICE), causing a split-brain: half the TP - # group uses the MNNVL-based TRTLLM allreduce path while the other half - # falls back to standard NCCL. They never rendezvous → permanent hang. - fuse_allreduce_rms: false - - colocated: - enabled: false - resources: - gpus_per_node: 4 - num_nodes: 26 # Overridden by launch script - -# ============================================================================= -# Data -# ============================================================================= -data: - max_input_seq_length: 16384 # NOTE: no-op for NemoGymDataset (processor ignores max_seq_length) - 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 + 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/over_refusal_detection/configs/over_refusal_detection_nemotron_tp8.yaml - - resources_servers/indirect_prompt_injection/configs/indirect_prompt_injection.yaml - - # Increase the num workers to help with the serialization/deserialization overhead of token IDs - policy_model: - responses_api_models: - vllm_model: - num_workers: 16 - num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} - policy_model_reasoning_off: - responses_api_models: - vllm_model: - num_workers: 4 - num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} - - 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, so TP=1, so fits single GPU - 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 (or omit for no safety judge) - return_token_id_information: false - uses_reasoning_parser: false - debug: true - show_vllm_engine_stats: true - vllm_serve_env_vars: - VLLM_RAY_DP_PACK_STRATEGY: strict - - vllm_serve_kwargs: - attention_backend: TRITON_ATTN - tensor_parallel_size: 4 - data_parallel_size: 1 - data_parallel_size_local: 1 - pipeline_parallel_size: 1 - gpu_memory_utilization: 0.85 - max_model_len: 96000 - 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 - show_vllm_engine_stats: true - vllm_serve_env_vars: - VLLM_RAY_DP_PACK_STRATEGY: strict - - vllm_serve_kwargs: - attention_backend: FLASH_ATTN - tensor_parallel_size: 4 - data_parallel_size: 2 - 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 - 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: 16384 - temperature: 0.6 - top_p: 0.95 - comparison_strategy: "circular" - 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 - genrm_parse_retries: 0 - group_reasoning_length_penalty_coeff: 0.1 - group_answer_length_penalty_coeff: 0.2 - group_style_penalty_coeff: 0.0 - default_score: 3.0 - default_ranking: 3.5 - - genrm_model: - responses_api_models: - genrm_model: - entrypoint: app.py - api_key: dummy_key - model: null # Set by launch script - uses_reasoning_parser: true - return_token_id_information: false - debug: true - show_vllm_engine_stats: true - ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python - vllm_serve_env_vars: - VLLM_RAY_DP_PACK_STRATEGY: strict - - vllm_serve_kwargs: - tensor_parallel_size: 4 - data_parallel_size: 2 - data_parallel_size_local: 1 - pipeline_parallel_size: 1 - reasoning_parser: deepseek_r1 - gpu_memory_utilization: 0.85 - max_model_len: 60000 - enable_prefix_caching: true - enable_chunked_prefill: true - model_loader_extra_config: - enable_multithread_load: true - num_threads: 112 - load_format: instanttensor - compilation_config: - cudagraph_capture_sizes: [1,2,4,8,16,32] - server_env: - NCCL_MNNVL_ENABLE: "0" - - 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: 1024 - unit_test_timeout_secs: 10 - debug: false - - math_formal_lean_refinement_agent: - responses_api_agents: - proof_refinement_agent: - max_correction_turns: 0 - -# ============================================================================= -# Logger -# ============================================================================= -logger: - log_dir: "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-pipeclean" - name: "grpo-ultra-v3-pipeclean" - tensorboard: {} - mlflow: - experiment_name: "grpo-ultra-v3-pipeclean" - run_name: "grpo-ultra-v3-pipeclean" - gpu_monitoring: - collection_interval: 10 - flush_interval: 10 - -# ============================================================================= -# Effort Levels -# ============================================================================= -effort_levels: - low_string: "{reasoning effort: low}" - low_weight: 0.2 - low_penalty: 1 - low_ub: 3000 diff --git a/examples/nemo_gym/nemotron-3-ultra/length_adj_pipeclean_launch.sh b/examples/nemo_gym/nemotron-3-ultra/length_adj_pipeclean_launch.sh deleted file mode 100755 index 875e3ee8d68..00000000000 --- a/examples/nemo_gym/nemotron-3-ultra/length_adj_pipeclean_launch.sh +++ /dev/null @@ -1,35 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -# Convenience launcher for the length-adjusted Ultra pipeclean recipe. -# -# This intentionally follows examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh: -# callers provide cluster, container, model, data, cache, and optional judge -# model paths through environment variables. This wrapper only selects the -# length-adjusted config by default. -# -# Required by ultra_launch.sh: -# EXP_NAME -# MODEL_PATH -# TRAIN_PATH -# VAL_PATH -# CONTAINER -# SANDBOX_CONTAINER -# PERSISTENT_CACHE -# SLURM_PARTITION -# SLURM_ACCOUNT -# -# Optional: -# CONFIG_PATH Override the recipe config. -# WANDB_PROJ W&B project name. -# GENRM_MODEL GenRM model path or HF id. -# NL2BASH_JUDGE_MODEL -# SAFETY_JUDGE_MODEL -# EXTRA_MOUNTS Comma-separated host:container mount pairs. - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" - -export CONFIG_PATH="${CONFIG_PATH:-examples/configs/grpo_ultra_64n4g_length_adj_pipeclean.yaml}" -export WANDB_PROJ="${WANDB_PROJ:-nemotron-3-ultra-length-adjusted}" - -exec "${SCRIPT_DIR}/ultra_launch.sh" "$@" diff --git a/scripts/profile_band.yaml b/scripts/profile_band.yaml deleted file mode 100644 index 0538bff5926..00000000000 --- a/scripts/profile_band.yaml +++ /dev/null @@ -1,28 +0,0 @@ -# Config for build_profile_band_dataset.py -# -# a, b are derived from the per-prompt profiled length distribution: -# reference set = passing rollouts (reward > 0); fall back to all -# profiled rollouts if fewer than ``min_passing`` pass. -# a = mean(reference) -# b = mean(reference) + n_std * std(reference) -# -# f is selected per-row by exact pass_rate match in f_table. -# Pass-rates not in the table → row gets no profile_band block (mechanism -# is a no-op for that prompt). -# -# Channels with std==0 (cap-clamped) or fewer than 2 samples are -# omitted from a row's profile_band block but other channels still apply. - -n_std: 2.0 -min_passing: 2 - -channels: - - total - - reasoning - - answer - -f_table: - - {pass_rate: 1.000, f: 0.6} - - {pass_rate: 0.875, f: 0.7} - - {pass_rate: 0.750, f: 0.8} - - {pass_rate: 0.625, f: 0.9} From 16e4d84c8d293f59fe6c50fe83d7dc333b56023b Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 15:46:04 -0700 Subject: [PATCH 10/27] fix: profiled_length_penalty requires min_samples PASSING profiled rollouts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously, when fewer than profiled_length_min_samples profiled rollouts passed, the threshold fell back to ALL profiled lengths — deriving a length budget from a profile that failed the problem, then applying a flat penalty to every current rollout over it (correct ones included, since this algorithm has no reward gate on the receiving side). A problem the profiled model could not solve likely needs a larger thinking budget, not a penalty computed from failed attempts. Now: fewer than min_samples passing profiled rollouts -> no profiled penalty for that prompt. min_samples semantics change from 'samples needed to compute the threshold (any)' to 'passing samples required to apply the penalty'. Set profiled_length_min_samples: 1 to opt in to single-pass thresholds. Adds 4 unit tests; updates the doc. Co-Authored-By: Claude Fable 5 --- nemo_rl/utils/length_adjustments.md | 11 ++-- nemo_rl/utils/length_adjustments.py | 6 +- tests/unit/utils/test_length_adjustments.py | 63 +++++++++++++++++++++ 3 files changed, 72 insertions(+), 8 deletions(-) diff --git a/nemo_rl/utils/length_adjustments.md b/nemo_rl/utils/length_adjustments.md index ed2e1b55759..319ecf01212 100644 --- a/nemo_rl/utils/length_adjustments.md +++ b/nemo_rl/utils/length_adjustments.md @@ -113,7 +113,7 @@ grpo: | `total_zmad_penalty` | Flat penalty applied to total lengths above the zMAD threshold. | | `profiled_length_penalty` | Flat penalty for rollouts longer than a per-prompt profiled-length threshold. | | `profiled_length_n_std` | Number of standard deviations used in `mean + n_std * std` for profiled-length thresholding. | -| `profiled_length_min_samples` | Minimum profiled length samples needed before computing the profiled threshold. | +| `profiled_length_min_samples` | Minimum PASSING profiled rollouts required; below this no profiled penalty is applied. | | `profile_band_total` | Enables per-prompt `{a,b,f}` multiplier on total length for correct rollouts. | | `profile_band_reasoning` | Enables per-prompt `{a,b,f}` multiplier on reasoning length for correct rollouts. | | `profile_band_answer` | Enables per-prompt `{a,b,f}` multiplier on answer length for correct rollouts. | @@ -384,10 +384,11 @@ This uses per-prompt profiling metadata: For each prompt group: -1. Prefer profiled lengths from passing rollouts. -2. If there are fewer than `profiled_length_min_samples` passing rollouts, fall back to all - profiled lengths. -3. Compute: +1. Select profiled lengths from passing rollouts (`profiled_rewards > 0`) only. +2. If there are fewer than `profiled_length_min_samples` passing rollouts, apply **no penalty** + for this prompt. A profile that mostly failed carries no signal about the right thinking + budget — the problem may simply need more than the profiled model had. +3. Otherwise compute: ```text threshold = mean(profiled_lengths) + profiled_length_n_std * std(profiled_lengths) diff --git a/nemo_rl/utils/length_adjustments.py b/nemo_rl/utils/length_adjustments.py index 1f818ab6301..0a86bce3ff6 100644 --- a/nemo_rl/utils/length_adjustments.py +++ b/nemo_rl/utils/length_adjustments.py @@ -422,7 +422,9 @@ def apply_group_length_adjustments( groups_adjusted += 1 # Profiled length penalty: penalize rollouts longer than mean + n_std of - # passing profiled lengths for this prompt. + # passing profiled lengths for this prompt. If fewer than min_samples + # profiled rollouts passed, the profiled model found the problem hard + # and its lengths carry no budget signal — apply no penalty at all. plp = params.get("profiled_length_penalty", 0.0) if plp > 0.0: p_rewards = results[g].get("profiled_rewards") @@ -430,8 +432,6 @@ def apply_group_length_adjustments( if p_rewards is not None and p_lengths is not None: min_samples = int(params.get("profiled_length_min_samples", 2)) passing = [l for r, l in zip(p_rewards, p_lengths) if r > 0] - if len(passing) < min_samples: - passing = list(p_lengths) if len(passing) >= min_samples: mean_l = statistics.mean(passing) std_l = statistics.stdev(passing) if len(passing) >= 2 else 0.0 diff --git a/tests/unit/utils/test_length_adjustments.py b/tests/unit/utils/test_length_adjustments.py index 6c9c286ba0a..ffe0e23ba7e 100644 --- a/tests/unit/utils/test_length_adjustments.py +++ b/tests/unit/utils/test_length_adjustments.py @@ -318,3 +318,66 @@ def test_agent_override_disables_for_agent(self): ] apply_group_length_adjustments(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + +class TestProfiledLengthPenalty: + """Profiled-length threshold penalty and its passing-samples requirement.""" + + @staticmethod + def make_profiled(reasoning, answer, reward, p_rewards, p_lengths): + r = make_result(reasoning, answer, reward) + r["profiled_rewards"] = p_rewards + r["profiled_output_lengths"] = p_lengths + return r + + def cfg(self, min_samples=2): + return make_config( + default={ + "enabled": True, + "profiled_length_penalty": 0.3, + "profiled_length_n_std": 1.0, + "profiled_length_min_samples": min_samples, + } + ) + + def test_enough_passes_penalizes_over_threshold(self): + # Passing profiled lengths 10 and 14: threshold = 12 + 1*std(=~2.83) ~ 14.83. + p_rewards, p_lengths = [1, 1, 0], [10, 14, 100] + results = [ + self.make_profiled("12345", "12345", 1.0, p_rewards, p_lengths), # 10 < thr + self.make_profiled("1234567890", "1234567890", 1.0, p_rewards, p_lengths), # 20 >= thr + ] + apply_group_length_adjustments(results, self.cfg()) + assert rewards_of(results) == pytest.approx([1.0, 0.7]) + # The failing profiled length (100) must not have entered the threshold: + # with it, mean+std would exceed 20 and nothing would be penalized. + + def test_one_pass_below_min_samples_no_penalty(self): + # Only 1 passing profiled rollout with min_samples=2: no penalty for + # anyone — no fallback to failing profiled lengths. + p_rewards, p_lengths = [1, 0, 0], [10, 100, 120] + results = [ + self.make_profiled("1234567890", "1234567890", 1.0, p_rewards, p_lengths), + self.make_profiled("1" * 50, "1" * 50, 1.0, p_rewards, p_lengths), + ] + apply_group_length_adjustments(results, self.cfg()) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_zero_passes_no_penalty(self): + p_rewards, p_lengths = [0, 0, 0], [10, 12, 14] + results = [ + self.make_profiled("1234567890", "1234567890", 1.0, p_rewards, p_lengths), + self.make_profiled("1" * 50, "1" * 50, 1.0, p_rewards, p_lengths), + ] + apply_group_length_adjustments(results, self.cfg()) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_one_pass_allowed_when_min_samples_is_one(self): + # min_samples=1 opts in to single-pass thresholds: threshold = 10 + 0. + p_rewards, p_lengths = [1, 0], [10, 100] + results = [ + self.make_profiled("1234", "1234", 1.0, p_rewards, p_lengths), # 8 < 10 + self.make_profiled("1234567890", "1234567890", 1.0, p_rewards, p_lengths), # 20 >= 10 + ] + apply_group_length_adjustments(results, self.cfg(min_samples=1)) + assert rewards_of(results) == pytest.approx([1.0, 0.7]) From f02100ff034052495b5d4b6b77f695144bad8f1c Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 19:53:07 -0700 Subject: [PATCH 11/27] revert: restore should_log_nemo_gym_responses default in grpo_nanov3.yaml Accidental passenger from the original port; the flag is deliberately managed per-recipe (review feedback on #3852). Co-Authored-By: Claude Fable 5 --- examples/nemo_gym/grpo_nanov3.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/nemo_gym/grpo_nanov3.yaml b/examples/nemo_gym/grpo_nanov3.yaml index b7fced772f1..6bb2f6963a6 100644 --- a/examples/nemo_gym/grpo_nanov3.yaml +++ b/examples/nemo_gym/grpo_nanov3.yaml @@ -312,7 +312,7 @@ data: env: should_use_nemo_gym: true # true: skip expensive train_data_step*.jsonl; false: write full jsonl. - should_log_nemo_gym_responses: false + should_log_nemo_gym_responses: true # false: ignore env mask_sample flags so the loss trains on every sample. should_mask_flagged_samples: true nemo_gym: From 541be20105b82fdd89a92ad1bad51aa411d5b524 Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 20:09:55 -0700 Subject: [PATCH 12/27] refactor: strip GDPO feature mode; length adjustments are reward-mutation only Review feedback on #3852: the feature-mode layer dispatched on batch key presence (hijacking existing GDPO reward/* runs) and its documented config shape parsed to zero params. Remove the layer entirely rather than patch it under review; it can return as its own properly-gated PR. - advantage_estimator.py, test_grpo.py, test_reward_penalties.py: reverted to the base branch versions (existing GDPO estimator untouched) - rollouts.py: drop gdpo_reward_features recording/metrics/batch key and the think_count_delta feature; keep length_adjustment_config threading - length_adjustments.py: drop reward_features merge and feature recording; grpo.length_bonus reward-mutation mode is the single code path - length_adjustments.md: remove GDPO feature mode sections Co-Authored-By: Claude Fable 5 --- nemo_rl/algorithms/advantage_estimator.py | 246 --------------- nemo_rl/experience/rollouts.py | 148 --------- nemo_rl/utils/length_adjustments.md | 129 +------- nemo_rl/utils/length_adjustments.py | 291 +----------------- tests/unit/algorithms/test_grpo.py | 49 --- .../unit/experience/test_reward_penalties.py | 120 -------- 6 files changed, 11 insertions(+), 972 deletions(-) diff --git a/nemo_rl/algorithms/advantage_estimator.py b/nemo_rl/algorithms/advantage_estimator.py index 293a9f0f485..491de0d9414 100644 --- a/nemo_rl/algorithms/advantage_estimator.py +++ b/nemo_rl/algorithms/advantage_estimator.py @@ -28,8 +28,6 @@ - MOPD: https://arxiv.org/abs/2601.02780 """ -from typing import Any - import torch from pydantic import BaseModel @@ -52,89 +50,10 @@ class AdvEstimatorConfig(BaseModel, extra="allow"): use_leave_one_out_baseline: bool = True # GDPO specific: optional per-component weights w_n for the aggregation. reward_weights: list[float] | None = None - reward_features: list[str] | dict[str, Any] | None = None - verbose: bool = False # Reinforce++ specific minus_baseline: bool = True -def _feature_reward(feature_entry: Any) -> float: - if feature_entry is None: - return 0.0 - if isinstance(feature_entry, dict): - return float(feature_entry.get("reward", 0.0)) - return float(feature_entry) - - -def _format_gdpo_feature_for_log( - name: str, feature_entry: Any, feature_weight: float = 1.0 -) -> str: - if feature_entry is None: - return f"{name}=MISSING" - if not isinstance(feature_entry, dict): - return f"{name}={float(feature_entry):+.4f}" - - details = [] - adjustment = feature_entry.get("adjustment") - if adjustment is not None: - details.append(str(adjustment)) - multiplier = feature_entry.get("multiplier") - if multiplier is not None: - details.append(f"mult={float(multiplier):.4f}") - close_thinks = feature_entry.get("num_close_think_tags") - if close_thinks is not None: - details.append(f"close_thinks={int(close_thinks)}") - if feature_weight != 1.0: - details.append(f"weight={feature_weight:.4f}") - - suffix = f"({','.join(details)})" if details else "" - return f"{name}={_feature_reward(feature_entry):+.4f}{suffix}" - - -def _resolve_gdpo_reward_feature_selection( - estimator_config: AdvEstimatorConfig, agent_name: str | None -) -> list[str] | dict[str, Any]: - cfg = estimator_config.reward_features or ["env_reward"] - if isinstance(cfg, (list, tuple)): - return list(cfg) - if not isinstance(cfg, dict): - raise TypeError("gdpo reward_features must be a list or a dict") - - if "default" not in cfg and "agent_overrides" not in cfg: - return cfg - - default = cfg.get("default", ["env_reward"]) - overrides = cfg.get("agent_overrides", {}) - return overrides.get(agent_name, default) - - -def _resolve_gdpo_reward_features( - estimator_config: AdvEstimatorConfig, agent_name: str | None -) -> list[str]: - selected = _resolve_gdpo_reward_feature_selection(estimator_config, agent_name) - if isinstance(selected, dict): - return list(selected) - return list(selected) - - -def _resolve_gdpo_reward_feature_weights( - estimator_config: AdvEstimatorConfig, agent_name: str | None -) -> dict[str, float]: - selected = _resolve_gdpo_reward_feature_selection(estimator_config, agent_name) - if not isinstance(selected, dict): - return {name: 1.0 for name in selected} - - weights = {} - for name, feature_cfg in selected.items(): - if isinstance(feature_cfg, dict): - weights[name] = float(feature_cfg.get("weight", 1.0)) - elif isinstance(feature_cfg, (int, float)): - weights[name] = float(feature_cfg) - else: - weights[name] = 1.0 - return weights - - class GRPOAdvantageEstimator: """GRPO-style advantage estimator with leave-one-out baseline. @@ -193,8 +112,6 @@ def __init__( # Optional per-reward weights w_n for the aggregation A = sum_n w_n * A_n # (paper: https://arxiv.org/abs/2601.05242). None => equal weights (all 1.0). self.reward_weights = estimator_config.reward_weights - self.estimator_config = estimator_config - self.verbose = estimator_config.verbose def compute_advantage( self, @@ -216,14 +133,6 @@ def compute_advantage( Returns: Advantages tensor of shape [batch_size, seq_len]. """ - if "gdpo_reward_features" in repeated_batch: - return self._compute_feature_advantage( - prompt_ids=prompt_ids, - rewards=rewards, - mask=mask, - repeated_batch=repeated_batch, - ) - reward_component_keys = get_gdpo_reward_component_keys(repeated_batch) if len(reward_component_keys) < 2: raise ValueError( @@ -280,161 +189,6 @@ def compute_advantage( return advantages.expand(mask.shape) - def _compute_feature_advantage( - self, - prompt_ids, - rewards, - mask, - repeated_batch, - ): - gdpo_features = repeated_batch["gdpo_reward_features"] - agent_refs = repeated_batch.get("agent_ref", [{} for _ in gdpo_features]) - if len(gdpo_features) != prompt_ids.shape[0]: - raise ValueError( - "gdpo_reward_features must match batch size; " - f"got {len(gdpo_features)} vs {prompt_ids.shape[0]}" - ) - - resolved_features: list[list[str]] = [] - resolved_feature_weights: list[dict[str, float]] = [] - for agent_ref in agent_refs: - agent_name = agent_ref.get("name") if isinstance(agent_ref, dict) else None - resolved_features.append( - _resolve_gdpo_reward_features(self.estimator_config, agent_name) - ) - resolved_feature_weights.append( - _resolve_gdpo_reward_feature_weights( - self.estimator_config, agent_name - ) - ) - self._validate_prompt_group_features( - prompt_ids, resolved_features, resolved_feature_weights - ) - - feature_names: list[str] = [] - seen = set() - for features in resolved_features: - for name in features: - if name not in seen: - seen.add(name) - feature_names.append(name) - if not feature_names: - raise ValueError("GDPO requires at least one reward feature") - - if self.verbose: - self._print_reward_feature_summary( - prompt_ids=prompt_ids, - rewards=rewards, - gdpo_features=gdpo_features, - resolved_features=resolved_features, - resolved_feature_weights=resolved_feature_weights, - ) - - advantage_parts = [] - reward_device = prompt_ids.device - reward_dtype = rewards.dtype if rewards.is_floating_point() else torch.float32 - for name in feature_names: - vals = [] - weights = [] - for selected, feature_dict, feature_weight_dict in zip( - resolved_features, gdpo_features, resolved_feature_weights - ): - if name not in selected: - vals.append(0.0) - weights.append(0.0) - continue - - entry = feature_dict.get(name) if isinstance(feature_dict, dict) else None - vals.append(_feature_reward(entry)) - weights.append(feature_weight_dict.get(name, 1.0)) - - feature_rewards = torch.tensor( - vals, dtype=reward_dtype, device=reward_device - ) - feature_weights = torch.tensor( - weights, dtype=reward_dtype, device=reward_device - ) - baseline, std = calculate_baseline_and_std_per_prompt( - prompt_ids, - feature_rewards, - torch.ones_like(feature_rewards), - leave_one_out_baseline=self.use_leave_one_out_baseline, - ) - adv = (feature_rewards - baseline).unsqueeze(-1) - - if self.normalize_rewards: - epsilon = 1e-6 - non_zero_std_mask = std > 0 - adv[non_zero_std_mask] = adv[non_zero_std_mask] / ( - std.unsqueeze(-1)[non_zero_std_mask] + epsilon - ) - adv = adv * feature_weights.unsqueeze(-1) - advantage_parts.append(adv) - - advantages = sum(advantage_parts) - adv_std = advantages.std() - if adv_std > 0: - advantages = (advantages - advantages.mean()) / adv_std - else: - advantages = advantages - advantages.mean() - - return advantages.to(mask.device).expand(mask.shape) - - @staticmethod - def _validate_prompt_group_features( - prompt_ids, - resolved_features: list[list[str]], - resolved_feature_weights: list[dict[str, float]], - ) -> None: - _, inverse = torch.unique(prompt_ids, dim=0, return_inverse=True) - for group_idx in torch.unique(inverse).cpu().tolist(): - indices = (inverse == group_idx).nonzero(as_tuple=True)[0].cpu().tolist() - group_feature_lists = { - tuple( - (name, resolved_feature_weights[i].get(name, 1.0)) - for name in resolved_features[i] - ) - for i in indices - } - if len(group_feature_lists) > 1: - raise ValueError( - "GDPO reward_features must resolve to one feature spec per prompt group; " - f"got {sorted(group_feature_lists)}" - ) - - @staticmethod - def _print_reward_feature_summary( - prompt_ids, - rewards, - gdpo_features, - resolved_features: list[list[str]], - resolved_feature_weights: list[dict[str, float]], - ) -> None: - print(f"\n{'=' * 70}", flush=True) - print("[Advantage] GDPO reward features", flush=True) - _, inverse = torch.unique(prompt_ids, dim=0, return_inverse=True) - scalar_rewards = rewards.detach().cpu().tolist() - for group_idx in torch.unique(inverse).cpu().tolist(): - indices = (inverse == group_idx).nonzero(as_tuple=True)[0].cpu().tolist() - print(f"\n group {group_idx}", flush=True) - for local_idx, batch_idx in enumerate(indices): - feature_dict = gdpo_features[batch_idx] - if not isinstance(feature_dict, dict): - feature_dict = {} - parts = [ - f" [{local_idx}] scalar_reward={float(scalar_rewards[batch_idx]):.4f}" - ] - for name in resolved_features[batch_idx]: - parts.append( - _format_gdpo_feature_for_log( - name, - feature_dict.get(name), - resolved_feature_weights[batch_idx].get(name, 1.0), - ) - ) - print(" ".join(parts), flush=True) - print(f"{'=' * 70}\n", flush=True) - class ReinforcePlusPlusAdvantageEstimator: """Reinforce++ advantage estimator with optional baseline subtraction and KL penalty in reward. diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 3b9e4f45cf7..3a4955ea644 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -2220,146 +2220,6 @@ def apply_reward_penalties( return counts -def _as_token_id_list(token_ids: Any) -> list[int]: - if token_ids is None: - return [] - if isinstance(token_ids, torch.Tensor): - return [int(x) for x in token_ids.detach().cpu().flatten().tolist()] - if isinstance(token_ids, (list, tuple)): - return [int(x) for x in token_ids] - return [] - - -def _assistant_generated_token_ids(result: dict) -> list[int]: - ids: list[int] = [] - for msg in result.get("message_log", []): - if msg.get("role") == "assistant": - ids.extend(_as_token_id_list(msg.get("token_ids"))) - return ids - - -def _output_generation_text(result: dict) -> str: - chunks: list[str] = [] - output_items = result["full_result"].get("response", {}).get("output", []) - for item in output_items: - gen_str = item.get("generation_str", "") - if isinstance(gen_str, str) and gen_str: - chunks.append(gen_str) - return "".join(chunks) - - -def _count_close_think_tags( - result: dict, - reward_penalty_config: dict[str, Any] | BaseModel | None, - tokenizer: TokenizerType | None = None, -) -> dict[str, Any]: - # Token-based counting requires an explicit think_close token id from the - # reward-penalty config. Token ids are tokenizer-specific, so guessing a - # default here would count unrelated tokens (e.g. newlines) on other - # tokenizers and silently inject large spurious penalties. - think_close_token_id = _get_reward_penalty_token_id( - reward_penalty_config or {}, "think_close" - ) - - assistant_ids = _assistant_generated_token_ids(result) - token_count = ( - sum(1 for token_id in assistant_ids if token_id == think_close_token_id) - if think_close_token_id is not None - else None - ) - - decoded_count = None - if tokenizer is not None and assistant_ids: - try: - decoded = tokenizer.decode(assistant_ids, skip_special_tokens=False) - decoded_count = decoded.count("") - except Exception: - decoded_count = None - - generation_text = _output_generation_text(result) - generation_str_count = generation_text.count("") if generation_text else 0 - - if decoded_count is not None: - close_count = decoded_count - source = "decoded_assistant_tokens" - elif generation_text: - close_count = ( - max(token_count, generation_str_count) - if token_count is not None - else generation_str_count - ) - source = "max_token_or_generation_str" - elif token_count is not None: - close_count = token_count - source = "assistant_token_ids" - else: - # No way to count: no tokenizer decode, no generation text, and no - # configured think_close token id. Report one close tag so the - # think_count_delta feature stays neutral rather than penalizing. - close_count = 1 - source = "unavailable" - - return { - "count": int(close_count), - "source": source, - "token_count": int(token_count) if token_count is not None else None, - "generation_str_count": int(generation_str_count), - "decoded_count": decoded_count, - } - - -def _record_gdpo_think_count_features( - results: list[dict], - reward_penalty_config: dict[str, Any] | BaseModel | None, - tokenizer: TokenizerType | None = None, -) -> None: - for result in results: - close_info = _count_close_think_tags(result, reward_penalty_config, tokenizer) - close_count = close_info["count"] - features = result["full_result"].setdefault("gdpo_reward_features", {}) - features["think_count_delta"] = { - "reward": float(-abs(close_count - 1)), - "adjustment": "format", - "num_close_think_tags": close_count, - "count_source": close_info["source"], - "token_close_think_tags": close_info["token_count"], - "generation_str_close_think_tags": close_info["generation_str_count"], - "decoded_close_think_tags": close_info["decoded_count"], - } - - -def _record_gdpo_env_reward_features(results: list[dict]) -> None: - for result in results: - features = result["full_result"].setdefault("gdpo_reward_features", {}) - reward = float(result["full_result"]["reward"]) - features["env_reward"] = {"reward": reward, "adjustment": None} - features.setdefault( - "length_adjusted_reward", - {"reward": reward, "adjustment": "combined"}, - ) - - -def _calculate_gdpo_reward_feature_metrics(results: list[dict]) -> dict[str, float]: - values_by_feature: dict[str, list[float]] = defaultdict(list) - for result in results: - features = result["full_result"].get("gdpo_reward_features", {}) - if not isinstance(features, dict): - continue - for name, entry in features.items(): - if isinstance(entry, dict) and "reward" in entry: - values_by_feature[name].append(float(entry["reward"])) - - metrics: dict[str, float] = {} - for name, values in values_by_feature.items(): - if not values: - continue - metric_prefix = f"gdpo_{name}" - metrics[f"{metric_prefix}/mean"] = sum(values) / len(values) - metrics[f"{metric_prefix}/min"] = min(values) - metrics[f"{metric_prefix}/max"] = max(values) - return metrics - - def _prepare_nemo_gym_rows( rows: list[dict], generation_config: GenerationConfig, @@ -2765,11 +2625,7 @@ def _postprocess_single_nemo_gym_group( resolved_reward_penalty_config = resolve_reward_penalty_config( reward_penalty_config, tokenizer, thinking_tags=thinking_tags ) - _record_gdpo_think_count_features( - results, resolved_reward_penalty_config, tokenizer - ) penalty_counts = apply_reward_penalties(results, resolved_reward_penalty_config) - _record_gdpo_env_reward_features(results) for nemo_gym_row, result in zip(nemo_gym_rows, results): result["agent_ref"] = nemo_gym_row["agent_ref"] @@ -2923,7 +2779,6 @@ def _postprocess_single_nemo_gym_group( ) rollout_metrics.update(per_agent_metrics) - rollout_metrics.update(_calculate_gdpo_reward_feature_metrics(results)) # Necessary for downstream nemo rl logging/printing. rollout_metrics["mean_gen_tokens_per_sample"] = rollout_metrics[ @@ -2958,9 +2813,6 @@ def _postprocess_single_nemo_gym_group( # stop_strings: NotRequired[list[str]] # Optional stop strings for generation # Extra information not in the DatumSpec used by the GRPO algorithm "total_reward": torch.tensor([r["full_result"]["reward"] for r in results]), - "gdpo_reward_features": [ - r["full_result"].get("gdpo_reward_features", {}) for r in results - ], # Add truncated field to match other rollout paths (reusing hit_max_tokens logic) "truncated": torch.tensor( [m["hit_max_tokens"] for m in all_sample_metrics], dtype=torch.bool diff --git a/nemo_rl/utils/length_adjustments.md b/nemo_rl/utils/length_adjustments.md index 319ecf01212..5099e569b90 100644 --- a/nemo_rl/utils/length_adjustments.md +++ b/nemo_rl/utils/length_adjustments.md @@ -4,16 +4,8 @@ This file documents the length-penalty and length-bonus algorithms implemented i `nemo_rl/utils/length_adjustments.py` -The code supports two usage modes: - -1. **Reward mutation mode** - Configure `grpo.length_bonus`. The length adjustment mutates `full_result["reward"]`. - -2. **GDPO feature mode** - Configure `grpo.adv_estimator.name: gdpo` and put length feature knobs under - `grpo.adv_estimator.reward_features`. The same calculations are recorded in - `full_result["gdpo_reward_features"]` and consumed by GDPO without mutating the scalar - environment reward. +Configure `grpo.length_bonus`. The length adjustments mutate `full_result["reward"]` in +place during rollout postprocessing. All algorithms are resolved per prompt group. Unless otherwise stated, only rollouts with `reward > 0` participate in length comparisons and receive length-based adjustments. @@ -122,45 +114,6 @@ grpo: | `group_length_penalty_profile_gate_field` | Selects which field from the chosen profile-band channel to use as the gate threshold, usually `a`. | | `group_length_penalty_profile_gate_positive_only` | If true, computes the gate mean using only `reward > 0` rollouts; if false, uses all rollouts. | -## GDPO Feature Mode - -In GDPO feature mode, the same feature names can be selected under `reward_features`. - -```yaml -grpo: - adv_estimator: - name: gdpo - reward_features: - default: - env_reward: 1.0 - length_adjusted_reward: - group_total_length_penalty_coeff: 0.1 - think_count_delta: 1.0 -``` - -Feature entries can also be weighted: - -```yaml -grpo: - adv_estimator: - name: gdpo - reward_features: - default: - env_reward: 1.0 - length_adjusted_reward: - group_total_length_penalty_coeff: 0.1 - think_count_delta: 0.5 -``` - -The rollout code records feature metrics to WandB using names like: - -```text -train/gdpo_length_adjusted_reward/mean -train/gdpo_length_adjusted_reward/min -train/gdpo_length_adjusted_reward/max -train/gdpo_think_count_delta/mean -``` - ## Per-Prompt Data Format Some algorithms depend on metadata stored on each training-data row. The rollout code copies @@ -561,82 +514,6 @@ grpo: group_length_penalty_profile_gate_positive_only: true ``` -## Recorded GDPO Feature Names - -The implementation records these feature names in `full_result["gdpo_reward_features"]`: - -- `env_reward` -- `reasoning_bonus` -- `answer_bonus` -- `total_bonus` -- `longest_reasoning_penalty` -- `longest_answer_penalty` -- `longest_total_penalty` -- `group_reasoning_length_penalty_coeff` -- `group_answer_length_penalty_coeff` -- `group_total_length_penalty_coeff` -- `reasoning_zmad_penalty` -- `answer_zmad_penalty` -- `total_zmad_penalty` -- `profiled_length_penalty` -- `profile_band_total` -- `profile_band_reasoning` -- `profile_band_answer` -- `profile_band_delta` -- `length_additive_delta` -- `length_total_delta` -- `length_adjusted_reward` - -`length_adjusted_reward` is the combined length-adjusted scalar that GDPO can use as one reward -feature. `length_additive_delta`, `profile_band_delta`, and `length_total_delta` are derived -summary features. - -## Formatting Feature: think_count_delta - -The GDPO branch also supports a formatting feature: - -```text -think_count_delta = -abs(num_close_think_tags - 1) -``` - -This is not a length-penalty algorithm, but it can be selected alongside length features in GDPO: - -```yaml -grpo: - adv_estimator: - name: gdpo - reward_features: - default: - env_reward: 1.0 - length_adjusted_reward: - group_total_length_penalty_coeff: 0.1 - think_count_delta: 1.0 -``` - -For a format-only GDPO setup, use only the environment reward and the malformed-format feature: - -```yaml -grpo: - adv_estimator: - name: gdpo - reward_features: - default: - env_reward: 1.0 - think_count_delta: 1.0 -``` - -To make the malformed-format feature weaker than the task reward, lower its weight: - -```yaml -grpo: - adv_estimator: - name: gdpo - reward_features: - default: - env_reward: 1.0 - think_count_delta: 0.5 -``` - ## Practical Notes - Most length algorithms act only on positive rollouts (`reward > 0`). @@ -644,8 +521,6 @@ grpo: - Group-relative scaling can reduce average length aggressively because it gives dense per-group pressure. - zMAD is more selective: it only hits high-side outliers. -- GDPO feature mode is useful when you want length or formatting behavior to be represented as a - separate feature instead of mixing it into the scalar environment reward. ## Final Recommendations diff --git a/nemo_rl/utils/length_adjustments.py b/nemo_rl/utils/length_adjustments.py index 0a86bce3ff6..54c4f7b8691 100644 --- a/nemo_rl/utils/length_adjustments.py +++ b/nemo_rl/utils/length_adjustments.py @@ -87,64 +87,6 @@ "group_length_penalty_profile_gate_field", }) -_GDPO_LENGTH_FEATURE_PARAM_KEYS = frozenset({ - "reasoning_bonus", - "answer_bonus", - "total_bonus", - "longest_reasoning_penalty", - "longest_answer_penalty", - "longest_total_penalty", - "group_reasoning_length_penalty_coeff", - "group_answer_length_penalty_coeff", - "group_total_length_penalty_coeff", - "reasoning_zmad_penalty", - "answer_zmad_penalty", - "total_zmad_penalty", - "profiled_length_penalty", - "profile_band_total", - "profile_band_reasoning", - "profile_band_answer", - "group_length_penalty_profile_gate", - "group_length_penalty_profile_gate_channel", - "group_length_penalty_profile_gate_field", - "group_length_penalty_profile_gate_positive_only", -}) - - -def _resolve_gdpo_feature_log_names_for_agent( - adv_cfg: dict[str, Any], agent_name: str -) -> list[str]: - reward_features = adv_cfg.get("reward_features", ["env_reward"]) - if isinstance(reward_features, (list, tuple)): - return list(reward_features) - if not isinstance(reward_features, dict): - return ["env_reward"] - if "default" not in reward_features and "agent_overrides" not in reward_features: - return list(reward_features) - - selected = reward_features.get("agent_overrides", {}).get( - agent_name, reward_features.get("default", ["env_reward"]) - ) - if isinstance(selected, dict): - return list(selected) - if isinstance(selected, (list, tuple)): - return list(selected) - return ["env_reward"] - - -def _set_gdpo_reward_feature( - result: dict[str, Any], - name: str, - reward: float, - adjustment: str | None, - **metadata: Any, -) -> None: - features = result["full_result"].setdefault("gdpo_reward_features", {}) - entry = {"reward": float(reward), "adjustment": adjustment} - entry.update(metadata) - features[name] = entry - - def _extract_reasoning_and_answer_text(result: dict[str, Any]) -> tuple[str, str]: """Extract reasoning and answer text from the Response API output items.""" fr = result.get("full_result", {}) @@ -176,78 +118,6 @@ def _extract_reasoning_and_answer_text(result: dict[str, Any]) -> tuple[str, str return reasoning_text, answer_text -def _extract_gdpo_length_feature_params(feature_cfg: Any) -> dict[str, Any]: - if not isinstance(feature_cfg, dict): - return {} - params = {} - for key, value in feature_cfg.items(): - if key not in _GDPO_LENGTH_FEATURE_PARAM_KEYS: - continue - if key in _BOOL_PARAM_KEYS: - if isinstance(value, bool): - params[key] = value - continue - if key in _STR_PARAM_KEYS: - if isinstance(value, str): - params[key] = value - continue - if isinstance(value, bool) or not isinstance(value, (int, float)): - continue - params[key] = float(value) - return params - - -def _merge_gdpo_reward_features_into_length_cfg( - grpo_config: dict[str, Any], -) -> dict[str, Any]: - length_cfg = dict(grpo_config.get("length_bonus", {}) or {}) - has_explicit_length_cfg = bool(length_cfg) - adv_cfg = grpo_config.get("adv_estimator", {}) or {} - if adv_cfg.get("name") != "gdpo": - return length_cfg - - reward_features = adv_cfg.get("reward_features") - if not isinstance(reward_features, dict): - return length_cfg - - if "default" in reward_features or "agent_overrides" in reward_features: - default_features = reward_features.get("default", {}) - agent_features = reward_features.get("agent_overrides", {}) - else: - default_features = reward_features - agent_features = {} - - default_params = _extract_gdpo_length_feature_params(default_features) - agent_params = { - agent_name: _extract_gdpo_length_feature_params(features) - for agent_name, features in agent_features.items() - } - agent_params = {k: v for k, v in agent_params.items() if v} - if not default_params and not agent_params: - return length_cfg - - if not has_explicit_length_cfg: - length_cfg["_gdpo_feature_only"] = True - - default_cfg = dict(length_cfg.get("default", {}) or {}) - default_cfg.update(default_params) - default_cfg.setdefault("enabled", True) - length_cfg["default"] = default_cfg - - overrides = dict(length_cfg.get("agent_overrides", {}) or {}) - for agent_name, params in agent_params.items(): - agent_cfg = dict(overrides.get(agent_name, {}) or {}) - agent_cfg.update(params) - agent_cfg.setdefault("enabled", True) - overrides[agent_name] = agent_cfg - if overrides: - length_cfg["agent_overrides"] = overrides - - if adv_cfg.get("verbose", False): - length_cfg["_gdpo_feature_verbose"] = True - return length_cfg - - def apply_group_length_adjustments( results: list[dict[str, Any]], master_config: dict[str, Any], @@ -255,12 +125,9 @@ def apply_group_length_adjustments( ) -> None: """Apply per-prompt-group length bonuses/penalties. - Reads ``grpo.length_bonus`` and GDPO ``reward_features`` for configuration. - No-ops when no length-adjustment feature is enabled. - - Calculates all adjustments first and records GDPO reward features. Explicit - ``grpo.length_bonus`` configs mutate scalar rewards; GDPO-only feature - configs do not. + Reads ``grpo.length_bonus`` for configuration and mutates + ``full_result["reward"]`` in place. No-ops when no length-adjustment + feature is enabled. Args: results: List of per-generation result dicts. @@ -268,20 +135,14 @@ def apply_group_length_adjustments( tokenizer: Tokenizer for computing reasoning/answer token counts. """ grpo_config = master_config.get("grpo", {}) - length_cfg = _merge_gdpo_reward_features_into_length_cfg(grpo_config) + length_cfg = dict(grpo_config.get("length_bonus", {}) or {}) if not length_cfg: return default_cfg = length_cfg.get("default", {}) agents_cfg = length_cfg.get("agent_overrides") global_band = _resolve_global_profile_band(length_cfg.get("profile_band")) - gdpo_feature_only = bool(length_cfg.get("_gdpo_feature_only", False)) - verbose = bool(length_cfg.get("verbose", False)) and not gdpo_feature_only - gdpo_feature_verbose = bool( - length_cfg.get("_gdpo_feature_verbose", False) - or (length_cfg.get("verbose", False) and gdpo_feature_only) - ) - should_mutate_reward = not gdpo_feature_only + verbose = bool(length_cfg.get("verbose", False)) if not default_cfg.get("enabled", False) and not agents_cfg and not global_band: return @@ -546,7 +407,7 @@ def apply_group_length_adjustments( print(f"{'=' * 70}\n", flush=True) - # Phase 3: apply additive adjustments and record GDPO reward features + # Phase 3: apply additive adjustments additive_base_rewards = [0.0] * n for i, r in enumerate(results): # The profiled-length penalty stacks only on rollouts whose group @@ -556,48 +417,7 @@ def apply_group_length_adjustments( profiled_adj = all_profiled_length_adj[i] if all_adjustments[i] >= 0 else 0.0 additive_delta = all_adjustments[i] + profiled_adj additive_base_rewards[i] = original_rewards[i] + additive_delta - if "env_reward" not in r["full_result"].setdefault( - "gdpo_reward_features", {} - ): - _set_gdpo_reward_feature(r, "env_reward", original_rewards[i], None) - - for name, value in ( - ("reasoning_bonus", all_reasoning_bonus[i]), - ("answer_bonus", all_answer_bonus[i]), - ("total_bonus", all_total_bonus[i]), - ("longest_reasoning_penalty", all_reasoning_longest_pen[i]), - ("longest_answer_penalty", all_answer_longest_pen[i]), - ("longest_total_penalty", all_total_longest_pen[i]), - ("group_reasoning_length_penalty_coeff", all_reasoning_adj[i]), - ("group_answer_length_penalty_coeff", all_answer_adj[i]), - ("group_total_length_penalty_coeff", all_total_adj[i]), - ("reasoning_zmad_penalty", all_zmad_reasoning_adj[i]), - ("answer_zmad_penalty", all_zmad_answer_adj[i]), - ("total_zmad_penalty", all_zmad_total_adj[i]), - ("profiled_length_penalty", profiled_adj), - ): - _set_gdpo_reward_feature(r, name, value, "additive") - - _set_gdpo_reward_feature( - r, "profile_band_total", 0.0, "multiplicative", multiplier=1.0 - ) - _set_gdpo_reward_feature( - r, "profile_band_reasoning", 0.0, "multiplicative", multiplier=1.0 - ) - _set_gdpo_reward_feature( - r, "profile_band_answer", 0.0, "multiplicative", multiplier=1.0 - ) - _set_gdpo_reward_feature(r, "profile_band_delta", 0.0, "derived_sum") - _set_gdpo_reward_feature( - r, "length_additive_delta", additive_delta, "derived_sum" - ) - _set_gdpo_reward_feature(r, "length_total_delta", additive_delta, "derived_sum") - _set_gdpo_reward_feature( - r, "length_adjusted_reward", additive_base_rewards[i], "combined" - ) - - if should_mutate_reward: - r["full_result"]["reward"] = additive_base_rewards[i] + r["full_result"]["reward"] = additive_base_rewards[i] # Phase 4: apply per-prompt profile_band multipliers (correct rollouts only). _apply_profile_band_multipliers( @@ -611,64 +431,9 @@ def apply_group_length_adjustments( agents_cfg=agents_cfg, defaults=defaults, num_gens=num_gens, - should_mutate_reward=should_mutate_reward, global_band=global_band, ) - if verbose or gdpo_feature_verbose: - _print_gdpo_reward_feature_summary( - results=results, - agent_names=agent_names, - original_rewards=original_rewards, - num_gens=num_gens, - adv_cfg=grpo_config.get("adv_estimator", {}) or {}, - ) - - -def _print_gdpo_reward_feature_summary( - results: list[dict[str, Any]], - agent_names: list[str], - original_rewards: list[float], - num_gens: int, - adv_cfg: dict[str, Any], -) -> None: - print(f"\n{'=' * 70}", flush=True) - print("[Rollout] GDPO reward features", flush=True) - n = len(results) - for g in range(0, n, num_gens): - group_size = min(num_gens, n - g) - agent_name = agent_names[g] - print( - f"\n group {g // num_gens} agent={agent_name}", - flush=True, - ) - feature_names = _resolve_gdpo_feature_log_names_for_agent(adv_cfg, agent_name) - for k in range(group_size): - idx = g + k - fr = results[idx]["full_result"] - features = fr.get("gdpo_reward_features", {}) - parts = [ - f" [{k}] original_reward={float(original_rewards[idx]):.4f}", - f"final_reward={float(fr['reward']):.4f}", - ] - for name in feature_names: - entry = features.get(name) - if not isinstance(entry, dict): - continue - reward = float(entry.get("reward", 0.0)) - adjustment = entry.get("adjustment") - multiplier = entry.get("multiplier") - suffix = ( - f":{float(multiplier):.4f}" - if multiplier is not None - else "" - ) - parts.append( - f"{name}={reward:+.4f}" - f"({adjustment}{suffix})" - ) - print(" ".join(parts), flush=True) - print(f"{'=' * 70}\n", flush=True) def _resolve_global_profile_band(pb_cfg: Any) -> dict[str, dict[str, Any]]: @@ -724,16 +489,13 @@ def _apply_profile_band_multipliers( agents_cfg: dict[str, Any] | None, defaults: dict[str, Any], num_gens: int, - should_mutate_reward: bool, global_band: dict[str, dict[str, Any]] | None = None, ) -> None: """Apply per-channel profile_band multipliers to correct rollouts. Each enabled channel contributes a multiplier in [0.0, 1.0] derived from the per-row {a, b, f} block, falling back to ``length_bonus.profile_band.defaults`` - for channels the row does not provide. Records multiplicative deltas as - additive-equivalent GDPO reward features, and mutates scalar rewards for - length-bonus configs. + for channels the row does not provide. Mutates scalar rewards in place. Skips any group where the low-effort bypass already replaced the reward (parity with Phase 1 of ``apply_group_length_adjustments``). @@ -778,42 +540,7 @@ def _apply_profile_band_multipliers( answer_delta = current_reward * answer_m - current_reward current_reward += answer_delta - profile_band_delta = current_reward - base_rewards[idx] - _set_gdpo_reward_feature( - results[idx], - "profile_band_total", - total_delta, - "multiplicative", - multiplier=total_m, - ) - _set_gdpo_reward_feature( - results[idx], - "profile_band_reasoning", - reasoning_delta, - "multiplicative", - multiplier=reasoning_m, - ) - _set_gdpo_reward_feature( - results[idx], - "profile_band_answer", - answer_delta, - "multiplicative", - multiplier=answer_m, - ) - _set_gdpo_reward_feature( - results[idx], "profile_band_delta", profile_band_delta, "derived_sum" - ) - _set_gdpo_reward_feature( - results[idx], - "length_total_delta", - current_reward - original_rewards[idx], - "derived_sum", - ) - _set_gdpo_reward_feature( - results[idx], "length_adjusted_reward", current_reward, "combined" - ) - if should_mutate_reward: - results[idx]["full_result"]["reward"] = current_reward + results[idx]["full_result"]["reward"] = current_reward def _band_multiplier(rl: int, ch: dict[str, Any] | None) -> float: diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index b26f8502681..44debcd1273 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -4719,55 +4719,6 @@ def test_grpo_advantage_estimator_small_nonzero_std(): assert result[0, 0] * result[1, 0] < 0 -def test_gdpo_advantage_estimator_feature_weight_changes_advantage(): - """Feature weights are applied after per-feature normalization.""" - - def compute_advantage(think_weight): - estimator_config = AdvEstimatorConfig( - use_leave_one_out_baseline=False, - normalize_rewards=True, - reward_features={ - "default": { - "length_adjusted_reward": {}, - "think_count_delta": {"weight": think_weight}, - } - }, - ) - estimator = GDPOAdvantageEstimator(estimator_config, ClippedPGLossConfig()) - prompt_ids = torch.tensor([[0], [0], [0], [0]]) - rewards = torch.zeros(4) - mask = torch.ones(4, 1) - repeated_batch = { - "agent_ref": [{"name": "math_with_judge_simple_agent"} for _ in range(4)], - "gdpo_reward_features": [ - { - "length_adjusted_reward": {"reward": 1.0}, - "think_count_delta": {"reward": 0.0}, - }, - { - "length_adjusted_reward": {"reward": 1.0}, - "think_count_delta": {"reward": -1.0}, - }, - { - "length_adjusted_reward": {"reward": 0.0}, - "think_count_delta": {"reward": 0.0}, - }, - { - "length_adjusted_reward": {"reward": 0.0}, - "think_count_delta": {"reward": -1.0}, - }, - ], - } - return estimator.compute_advantage( - prompt_ids, rewards, mask, repeated_batch=repeated_batch - )[:, 0] - - without_think_feature = compute_advantage(0.0) - with_think_feature = compute_advantage(1.0) - - assert not torch.allclose(without_think_feature, with_think_feature) - - # ============================================================================ # Tests for ReinforcePlusPlusAdvantageEstimator class # ============================================================================ diff --git a/tests/unit/experience/test_reward_penalties.py b/tests/unit/experience/test_reward_penalties.py index 70de89e2def..ea5f4921519 100644 --- a/tests/unit/experience/test_reward_penalties.py +++ b/tests/unit/experience/test_reward_penalties.py @@ -19,10 +19,8 @@ from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.rollouts import ( - _calculate_gdpo_reward_feature_metrics, _extract_mask_sample_flags, _postprocess_single_nemo_gym_group, - _record_gdpo_think_count_features, apply_reward_penalties, resolve_reward_penalty_config, should_mask_flagged_samples, @@ -829,123 +827,6 @@ def test_no_generation_str_skipped(self): assert result["full_result"]["reward"] == 1.0 -# ===================================================================== -# GDPO feature: think_count_delta -# ===================================================================== - - -class TestGdpoThinkCountFeature: - CFG = {"token_ids": {"think_close": 13}} - - def test_one_close_think_has_zero_delta(self): - result = _make_result( - message_log=[ - _msg("user", [100, 12]), - _msg("assistant", [300, 13, 400]), - ] - ) - _record_gdpo_think_count_features([result], self.CFG) - entry = result["full_result"]["gdpo_reward_features"]["think_count_delta"] - assert entry["reward"] == 0.0 - assert entry["num_close_think_tags"] == 1 - - def test_missing_close_think_has_negative_delta(self): - result = _make_result( - message_log=[ - _msg("user", [100, 12]), - _msg("assistant", [300, 400]), - ] - ) - _record_gdpo_think_count_features([result], self.CFG) - entry = result["full_result"]["gdpo_reward_features"]["think_count_delta"] - assert entry["reward"] == -1.0 - assert entry["num_close_think_tags"] == 0 - - def test_double_close_think_has_negative_delta(self): - result = _make_result( - message_log=[ - _msg("user", [100, 12]), - _msg("assistant", [300, 13, 400, 13]), - ] - ) - _record_gdpo_think_count_features([result], self.CFG) - entry = result["full_result"]["gdpo_reward_features"]["think_count_delta"] - assert entry["reward"] == -1.0 - assert entry["num_close_think_tags"] == 2 - - def test_generation_str_fallback_counts_close_thinks(self): - result = _make_result( - output_items=[ - _message_item("answer", generation_str="a b ") - ], - message_log=[ - _msg("user", [100, 12]), - _msg("assistant", [300, 400]), - ], - ) - _record_gdpo_think_count_features([result], self.CFG) - entry = result["full_result"]["gdpo_reward_features"]["think_count_delta"] - assert entry["reward"] == -1.0 - assert entry["num_close_think_tags"] == 2 - - def test_unconfigured_token_id_is_neutral_without_text(self): - # No think_close token id configured, no tokenizer, no generation - # text: counting is impossible, so the feature must stay neutral - # (count=1 -> delta=0) instead of guessing a token id. - result = _make_result( - message_log=[ - _msg("user", [100, 12]), - _msg("assistant", [300, 13, 400, 13]), - ] - ) - _record_gdpo_think_count_features([result], {}) - entry = result["full_result"]["gdpo_reward_features"]["think_count_delta"] - assert entry["reward"] == 0.0 - assert entry["count_source"] == "unavailable" - - def test_unconfigured_token_id_uses_generation_str(self): - # Without a configured token id the generation string count is used - # alone (token ids of other tokenizers must not leak into the count). - result = _make_result( - output_items=[ - _message_item("answer", generation_str="a b") - ], - message_log=[ - _msg("user", [100, 12]), - _msg("assistant", [300, 13, 400, 13]), - ], - ) - _record_gdpo_think_count_features([result], {}) - entry = result["full_result"]["gdpo_reward_features"]["think_count_delta"] - assert entry["reward"] == 0.0 - assert entry["num_close_think_tags"] == 1 - - def test_generic_gdpo_feature_metrics(self): - r1 = _make_result() - r1["full_result"]["gdpo_reward_features"] = { - "length_adjusted_reward": {"reward": 1.0}, - "think_count_delta": {"reward": 0.0}, - } - r2 = _make_result() - r2["full_result"]["gdpo_reward_features"] = { - "length_adjusted_reward": {"reward": -0.5}, - "think_count_delta": {"reward": -2.0}, - } - r3 = _make_result() - r3["full_result"]["gdpo_reward_features"] = { - "length_adjusted_reward": {"reward": 0.5}, - } - - metrics = _calculate_gdpo_reward_feature_metrics([r1, r2, r3]) - - assert metrics["gdpo_length_adjusted_reward/min"] == -0.5 - assert metrics["gdpo_length_adjusted_reward/max"] == 1.0 - assert metrics["gdpo_length_adjusted_reward/mean"] == (1.0 - 0.5 + 0.5) / 3 - assert metrics["gdpo_think_count_delta/min"] == -2.0 - assert metrics["gdpo_think_count_delta/max"] == 0.0 - assert metrics["gdpo_think_count_delta/mean"] == -1.0 - - # ===================================================================== # Cross-cutting: multiple penalties, config gating, batch behavior # ===================================================================== @@ -1021,7 +902,6 @@ def test_batch_of_results_mixed(self): TestPenalizeEmptyFinalAnswer, TestPenalizeUnwantedTokens, TestPenalizeMultiEndThink, - TestGdpoThinkCountFeature, TestCrossCutting, ] From 87bea4fc8ba109878c3f47896714179ab4f9286b Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 20:11:49 -0700 Subject: [PATCH 13/27] fix: keep validation metrics on raw env reward MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The length-adjustment code groups by the training stride (num_generations_per_prompt), which mismatches val_num_generations_per_prompt — groups straddled distinct prompts and shaped rewards corrupted accuracy/pass_k (review finding on #3852). Co-Authored-By: Claude Fable 5 --- nemo_rl/algorithms/grpo.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index f297bf304b7..e971fadd4f4 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -3939,7 +3939,11 @@ def validate( greedy=False, effort_config=_get_effort_config(master_config), reward_penalty_config=master_config.reward_penalties, - length_adjustment_config=master_config.grpo.model_dump(), + # No length_adjustment_config here: validation metrics + # (accuracy/pass_k) must reflect the raw env reward, and the + # adjustment code groups by the TRAINING stride + # (num_generations_per_prompt), which does not match + # val_num_generations_per_prompt. thinking_tags=get_nemo_gym_thinking_tags(master_config.env), mask_env_flagged_samples=should_mask_flagged_samples( master_config.env From bbcc4d9df5ca2ec02dfc9269bc8a0ee40260c32f Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 20:11:49 -0700 Subject: [PATCH 14/27] fix: band multiplier negative-base guard, config validation, top_percentile default Review findings on #3852: - profile_band multipliers now skip rollouts whose additive-adjusted base is <= 0: scaling a negative base by m < 1 RAISED rewards for longer rollouts (incentive inversion) - unknown keys anywhere under grpo.length_bonus now raise ValueError instead of being silently ignored - top_percentile gets one canonical default (0.5, matching the docs); the previous generic 0.0 default made the later setdefault(0.2) dead code - global profile_band defaults no longer override an explicitly configured profile_band_: false Adds 4 regression tests (30 total in the file). Co-Authored-By: Claude Fable 5 --- nemo_rl/utils/length_adjustments.py | 56 +++++++++++++++- tests/unit/utils/test_length_adjustments.py | 74 +++++++++++++++++++++ 2 files changed, 127 insertions(+), 3 deletions(-) diff --git a/nemo_rl/utils/length_adjustments.py b/nemo_rl/utils/length_adjustments.py index 54c4f7b8691..0ce9cc1975d 100644 --- a/nemo_rl/utils/length_adjustments.py +++ b/nemo_rl/utils/length_adjustments.py @@ -118,6 +118,47 @@ def _extract_reasoning_and_answer_text(result: dict[str, Any]) -> tuple[str, str return reasoning_text, answer_text +_TOP_LEVEL_LENGTH_BONUS_KEYS = frozenset( + {"verbose", "default", "agent_overrides", "profile_band"} +) +_PROFILE_BAND_BLOCK_KEYS = frozenset({"enabled", "defaults"}) +_PROFILE_BAND_CHANNELS = frozenset({"total", "reasoning", "answer"}) + + +def _reject_unknown_length_bonus_keys(length_cfg: dict[str, Any]) -> None: + """Raise on unknown config keys instead of silently ignoring them.""" + + def _check(block: Any, allowed: frozenset, where: str) -> None: + if not isinstance(block, dict): + return + unknown = sorted(set(block) - allowed) + if unknown: + raise ValueError( + f"Unknown key(s) {unknown} in {where}; allowed: {sorted(allowed)}" + ) + + _check(length_cfg, _TOP_LEVEL_LENGTH_BONUS_KEYS, "grpo.length_bonus") + _check( + length_cfg.get("default"), frozenset(_PARAM_KEYS), "grpo.length_bonus.default" + ) + agents_cfg = length_cfg.get("agent_overrides") + if isinstance(agents_cfg, dict): + for agent_name, overrides in agents_cfg.items(): + _check( + overrides, + frozenset(_PARAM_KEYS), + f"grpo.length_bonus.agent_overrides.{agent_name}", + ) + pb_cfg = length_cfg.get("profile_band") + if isinstance(pb_cfg, dict): + _check(pb_cfg, _PROFILE_BAND_BLOCK_KEYS, "grpo.length_bonus.profile_band") + _check( + pb_cfg.get("defaults"), + _PROFILE_BAND_CHANNELS, + "grpo.length_bonus.profile_band.defaults", + ) + + def apply_group_length_adjustments( results: list[dict[str, Any]], master_config: dict[str, Any], @@ -139,6 +180,7 @@ def apply_group_length_adjustments( if not length_cfg: return + _reject_unknown_length_bonus_keys(length_cfg) default_cfg = length_cfg.get("default", {}) agents_cfg = length_cfg.get("agent_overrides") global_band = _resolve_global_profile_band(length_cfg.get("profile_band")) @@ -165,13 +207,16 @@ def apply_group_length_adjustments( defaults[k] = default_cfg.get(k, 2) elif k == "profiled_length_n_std": defaults[k] = default_cfg.get(k, 1.0) + elif k == "top_percentile": + defaults[k] = default_cfg.get(k, 0.5) else: defaults[k] = default_cfg.get(k, 0.0) - defaults.setdefault("top_percentile", 0.2) # Channels listed under length_bonus.profile_band.defaults are implicitly - # enabled; per-agent overrides can still disable them. + # enabled — unless the user explicitly configured the channel flag, which + # always wins (e.g. profile_band_total: false stays false). for _ch in global_band: - defaults[f"profile_band_{_ch}"] = True + if f"profile_band_{_ch}" not in default_cfg: + defaults[f"profile_band_{_ch}"] = True n = len(results) original_rewards = [r["full_result"]["reward"] for r in results] @@ -526,6 +571,11 @@ def _apply_profile_band_multipliers( # Gate on the env reward (correct rollouts only). if original_rewards[idx] <= 0: continue + # Additive penalties can push the base below zero; multiplying a + # negative base by m < 1 would RAISE the reward for longer + # rollouts. Scale only the non-negative part. + if base_rewards[idx] <= 0: + continue current_reward = base_rewards[idx] total_m = _band_multiplier(total_lengths[idx], ch_total) diff --git a/tests/unit/utils/test_length_adjustments.py b/tests/unit/utils/test_length_adjustments.py index ffe0e23ba7e..69097948d0b 100644 --- a/tests/unit/utils/test_length_adjustments.py +++ b/tests/unit/utils/test_length_adjustments.py @@ -381,3 +381,77 @@ def test_one_pass_allowed_when_min_samples_is_one(self): ] apply_group_length_adjustments(results, self.cfg(min_samples=1)) assert rewards_of(results) == pytest.approx([1.0, 0.7]) + + +class TestReviewFixes: + """Regression tests for review findings on PR #3852.""" + + def test_band_multiplier_skips_negative_base(self): + # A flat additive penalty pushes the correct-but-long rollout's base + # negative; the band multiplier must NOT scale it (base * m - base > 0 + # for base < 0 would make longer rollouts score HIGHER). + band = {"total": {"a": 1, "b": 2, "f": 0.1}} # length 20 -> m = 0.1 + results = [ + make_result("12345", "12345", 1.0, band=band), + make_result("1234567890", "1234567890", 1.0, band=band), + ] + # zMAD-free setup: use profiled penalty to force one base negative. + results[1]["profiled_rewards"] = [1, 1] + results[1]["profiled_output_lengths"] = [5, 5] + results[0]["profiled_rewards"] = [1, 1] + results[0]["profiled_output_lengths"] = [5, 5] + cfg = make_config( + default={ + "enabled": True, + "profiled_length_penalty": 1.5, # threshold 5: len 20 -> base -0.5 + "profile_band_total": True, + } + ) + apply_group_length_adjustments(results, cfg) + r_short, r_long = rewards_of(results) + # Short rollout (len 10 >= 5) also penalized to -0.5; both bases are + # negative, so the band phase must leave both untouched. + assert r_short == pytest.approx(-0.5) + assert r_long == pytest.approx(-0.5) + assert r_long <= r_short # longer must never beat shorter + + def test_explicit_false_channel_not_overridden_by_global_defaults(self): + cfg = make_config( + default={"enabled": True, "profile_band_total": False}, + profile_band={ + "enabled": True, + "defaults": {"total": {"a": 10, "b": 20, "f": 0.5}}, + }, + ) + results = [ + make_result("1234567890123456789012345", "12345", 1.0), # total 30 + make_result("12345", "12345", 1.0), + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_unknown_key_raises(self): + cfg = make_config(default={"enabled": True, "group_total_length_coeff": 0.1}) + results = [make_result("12345", "12345", 1.0), make_result("123", "123", 1.0)] + try: + apply_group_length_adjustments(results, cfg) + except ValueError as e: + assert "group_total_length_coeff" in str(e) + else: + raise AssertionError("expected ValueError for unknown config key") + + def test_top_percentile_default_is_half(self): + # With 2 positive scorers (rewards 1.0 and 0.9) and default + # top_percentile=0.5, only the best scorer is a "top scorer" — the + # longest-total penalty needs >= 2 eligible top scorers, so it no-ops. + # With the old effective default of 0.0 the behavior was the same, but + # an explicit 1.0 must include both and penalize the longer one. + results = [ + make_result("12345", "12345", 1.0), + make_result("1234567890", "1234567890", 0.9), + ] + cfg = make_config( + default={"enabled": True, "longest_total_penalty": 0.2, "top_percentile": 1.0} + ) + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 0.7]) From 7a4d88897ebfa26458ce48912b0461d727fca411 Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 20:33:52 -0700 Subject: [PATCH 15/27] feat: clamp length-adjusted rewards at 0 for correct rollouts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adopt the clamping alternative from review (#3852): stacked flat penalties exceeding the reward now produce 0, never a negative value — length penalties may wipe a correct rollout's reward out but not flip its sign. The clamp applies only to originally-positive rollouts, so an env's own negative rewards pass through untouched. The base<=0 gate in the band phase stays as a second line of defense against the multiplier inversion. Co-Authored-By: Claude Fable 5 --- nemo_rl/utils/length_adjustments.md | 4 ++ nemo_rl/utils/length_adjustments.py | 6 +++ tests/unit/utils/test_length_adjustments.py | 55 +++++++++++++++------ 3 files changed, 50 insertions(+), 15 deletions(-) diff --git a/nemo_rl/utils/length_adjustments.md b/nemo_rl/utils/length_adjustments.md index 5099e569b90..1fa6bc8667a 100644 --- a/nemo_rl/utils/length_adjustments.md +++ b/nemo_rl/utils/length_adjustments.md @@ -517,6 +517,10 @@ grpo: ## Practical Notes - Most length algorithms act only on positive rollouts (`reward > 0`). +- Additive penalties can wipe a correct rollout's reward out but never flip its sign: the + adjusted reward is clamped at 0 for originally-positive rollouts (stacked flat penalties + exceeding the reward produce 0, not a negative value). Environments' own negative rewards + pass through untouched. - `profile_band_*` multipliers also apply only to originally correct rollouts. - Group-relative scaling can reduce average length aggressively because it gives dense per-group pressure. diff --git a/nemo_rl/utils/length_adjustments.py b/nemo_rl/utils/length_adjustments.py index 0ce9cc1975d..4877b92e4a5 100644 --- a/nemo_rl/utils/length_adjustments.py +++ b/nemo_rl/utils/length_adjustments.py @@ -462,6 +462,12 @@ def apply_group_length_adjustments( profiled_adj = all_profiled_length_adj[i] if all_adjustments[i] >= 0 else 0.0 additive_delta = all_adjustments[i] + profiled_adj additive_base_rewards[i] = original_rewards[i] + additive_delta + # Stacked flat penalties can exceed the reward itself. Length penalties + # may wipe a correct rollout's reward out, but never flip its sign — + # clamp at 0. Only for originally-positive rollouts: an env's own + # negative reward must pass through untouched. + if original_rewards[i] > 0 and additive_base_rewards[i] < 0: + additive_base_rewards[i] = 0.0 r["full_result"]["reward"] = additive_base_rewards[i] # Phase 4: apply per-prompt profile_band multipliers (correct rollouts only). diff --git a/tests/unit/utils/test_length_adjustments.py b/tests/unit/utils/test_length_adjustments.py index 69097948d0b..502c8fedd55 100644 --- a/tests/unit/utils/test_length_adjustments.py +++ b/tests/unit/utils/test_length_adjustments.py @@ -386,33 +386,58 @@ def test_one_pass_allowed_when_min_samples_is_one(self): class TestReviewFixes: """Regression tests for review findings on PR #3852.""" - def test_band_multiplier_skips_negative_base(self): - # A flat additive penalty pushes the correct-but-long rollout's base - # negative; the band multiplier must NOT scale it (base * m - base > 0 - # for base < 0 would make longer rollouts score HIGHER). - band = {"total": {"a": 1, "b": 2, "f": 0.1}} # length 20 -> m = 0.1 + def test_flat_penalty_exceeding_reward_clamps_at_zero(self): + # Stacked/oversized flat penalties may wipe a correct rollout's reward + # out but never flip its sign: 1.0 - 1.5 clamps to 0.0, not -0.5. + results = [ + make_result("1234", "1234", 1.0), # len 8 < threshold 10: untouched + make_result("1234567890", "1234567890", 1.0), # len 20 >= 10: clamped + ] + for r in results: + r["profiled_rewards"] = [1, 1] + r["profiled_output_lengths"] = [10, 10] + cfg = make_config( + default={"enabled": True, "profiled_length_penalty": 1.5} + ) + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 0.0]) + + def test_negative_env_reward_not_raised_by_clamp(self): + # The clamp applies only to originally-positive rollouts: an env's own + # negative reward must pass through untouched. + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.1} + ) + results = [ + make_result("12345", "12345", -1.0), + make_result("1234567890", "1234567890", 1.0), + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results)[0] == pytest.approx(-1.0) + + def test_band_multiplier_never_rewards_length_on_penalized_base(self): + # Flat penalty exceeds the reward, so bases clamp to 0; the band phase + # must leave them at 0 (base * m - base > 0 for base < 0 would have + # made longer rollouts score HIGHER pre-clamp). + band = {"total": {"a": 1, "b": 2, "f": 0.1}} results = [ make_result("12345", "12345", 1.0, band=band), make_result("1234567890", "1234567890", 1.0, band=band), ] - # zMAD-free setup: use profiled penalty to force one base negative. - results[1]["profiled_rewards"] = [1, 1] - results[1]["profiled_output_lengths"] = [5, 5] - results[0]["profiled_rewards"] = [1, 1] - results[0]["profiled_output_lengths"] = [5, 5] + for r in results: + r["profiled_rewards"] = [1, 1] + r["profiled_output_lengths"] = [5, 5] cfg = make_config( default={ "enabled": True, - "profiled_length_penalty": 1.5, # threshold 5: len 20 -> base -0.5 + "profiled_length_penalty": 1.5, # threshold 5: both clamped to 0 "profile_band_total": True, } ) apply_group_length_adjustments(results, cfg) r_short, r_long = rewards_of(results) - # Short rollout (len 10 >= 5) also penalized to -0.5; both bases are - # negative, so the band phase must leave both untouched. - assert r_short == pytest.approx(-0.5) - assert r_long == pytest.approx(-0.5) + assert r_short == pytest.approx(0.0) + assert r_long == pytest.approx(0.0) assert r_long <= r_short # longer must never beat shorter def test_explicit_false_channel_not_overridden_by_global_defaults(self): From b4d4b73c1e8b0c3c5e2241047c429679c67e428f Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 21:45:30 -0700 Subject: [PATCH 16/27] refactor: gate per-row metadata copying on an actual length_bonus config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The copy loop ran for every NeMo-Gym rollout; now the postprocess block is a strict no-op unless grpo.length_bonus is configured. Also drop the profiled_reasoning_lengths / profiled_answer_lengths copies — nothing in the runtime consumes them (only the offline dataset builder reads those names, directly from the JSONL). Co-Authored-By: Claude Fable 5 --- nemo_rl/experience/rollouts.py | 31 ++++++++++++++----------------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 3a4955ea644..3fe07d52675 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -2627,29 +2627,26 @@ def _postprocess_single_nemo_gym_group( ) penalty_counts = apply_reward_penalties(results, resolved_reward_penalty_config) - for nemo_gym_row, result in zip(nemo_gym_rows, results): - result["agent_ref"] = nemo_gym_row["agent_ref"] - result["profiled_rewards"] = nemo_gym_row.get("profiled_rewards") - result["profiled_output_lengths"] = nemo_gym_row.get( - "profiled_output_lengths" - ) - result["profiled_reasoning_lengths"] = nemo_gym_row.get( - "profiled_reasoning_lengths" - ) - result["profiled_answer_lengths"] = nemo_gym_row.get( - "profiled_answer_lengths" - ) - result["profile_band"] = nemo_gym_row.get("profile_band") - if length_adjustment_config is not None: grpo_config = ( length_adjustment_config.model_dump() if isinstance(length_adjustment_config, BaseModel) else dict(length_adjustment_config) ) - apply_group_length_adjustments( - results, {"grpo": grpo_config}, tokenizer=tokenizer - ) + # Callers pass the whole grpo config block; runs without a + # grpo.length_bonus section are untouched by this block. + if grpo_config.get("length_bonus"): + # Copy the per-row fields the length adjustments consume. + for nemo_gym_row, result in zip(nemo_gym_rows, results): + result["agent_ref"] = nemo_gym_row["agent_ref"] + result["profiled_rewards"] = nemo_gym_row.get("profiled_rewards") + result["profiled_output_lengths"] = nemo_gym_row.get( + "profiled_output_lengths" + ) + result["profile_band"] = nemo_gym_row.get("profile_band") + apply_group_length_adjustments( + results, {"grpo": grpo_config}, tokenizer=tokenizer + ) # Prepare for the rollout metrics calculation below. Not strictly necessary here, but good to have parity with `run_async_multi_turn_rollout` with timer.time(f"{timer_prefix}/prepare_for_metrics_calculation"): From ded553a78a15962bb87fa51e91cd46b3ec04c0e8 Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 22:01:28 -0700 Subject: [PATCH 17/27] feat: restrict length adjustments to binary (0/1) env rewards MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every algorithm and the final clamp assume binary rewards; enforcing it removes the graded/negative-reward edge cases wholesale. A prompt group containing any non-binary reward is skipped untouched (no adjustment, no clamp) with a warning logged once per agent — a graded judge agent in a mixed batch degrades loudly instead of killing the run. With binariness guaranteed the phase-3 clamp is now unconditional, which also closes the profiled_length_penalty asymmetry from review follow-up: penalty-created negatives on wrong rollouts floor back to 0, so an all-wrong group stays variance-free and contributes no GRPO gradient, matching vanilla behavior. Adds binary-skip, graded-skip, and all-wrong-variance-free tests (34 total); documents the requirement and the top_percentile inertness under binary rewards. Co-Authored-By: Claude Fable 5 --- nemo_rl/utils/length_adjustments.md | 25 ++++++++-- nemo_rl/utils/length_adjustments.py | 55 +++++++++++++++++---- tests/unit/utils/test_length_adjustments.py | 53 +++++++++++++++----- 3 files changed, 107 insertions(+), 26 deletions(-) diff --git a/nemo_rl/utils/length_adjustments.md b/nemo_rl/utils/length_adjustments.md index 1fa6bc8667a..babcf1b8bc8 100644 --- a/nemo_rl/utils/length_adjustments.md +++ b/nemo_rl/utils/length_adjustments.md @@ -10,6 +10,23 @@ place during rollout postprocessing. All algorithms are resolved per prompt group. Unless otherwise stated, only rollouts with `reward > 0` participate in length comparisons and receive length-based adjustments. +## Binary Rewards Requirement + +Length adjustments are defined for **binary (0/1) environment rewards** only. Every algorithm +and the final clamp assume it. At rollout time each prompt group's rewards are checked: a group +containing any graded or negative reward is skipped entirely — no adjustment, no clamp, rewards +pass through untouched — and a warning is logged once per agent. To silence the warning for a +deliberately graded agent (e.g. a genrm judge), disable it explicitly under `agent_overrides` +with `enabled: false`. + +Consequences of binariness worth knowing: + +- The adjusted reward is clamped at 0, so penalties can wipe a correct rollout's reward out but + never flip its sign — and an all-wrong group stays all-zero (no within-group variance, hence + no GRPO gradient, matching vanilla behavior on groups with no correctness signal). +- `top_percentile` is effectively inert: all correct rollouts tie at the top score, so every + positive rollout is always a "top scorer". + ## Common Config ```yaml @@ -517,10 +534,10 @@ grpo: ## Practical Notes - Most length algorithms act only on positive rollouts (`reward > 0`). -- Additive penalties can wipe a correct rollout's reward out but never flip its sign: the - adjusted reward is clamped at 0 for originally-positive rollouts (stacked flat penalties - exceeding the reward produce 0, not a negative value). Environments' own negative rewards - pass through untouched. +- Additive penalties can wipe a rollout's reward out but never flip its sign: the adjusted + reward is clamped at 0 (stacked flat penalties exceeding the reward produce 0, not a negative + value). Non-binary rewards never reach the clamp — their groups are skipped wholesale (see + Binary Rewards Requirement). - `profile_band_*` multipliers also apply only to originally correct rollouts. - Group-relative scaling can reduce average length aggressively because it gives dense per-group pressure. diff --git a/nemo_rl/utils/length_adjustments.py b/nemo_rl/utils/length_adjustments.py index 4877b92e4a5..21304003970 100644 --- a/nemo_rl/utils/length_adjustments.py +++ b/nemo_rl/utils/length_adjustments.py @@ -87,6 +87,17 @@ "group_length_penalty_profile_gate_field", }) +# Length adjustments are defined for binary (0/1) env rewards only. Agents +# already warned about non-binary rewards (warn once per agent, then skip +# their prompt groups). +_NON_BINARY_WARNED_AGENTS: set[str] = set() +_BINARY_REWARD_TOL = 1e-6 + + +def _is_binary_reward(value: float) -> bool: + v = float(value) + return abs(v) <= _BINARY_REWARD_TOL or abs(v - 1.0) <= _BINARY_REWARD_TOL + def _extract_reasoning_and_answer_text(result: dict[str, Any]) -> tuple[str, str]: """Extract reasoning and answer text from the Response API output items.""" fr = result.get("full_result", {}) @@ -247,10 +258,29 @@ def apply_group_length_adjustments( total_lengths = [0] * n groups_adjusted = 0 group_gate_infos: dict[int, dict[str, Any]] = {} + # Rows whose group passed the binary-rewards check; all other rows are + # left completely untouched (no adjustment, no clamp, no reward writeback). + binary_ok = [False] * n for g in range(0, n, num_gens): agent_name = agent_names[g] group_size = min(num_gens, n - g) + # Length adjustments are defined for binary (0/1) env rewards only: + # every algorithm and the phase-3 clamp assume it. Skip (and warn once + # per agent) on graded or negative rewards. + if any( + not _is_binary_reward(original_rewards[g + k]) for k in range(group_size) + ): + if agent_name not in _NON_BINARY_WARNED_AGENTS: + _NON_BINARY_WARNED_AGENTS.add(agent_name) + logger.warning( + f"length adjustments require binary (0/1) env rewards; agent " + f"{agent_name} produced non-binary rewards — skipping length " + f"adjustments for its prompt groups" + ) + continue + for k in range(group_size): + binary_ok[g + k] = True if any(results[g + k].get("low_effort_applied") for k in range(group_size)): continue params = _resolve_agent_params(agent_name, agents_cfg, defaults) @@ -452,9 +482,11 @@ def apply_group_length_adjustments( print(f"{'=' * 70}\n", flush=True) - # Phase 3: apply additive adjustments + # Phase 3: apply additive adjustments (binary-verified rows only) additive_base_rewards = [0.0] * n for i, r in enumerate(results): + if not binary_ok[i]: + continue # The profiled-length penalty stacks only on rollouts whose group # adjustments are non-negative: a rollout already penalized by the # group-relative channels should not be double-penalized for the same @@ -462,11 +494,11 @@ def apply_group_length_adjustments( profiled_adj = all_profiled_length_adj[i] if all_adjustments[i] >= 0 else 0.0 additive_delta = all_adjustments[i] + profiled_adj additive_base_rewards[i] = original_rewards[i] + additive_delta - # Stacked flat penalties can exceed the reward itself. Length penalties - # may wipe a correct rollout's reward out, but never flip its sign — - # clamp at 0. Only for originally-positive rollouts: an env's own - # negative reward must pass through untouched. - if original_rewards[i] > 0 and additive_base_rewards[i] < 0: + # Rewards here are binary, so any negative value is penalty-created. + # Length penalties may wipe a reward out but never flip its sign; this + # also keeps all-wrong groups variance-free (no gradient from a group + # with no correctness signal). + if additive_base_rewards[i] < 0: additive_base_rewards[i] = 0.0 r["full_result"]["reward"] = additive_base_rewards[i] @@ -483,6 +515,7 @@ def apply_group_length_adjustments( defaults=defaults, num_gens=num_gens, global_band=global_band, + binary_ok=binary_ok, ) @@ -541,6 +574,7 @@ def _apply_profile_band_multipliers( defaults: dict[str, Any], num_gens: int, global_band: dict[str, dict[str, Any]] | None = None, + binary_ok: list[bool] | None = None, ) -> None: """Apply per-channel profile_band multipliers to correct rollouts. @@ -574,12 +608,15 @@ def _apply_profile_band_multipliers( ch_ans = band.get("answer") if use_ans else None for k in range(group_size): idx = g + k + # Only rows whose group passed the binary-rewards check. + if binary_ok is not None and not binary_ok[idx]: + continue # Gate on the env reward (correct rollouts only). if original_rewards[idx] <= 0: continue - # Additive penalties can push the base below zero; multiplying a - # negative base by m < 1 would RAISE the reward for longer - # rollouts. Scale only the non-negative part. + # The phase-3 clamp floors the base at 0; multiplying a negative + # base by m < 1 would RAISE the reward for longer rollouts, so + # scale only strictly-positive bases. if base_rewards[idx] <= 0: continue current_reward = base_rewards[idx] diff --git a/tests/unit/utils/test_length_adjustments.py b/tests/unit/utils/test_length_adjustments.py index 502c8fedd55..7a38e2be69e 100644 --- a/tests/unit/utils/test_length_adjustments.py +++ b/tests/unit/utils/test_length_adjustments.py @@ -402,9 +402,9 @@ def test_flat_penalty_exceeding_reward_clamps_at_zero(self): apply_group_length_adjustments(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 0.0]) - def test_negative_env_reward_not_raised_by_clamp(self): - # The clamp applies only to originally-positive rollouts: an env's own - # negative reward must pass through untouched. + def test_negative_env_reward_group_skipped(self): + # Length adjustments require binary rewards: a group containing a + # negative env reward is skipped wholesale — no adjustment, no clamp. cfg = make_config( default={"enabled": True, "group_total_length_penalty_coeff": 0.1} ) @@ -413,7 +413,36 @@ def test_negative_env_reward_not_raised_by_clamp(self): make_result("1234567890", "1234567890", 1.0), ] apply_group_length_adjustments(results, cfg) - assert rewards_of(results)[0] == pytest.approx(-1.0) + assert rewards_of(results) == pytest.approx([-1.0, 1.0]) + + def test_graded_rewards_group_skipped(self): + # Graded (non-binary) rewards also skip the group untouched. + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.1} + ) + results = [ + make_result("12345", "12345", 0.5), + make_result("1234567890", "1234567890", 1.0), + ] + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([0.5, 1.0]) + + def test_all_wrong_group_stays_variance_free(self): + # profiled_length_penalty has no positive-reward gate, but the binary + # clamp floors penalty-created negatives at 0: an all-wrong group must + # stay all-zero (no within-group variance -> no GRPO gradient). + results = [ + make_result("1234", "1234", 0.0), # len 8, under threshold + make_result("1234567890", "1234567890", 0.0), # len 20, over + ] + for r in results: + r["profiled_rewards"] = [1, 1] + r["profiled_output_lengths"] = [10, 10] + cfg = make_config( + default={"enabled": True, "profiled_length_penalty": 0.3} + ) + apply_group_length_adjustments(results, cfg) + assert rewards_of(results) == pytest.approx([0.0, 0.0]) def test_band_multiplier_never_rewards_length_on_penalized_base(self): # Flat penalty exceeds the reward, so bases clamp to 0; the band phase @@ -465,18 +494,16 @@ def test_unknown_key_raises(self): else: raise AssertionError("expected ValueError for unknown config key") - def test_top_percentile_default_is_half(self): - # With 2 positive scorers (rewards 1.0 and 0.9) and default - # top_percentile=0.5, only the best scorer is a "top scorer" — the - # longest-total penalty needs >= 2 eligible top scorers, so it no-ops. - # With the old effective default of 0.0 the behavior was the same, but - # an explicit 1.0 must include both and penalize the longer one. + def test_longest_penalty_under_binary_rewards(self): + # Under binary rewards all correct rollouts tie at the top score, so + # top_percentile is inert and both rollouts are eligible top scorers; + # the longest one takes the flat penalty (default top_percentile 0.5). results = [ make_result("12345", "12345", 1.0), - make_result("1234567890", "1234567890", 0.9), + make_result("1234567890", "1234567890", 1.0), ] cfg = make_config( - default={"enabled": True, "longest_total_penalty": 0.2, "top_percentile": 1.0} + default={"enabled": True, "longest_total_penalty": 0.2} ) apply_group_length_adjustments(results, cfg) - assert rewards_of(results) == pytest.approx([1.0, 0.7]) + assert rewards_of(results) == pytest.approx([1.0, 0.8]) From aa2a1a84a3db932f29866b1c80a571b4494fa4f0 Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 22:15:36 -0700 Subject: [PATCH 18/27] refactor: rename length_bonus -> length_penalty (config key, module, docs) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rename before anyone depends on the key — the feature is new to this PR: - config: grpo.length_bonus -> grpo.length_penalty - module: nemo_rl/utils/length_adjustments.{py,md} -> length_penalty.{py,md} - entry point: apply_group_length_adjustments -> apply_group_length_penalties - threaded kwarg: length_adjustment_config -> length_penalty_config - tests: test_length_adjustments.py -> test_length_penalty.py Pure mechanical rename; no behavior change (34/34 tests unchanged). Co-Authored-By: Claude Fable 5 --- .../async_utils/trajectory_collector.py | 2 +- nemo_rl/algorithms/grpo.py | 4 +- nemo_rl/experience/rollouts.py | 28 +++---- ...ength_adjustments.md => length_penalty.md} | 36 ++++----- ...ength_adjustments.py => length_penalty.py} | 42 +++++----- scripts/build_profile_band_dataset.py | 2 +- ..._adjustments.py => test_length_penalty.py} | 78 +++++++++---------- 7 files changed, 96 insertions(+), 96 deletions(-) rename nemo_rl/utils/{length_adjustments.md => length_penalty.md} (96%) rename nemo_rl/utils/{length_adjustments.py => length_penalty.py} (96%) rename tests/unit/utils/{test_length_adjustments.py => test_length_penalty.py} (90%) diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index 4bf824c1777..dd0ff5bcd94 100644 --- a/nemo_rl/algorithms/async_utils/trajectory_collector.py +++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py @@ -1291,7 +1291,7 @@ async def _iter_rollout_groups( max_rollout_turns=None, greedy=False, reward_penalty_config=self.master_config.reward_penalties, - length_adjustment_config=self.master_config.grpo.model_dump(), + length_penalty_config=self.master_config.grpo.model_dump(), thinking_tags=get_nemo_gym_thinking_tags(self.master_config.env), mask_env_flagged_samples=should_mask_flagged_samples( self.master_config.env diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index e971fadd4f4..bff55e9c385 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -3015,7 +3015,7 @@ def grpo_train( greedy=False, effort_config=_get_effort_config(master_config), reward_penalty_config=master_config.reward_penalties, - length_adjustment_config=master_config.grpo.model_dump(), + length_penalty_config=master_config.grpo.model_dump(), thinking_tags=get_nemo_gym_thinking_tags(master_config.env), mask_env_flagged_samples=should_mask_flagged_samples( master_config.env @@ -3939,7 +3939,7 @@ def validate( greedy=False, effort_config=_get_effort_config(master_config), reward_penalty_config=master_config.reward_penalties, - # No length_adjustment_config here: validation metrics + # No length_penalty_config here: validation metrics # (accuracy/pass_k) must reflect the raw env reward, and the # adjustment code groups by the TRAINING stride # (num_generations_per_prompt), which does not match diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 3fe07d52675..25218687942 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -72,7 +72,7 @@ collect_multimodal_payload_metrics, print_multimodal_payload_metrics, ) -from nemo_rl.utils.length_adjustments import apply_group_length_adjustments +from nemo_rl.utils.length_penalty import apply_group_length_penalties from nemo_rl.utils.timer import Timer TokenizerType = PreTrainedTokenizerBase @@ -2272,7 +2272,7 @@ async def run_async_nemo_gym_rollout( greedy: bool = False, effort_config: Optional[EffortLevelsConfig] = None, reward_penalty_config: dict[str, Any] | BaseModel | None = None, - length_adjustment_config: dict[str, Any] | BaseModel | None = None, + length_penalty_config: dict[str, Any] | BaseModel | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, mask_env_flagged_samples: bool = True, returns_entire_batch: bool = False, @@ -2304,7 +2304,7 @@ async def run_async_nemo_gym_rollout( greedy: Must be ``False`` because this path does not support greedy mode. effort_config: Optional configuration for effort-based reward shaping. reward_penalty_config: Optional reward-penalty configuration. - length_adjustment_config: Optional GRPO config block for length adjustments. + length_penalty_config: Optional GRPO config block for length adjustments. thinking_tags: Optional opening and closing tags used by thinking penalties. mask_env_flagged_samples: Whether to carry env-driven ``mask_sample`` flags in the rollout batch for loss masking. @@ -2480,7 +2480,7 @@ async def run_async_nemo_gym_rollout( log_full_result_tables=log_full_result_tables, effort_config=effort_config, reward_penalty_config=reward_penalty_config, - length_adjustment_config=length_adjustment_config, + length_penalty_config=length_penalty_config, thinking_tags=thinking_tags, mask_env_flagged_samples=mask_env_flagged_samples, ) @@ -2518,7 +2518,7 @@ def run_nemo_gym_rollout_sync( greedy: bool = False, effort_config: Optional[EffortLevelsConfig] = None, reward_penalty_config: dict[str, Any] | BaseModel | None = None, - length_adjustment_config: dict[str, Any] | BaseModel | None = None, + length_penalty_config: dict[str, Any] | BaseModel | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, sampling_params: Optional[GenerationSamplingParams] = None, mask_env_flagged_samples: bool = True, @@ -2583,7 +2583,7 @@ async def _consume_rollout() -> NemoGymRolloutResult: greedy=greedy, effort_config=effort_config, reward_penalty_config=reward_penalty_config, - length_adjustment_config=length_adjustment_config, + length_penalty_config=length_penalty_config, thinking_tags=thinking_tags, mask_env_flagged_samples=mask_env_flagged_samples, returns_entire_batch=True, @@ -2610,7 +2610,7 @@ def _postprocess_single_nemo_gym_group( log_full_result_tables: bool, effort_config: Optional[EffortLevelsConfig] = None, reward_penalty_config: dict[str, Any] | BaseModel | None = None, - length_adjustment_config: dict[str, Any] | BaseModel | None = None, + length_penalty_config: dict[str, Any] | BaseModel | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, mask_env_flagged_samples: bool = True, ) -> NemoGymRolloutResult: @@ -2627,15 +2627,15 @@ def _postprocess_single_nemo_gym_group( ) penalty_counts = apply_reward_penalties(results, resolved_reward_penalty_config) - if length_adjustment_config is not None: + if length_penalty_config is not None: grpo_config = ( - length_adjustment_config.model_dump() - if isinstance(length_adjustment_config, BaseModel) - else dict(length_adjustment_config) + length_penalty_config.model_dump() + if isinstance(length_penalty_config, BaseModel) + else dict(length_penalty_config) ) # Callers pass the whole grpo config block; runs without a - # grpo.length_bonus section are untouched by this block. - if grpo_config.get("length_bonus"): + # grpo.length_penalty section are untouched by this block. + if grpo_config.get("length_penalty"): # Copy the per-row fields the length adjustments consume. for nemo_gym_row, result in zip(nemo_gym_rows, results): result["agent_ref"] = nemo_gym_row["agent_ref"] @@ -2644,7 +2644,7 @@ def _postprocess_single_nemo_gym_group( "profiled_output_lengths" ) result["profile_band"] = nemo_gym_row.get("profile_band") - apply_group_length_adjustments( + apply_group_length_penalties( results, {"grpo": grpo_config}, tokenizer=tokenizer ) diff --git a/nemo_rl/utils/length_adjustments.md b/nemo_rl/utils/length_penalty.md similarity index 96% rename from nemo_rl/utils/length_adjustments.md rename to nemo_rl/utils/length_penalty.md index babcf1b8bc8..869fb5dd057 100644 --- a/nemo_rl/utils/length_adjustments.md +++ b/nemo_rl/utils/length_penalty.md @@ -1,10 +1,10 @@ -# Length Adjustment Algorithms +# Length Penalty Algorithms This file documents the length-penalty and length-bonus algorithms implemented in: -`nemo_rl/utils/length_adjustments.py` +`nemo_rl/utils/length_penalty.py` -Configure `grpo.length_bonus`. The length adjustments mutate `full_result["reward"]` in +Configure `grpo.length_penalty`. The length adjustments mutate `full_result["reward"]` in place during rollout postprocessing. All algorithms are resolved per prompt group. Unless otherwise stated, only rollouts with @@ -31,7 +31,7 @@ Consequences of binariness worth knowing: ```yaml grpo: - length_bonus: + length_penalty: verbose: true default: enabled: true @@ -84,7 +84,7 @@ agent: ```yaml grpo: - length_bonus: + length_penalty: default: enabled: true group_total_length_penalty_coeff: 0.1 @@ -226,7 +226,7 @@ Example: ```yaml grpo: - length_bonus: + length_penalty: default: enabled: true total_bonus: 0.1 @@ -256,7 +256,7 @@ Example: ```yaml grpo: - length_bonus: + length_penalty: default: enabled: true top_percentile: 0.5 @@ -292,7 +292,7 @@ Example: ```yaml grpo: - length_bonus: + length_penalty: default: enabled: true group_total_length_penalty_coeff: 0.1 @@ -332,7 +332,7 @@ Example: ```yaml grpo: - length_bonus: + length_penalty: default: enabled: true total_zmad_threshold: 2.5 @@ -370,7 +370,7 @@ Example: ```yaml grpo: - length_bonus: + length_penalty: default: enabled: true profiled_length_penalty: 0.1 @@ -416,7 +416,7 @@ Example: ```yaml grpo: - length_bonus: + length_penalty: default: enabled: true profile_band_total: true @@ -425,12 +425,12 @@ grpo: #### Global Defaults (dataset without per-prompt bands) When the dataset has no per-prompt `profile_band` metadata, global `{a, b, f}` values can be -set directly in the config under `length_bonus.profile_band`. Only the channels listed under +set directly in the config under `length_penalty.profile_band`. Only the channels listed under `defaults` are activated: ```yaml grpo: - length_bonus: + length_penalty: profile_band: enabled: true defaults: @@ -439,7 +439,7 @@ grpo: ```yaml grpo: - length_bonus: + length_penalty: profile_band: enabled: true defaults: @@ -448,7 +448,7 @@ grpo: ```yaml grpo: - length_bonus: + length_penalty: profile_band: enabled: true defaults: @@ -461,7 +461,7 @@ only, and the last on answer length only. Multiple channels may be listed togeth Semantics: - Channels under `defaults` are implicitly enabled — no need to also set - `profile_band_total/reasoning/answer: true` under `length_bonus.default`. Per-agent + `profile_band_total/reasoning/answer: true` under `length_penalty.default`. Per-agent `agent_overrides` can still disable a channel (e.g. `profile_band_total: false`). - Per-prompt `profile_band` metadata, when present on a row, takes precedence over the global defaults on a per-channel basis (a row that only provides `total` still falls back to the @@ -493,7 +493,7 @@ For each prompt group: 1. Read a threshold from `profile_band[channel][field]`, for example `profile_band["total"]["a"]`. The band is the row's `profile_band` metadata merged over the global - `length_bonus.profile_band.defaults` (row channels win), so the gate also works on datasets + `length_penalty.profile_band.defaults` (row channels win), so the gate also works on datasets without per-prompt bands when global defaults are configured. 2. Compute the mean rollout length for the selected channel. 3. If `group_length_penalty_profile_gate_positive_only` is true, use only positive rollouts in @@ -521,7 +521,7 @@ Example: ```yaml grpo: - length_bonus: + length_penalty: default: enabled: true group_total_length_penalty_coeff: 0.1 diff --git a/nemo_rl/utils/length_adjustments.py b/nemo_rl/utils/length_penalty.py similarity index 96% rename from nemo_rl/utils/length_adjustments.py rename to nemo_rl/utils/length_penalty.py index 21304003970..c416b82f75c 100644 --- a/nemo_rl/utils/length_adjustments.py +++ b/nemo_rl/utils/length_penalty.py @@ -129,14 +129,14 @@ def _extract_reasoning_and_answer_text(result: dict[str, Any]) -> tuple[str, str return reasoning_text, answer_text -_TOP_LEVEL_LENGTH_BONUS_KEYS = frozenset( +_TOP_LEVEL_LENGTH_PENALTY_KEYS = frozenset( {"verbose", "default", "agent_overrides", "profile_band"} ) _PROFILE_BAND_BLOCK_KEYS = frozenset({"enabled", "defaults"}) _PROFILE_BAND_CHANNELS = frozenset({"total", "reasoning", "answer"}) -def _reject_unknown_length_bonus_keys(length_cfg: dict[str, Any]) -> None: +def _reject_unknown_length_penalty_keys(length_cfg: dict[str, Any]) -> None: """Raise on unknown config keys instead of silently ignoring them.""" def _check(block: Any, allowed: frozenset, where: str) -> None: @@ -148,9 +148,9 @@ def _check(block: Any, allowed: frozenset, where: str) -> None: f"Unknown key(s) {unknown} in {where}; allowed: {sorted(allowed)}" ) - _check(length_cfg, _TOP_LEVEL_LENGTH_BONUS_KEYS, "grpo.length_bonus") + _check(length_cfg, _TOP_LEVEL_LENGTH_PENALTY_KEYS, "grpo.length_penalty") _check( - length_cfg.get("default"), frozenset(_PARAM_KEYS), "grpo.length_bonus.default" + length_cfg.get("default"), frozenset(_PARAM_KEYS), "grpo.length_penalty.default" ) agents_cfg = length_cfg.get("agent_overrides") if isinstance(agents_cfg, dict): @@ -158,26 +158,26 @@ def _check(block: Any, allowed: frozenset, where: str) -> None: _check( overrides, frozenset(_PARAM_KEYS), - f"grpo.length_bonus.agent_overrides.{agent_name}", + f"grpo.length_penalty.agent_overrides.{agent_name}", ) pb_cfg = length_cfg.get("profile_band") if isinstance(pb_cfg, dict): - _check(pb_cfg, _PROFILE_BAND_BLOCK_KEYS, "grpo.length_bonus.profile_band") + _check(pb_cfg, _PROFILE_BAND_BLOCK_KEYS, "grpo.length_penalty.profile_band") _check( pb_cfg.get("defaults"), _PROFILE_BAND_CHANNELS, - "grpo.length_bonus.profile_band.defaults", + "grpo.length_penalty.profile_band.defaults", ) -def apply_group_length_adjustments( +def apply_group_length_penalties( results: list[dict[str, Any]], master_config: dict[str, Any], tokenizer: Any = None, ) -> None: """Apply per-prompt-group length bonuses/penalties. - Reads ``grpo.length_bonus`` for configuration and mutates + Reads ``grpo.length_penalty`` for configuration and mutates ``full_result["reward"]`` in place. No-ops when no length-adjustment feature is enabled. @@ -187,11 +187,11 @@ def apply_group_length_adjustments( tokenizer: Tokenizer for computing reasoning/answer token counts. """ grpo_config = master_config.get("grpo", {}) - length_cfg = dict(grpo_config.get("length_bonus", {}) or {}) + length_cfg = dict(grpo_config.get("length_penalty", {}) or {}) if not length_cfg: return - _reject_unknown_length_bonus_keys(length_cfg) + _reject_unknown_length_penalty_keys(length_cfg) default_cfg = length_cfg.get("default", {}) agents_cfg = length_cfg.get("agent_overrides") global_band = _resolve_global_profile_band(length_cfg.get("profile_band")) @@ -222,7 +222,7 @@ def apply_group_length_adjustments( defaults[k] = default_cfg.get(k, 0.5) else: defaults[k] = default_cfg.get(k, 0.0) - # Channels listed under length_bonus.profile_band.defaults are implicitly + # Channels listed under length_penalty.profile_band.defaults are implicitly # enabled — unless the user explicitly configured the channel flag, which # always wins (e.g. profile_band_total: false stays false). for _ch in global_band: @@ -274,9 +274,9 @@ def apply_group_length_adjustments( if agent_name not in _NON_BINARY_WARNED_AGENTS: _NON_BINARY_WARNED_AGENTS.add(agent_name) logger.warning( - f"length adjustments require binary (0/1) env rewards; agent " + f"length penalties require binary (0/1) env rewards; agent " f"{agent_name} produced non-binary rewards — skipping length " - f"adjustments for its prompt groups" + f"penalties for its prompt groups" ) continue for k in range(group_size): @@ -335,7 +335,7 @@ def apply_group_length_adjustments( zmad_r_adj, zmad_a_adj, zmad_t_adj, - ) = _apply_length_bonuses_and_penalties( + ) = _apply_length_penaltyes_and_penalties( group_rewards, group_reasoning, group_answer, group_total, **params ) @@ -521,7 +521,7 @@ def apply_group_length_adjustments( def _resolve_global_profile_band(pb_cfg: Any) -> dict[str, dict[str, Any]]: - """Parse ``length_bonus.profile_band`` into per-channel {a, b, f} blocks. + """Parse ``length_penalty.profile_band`` into per-channel {a, b, f} blocks. Returns only channels ("total", "reasoning", "answer") present under ``defaults`` with a complete, well-formed block. Empty dict when the @@ -540,7 +540,7 @@ def _resolve_global_profile_band(pb_cfg: Any) -> dict[str, dict[str, Any]]: a, b, f = ch_cfg.get("a"), ch_cfg.get("b"), ch_cfg.get("f") if a is None or b is None or f is None or b <= a: logger.warning( - f"length_bonus.profile_band.defaults.{ch} is malformed " + f"length_penalty.profile_band.defaults.{ch} is malformed " f"(a={a}, b={b}, f={f}); ignoring this channel" ) continue @@ -579,11 +579,11 @@ def _apply_profile_band_multipliers( """Apply per-channel profile_band multipliers to correct rollouts. Each enabled channel contributes a multiplier in [0.0, 1.0] derived from the - per-row {a, b, f} block, falling back to ``length_bonus.profile_band.defaults`` + per-row {a, b, f} block, falling back to ``length_penalty.profile_band.defaults`` for channels the row does not provide. Mutates scalar rewards in place. Skips any group where the low-effort bypass already replaced the reward - (parity with Phase 1 of ``apply_group_length_adjustments``). + (parity with Phase 1 of ``apply_group_length_penalties``). """ n = len(results) global_band = global_band or {} @@ -751,7 +751,7 @@ def _resolve_agent_params( if agent_name not in agents_cfg: print( - f"[length_adjustments] WARNING: agent '{agent_name}' not found in " + f"[length_penalty] WARNING: agent '{agent_name}' not found in " f"agent_overrides, falling back to defaults", flush=True, ) @@ -794,7 +794,7 @@ def _zmad_local_outliers( return out -def _apply_length_bonuses_and_penalties( +def _apply_length_penaltyes_and_penalties( rewards: list[float], reasoning_lengths: list[int], answer_lengths: list[int], diff --git a/scripts/build_profile_band_dataset.py b/scripts/build_profile_band_dataset.py index d6352bf38cf..b4f5a4721fb 100644 --- a/scripts/build_profile_band_dataset.py +++ b/scripts/build_profile_band_dataset.py @@ -19,7 +19,7 @@ ``profiled_reasoning_lengths``, ``profiled_answer_lengths``, and ``pass_rate``) and writes a new JSONL where each row has an additional ``profile_band`` field consumed by Mechanism 6 in -``nemo_rl/utils/length_adjustments.py``:: +``nemo_rl/utils/length_penalty.py``:: profile_band: total: {a, b, f} # only present if data is non-degenerate diff --git a/tests/unit/utils/test_length_adjustments.py b/tests/unit/utils/test_length_penalty.py similarity index 90% rename from tests/unit/utils/test_length_adjustments.py rename to tests/unit/utils/test_length_penalty.py index 7a38e2be69e..e0d9bbf5e49 100644 --- a/tests/unit/utils/test_length_adjustments.py +++ b/tests/unit/utils/test_length_penalty.py @@ -13,13 +13,13 @@ # limitations under the License. """Unit tests for the profile_band multiplier and group relative-length -scaling algorithms in nemo_rl/utils/length_adjustments.py.""" +scaling algorithms in nemo_rl/utils/length_penalty.py.""" import pytest -from nemo_rl.utils.length_adjustments import ( +from nemo_rl.utils.length_penalty import ( _band_multiplier, - apply_group_length_adjustments, + apply_group_length_penalties, ) AGENT = "math_with_judge_simple_agent" @@ -45,15 +45,15 @@ def make_result(reasoning: str, answer: str, reward: float, band=None): def make_config(default=None, profile_band=None, num_gens=2): - length_bonus = {} + length_penalty = {} if default is not None: - length_bonus["default"] = {"length_type": "chars", **default} + length_penalty["default"] = {"length_type": "chars", **default} if profile_band is not None: - length_bonus["profile_band"] = profile_band + length_penalty["profile_band"] = profile_band return { "grpo": { "num_generations_per_prompt": num_gens, - "length_bonus": length_bonus, + "length_penalty": length_penalty, } } @@ -100,7 +100,7 @@ def test_total_channel_scales_correct_rollouts(self): make_result("1234567890", "1234567890", 1.0, band=band), # 20 -> x0.5 ] cfg = make_config(default={"enabled": True, "profile_band_total": True}) - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 0.5]) def test_zero_reward_rollouts_untouched(self): @@ -110,7 +110,7 @@ def test_zero_reward_rollouts_untouched(self): make_result("1234567890", "1234567890", 1.0, band=band), ] cfg = make_config(default={"enabled": True, "profile_band_total": True}) - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([0.0, 0.5]) def test_reasoning_channel_ignores_answer_length(self): @@ -121,7 +121,7 @@ def test_reasoning_channel_ignores_answer_length(self): make_result("123456789012345", long_answer, 1.0, band=band), # 15 -> x0.75 ] cfg = make_config(default={"enabled": True, "profile_band_reasoning": True}) - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 0.75]) def test_missing_row_band_is_noop(self): @@ -130,7 +130,7 @@ def test_missing_row_band_is_noop(self): make_result("1234567890123456789012345", "12345", 1.0), ] cfg = make_config(default={"enabled": True, "profile_band_total": True}) - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 1.0]) def test_channel_not_enabled_in_config_is_noop(self): @@ -140,12 +140,12 @@ def test_channel_not_enabled_in_config_is_noop(self): make_result("12345", "12345", 1.0, band=band), ] cfg = make_config(default={"enabled": True}) # no profile_band_* flag - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 1.0]) class TestProfileBandGlobalDefaults: - """profile_band driven by config-level length_bonus.profile_band defaults.""" + """profile_band driven by config-level length_penalty.profile_band defaults.""" def test_global_total_only(self): cfg = make_config( @@ -159,11 +159,11 @@ def test_global_total_only(self): make_result("12345", "12345", 1.0), # total 10 -> x1.0 make_result("1234567890123456789012345", "12345", 1.0), # 30 -> x0.5 ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 0.5]) def test_global_works_without_default_block(self): - # Channels under defaults are implicitly enabled; no length_bonus.default + # Channels under defaults are implicitly enabled; no length_penalty.default # `enabled` or profile_band_* booleans required. cfg = make_config( default={}, # only length_type @@ -176,7 +176,7 @@ def test_global_works_without_default_block(self): make_result("123456789012345", "xx", 1.0), # reasoning 15 -> x0.75 make_result("12345", "xx", 1.0), # 5 -> x1.0 ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([0.75, 1.0]) def test_row_band_wins_over_global(self): @@ -192,7 +192,7 @@ def test_row_band_wins_over_global(self): make_result("12345", "12345", 1.0, band=generous), make_result("1234567890123456789012345", "12345", 1.0, band=generous), ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 1.0]) def test_disabled_block_is_noop(self): @@ -207,7 +207,7 @@ def test_disabled_block_is_noop(self): make_result("1234567890123456789012345", "12345", 1.0), make_result("12345", "12345", 1.0), ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 1.0]) def test_malformed_global_channel_ignored(self): @@ -222,7 +222,7 @@ def test_malformed_global_channel_ignored(self): make_result("1234567890123456789012345", "12345", 1.0), make_result("12345", "12345", 1.0), ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 1.0]) @@ -238,7 +238,7 @@ def test_two_rollouts_symmetric_adjustment(self): make_result("12345", "12345", 1.0), # total 10 -> +0.05 make_result("1234567890123456789012345", "12345", 1.0), # 30 -> -0.05 ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.05, 0.95]) def test_three_rollouts_zero_centered(self): @@ -252,7 +252,7 @@ def test_three_rollouts_zero_centered(self): make_result("1234567890", "1234567890", 1.0), # 20 make_result("123456789012345", "123456789012345", 1.0), # 30 ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.05, 1.0, 0.95]) # Zero-centered: the group's mean reward is unchanged by the adjustment. assert sum(rewards_of(results)) == pytest.approx(3.0) @@ -265,7 +265,7 @@ def test_equal_lengths_no_adjustment(self): make_result("12345", "12345", 1.0), make_result("12345", "12345", 1.0), ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 1.0]) def test_only_positive_rollouts_participate(self): @@ -280,7 +280,7 @@ def test_only_positive_rollouts_participate(self): make_result("1" * 1000, "1" * 1000, 0.0), # untouched, excluded make_result("1234567890123456789012345", "12345", 1.0), # 30 -> -0.05 ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.05, 0.0, 0.95]) def test_reasoning_channel_uses_reasoning_length_only(self): @@ -293,7 +293,7 @@ def test_reasoning_channel_uses_reasoning_length_only(self): make_result("12345", "123456789012345", 1.0), # reasoning 5 -> +0.05 make_result("123456789012345", "12345", 1.0), # reasoning 15 -> -0.05 ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.05, 0.95]) def test_zero_coefficient_is_noop(self): @@ -304,19 +304,19 @@ def test_zero_coefficient_is_noop(self): make_result("12345", "12345", 1.0), make_result("1234567890123456789012345", "12345", 1.0), ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 1.0]) def test_agent_override_disables_for_agent(self): cfg = make_config( default={"enabled": True, "group_total_length_penalty_coeff": 0.1} ) - cfg["grpo"]["length_bonus"]["agent_overrides"] = {AGENT: {"enabled": False}} + cfg["grpo"]["length_penalty"]["agent_overrides"] = {AGENT: {"enabled": False}} results = [ make_result("12345", "12345", 1.0), make_result("1234567890123456789012345", "12345", 1.0), ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 1.0]) @@ -347,7 +347,7 @@ def test_enough_passes_penalizes_over_threshold(self): self.make_profiled("12345", "12345", 1.0, p_rewards, p_lengths), # 10 < thr self.make_profiled("1234567890", "1234567890", 1.0, p_rewards, p_lengths), # 20 >= thr ] - apply_group_length_adjustments(results, self.cfg()) + apply_group_length_penalties(results, self.cfg()) assert rewards_of(results) == pytest.approx([1.0, 0.7]) # The failing profiled length (100) must not have entered the threshold: # with it, mean+std would exceed 20 and nothing would be penalized. @@ -360,7 +360,7 @@ def test_one_pass_below_min_samples_no_penalty(self): self.make_profiled("1234567890", "1234567890", 1.0, p_rewards, p_lengths), self.make_profiled("1" * 50, "1" * 50, 1.0, p_rewards, p_lengths), ] - apply_group_length_adjustments(results, self.cfg()) + apply_group_length_penalties(results, self.cfg()) assert rewards_of(results) == pytest.approx([1.0, 1.0]) def test_zero_passes_no_penalty(self): @@ -369,7 +369,7 @@ def test_zero_passes_no_penalty(self): self.make_profiled("1234567890", "1234567890", 1.0, p_rewards, p_lengths), self.make_profiled("1" * 50, "1" * 50, 1.0, p_rewards, p_lengths), ] - apply_group_length_adjustments(results, self.cfg()) + apply_group_length_penalties(results, self.cfg()) assert rewards_of(results) == pytest.approx([1.0, 1.0]) def test_one_pass_allowed_when_min_samples_is_one(self): @@ -379,7 +379,7 @@ def test_one_pass_allowed_when_min_samples_is_one(self): self.make_profiled("1234", "1234", 1.0, p_rewards, p_lengths), # 8 < 10 self.make_profiled("1234567890", "1234567890", 1.0, p_rewards, p_lengths), # 20 >= 10 ] - apply_group_length_adjustments(results, self.cfg(min_samples=1)) + apply_group_length_penalties(results, self.cfg(min_samples=1)) assert rewards_of(results) == pytest.approx([1.0, 0.7]) @@ -399,7 +399,7 @@ def test_flat_penalty_exceeding_reward_clamps_at_zero(self): cfg = make_config( default={"enabled": True, "profiled_length_penalty": 1.5} ) - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 0.0]) def test_negative_env_reward_group_skipped(self): @@ -412,7 +412,7 @@ def test_negative_env_reward_group_skipped(self): make_result("12345", "12345", -1.0), make_result("1234567890", "1234567890", 1.0), ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([-1.0, 1.0]) def test_graded_rewards_group_skipped(self): @@ -424,7 +424,7 @@ def test_graded_rewards_group_skipped(self): make_result("12345", "12345", 0.5), make_result("1234567890", "1234567890", 1.0), ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([0.5, 1.0]) def test_all_wrong_group_stays_variance_free(self): @@ -441,7 +441,7 @@ def test_all_wrong_group_stays_variance_free(self): cfg = make_config( default={"enabled": True, "profiled_length_penalty": 0.3} ) - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([0.0, 0.0]) def test_band_multiplier_never_rewards_length_on_penalized_base(self): @@ -463,7 +463,7 @@ def test_band_multiplier_never_rewards_length_on_penalized_base(self): "profile_band_total": True, } ) - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) r_short, r_long = rewards_of(results) assert r_short == pytest.approx(0.0) assert r_long == pytest.approx(0.0) @@ -481,14 +481,14 @@ def test_explicit_false_channel_not_overridden_by_global_defaults(self): make_result("1234567890123456789012345", "12345", 1.0), # total 30 make_result("12345", "12345", 1.0), ] - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 1.0]) def test_unknown_key_raises(self): cfg = make_config(default={"enabled": True, "group_total_length_coeff": 0.1}) results = [make_result("12345", "12345", 1.0), make_result("123", "123", 1.0)] try: - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) except ValueError as e: assert "group_total_length_coeff" in str(e) else: @@ -505,5 +505,5 @@ def test_longest_penalty_under_binary_rewards(self): cfg = make_config( default={"enabled": True, "longest_total_penalty": 0.2} ) - apply_group_length_adjustments(results, cfg) + apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 0.8]) From 44de07de9a646454d64392b121d0cff7982c1ab9 Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 22:32:18 -0700 Subject: [PATCH 19/27] feat: pass-rate-scaled length penalty (MAI-style) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements -w_len * R_len(y_i) with R_len = rho_q * |y_i| / l_max, where rho_q is the prompt group's pass rate and l_max is the group's longest total rollout length (not a hyperparameter — self-normalizing per group; the longest correct rollout loses exactly w * rho_q). Dense pressure scaled by difficulty: easy prompts get strong shortening pressure, hard prompts little, and all-wrong groups (rho_q = 0) exactly none — preserving the no-gradient-without-correctness-signal invariant by construction. Applied to correct rollouts only, consistent with the module; single knob pass_rate_length_penalty_weight. Adds 5 unit tests (39 total) and the doc section. Co-Authored-By: Claude Fable 5 --- nemo_rl/utils/length_penalty.md | 39 +++++++++++++++++ nemo_rl/utils/length_penalty.py | 36 +++++++++++++++- tests/unit/utils/test_length_penalty.py | 56 +++++++++++++++++++++++++ 3 files changed, 129 insertions(+), 2 deletions(-) diff --git a/nemo_rl/utils/length_penalty.md b/nemo_rl/utils/length_penalty.md index 869fb5dd057..90d47d260f5 100644 --- a/nemo_rl/utils/length_penalty.md +++ b/nemo_rl/utils/length_penalty.md @@ -123,6 +123,7 @@ grpo: | `profiled_length_penalty` | Flat penalty for rollouts longer than a per-prompt profiled-length threshold. | | `profiled_length_n_std` | Number of standard deviations used in `mean + n_std * std` for profiled-length thresholding. | | `profiled_length_min_samples` | Minimum PASSING profiled rollouts required; below this no profiled penalty is applied. | +| `pass_rate_length_penalty_weight` | Weight of the pass-rate-scaled (MAI-style) dense length penalty; 0 disables. | | `profile_band_total` | Enables per-prompt `{a,b,f}` multiplier on total length for correct rollouts. | | `profile_band_reasoning` | Enables per-prompt `{a,b,f}` multiplier on reasoning length for correct rollouts. | | `profile_band_answer` | Enables per-prompt `{a,b,f}` multiplier on answer length for correct rollouts. | @@ -531,6 +532,44 @@ grpo: group_length_penalty_profile_gate_positive_only: true ``` +### 8. Pass-Rate-Scaled Length Penalty (MAI-style) + +Config keys: + +- `pass_rate_length_penalty_weight` + +Implements the MAI-paper length penalty `- w_len * R_len(y_i)` with + +```text +R_len(y_i) = rho_q * |y_i| / l_max +``` + +where `rho_q` is the prompt group's pass rate (fraction of rollouts with `reward > 0`), +`|y_i|` is the rollout's total generated length (reasoning + answer, in `length_type` units), +and `l_max` is the longest total length in the group (not a hyperparameter — the penalty is +self-normalizing per group). + +Properties: + +- Dense pressure scaled by difficulty: easy prompts (high pass rate) get strong shortening + pressure; hard prompts (low pass rate) get little. +- All-wrong groups (`rho_q = 0`) receive exactly zero penalty by construction, preserving the + no-gradient-without-correctness-signal invariant. +- Applied to correct rollouts only (consistent with the rest of this module; under binary + rewards plus the 0-clamp this matches the paper's behavior for wrong rollouts). +- The longest correct rollout in a group loses exactly `w_len * rho_q`; shorter ones lose + proportionally less. + +Example: + +```yaml +grpo: + length_penalty: + default: + enabled: true + pass_rate_length_penalty_weight: 0.2 +``` + ## Practical Notes - Most length algorithms act only on positive rollouts (`reward > 0`). diff --git a/nemo_rl/utils/length_penalty.py b/nemo_rl/utils/length_penalty.py index c416b82f75c..395d7c6c4fc 100644 --- a/nemo_rl/utils/length_penalty.py +++ b/nemo_rl/utils/length_penalty.py @@ -62,6 +62,7 @@ "profiled_length_penalty", "profiled_length_n_std", "profiled_length_min_samples", + "pass_rate_length_penalty_weight", "profile_band_total", "profile_band_reasoning", "profile_band_answer", @@ -253,6 +254,7 @@ def apply_group_length_penalties( all_zmad_answer_adj = [0.0] * n all_zmad_total_adj = [0.0] * n all_profiled_length_adj = [0.0] * n + all_pass_rate_len_adj = [0.0] * n reasoning_lengths = [0] * n answer_lengths = [0] * n total_lengths = [0] * n @@ -378,6 +380,29 @@ def apply_group_length_penalties( if total_lengths[idx] >= threshold: all_profiled_length_adj[idx] = -plp + # Pass-rate-scaled length penalty (MAI): -w * rho_q * |y_i| / l_max on + # correct rollouts, where rho_q is the group's pass rate. Easy prompts + # (high pass rate) get strong shortening pressure; hard prompts get + # little; all-wrong groups (rho_q = 0) get exactly none. + prlp_w = params.get("pass_rate_length_penalty_weight", 0.0) + if prlp_w > 0.0: + pass_rate = sum( + 1 for k in range(group_size) if original_rewards[g + k] > 0 + ) / float(group_size) + if pass_rate > 0.0: + # l_max is the group's longest rollout, so the penalty is + # self-normalizing: the longest correct rollout loses exactly + # w * rho_q, shorter ones proportionally less. + l_max = float(max(group_total[:group_size])) + if l_max > 0.0: + for k in range(group_size): + idx = g + k + if original_rewards[idx] <= 0: + continue + all_pass_rate_len_adj[idx] = ( + -prlp_w * pass_rate * total_lengths[idx] / l_max + ) + # Phase 2: debug print (only when verbose flag is set) if verbose: num_groups = n // num_gens if num_gens > 0 else 0 @@ -457,7 +482,13 @@ def apply_group_length_penalties( idx = g + k orig = original_rewards[idx] profiled_adj = all_profiled_length_adj[idx] if all_adjustments[idx] >= 0 else 0.0 - final = orig + all_adjustments[idx] + profiled_adj + final = max( + 0.0, + orig + + all_adjustments[idx] + + profiled_adj + + all_pass_rate_len_adj[idx], + ) print( f" [{k}] reward={orig:.4f}" f" reasoning_len={reasoning_lengths[idx]}" @@ -476,6 +507,7 @@ def apply_group_length_penalties( f" zmad_a={all_zmad_answer_adj[idx]:+.4f}" f" zmad_t={all_zmad_total_adj[idx]:+.4f}" f" profiled_len_adj={all_profiled_length_adj[idx]:+.4f}" + f" pass_rate_len_adj={all_pass_rate_len_adj[idx]:+.4f}" f" final_reward={final:.4f}", flush=True, ) @@ -492,7 +524,7 @@ def apply_group_length_penalties( # group-relative channels should not be double-penalized for the same # excess length. profiled_adj = all_profiled_length_adj[i] if all_adjustments[i] >= 0 else 0.0 - additive_delta = all_adjustments[i] + profiled_adj + additive_delta = all_adjustments[i] + profiled_adj + all_pass_rate_len_adj[i] additive_base_rewards[i] = original_rewards[i] + additive_delta # Rewards here are binary, so any negative value is penalty-created. # Length penalties may wipe a reward out but never flip its sign; this diff --git a/tests/unit/utils/test_length_penalty.py b/tests/unit/utils/test_length_penalty.py index e0d9bbf5e49..e1a42bc1bac 100644 --- a/tests/unit/utils/test_length_penalty.py +++ b/tests/unit/utils/test_length_penalty.py @@ -383,6 +383,62 @@ def test_one_pass_allowed_when_min_samples_is_one(self): assert rewards_of(results) == pytest.approx([1.0, 0.7]) +class TestPassRateLengthPenalty: + """MAI-style penalty: -w * rho_q * |y_i| / l_max on correct rollouts.""" + + def cfg(self, w=0.2, num_gens=2): + return make_config( + default={"enabled": True, "pass_rate_length_penalty_weight": w}, + num_gens=num_gens, + ) + + def test_all_correct_group_scales_by_relative_length(self): + # rho_q = 1.0, l_max = 20: penalties are w*10/20 and w*20/20. + results = [ + make_result("12345", "12345", 1.0), # total 10 -> -0.1 + make_result("1234567890", "1234567890", 1.0), # total 20 -> -0.2 + ] + apply_group_length_penalties(results, self.cfg(w=0.2)) + assert rewards_of(results) == pytest.approx([0.9, 0.8]) + + def test_pass_rate_scales_penalty(self): + # 1 of 2 correct -> rho_q = 0.5; only the correct rollout is penalized: + # 1.0 - 0.2 * 0.5 * 20/20 = 0.9. The wrong rollout stays at 0. + results = [ + make_result("1234567890", "1234567890", 1.0), # l_max contributor + make_result("12345", "12345", 0.0), + ] + apply_group_length_penalties(results, self.cfg(w=0.2)) + assert rewards_of(results) == pytest.approx([0.9, 0.0]) + + def test_all_wrong_group_gets_zero_penalty(self): + # rho_q = 0 -> no penalty at all, group stays variance-free. + results = [ + make_result("12345", "12345", 0.0), + make_result("1234567890", "1234567890", 0.0), + ] + apply_group_length_penalties(results, self.cfg(w=0.2)) + assert rewards_of(results) == pytest.approx([0.0, 0.0]) + + def test_zero_weight_is_noop(self): + results = [ + make_result("12345", "12345", 1.0), + make_result("1234567890", "1234567890", 1.0), + ] + apply_group_length_penalties(results, self.cfg(w=0.0)) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_wrong_rollout_length_still_sets_normalizer(self): + # l_max comes from the whole group (longest is a wrong rollout with + # total 40): the correct rollout (total 20) loses w * rho * 20/40. + results = [ + make_result("1234567890", "1234567890", 1.0), + make_result("1" * 20, "1" * 20, 0.0), + ] + apply_group_length_penalties(results, self.cfg(w=0.2)) + assert rewards_of(results) == pytest.approx([1.0 - 0.2 * 0.5 * 0.5, 0.0]) + + class TestReviewFixes: """Regression tests for review findings on PR #3852.""" From 89251d457542ddc6121c87a1277aa70af0270409 Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 22:47:47 -0700 Subject: [PATCH 20/27] fix: pass-rate penalty normalizes by longest CORRECT rollout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit l_max over the whole group let long wrong rambles dilute the penalty spread between correct rollouts — the only component GRPO's mean-centered advantage actually sees — and made 'longest correct loses exactly w*rho' hold only when the longest rollout happened to be correct. Normalizing over correct rollouts only decouples the pressure from wrong-rollout noise and makes the invariant exact. Co-Authored-By: Claude Fable 5 --- nemo_rl/utils/length_penalty.md | 4 ++-- nemo_rl/utils/length_penalty.py | 16 ++++++++++++---- tests/unit/utils/test_length_penalty.py | 10 ++++++---- 3 files changed, 20 insertions(+), 10 deletions(-) diff --git a/nemo_rl/utils/length_penalty.md b/nemo_rl/utils/length_penalty.md index 90d47d260f5..60ffdc469b3 100644 --- a/nemo_rl/utils/length_penalty.md +++ b/nemo_rl/utils/length_penalty.md @@ -546,8 +546,8 @@ R_len(y_i) = rho_q * |y_i| / l_max where `rho_q` is the prompt group's pass rate (fraction of rollouts with `reward > 0`), `|y_i|` is the rollout's total generated length (reasoning + answer, in `length_type` units), -and `l_max` is the longest total length in the group (not a hyperparameter — the penalty is -self-normalizing per group). +and `l_max` is the longest total length among the group's CORRECT rollouts (not a hyperparameter — +the penalty is self-normalizing over the set it applies to; long wrong rollouts do not dilute it). Properties: diff --git a/nemo_rl/utils/length_penalty.py b/nemo_rl/utils/length_penalty.py index 395d7c6c4fc..f66d5574dfd 100644 --- a/nemo_rl/utils/length_penalty.py +++ b/nemo_rl/utils/length_penalty.py @@ -390,10 +390,18 @@ def apply_group_length_penalties( 1 for k in range(group_size) if original_rewards[g + k] > 0 ) / float(group_size) if pass_rate > 0.0: - # l_max is the group's longest rollout, so the penalty is - # self-normalizing: the longest correct rollout loses exactly - # w * rho_q, shorter ones proportionally less. - l_max = float(max(group_total[:group_size])) + # l_max is the longest CORRECT rollout, so the penalty is + # self-normalizing over the set it applies to: the longest + # correct rollout loses exactly w * rho_q, shorter ones + # proportionally less — decoupled from wrong-rollout lengths + # (long wrong rambles must not dilute the pressure). + l_max = float( + max( + total_lengths[g + k] + for k in range(group_size) + if original_rewards[g + k] > 0 + ) + ) if l_max > 0.0: for k in range(group_size): idx = g + k diff --git a/tests/unit/utils/test_length_penalty.py b/tests/unit/utils/test_length_penalty.py index e1a42bc1bac..fa901ceb3e3 100644 --- a/tests/unit/utils/test_length_penalty.py +++ b/tests/unit/utils/test_length_penalty.py @@ -428,15 +428,17 @@ def test_zero_weight_is_noop(self): apply_group_length_penalties(results, self.cfg(w=0.0)) assert rewards_of(results) == pytest.approx([1.0, 1.0]) - def test_wrong_rollout_length_still_sets_normalizer(self): - # l_max comes from the whole group (longest is a wrong rollout with - # total 40): the correct rollout (total 20) loses w * rho * 20/40. + def test_wrong_rollout_length_does_not_set_normalizer(self): + # l_max comes from CORRECT rollouts only: the wrong rollout's total of + # 40 is ignored, the correct rollout (total 20) is its own max and + # loses exactly w * rho = 0.2 * 0.5. A long wrong ramble must not + # dilute the pressure on correct rollouts. results = [ make_result("1234567890", "1234567890", 1.0), make_result("1" * 20, "1" * 20, 0.0), ] apply_group_length_penalties(results, self.cfg(w=0.2)) - assert rewards_of(results) == pytest.approx([1.0 - 0.2 * 0.5 * 0.5, 0.0]) + assert rewards_of(results) == pytest.approx([1.0 - 0.2 * 0.5, 0.0]) class TestReviewFixes: From c89b83f7169d3ac3bb6616ed9823742aa40682f9 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Thu, 27 Aug 2026 11:32:57 -0700 Subject: [PATCH 21/27] style: satisfy pre-commit hooks (rename doc, ruff format/isort, F541) CI runs `pre-commit run --all-files` (.github/workflows/cicd-main.yml:387), which failed on four hooks for this branch: - no-underscore-md: nemo_rl/utils/length_penalty.md was the sole offender. This hook is a hard exit 1 and is not auto-fixable. Renamed to length-penalty.md, matching the hyphenated sibling docs under nemo_rl/. No references to the old filename exist anywhere in the tree. - ruff (F541): dropped a needless f-string prefix in scripts/build_profile_band_dataset.py. - ruff check --select I: reordered the length_penalty import in nemo_rl/experience/rollouts.py. - ruff-format: reflowed 3 files. Verified: `pre-commit run --all-files` now exits 0; the 4 touched files AST-compare identical to their previous versions apart from the intended F541 fix and the import reorder; tests/unit/utils/test_length_penalty.py is 39 passed. Co-Authored-By: Claude Opus 5 (1M context) --- nemo_rl/experience/rollouts.py | 2 +- .../{length_penalty.md => length-penalty.md} | 0 nemo_rl/utils/length_penalty.py | 122 +++++++++++++----- scripts/build_profile_band_dataset.py | 24 ++-- tests/unit/utils/test_length_penalty.py | 20 ++- 5 files changed, 118 insertions(+), 50 deletions(-) rename nemo_rl/utils/{length_penalty.md => length-penalty.md} (100%) diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 25218687942..b920aebc1b1 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -68,11 +68,11 @@ GenerationOutputSpec, GenerationSamplingParams, ) +from nemo_rl.utils.length_penalty import apply_group_length_penalties from nemo_rl.utils.multimodal_payload_metrics import ( collect_multimodal_payload_metrics, print_multimodal_payload_metrics, ) -from nemo_rl.utils.length_penalty import apply_group_length_penalties from nemo_rl.utils.timer import Timer TokenizerType = PreTrainedTokenizerBase diff --git a/nemo_rl/utils/length_penalty.md b/nemo_rl/utils/length-penalty.md similarity index 100% rename from nemo_rl/utils/length_penalty.md rename to nemo_rl/utils/length-penalty.md diff --git a/nemo_rl/utils/length_penalty.py b/nemo_rl/utils/length_penalty.py index f66d5574dfd..9d87a4b5175 100644 --- a/nemo_rl/utils/length_penalty.py +++ b/nemo_rl/utils/length_penalty.py @@ -73,20 +73,24 @@ ) # Param keys that should be merged as bools rather than floats. -_BOOL_PARAM_KEYS = frozenset({ - "enabled", - "profile_band_total", - "profile_band_reasoning", - "profile_band_answer", - "group_length_penalty_profile_gate", - "group_length_penalty_profile_gate_positive_only", -}) +_BOOL_PARAM_KEYS = frozenset( + { + "enabled", + "profile_band_total", + "profile_band_reasoning", + "profile_band_answer", + "group_length_penalty_profile_gate", + "group_length_penalty_profile_gate_positive_only", + } +) -_STR_PARAM_KEYS = frozenset({ - "length_type", - "group_length_penalty_profile_gate_channel", - "group_length_penalty_profile_gate_field", -}) +_STR_PARAM_KEYS = frozenset( + { + "length_type", + "group_length_penalty_profile_gate_channel", + "group_length_penalty_profile_gate_field", + } +) # Length adjustments are defined for binary (0/1) env rewards only. Agents # already warned about non-binary rewards (warn once per agent, then skip @@ -99,6 +103,7 @@ def _is_binary_reward(value: float) -> bool: v = float(value) return abs(v) <= _BINARY_REWARD_TOL or abs(v - 1.0) <= _BINARY_REWARD_TOL + def _extract_reasoning_and_answer_text(result: dict[str, Any]) -> tuple[str, str]: """Extract reasoning and answer text from the Response API output items.""" fr = result.get("full_result", {}) @@ -112,17 +117,33 @@ def _extract_reasoning_and_answer_text(result: dict[str, Any]) -> tuple[str, str reasoning_text = "" answer_text = "" for item in output_items: - item_type = item.get("type", "") if isinstance(item, dict) else getattr(item, "type", "") + item_type = ( + item.get("type", "") + if isinstance(item, dict) + else getattr(item, "type", "") + ) if item_type == "reasoning": - summaries = item.get("summary", []) if isinstance(item, dict) else getattr(item, "summary", []) + summaries = ( + item.get("summary", []) + if isinstance(item, dict) + else getattr(item, "summary", []) + ) for s in summaries: t = s.get("text", "") if isinstance(s, dict) else getattr(s, "text", "") reasoning_text += t elif item_type == "message": - content = item.get("content", []) if isinstance(item, dict) else getattr(item, "content", []) + content = ( + item.get("content", []) + if isinstance(item, dict) + else getattr(item, "content", []) + ) if isinstance(content, list): for c in content: - t = c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") + t = ( + c.get("text", "") + if isinstance(c, dict) + else getattr(c, "text", "") + ) answer_text += t elif isinstance(content, str): answer_text += content @@ -298,8 +319,16 @@ def apply_group_length_penalties( idx = g + k r_text, a_text = texts[idx] if use_tokens and tokenizer is not None: - reasoning_lengths[idx] = len(tokenizer.encode(r_text, add_special_tokens=False)) if r_text else 0 - answer_lengths[idx] = len(tokenizer.encode(a_text, add_special_tokens=False)) if a_text else 0 + reasoning_lengths[idx] = ( + len(tokenizer.encode(r_text, add_special_tokens=False)) + if r_text + else 0 + ) + answer_lengths[idx] = ( + len(tokenizer.encode(a_text, add_special_tokens=False)) + if a_text + else 0 + ) else: reasoning_lengths[idx] = len(r_text) answer_lengths[idx] = len(a_text) @@ -425,7 +454,9 @@ def apply_group_length_penalties( for g in range(0, n, num_gens): agent_name = agent_names[g] group_size = min(num_gens, n - g) - low_effort = any(results[g + k].get("low_effort_applied") for k in range(group_size)) + low_effort = any( + results[g + k].get("low_effort_applied") for k in range(group_size) + ) params = _resolve_agent_params(agent_name, agents_cfg, defaults) skipped = params is None disabled = params is not None and not params.get("enabled", True) @@ -489,7 +520,9 @@ def apply_group_length_penalties( for k in range(group_size): idx = g + k orig = original_rewards[idx] - profiled_adj = all_profiled_length_adj[idx] if all_adjustments[idx] >= 0 else 0.0 + profiled_adj = ( + all_profiled_length_adj[idx] if all_adjustments[idx] >= 0 else 0.0 + ) final = max( 0.0, orig @@ -559,7 +592,6 @@ def apply_group_length_penalties( ) - def _resolve_global_profile_band(pb_cfg: Any) -> dict[str, dict[str, Any]]: """Parse ``length_penalty.profile_band`` into per-channel {a, b, f} blocks. @@ -716,7 +748,9 @@ def _group_length_profile_gate_info( enabled = bool(params.get("group_length_penalty_profile_gate", False)) channel = str(params.get("group_length_penalty_profile_gate_channel", "total")) field = str(params.get("group_length_penalty_profile_gate_field", "a")) - positive_only = bool(params.get("group_length_penalty_profile_gate_positive_only", True)) + positive_only = bool( + params.get("group_length_penalty_profile_gate_positive_only", True) + ) info = { "enabled": enabled, "open": True, @@ -954,7 +988,11 @@ def _apply_length_penaltyes_and_penalties( # Reasoning bonus: shortest non-empty reasoning among top scorers if reasoning_bonus > 0: - valid = [(pi, pos_reasoning[k]) for k, pi in enumerate(positive_indices) if pos_reasoning[k] > 0] + valid = [ + (pi, pos_reasoning[k]) + for k, pi in enumerate(positive_indices) + if pos_reasoning[k] > 0 + ] if valid: shortest_pi, _ = min(valid, key=lambda x: x[1]) if adjusted[shortest_pi] >= top_threshold: @@ -964,7 +1002,11 @@ def _apply_length_penaltyes_and_penalties( # Answer bonus: shortest non-empty answer among top scorers if answer_bonus > 0: - valid = [(pi, pos_answer[k]) for k, pi in enumerate(positive_indices) if pos_answer[k] > 0] + valid = [ + (pi, pos_answer[k]) + for k, pi in enumerate(positive_indices) + if pos_answer[k] > 0 + ] if valid: shortest_pi, _ = min(valid, key=lambda x: x[1]) if adjusted[shortest_pi] >= top_threshold: @@ -974,7 +1016,11 @@ def _apply_length_penaltyes_and_penalties( # Total bonus: shortest combined (reasoning + answer) among top scorers if total_bonus > 0: - valid = [(pi, pos_total[k]) for k, pi in enumerate(positive_indices) if pos_total[k] > 0] + valid = [ + (pi, pos_total[k]) + for k, pi in enumerate(positive_indices) + if pos_total[k] > 0 + ] if valid: shortest_pi, _ = min(valid, key=lambda x: x[1]) if adjusted[shortest_pi] >= top_threshold: @@ -984,7 +1030,11 @@ def _apply_length_penaltyes_and_penalties( # Longest reasoning penalty: longest among top-percentile scorers only if longest_reasoning_penalty > 0: - valid = [(pi, reasoning_lengths[pi]) for pi in top_scorer_indices if reasoning_lengths[pi] > 0] + valid = [ + (pi, reasoning_lengths[pi]) + for pi in top_scorer_indices + if reasoning_lengths[pi] > 0 + ] if len(valid) >= 2: longest_pi, _ = max(valid, key=lambda x: x[1]) pen = -longest_reasoning_penalty @@ -994,7 +1044,11 @@ def _apply_length_penaltyes_and_penalties( # Longest answer penalty: longest among top-percentile scorers only if longest_answer_penalty > 0: - valid = [(pi, answer_lengths[pi]) for pi in top_scorer_indices if answer_lengths[pi] > 0] + valid = [ + (pi, answer_lengths[pi]) + for pi in top_scorer_indices + if answer_lengths[pi] > 0 + ] if len(valid) >= 2: longest_pi, _ = max(valid, key=lambda x: x[1]) pen = -longest_answer_penalty @@ -1004,7 +1058,11 @@ def _apply_length_penaltyes_and_penalties( # Longest total penalty: longest combined length among top-percentile scorers only if longest_total_penalty > 0: - valid = [(pi, total_lengths[pi]) for pi in top_scorer_indices if total_lengths[pi] > 0] + valid = [ + (pi, total_lengths[pi]) + for pi in top_scorer_indices + if total_lengths[pi] > 0 + ] if len(valid) >= 2: longest_pi, _ = max(valid, key=lambda x: x[1]) pen = -longest_total_penalty @@ -1013,7 +1071,11 @@ def _apply_length_penaltyes_and_penalties( t_longest_pen_per[longest_pi] = pen # Independent reasoning, answer, and total length penalties (zero-centered) - if group_reasoning_length_penalty_coeff > 0 or group_answer_length_penalty_coeff > 0 or group_total_length_penalty_coeff > 0: + if ( + group_reasoning_length_penalty_coeff > 0 + or group_answer_length_penalty_coeff > 0 + or group_total_length_penalty_coeff > 0 + ): reasoning_weights = _compute_length_weights(pos_reasoning) answer_weights = _compute_length_weights(pos_answer) total_weights = _compute_length_weights(pos_total) diff --git a/scripts/build_profile_band_dataset.py b/scripts/build_profile_band_dataset.py index b4f5a4721fb..afd4e94bdf1 100644 --- a/scripts/build_profile_band_dataset.py +++ b/scripts/build_profile_band_dataset.py @@ -72,18 +72,22 @@ yaml = None # type: ignore CHANNEL_TO_LENGTHS_KEY = { - "total": "profiled_output_lengths", + "total": "profiled_output_lengths", "reasoning": "profiled_reasoning_lengths", - "answer": "profiled_answer_lengths", + "answer": "profiled_answer_lengths", } def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) p.add_argument("--input", required=True, help="Path to profiled JSONL (input).") p.add_argument("--config", required=True, help="Path to profile_band yaml config.") p.add_argument("--output", required=True, help="Path to write augmented JSONL.") - p.add_argument("--quiet", action="store_true", help="Suppress per-row diagnostics summary.") + p.add_argument( + "--quiet", action="store_true", help="Suppress per-row diagnostics summary." + ) return p.parse_args() @@ -106,7 +110,9 @@ def load_config(path: str) -> dict[str, Any]: cfg["_f_table"] = f_table bad_channels = [c for c in cfg["channels"] if c not in CHANNEL_TO_LENGTHS_KEY] if bad_channels: - sys.exit(f"unknown channels in config: {bad_channels}; valid: {sorted(CHANNEL_TO_LENGTHS_KEY)}") + sys.exit( + f"unknown channels in config: {bad_channels}; valid: {sorted(CHANNEL_TO_LENGTHS_KEY)}" + ) return cfg @@ -197,7 +203,9 @@ def main() -> None: print(f"input : {in_path}") print(f"output : {out_path}") print(f"rows in: {n_total}") - print(f"rows w/ profile_band: {n_with_band} ({100*n_with_band/n_total:.1f}%)") + print( + f"rows w/ profile_band: {n_with_band} ({100 * n_with_band / n_total:.1f}%)" + ) print(f"rows skipped (pass_rate not in f_table): {n_pr_skipped}") if n_pr_skipped_by_passrate: print(" by pass_rate:") @@ -206,8 +214,8 @@ def main() -> None: print("channels emitted (per row, summed):") for ch in cfg["channels"]: c = n_channels_emitted.get(ch, 0) - print(f" {ch:>10}: {c} ({100*c/n_total:.1f}% of rows)") - print(f"f_table (rounded pass_rate -> f):") + print(f" {ch:>10}: {c} ({100 * c / n_total:.1f}% of rows)") + print("f_table (rounded pass_rate -> f):") for pr in sorted(cfg["_f_table"]): print(f" {pr:.3f} -> {cfg['_f_table'][pr]}") print(f"n_std={cfg['n_std']} min_passing={cfg['min_passing']}") diff --git a/tests/unit/utils/test_length_penalty.py b/tests/unit/utils/test_length_penalty.py index fa901ceb3e3..80a654b28d3 100644 --- a/tests/unit/utils/test_length_penalty.py +++ b/tests/unit/utils/test_length_penalty.py @@ -345,7 +345,9 @@ def test_enough_passes_penalizes_over_threshold(self): p_rewards, p_lengths = [1, 1, 0], [10, 14, 100] results = [ self.make_profiled("12345", "12345", 1.0, p_rewards, p_lengths), # 10 < thr - self.make_profiled("1234567890", "1234567890", 1.0, p_rewards, p_lengths), # 20 >= thr + self.make_profiled( + "1234567890", "1234567890", 1.0, p_rewards, p_lengths + ), # 20 >= thr ] apply_group_length_penalties(results, self.cfg()) assert rewards_of(results) == pytest.approx([1.0, 0.7]) @@ -377,7 +379,9 @@ def test_one_pass_allowed_when_min_samples_is_one(self): p_rewards, p_lengths = [1, 0], [10, 100] results = [ self.make_profiled("1234", "1234", 1.0, p_rewards, p_lengths), # 8 < 10 - self.make_profiled("1234567890", "1234567890", 1.0, p_rewards, p_lengths), # 20 >= 10 + self.make_profiled( + "1234567890", "1234567890", 1.0, p_rewards, p_lengths + ), # 20 >= 10 ] apply_group_length_penalties(results, self.cfg(min_samples=1)) assert rewards_of(results) == pytest.approx([1.0, 0.7]) @@ -454,9 +458,7 @@ def test_flat_penalty_exceeding_reward_clamps_at_zero(self): for r in results: r["profiled_rewards"] = [1, 1] r["profiled_output_lengths"] = [10, 10] - cfg = make_config( - default={"enabled": True, "profiled_length_penalty": 1.5} - ) + cfg = make_config(default={"enabled": True, "profiled_length_penalty": 1.5}) apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 0.0]) @@ -496,9 +498,7 @@ def test_all_wrong_group_stays_variance_free(self): for r in results: r["profiled_rewards"] = [1, 1] r["profiled_output_lengths"] = [10, 10] - cfg = make_config( - default={"enabled": True, "profiled_length_penalty": 0.3} - ) + cfg = make_config(default={"enabled": True, "profiled_length_penalty": 0.3}) apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([0.0, 0.0]) @@ -560,8 +560,6 @@ def test_longest_penalty_under_binary_rewards(self): make_result("12345", "12345", 1.0), make_result("1234567890", "1234567890", 1.0), ] - cfg = make_config( - default={"enabled": True, "longest_total_penalty": 0.2} - ) + cfg = make_config(default={"enabled": True, "longest_total_penalty": 0.2}) apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 0.8]) From 3973211b866136731685722bb7fe86e4492a17f7 Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 23:12:22 -0700 Subject: [PATCH 22/27] feat(launcher): 4-GPU-node support, sbatch comment/qos/extra-dependency, overridable sandbox command --gres honors GPUS_PER_NODE (HSG GB200 has 4 GPUs/node; ray.sub already consumes the same variable), SLURM_COMMENT enables idle-reaper exemptions required by async-GRPO buffer fill, SLURM_QOS and SLURM_EXTRA_DEPENDENCY are optional passthroughs, and SANDBOX_COMMAND can be overridden for sandbox startup retry wrappers. Co-Authored-By: Claude Fable 5 --- examples/nemo_gym/nemotron-3-super/super_launch.sh | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/examples/nemo_gym/nemotron-3-super/super_launch.sh b/examples/nemo_gym/nemotron-3-super/super_launch.sh index c058cf2fa38..5fbcd59291d 100755 --- a/examples/nemo_gym/nemotron-3-super/super_launch.sh +++ b/examples/nemo_gym/nemotron-3-super/super_launch.sh @@ -134,7 +134,7 @@ export LISTEN_PORT=6000 export NGINX_PORT=6000 export NEMO_SKILLS_SANDBOX_PORT=6000 export SANDBOX_CONTAINER -export SANDBOX_COMMAND="/start-with-nginx.sh" +export SANDBOX_COMMAND="${SANDBOX_COMMAND:-/start-with-nginx.sh}" export SANDBOX_ENV_VARS="NEMO_SKILLS_SANDBOX_PORT=${NEMO_SKILLS_SANDBOX_PORT}" # ---- Build the run command ---- @@ -216,11 +216,17 @@ SBATCH_CMD=( --job-name="${WANDB_NAME}" --partition="${SLURM_PARTITION}" --time="${SLURM_TIME_LIMIT}" - --gres=gpu:8 + --gres=gpu:"${GPUS_PER_NODE:-8}" --exclusive - --dependency=singleton + --dependency=singleton${SLURM_EXTRA_DEPENDENCY:+,${SLURM_EXTRA_DEPENDENCY}} ray.sub ) +if [[ -n "${SLURM_QOS:-}" ]]; then + SBATCH_CMD=("${SBATCH_CMD[@]:0:1}" --qos="${SLURM_QOS}" "${SBATCH_CMD[@]:1}") +fi +if [[ -n "${SLURM_COMMENT:-}" ]]; then + SBATCH_CMD=("${SBATCH_CMD[@]:0:1}" --comment="${SLURM_COMMENT}" "${SBATCH_CMD[@]:1}") +fi if [[ "$DRY_RUN" == true ]]; then echo "" From d202b07ea7884aeca769addea9b7e3b0ec6e1223 Mon Sep 17 00:00:00 2001 From: adithyare Date: Wed, 26 Aug 2026 23:57:17 -0700 Subject: [PATCH 23/27] feat(launcher): optional stable RESULTS_DIR layout When RESULTS_DIR is set: checkpoints/ (stable dir for singleton auto-resume), runs//{logs,slurm} with a runs/latest symlink, ray_logs/ as ray.sub's BASE_LOG_DIR, sbatch --output/--error into the run's slurm/, and nemo_gym_log_dir under the run's logs. Unset keeps the legacy snapshot-relative results/ and logs/ dirs. Co-Authored-By: Claude Fable 5 --- .../nemo_gym/nemotron-3-super/super_launch.sh | 29 +++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/examples/nemo_gym/nemotron-3-super/super_launch.sh b/examples/nemo_gym/nemotron-3-super/super_launch.sh index 5fbcd59291d..e05499faffe 100755 --- a/examples/nemo_gym/nemotron-3-super/super_launch.sh +++ b/examples/nemo_gym/nemotron-3-super/super_launch.sh @@ -70,8 +70,26 @@ fi # ---- Derived paths ---- CODE_DIR=$(realpath "$PWD") WANDB_NAME="${EXP_NAME}" -CHECKPOINT_DIR="results/${EXP_NAME}" -LOG_DIR="logs/${EXP_NAME}" +# Optional stable results layout: set RESULTS_DIR to get +# $RESULTS_DIR/checkpoints (stable -> singleton auto-resume) +# $RESULTS_DIR/runs//{logs,slurm} (per-submission; runs/latest symlink) +# $RESULTS_DIR/ray_logs/-logs (ray.sub infra logs via BASE_LOG_DIR) +# Unset: legacy snapshot-relative results/ and logs/ dirs. +RESULTS_DIR="${RESULTS_DIR:-}" +if [[ -n "${RESULTS_DIR}" ]]; then + CHECKPOINT_DIR="${CHECKPOINT_DIR:-${RESULTS_DIR}/checkpoints}" + RUN_DIR="${RESULTS_DIR}/runs/$(date +%Y%m%d-%H%M)" + LOG_DIR="${RUN_DIR}/logs" + SLURM_LOG_DIR="${RUN_DIR}/slurm" + mkdir -p "${CHECKPOINT_DIR}" "${LOG_DIR}" "${SLURM_LOG_DIR}" + ln -sfn "$(realpath "${RUN_DIR}")" "${RESULTS_DIR}/runs/latest" + export BASE_LOG_DIR="${BASE_LOG_DIR:-${RESULTS_DIR}/ray_logs}" + mkdir -p "${BASE_LOG_DIR}" +else + CHECKPOINT_DIR="results/${EXP_NAME}" + LOG_DIR="logs/${EXP_NAME}" + SLURM_LOG_DIR="" +fi VLLM_CACHE_DIR="${PERSISTENT_CACHE}/vllm_compile_cache" FLASHINFER_CUBIN_CACHE="${PERSISTENT_CACHE}/flashinfer_cubins" @@ -171,6 +189,10 @@ export COMMAND="export HF_MODULES_CACHE=${HF_MODULES_CACHE_DIR} ; \ data.train.data_path=${TRAIN_PATH} \ data.validation.data_path=${VAL_PATH}" +if [[ -n "${RESULTS_DIR}" ]]; then + COMMAND="$COMMAND env.nemo_gym.nemo_gym_log_dir=${LOG_DIR}/nemo_gym" +fi + if [[ -n "$SIF_DIR" ]]; then COMMAND="$COMMAND sif_dir=${SIF_DIR}" fi @@ -227,6 +249,9 @@ fi if [[ -n "${SLURM_COMMENT:-}" ]]; then SBATCH_CMD=("${SBATCH_CMD[@]:0:1}" --comment="${SLURM_COMMENT}" "${SBATCH_CMD[@]:1}") fi +if [[ -n "${SLURM_LOG_DIR}" ]]; then + SBATCH_CMD=("${SBATCH_CMD[@]:0:1}" --output="${SLURM_LOG_DIR}/%j.out" --error="${SLURM_LOG_DIR}/%j.err" "${SBATCH_CMD[@]:1}") +fi if [[ "$DRY_RUN" == true ]]; then echo "" From 49746de179c5f913f40720d4d8ea89db673270ed Mon Sep 17 00:00:00 2001 From: adithyare Date: Thu, 27 Aug 2026 00:03:39 -0700 Subject: [PATCH 24/27] feat(launcher): optional --segment sbatch flag with divisibility guard Co-Authored-By: Claude Fable 5 --- examples/nemo_gym/nemotron-3-super/super_launch.sh | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/examples/nemo_gym/nemotron-3-super/super_launch.sh b/examples/nemo_gym/nemotron-3-super/super_launch.sh index e05499faffe..a90c6a83d49 100755 --- a/examples/nemo_gym/nemotron-3-super/super_launch.sh +++ b/examples/nemo_gym/nemotron-3-super/super_launch.sh @@ -252,6 +252,13 @@ fi if [[ -n "${SLURM_LOG_DIR}" ]]; then SBATCH_CMD=("${SBATCH_CMD[@]:0:1}" --output="${SLURM_LOG_DIR}/%j.out" --error="${SLURM_LOG_DIR}/%j.err" "${SBATCH_CMD[@]:1}") fi +if [[ -n "${SLURM_SEGMENT:-}" ]]; then + if (( SBATCH_NUM_NODES % SLURM_SEGMENT != 0 )); then + echo "Error: SBATCH_NUM_NODES=${SBATCH_NUM_NODES} not divisible by SLURM_SEGMENT=${SLURM_SEGMENT}" >&2 + exit 1 + fi + SBATCH_CMD=("${SBATCH_CMD[@]:0:1}" --segment="${SLURM_SEGMENT}" "${SBATCH_CMD[@]:1}") +fi if [[ "$DRY_RUN" == true ]]; then echo "" From 9901ae32fe27561e115345ef5b9e3eb9c8c3951a Mon Sep 17 00:00:00 2001 From: adithyare Date: Thu, 27 Aug 2026 01:09:23 -0700 Subject: [PATCH 25/27] =?UTF-8?q?fix(launcher):=20drop=20nemo=5Fgym=5Flog?= =?UTF-8?q?=5Fdir=20override=20=E2=80=94=20key=20not=20in=20this=20branch'?= =?UTF-8?q?s=20schema?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ported from the geshen/nano launcher but this branch's NemoGym config does not define it; OmegaConf struct mode rejects the override and both ablation primers (6580116/6580230) died at config load. Co-Authored-By: Claude Fable 5 --- examples/nemo_gym/nemotron-3-super/super_launch.sh | 4 ---- 1 file changed, 4 deletions(-) diff --git a/examples/nemo_gym/nemotron-3-super/super_launch.sh b/examples/nemo_gym/nemotron-3-super/super_launch.sh index a90c6a83d49..6b72177e722 100755 --- a/examples/nemo_gym/nemotron-3-super/super_launch.sh +++ b/examples/nemo_gym/nemotron-3-super/super_launch.sh @@ -189,10 +189,6 @@ export COMMAND="export HF_MODULES_CACHE=${HF_MODULES_CACHE_DIR} ; \ data.train.data_path=${TRAIN_PATH} \ data.validation.data_path=${VAL_PATH}" -if [[ -n "${RESULTS_DIR}" ]]; then - COMMAND="$COMMAND env.nemo_gym.nemo_gym_log_dir=${LOG_DIR}/nemo_gym" -fi - if [[ -n "$SIF_DIR" ]]; then COMMAND="$COMMAND sif_dir=${SIF_DIR}" fi From 79470b531394a29feba8c81f36e7f396024648b5 Mon Sep 17 00:00:00 2001 From: adithyare Date: Thu, 27 Aug 2026 01:26:13 -0700 Subject: [PATCH 26/27] =?UTF-8?q?fix(launcher):=20force-override=20(++)=20?= =?UTF-8?q?gym=20venv=20keys=20=E2=80=94=20absent=20from=20pipeline-derive?= =?UTF-8?q?d=20configs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit env.nemo_gym.uv_venv_dir / skip_venv_if_present exist in stage1-based configs but not pipeline-derived ones; plain overrides fail Hydra composition on the latter (killed ablation r1 and r2: 6580116/6580230, 6582890/6582899 — r1 was misattributed to nemo_gym_log_dir, which was merely later in the same list). '++' sets the key in both cases. Co-Authored-By: Claude Fable 5 --- examples/nemo_gym/nemotron-3-super/super_launch.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/nemo_gym/nemotron-3-super/super_launch.sh b/examples/nemo_gym/nemotron-3-super/super_launch.sh index 6b72177e722..ea71205cabd 100755 --- a/examples/nemo_gym/nemotron-3-super/super_launch.sh +++ b/examples/nemo_gym/nemotron-3-super/super_launch.sh @@ -178,8 +178,8 @@ export COMMAND="export HF_MODULES_CACHE=${HF_MODULES_CACHE_DIR} ; \ PYTHONPATH=${SNAPSHOT_DIR}:\${PYTHONPATH:-} \ python ./examples/nemo_gym/run_grpo_nemo_gym.py \ --config ${CONFIG_PATH} \ - env.nemo_gym.uv_venv_dir=${GYM_VENV_DIR} \ - env.nemo_gym.skip_venv_if_present=true \ + ++env.nemo_gym.uv_venv_dir=${GYM_VENV_DIR} \ + ++env.nemo_gym.skip_venv_if_present=true \ policy.model_name=${MODEL_PATH} \ checkpointing.checkpoint_dir=${CHECKPOINT_DIR} \ logger.log_dir=${LOG_DIR} \ From 1643600660037075d514749bba769a3cd01e4be7 Mon Sep 17 00:00:00 2001 From: adithyare Date: Thu, 27 Aug 2026 12:50:31 -0700 Subject: [PATCH 27/27] fix: enabled defaults True consistently in apply_group_length_penalties MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The early return read enabled with default False while per-group param resolution used default True — a default: block omitting enabled silently no-oped, and adding an unrelated empty agent_overrides entry flipped the same block active (review finding on #3852). A configured default: block is now intent-to-enable at both read sites; set enabled: false to opt out. Two regression tests cover the repro and the unrelated-key independence. Co-Authored-By: Claude Fable 5 --- nemo_rl/utils/length_penalty.py | 5 ++++- tests/unit/utils/test_length_penalty.py | 30 +++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/nemo_rl/utils/length_penalty.py b/nemo_rl/utils/length_penalty.py index 9d87a4b5175..5d9b64bef24 100644 --- a/nemo_rl/utils/length_penalty.py +++ b/nemo_rl/utils/length_penalty.py @@ -218,7 +218,10 @@ def apply_group_length_penalties( agents_cfg = length_cfg.get("agent_overrides") global_band = _resolve_global_profile_band(length_cfg.get("profile_band")) verbose = bool(length_cfg.get("verbose", False)) - if not default_cfg.get("enabled", False) and not agents_cfg and not global_band: + # `enabled` defaults True here AND in the per-group param resolution: a + # configured `default:` block is intent-to-enable; omitting `enabled` must + # not silently no-op (and must not depend on unrelated keys being present). + if not default_cfg.get("enabled", True) and not agents_cfg and not global_band: return num_gens = master_config["grpo"]["num_generations_per_prompt"] diff --git a/tests/unit/utils/test_length_penalty.py b/tests/unit/utils/test_length_penalty.py index 80a654b28d3..3bde11fd545 100644 --- a/tests/unit/utils/test_length_penalty.py +++ b/tests/unit/utils/test_length_penalty.py @@ -542,6 +542,36 @@ def test_explicit_false_channel_not_overridden_by_global_defaults(self): apply_group_length_penalties(results, cfg) assert rewards_of(results) == pytest.approx([1.0, 1.0]) + def test_default_block_without_enabled_is_active(self): + # `enabled` defaults True consistently: a default: block that omits it + # applies its penalties (previously the early return used default False + # and silently no-oped). + cfg = make_config(default={"group_total_length_penalty_coeff": 0.2}) + results = [ + make_result("12345", "12345", 1.0), + make_result("1234567890123456789012345", "12345", 1.0), + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.1, 0.9]) + + def test_behavior_independent_of_empty_agent_overrides(self): + # The same default: block must behave identically with or without an + # unrelated empty agent_overrides entry (previously it flipped the + # early return and changed behavior). + base = {"group_total_length_penalty_coeff": 0.2} + outs = [] + for extra_overrides in (None, {"some_agent": {}}): + cfg = make_config(default=dict(base)) + if extra_overrides is not None: + cfg["grpo"]["length_penalty"]["agent_overrides"] = extra_overrides + results = [ + make_result("12345", "12345", 1.0), + make_result("1234567890123456789012345", "12345", 1.0), + ] + apply_group_length_penalties(results, cfg) + outs.append(rewards_of(results)) + assert outs[0] == pytest.approx(outs[1]) + def test_unknown_key_raises(self): cfg = make_config(default={"enabled": True, "group_total_length_coeff": 0.1}) results = [make_result("12345", "12345", 1.0), make_result("123", "123", 1.0)]