|
| 1 | +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. |
| 2 | + |
| 3 | +import argparse |
| 4 | +from dataclasses import dataclass, field |
| 5 | +from typing import Literal |
| 6 | + |
| 7 | +SubmissionGranularity = Literal["R", "G", "B"] |
| 8 | +ConsumptionGranularity = Literal["G", "B"] |
| 9 | + |
| 10 | +_BOOL_OPTIONAL = {"argparse_meta": {"action": argparse.BooleanOptionalAction}} |
| 11 | + |
| 12 | + |
| 13 | +@dataclass(kw_only=True) |
| 14 | +class RLConfig: |
| 15 | + """Configuration settings for the built-in RL training loop (GRPO).""" |
| 16 | + |
| 17 | + perform_rl_step: bool = False |
| 18 | + """Use the RL training step.""" |
| 19 | + |
| 20 | + rl_prompts_per_eval: int = 32 |
| 21 | + """Number of prompts to evaluate for each RL task. This evaluation can be very expensive when |
| 22 | + using environments that evaluate pass@k so we default to a lower number.""" |
| 23 | + # TODO(rkirby): allow for "complete" evaluation when rl_prompts_per_eval is set to -1 |
| 24 | + |
| 25 | + grpo_prompts_per_step: int = 32 |
| 26 | + """Number of GRPO groups (G in the paper).""" |
| 27 | + |
| 28 | + grpo_group_size: int = 2 |
| 29 | + """Number of samples per GRPO group.""" |
| 30 | + |
| 31 | + rl_generation_lag: float | None = None |
| 32 | + """Number of trainer batches of rollout generation lag to allow. The number of in-flight trainer |
| 33 | + batches is this value plus one. May be fractional or negative; the minimum of -1 keeps a single |
| 34 | + unit of generation work in flight. If omitted, the lag is autotuned to the inference engine's |
| 35 | + request capacity when rl_partial_rollouts is set, and is 0 otherwise. Requires |
| 36 | + rl_partial_rollouts when greater than 0. Mutually exclusive with rl_max_inflight_requests; both |
| 37 | + are accepted only when they agree, which is the case when a normalized config is reloaded.""" |
| 38 | + |
| 39 | + rl_max_inflight_requests: int | None = None |
| 40 | + """Maximum number of inference requests RL generation may keep in flight: equivalent to |
| 41 | + (rl_generation_lag + 1) training batches of grpo_prompts_per_step * grpo_group_size requests |
| 42 | + each. Requires rl_partial_rollouts when above one training batch; mutually exclusive with |
| 43 | + rl_generation_lag.""" |
| 44 | + |
| 45 | + rl_submission_granularity: SubmissionGranularity = "B" |
| 46 | + """Granularity for submitting rollout generation work. R submits individual rollouts |
| 47 | + independently while still yielding complete rollout groups to training. G submits one rollout |
| 48 | + group at a time. B submits grpo_prompts_per_step rollout groups together.""" |
| 49 | + |
| 50 | + rl_consumption_granularity: ConsumptionGranularity = "B" |
| 51 | + """Granularity for consuming generated rollout groups. G consumes groups as they complete. |
| 52 | + B consumes complete trainer batches in submission order. R is not supported.""" |
| 53 | + |
| 54 | + rl_durable_rollout_bank: bool = False |
| 55 | + """Persist completed rollout groups to a durable, write-through ledger so they survive a SIGKILL |
| 56 | + (the SLURM time limit) and are restored at restart instead of regenerated. No-op when unset.""" |
| 57 | + |
| 58 | + rl_rollout_bank_dir: str | None = None |
| 59 | + """Directory for the durable rollout bank (on Lustre). Defaults to <save>/rollout_bank so the |
| 60 | + bank stays coupled to the checkpoint.""" |
| 61 | + |
| 62 | + rl_rollout_bank_max_bytes: int = 0 |
| 63 | + """Soft cap (bytes) on the rollout bank size; 0 = unbounded. On exceed, a warning is logged. |
| 64 | + Compaction occurs at the next checkpoint regardless of the cap and never blocks generation.""" |
| 65 | + |
| 66 | + grpo_iterations: int = 2 |
| 67 | + """Number of iterations per GRPO implementation.""" |
| 68 | + |
| 69 | + grpo_clamp_eps_lower: float = 0.01 |
| 70 | + """Lower GRPO clipping bound. As in DAPO the lower and upper bounds may differ; set them equal |
| 71 | + for vanilla GRPO.""" |
| 72 | + |
| 73 | + grpo_clamp_eps_upper: float = 0.01 |
| 74 | + """Upper GRPO clipping bound. In the vanilla implementation, equals the lower one.""" |
| 75 | + |
| 76 | + grpo_kl_beta: float = 0.001 |
| 77 | + """KL term weight in the GRPO loss.""" |
| 78 | + |
| 79 | + grpo_entropy_term_weight: float = 0.0 |
| 80 | + """Entropy term weight in the GRPO loss.""" |
| 81 | + |
| 82 | + grpo_filter_groups_with_same_reward: bool = False |
| 83 | + """Filter groups with same reward.""" |
| 84 | + |
| 85 | + langrl_env_config: str | None = None |
| 86 | + """Path to YAML config file for RL environment configuration.""" |
| 87 | + |
| 88 | + rl_default_temperature: float = 1.0 |
| 89 | + """Default temperature for model inference.""" |
| 90 | + |
| 91 | + rl_default_top_p: float = 0.0 |
| 92 | + """Default top-p for model inference.""" |
| 93 | + |
| 94 | + rl_default_top_k: int = -1 |
| 95 | + """Default top-k for model inference.""" |
| 96 | + |
| 97 | + rl_offload_optimizer_during_inference: bool = False |
| 98 | + """Offload optimizer state to CPU during inference/rollout to save GPU memory.""" |
| 99 | + |
| 100 | + rl_kv_cache_management_mode: Literal["persist", "offload", "recompute"] = "persist" |
| 101 | + """KV cache management mode during RL training. persist: leave KV cache in GPU memory (default). |
| 102 | + offload: offload KV cache to CPU during training. recompute: deallocate KV cache and recompute |
| 103 | + from scratch each cycle.""" |
| 104 | + |
| 105 | + rl_persist_cuda_graphs: bool = field(default=False, metadata=_BOOL_OPTIONAL) |
| 106 | + """Persist CUDA graphs when the inference engine is suspended. If False, CUDA graphs are deleted |
| 107 | + on suspend and re-captured on resume.""" |
| 108 | + |
| 109 | + rl_partial_rollouts: bool = field(default=False, metadata=_BOOL_OPTIONAL) |
| 110 | + """Allow inference to continue generating rollouts while training updates the policy weights. |
| 111 | + This enables off-policy training where rollouts may be generated with a stale version of the |
| 112 | + policy. Use rl_generation_lag to control the degree of staleness.""" |
| 113 | + |
| 114 | + rl_inference_logprobs_is_correction: bool = field(default=False, metadata=_BOOL_OPTIONAL) |
| 115 | + """If set, use inference logprobs in importance sampling correction of the loss.""" |
| 116 | + |
| 117 | + rl_importance_sampling_truncation_coef: float | None = None |
| 118 | + """If rl_inference_logprobs_is_correction is on and this coefficient is set, apply truncation |
| 119 | + for the IS correction at GRPO loss.""" |
| 120 | + |
| 121 | + rl_use_sequence_packing: bool = field(default=False, metadata=_BOOL_OPTIONAL) |
| 122 | + """Enable sequence packing.""" |
| 123 | + |
| 124 | + rl_sequence_packing_max_sequences_per_bin: int = 50 |
| 125 | + """Maximum number of sequences that can be packed into a single bin.""" |
| 126 | + |
| 127 | + rl_sequence_packing_algo: Literal["fifo", "round-robin"] = "fifo" |
| 128 | + """Algorithm for distributing packed bins across ranks. fifo: first-in-first-out sequential |
| 129 | + distribution. round-robin: distribute bins cyclically across ranks for better load balancing.""" |
| 130 | + |
| 131 | + rl_training_cuda_graphs: bool = field(default=False, metadata=_BOOL_OPTIONAL) |
| 132 | + """If set, do not toggle CUDA graphs on/off between inference and training phases.""" |
| 133 | + |
| 134 | + rl_inference_tensor_model_parallel_size: int | None = None |
| 135 | + """Degree of tensor model parallelism for inference for RL.""" |
| 136 | + |
| 137 | + rl_inference_pipeline_model_parallel_size: int | None = None |
| 138 | + """Degree of pipeline model parallelism for inference for RL.""" |
| 139 | + |
| 140 | + rl_inference_expert_model_parallel_size: int | None = None |
| 141 | + """Degree of expert model parallelism for inference for RL.""" |
| 142 | + |
| 143 | + rl_inference_expert_tensor_model_parallel_size: int | None = None |
| 144 | + """Degree of expert tensor model parallelism for inference for RL. For MoE models, this controls |
| 145 | + the TP size for expert layers specifically. Defaults to training expert_tensor_parallel_size if |
| 146 | + not specified.""" |
| 147 | + |
| 148 | + rl_inference_model_unified_memory_level: Literal[0, 1] = 0 |
| 149 | + """Allocate the separate RL inference model parameters from a unified virtual memory (UVM) CUDA |
| 150 | + mempool. Level 0 disables UVM (default). Level 1 enables UVM allocation so the inference model |
| 151 | + weights can be prefetched to CPU when idle while keeping CUDA-graph-safe device pointers.""" |
| 152 | + |
| 153 | + rl_offload_inference_model_weights: bool = field(default=False, metadata=_BOOL_OPTIONAL) |
| 154 | + """When using a separate RL inference model, offload its weights to CPU when not doing rollout |
| 155 | + inference, and restore to GPU right before inference. Works with two backends: 1) UVM (when |
| 156 | + rl_inference_model_unified_memory_level=1), or 2) torch_memory_saver (when UVM is not enabled; |
| 157 | + requires torch_memory_saver to be installed).""" |
| 158 | + |
| 159 | + refit_method: Literal["nccl", "nccl_m2n", "gloo", "nvshmem", "nixl"] = "gloo" |
| 160 | + """Method to refit model weights. nccl: use NCCLCopyService; nccl_m2n: use the official NCCL M2N |
| 161 | + API from a non-RL launcher such as the ReFIT benchmark; gloo: use GlooCopyService over CPU; |
| 162 | + nvshmem: use NVSHMEMCopyService; nixl: use NixlCopyService.""" |
| 163 | + |
| 164 | + refit_execution_batch_bytes: int | None = None |
| 165 | + """Optional soft per-rank byte limit for ReFIT execution staging. The default None preserves one |
| 166 | + model-wide generic submission and NCCL M2N's existing 256 MiB default.""" |
| 167 | + |
| 168 | + rl_verify_model_weights_swap: bool = field(default=False, metadata=_BOOL_OPTIONAL) |
| 169 | + """If set, verify that the model weights were correctly transferred by comparing forward pass |
| 170 | + outputs on the first swap of model weights.""" |
| 171 | + |
| 172 | + rl_skip_bos_token: bool = field(default=False, metadata=_BOOL_OPTIONAL) |
| 173 | + """Skip BOS token at the beginning of the sequences.""" |
| 174 | + |
| 175 | + rl_profile: bool = False |
| 176 | + """Enable RL profiling to collect detailed timer data (JSONL + CSV).""" |
| 177 | + |
| 178 | + rl_profile_dir: str | None = None |
| 179 | + """Directory to write RL profiling data. Defaults to {save}/profiles.""" |
| 180 | + |
| 181 | + rl_inference_parsers: list[str] = field( |
| 182 | + default_factory=list, metadata={"argparse_meta": {"nargs": "*"}} |
| 183 | + ) |
| 184 | + """List of response parsers to enable for RL inference (e.g. deepseek-r1-reasoning |
| 185 | + qwen3-coder-tool). qwen3-coder-tool-combined additionally treats <tool_call> as the end of an |
| 186 | + unterminated reasoning block, like vLLM's combined qwen3 parser.""" |
| 187 | + |
| 188 | + grpo_samples_per_iteration: int = field(init=False) |
| 189 | + """Derived: grpo_prompts_per_step * grpo_group_size, the samples generated per RL iteration.""" |
| 190 | + |
| 191 | + def __post_init__(self): |
| 192 | + self.grpo_samples_per_iteration = self.grpo_prompts_per_step * self.grpo_group_size |
| 193 | + |
| 194 | + if self.refit_execution_batch_bytes is not None: |
| 195 | + assert self.refit_execution_batch_bytes > 0, ( |
| 196 | + "--refit-execution-batch-bytes must be a positive integer" |
| 197 | + ) |
| 198 | + |
| 199 | + # The remaining checks describe the RL loop and only apply when it runs. |
| 200 | + if not self.perform_rl_step: |
| 201 | + return |
| 202 | + |
| 203 | + assert self.refit_method != "nccl_m2n", "nccl_m2n is unsupported by the built-in RL loop" |
| 204 | + |
| 205 | + if self.rl_max_inflight_requests is not None: |
| 206 | + assert self.rl_max_inflight_requests >= 1, ( |
| 207 | + f"--rl-max-inflight-requests ({self.rl_max_inflight_requests}) must be >= 1." |
| 208 | + ) |
| 209 | + if self.rl_max_inflight_requests > self.grpo_samples_per_iteration: |
| 210 | + assert self.rl_partial_rollouts, ( |
| 211 | + "--rl-max-inflight-requests above one training batch " |
| 212 | + f"({self.grpo_samples_per_iteration} requests) requires --rl-partial-rollouts." |
| 213 | + ) |
| 214 | + # Total in-flight requests = (lag + 1) trainer batches of P * G requests each. |
| 215 | + derived_lag = self.rl_max_inflight_requests / self.grpo_samples_per_iteration - 1 |
| 216 | + assert self.rl_generation_lag is None or self.rl_generation_lag == derived_lag, ( |
| 217 | + "--rl-generation-lag and --rl-max-inflight-requests are mutually exclusive." |
| 218 | + ) |
| 219 | + self.rl_generation_lag = derived_lag |
| 220 | + if self.rl_generation_lag is None: |
| 221 | + # With rl_partial_rollouts the lag is autotuned from engine capacity at inference |
| 222 | + # launch; otherwise generation is fully synchronous. |
| 223 | + if not self.rl_partial_rollouts: |
| 224 | + self.rl_generation_lag = 0.0 |
| 225 | + else: |
| 226 | + assert self.rl_generation_lag >= -1, ( |
| 227 | + f"--rl-generation-lag ({self.rl_generation_lag}) must be >= -1." |
| 228 | + ) |
| 229 | + if self.rl_generation_lag > 0: |
| 230 | + assert self.rl_partial_rollouts, ( |
| 231 | + "--rl-generation-lag requires --rl-partial-rollouts." |
| 232 | + ) |
| 233 | + |
| 234 | + assert self.rl_submission_granularity == "B" or self.rl_partial_rollouts, ( |
| 235 | + f"--rl-submission-granularity {self.rl_submission_granularity} requires " |
| 236 | + "--rl-partial-rollouts." |
| 237 | + ) |
| 238 | + assert self.rl_consumption_granularity != "R", ( |
| 239 | + "--rl-consumption-granularity R is not currently supported." |
| 240 | + ) |
| 241 | + assert not ( |
| 242 | + self.rl_submission_granularity == "B" and self.rl_consumption_granularity == "G" |
| 243 | + ), "--rl-submission-granularity B with --rl-consumption-granularity G is not supported." |
| 244 | + |
| 245 | + # KV cache offload requires CUDA graph persistence: recapturing CUDA graphs runs dummy |
| 246 | + # forward passes that corrupt the preserved KV data. |
| 247 | + assert self.rl_kv_cache_management_mode != "offload" or self.rl_persist_cuda_graphs, ( |
| 248 | + "--rl-kv-cache-management-mode=offload requires --rl-persist-cuda-graphs" |
| 249 | + ) |
| 250 | + |
| 251 | + if ( |
| 252 | + self.rl_offload_inference_model_weights |
| 253 | + and self.rl_inference_model_unified_memory_level != 1 |
| 254 | + ): |
| 255 | + # Not using UVM, so torch_memory_saver must provide the CPU backup. |
| 256 | + try: |
| 257 | + import torch_memory_saver # noqa: F401 |
| 258 | + except ImportError: |
| 259 | + raise AssertionError( |
| 260 | + "To use --rl-offload-inference-model-weights without UVM " |
| 261 | + "(--rl-inference-model-unified-memory-level=1), `torch_memory_saver` must be " |
| 262 | + "installed. See https://github.com/fzyzcjy/torch_memory_saver." |
| 263 | + ) |
0 commit comments