Add async non-colated PPO + connect to Gym rollout - #3185
Draft
HeyyyyyyG wants to merge 71 commits into
Draft
Conversation
… config Copy the GB200/aarch64 async-GRPO SWE launcher and its megatron config into examples/swe_bench for the oci-hsg-cs-001 cluster: - run_grpo_swe2_scale_gen_hsg.sh (single-knob NUM_VLLM_REPLICAS, ALIGN_BASELINE switch, LUSTRE_UV_CACHE_SEED default, 4 GPU/node geometry) - grpo_qwen3_30b_async_swe.yaml (incl. megatron_cfg.use_fused_weighted_squared_relu) Signed-off-by: ruit <ruit@nvidia.com>
…riton MoE) - rename config to grpo_qwen3_30b_async_swe_hsg.yaml with the GB200 fixes: megatron_cfg.use_fused_weighted_squared_relu + vllm_kwargs.moe_backend=triton (triton MoE backend avoids the weight-refit broadcast_weights_for_collective hang) - run_grpo_swe2_scale_gen_hsg.sh: point CONFIG_FILE at the hsg yaml and default CONTAINER to the baked nightly-063026-gymvenvs image (has /opt/gym_venvs prebuilt so gym spinup skips the concurrent venv build that deadlocks on the uv cache lock) Signed-off-by: ruit <ruit@nvidia.com>
… creds
Remove the hardcoded source of ${HOME}/script/export_env_vars.sh (a personal
file). Users export their own HF_HOME/HF_TOKEN/WANDB_API_KEY/etc. before running.
Signed-off-by: ruit <ruit@nvidia.com>
Set DEFAULT_MODEL_PATH to the SWE1 step_230_hf checkpoint and TRAIN/VAL_DATA_PATH to the r2e_gym_subset dataset (train=val) so the hsg launcher reproduces the baseline SWE2 run by default. All three remain overridable via env. Signed-off-by: ruit <ruit@nvidia.com>
…ic_val The r2e_gym_subset set has instances with no matching sif on this cluster (missing container -> 500, batch never fills -> step stalls). Revert the launcher defaults to the runnable balanced_language (train) + swe_public_val. Signed-off-by: ruit <ruit@nvidia.com>
…skip-training
- add examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh: single-knob
(NUM_VLLM_REPLICAS) launcher, 3 modes — ALIGN_BASELINE (32-node baseline,
CP16/EP32), default (scale train+gen from the validated CP2/EP8 2-node/DP
base), and SKIP_TRAINING (gen benchmark: pin train to 1 node, scale only gen).
Segment auto-picks the largest of {16,8,4,2,1} dividing each worker group;
wandb staging redirected off the coreai quota; 60min idle-reaper exemption.
- add examples/swe_bench/grpo_nano_v3_5_swe_hsg.yaml.
- gen-benchmark skip-training support (NRL_GEN_BENCHMARK_SKIP_TRAINING) in
grpo.py / lm_policy.py / base_policy_worker.py (per #2930): skip
optimizer build + policy.train(), NCCL-free keep-alive matmul to avoid the
idle-GPU reaper while training is skipped.
- MegatronConfig.mtp_num_layers accepts None (int | None) so MTP can be disabled
on hybrid (NemotronH) models where 0 still trips the mtp_num_layers>0 assertion.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: ruit <ruit@nvidia.com>
Signed-off-by: Jiaqi Zeng <jiaqiz@aws-cmh-slurm-1-vscode-02.cm.cluster>
Signed-off-by: Jiaqi Zeng <jiaqiz@aws-cmh-slurm-1-vscode-02.cm.cluster>
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
The multi-platform lock (aarch64 + x86_64) requires both to jointly resolve. Pinning a custom vLLM whose own transitive deps (e.g. flashinfer-python) differ per version can make the unused x86_64 split unsatisfiable and block the whole lock/sync, even though we only ever build/run on aarch64 (GB200/GB300). Per uv's own hint on this failure.
This reverts commit c68f226.
This reverts commit afce932.
This reverts commit c88955d.
…lds" This reverts commit b66cb73.
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
…diagnostics PPO previously required colocated generation and had no visibility into train/inference logprob mismatch, unlike GRPO. Ports both capabilities from GRPO's proven implementations: separate train/inference clusters with NCCL collective weight refit (SGLang guarded as unsupported non-colocated), and the sequence-level mismatch masking, W&B mismatch plot, and gen_kl_error logging GRPO already surfaces. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
Add async PPO training (async_ppo_train) mirroring async GRPO: an AsyncTrajectoryCollector + ReplayBuffer Ray-actor pipeline with a weight-version staleness window, critic-warmup trajectory banking (warmup_max_trajectory_age_steps), in-flight weight updates, and freeze-aware trajectory-age diagnostics. Also route async rollouts through validate(), fix replay-buffer and value-optimizer checkpoint/resume, and add unit + functional (ppo_async.sh) coverage. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
…n resume The partially-generated frontier target saved in the async replay buffer holds only the fast-completing (short) rollouts, so keeping and gap-filling it on resume trains the first post-resume step on a systematically shorter, higher-reward batch — a dip in mean_gen_tokens_per_sample and a reward spike after every resume. Add ppo.async_ppo.drop_incomplete_targets_on_restore (default false = historical gap-fill behavior). When true, incomplete restored targets are dropped and regenerated fresh (unbiased batch, at the cost of a one-target generation bubble at resume). Complete banked targets are always kept and replayed exactly. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
Finish the partial NeMo-Gym port for PPO so it reaches GRPO parity: - setup(): spin up the NeMo-Gym actor via the vLLM deferred-load path, bundling policy+value+vLLM as one GPU task while the gym overlaps (CPU/HTTP); add nemo_gym_actor to the return tuple. - ppo_train(): complete the gym rollout call (max_seq_len, effort_config, reward_penalty_config, thinking_tags, stop_token_ids/stop_strings=None). - validate(): add the gym branch (3-way dispatch) — closes the NemoGym.step() NotImplementedError crash gap. - async_ppo_train(): drop the gym-off guard; gym rollouts run inside the shared AsyncTrajectoryCollector, which now reads master_config.reward_penalties. - reward_penalties (reward-zeroing, pre-GAE) are wired functional and GAE-compatible; mask_sample is honored in both training loops (mirrors GRPO); the advantage-overwrite penalties (invalid_tool_call_advantage / malformed_thinking_advantage) are explicitly rejected for PPO (they'd break GAE advantage/return consistency). - Harden the collector gym branch to clear stop_token_ids/stop_strings by construction (avoids a swallowed assert -> silent buffer stall). Add examples/nemo_gym/run_ppo_nemo_gym.py + ppo_math_rlvr_nemo_gym.yaml (Math/RLVR), unit coverage for the validate gym dispatch, and functional tests ppo_nemo_gym.sh (sync) + ppo_async_gym.sh (async). Both smoke tests pass on GB200 (DAPO17k, Qwen2.5-0.5B) end to end. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
The AsyncTrajectoryCollector iterated the training dataloader with a single, non-cycling `for batch in self.dataloader`. A finite dataset is consumed once and the loop returns, stopping the collector and silently stalling the replay buffer (the driver waits on buffer_size=0 forever). This bites hardest on RESUME: the checkpoint saves the collector's real dataloader position, and a restored StatefulDataLoader whose samples_yielded is already at the epoch end yields ZERO batches on the first epoch — so the run hangs immediately after resume (observed at step 195 of a 218-batch DAPOMath17K run). Wrap the loop in `while self.running:` so it cycles epochs, mirroring sync ppo_train/grpo_train's `while epoch < max_num_epochs: for batch in dataloader`. StatefulDataLoader resume-once semantics (next_iter_state cleared after the first iterator) mean the resumed partial epoch yields its remainder and fresh full epochs follow; a two-empty-epoch guard fails loudly on a genuinely empty dataset instead of busy-spinning. Add unit tests for the resume-exhaustion recovery, multi-epoch cycling, and the empty-dataset guard. Shared with async GRPO. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
…ndition Calibration: EV tells you the critic ranks prefixes well, not that V is literally P(success | prefix). Add ECE (token-weighted |mean(V) - mean(R)| over 10 predicted-value bins) and signed bias per early/mid/late bucket, and wire the positional metrics into the async loop too. Miscalibration distorts advantage MAGNITUDES even when EV looks fine. Precondition: a missing/empty ground_truth silently degrades a sample to a blind critic (the grader note renders with an empty answer). Tolerate stragglers with a warning, but raise when the WHOLE batch lacks privilege — that means extra_env_info did not survive the data path and the run would measure a blind critic while labelled privileged. Also updates the positional-metrics docstring with the priv-vs-blind 7B DAPO finding: the answer acts as a verifier (wins at LATE tokens) rather than a forecaster of the policy's future behavior. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
- ppo.setup: a fabricated warm-start seed (a step_0 checkpoint holding only a pretrained value/ dir) carries no train_dataloader.pt; start the dataloader fresh with a warning instead of failing, mirroring how a missing policy/ falls back to base weights. - ppo.setup: segment topology only applies to multi-node clusters. In the single-node carve, train and inference SHARE the one node, so train_nodes + inference_nodes double-counted the roles and demanded 2 alive nodes on a 1-node cluster. NVLink-domain placement is meaningless there. - nemo_gym: when the agent fails before its first model request, the result has neither output items nor input messages and apply_chat_template([]) threw a bare IndexError that masked the diagnosis. Raise something actionable with the truncated gym result instead. - MegatronConfig.mtp_num_layers: type as int | None. For hybrid (mamba) models None is the only full off-switch — hybrid_model.py gates the MTP block on `is not None`, so 0 still enters it and asserts > 0. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Mirrors the existing moe_metrics passthrough so Multi-Token Prediction loss terms are visible in wandb under mtp/*, in sync ppo_train and async_ppo_train. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…atron - ppo_nano_v3_5_swe_cmh.yaml: the SWE PPO run config. Sets vllm disable_custom_all_reduce: jobs 37606/38016 each lost a generation engine ~15 min in, and the NCCL flight recorder showed all 4 TP ranks with their next collective enqueued-but-never-started — all four streams parked in the custom-allreduce kernel ahead of it (symmetric cross-GPU deadlock, different node each time). Costs a few % decode latency on TP-4. Also drops agent concurrency 3x -> 2x prompts*generations. - ppo_math_1B_megatron.yaml: activation_checkpointing + defer_fp32_logits. At long packed sequences the actor forward materializes full fp32 LM-head logits (~11GB over a 152k vocab) on top of unrecomputed activations, and the first policy-training step OOMs in the grad all-reduce. Memory-shape only. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
- grpo_ultra_256n4g_bf16.yaml: alltoall MoE dispatcher, lr 3e-6 -> 4e-6, sequence-length divisor accounts for context parallelism, generation lengths follow policy.max_total_sequence_length, wandb on - grpo/ppo_ultra_256n4g_bf16_pivotonly.yaml: pivot-only variants - ppo_math_1B.yaml: rollout dumps on (every 5 steps), seq logprob error threshold 2, drop incomplete async targets on restore, segment_size 2 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
67f7955 -> 5613d41: - upstream jiaqiz/ppo-dev (incl. responses_converter tool_choice fix #1925) - fix(rollout_collection): retry /run on transient connection failures - fix(swe_agents): re-bind /etc/resolv.conf into the apptainer sandbox Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
PPO's inner loop drove the critic and the actor off a single `ppo.ppo_epochs` count, so fitting the critic harder meant also pushing the policy further off the data that produced its advantages. Add `ppo.critic_ppo_epochs` (null/absent => coupled to ppo_epochs, the previous behavior; must be >= ppo_epochs). The surplus runs as critic-only passes before the shared critic/actor loop, which is left byte-identical -- the only change to existing code is hoisting the value_train_batch selection out of it. Ordering is free because every pass consumes the same returns/advantages, computed once per step and frozen before any update, so the critic never observes the actor's within-step updates. prepare_for_training/finish_training bracket the extra passes as a whole, so they cost one additional param+grad+optimizer onload cycle per step rather than one per pass. Each Megatron worker ticks its LR scheduler once per train() call, so the value model's train_iters budget now scales by critic_ppo_epochs instead of ppo_epochs; otherwise its decay schedule would end before training does. Note that lr_warmup_iters is counted in those same ticks, so a critic running N epochs warms up N times faster in PPO-step terms. Also add `ppo.log_post_update_critic_metrics` (default false). critic/* metrics come from a training pass's forward and therefore describe the critic BEFORE that pass's update; when enabled, one forward-only pass (eval_mode, so it skips optimizer.step and scheduler.step) re-scores the same batch after the final update and emits critic/explained_var_post_update and critic/loss_post_update. It runs while the value model is still resident and the policy is not, so it adds no co-residency -- just the forward. _resolve_critic_ppo_epochs also asserts ppo_epochs >= 1, which sync ppo_train never checked; ppo_epochs: 0 previously produced a run that trained nothing. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
PPO/VAPO port of grpo_nano_v3_5_swe_hsg.yaml for GB200 NVL72, run via examples/nemo_gym/run_ppo_nemo_gym.py. Mirrors the existing cmh variant, with the apptainer SWE image paths pointing at this cluster's image store and moe_hybridep_num_sms omitted (Megatron in the NCCL-2.30.4 container deprecates it in favour of moe_flex_dispatcher_num_sms, which this checkout does not plumb, and errors out when both deprecated knobs are set). Enables the per-token rollout dump every 5 steps and the post-update critic metrics for critic-quality debugging. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
decord is only reachable on the audio fallback path, but the top-level import made it a hard dependency of multimodal_utils, which every entrypoint pulls in via batched_data_dict. In containers whose prebaked venvs omit decord that killed the driver with ModuleNotFoundError before any text-only run could start. Move the import to its use site, matching what nemo_rl/data/datasets/utils.py already does. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A warm-start seed fabricated by prep_warm_start.sh carries critic weights from a stage-B run whose expert-parallel layout may differ from the resuming run's. Megatron's distributed optimizer state is sharded by (expert-)DP group index and cannot reshard across a changed EP, so the load died with "Missing key in checkpoint state_dict: chained_1.optimizer.distributed.dp_group_idx_N...". Model weights do reshard, so take those and rebuild Adam plus the LR schedule from scratch (value_lr_warmup_iters re-warms it). Gated on the warm_start_provenance.txt marker that only the fabricated seed carries, so ordinary step_N resumes keep restoring the optimizer as before. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Split the single 3600s swebench_agent_timeout into 2700s of agent time plus a dedicated 900s swebench_tests_timeout, leaving the 3600s total unchanged so a long test run can no longer consume the whole budget. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
With critic_ppo_epochs > 1 the loss-derived explained_var came from the last training pass's forward -- a critic that had already fit this batch several times -- so it read as batch-fit rather than critic quality. On run awk40z9y (critic_ppo_epochs=4) it sat at ~0.93 while the positional ev_early/mid/late diagnostics, computed from the rollout-time values, showed the critic's actual pre-update EV of ~0.33. Overwrite critic/explained_var at both metric sites with a pooled EV computed driver-side from train_data values/returns -- the exact tensors GAE consumed, i.e. the critic before ANY update this step, regardless of epoch count or gbs-vs-rollout microbatching. Masked like the value loss (token * sample) so it stays directly comparable to critic/explained_var_post_update, which remains the after-all-updates end of the bracket. critic/loss and critic/grad_norm still describe the last training pass. No extra compute: the tensors are already on the driver. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
normalize_advantages rescales advantages to unit std every step, so the observable advantage spread is 1.0 by construction and carries no signal. The pre-whitening spread does: at lambda=1 the advantage is R - V(s_t), so its std is the critic's residual scale and should shrink as the critic improves. GeneralizedAdvantageEstimator now records mean/std/abs_mean/max_abs (plus whiten_gain = 1/std) over valid tokens immediately before the whitening branch, exposed via last_metrics following the OPDAdvantageEstimator convention. Both ppo_train and async_ppo_train merge it into their metrics. Diagnostic only: no change to advantages, returns, or the loss. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The dump stores train_data["advantages"], which is post-whitening, so
offline analysis cannot recover the critic's residual scale. Reconstructing
it from values/reward instead requires knowing lambda_policy, which the
payload does not record.
Whitening is affine, so the two scalars suffice:
adv_raw = advantages * adv_raw_std + adv_raw_mean
Stored as scalars rather than a second packed tensor -- same information,
~13MB/dump cheaper. Threaded from the estimator's last_metrics at both dump
sites; absent for estimators that do not provide it.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Adds an opt-in turn MDP for long agentic rollouts, selected with
ppo.adv_estimator.name=turn_gae. One assistant message is one action:
d_k = r_k + g*V(s_{k+1}) - V(s_k), A_k = d_k + g*lam*A_{k+1}, G_k = A_k + V(s_k)
The motivation is not that the token-level critic is noisy within a turn
(measured: only 3.8% of V's variance is intra-turn, and at lambda=1
token- and turn-level baselines are near-identical). It is that lambda is
structurally unusable at the token level: its horizon is 1/(1-lambda)
TOKENS, so on a ~45k-token SWE rollout length_adaptive_alpha is forced to
drive lambda to 1-1.5e-5 and GAE degenerates to A_t = R - V(s_t) — a pure
baseline with no temporal credit assignment at all. Over ~92 turns,
lambda=0.97 is a 33-turn horizon: a usable knob. At lambda<1 the advantage
is built from TD increments, in which any constant per-trajectory offset
cancels — and the critic's measured weakness on this workload IS the
constant part (it barely reads task difficulty from the prompt), while
what survives, d_K = R - V(s_K), is its strongest signal.
Where V(s_k) is read: the value workers right-shift the value head, so
values[t] = V(state before token t). V(s_k) is therefore the FIRST token
of assistant message k, which is already inside the response mask. That
is what keeps the change small — no model, worker, sequence-packing or
context-parallel changes are needed.
Critic supervision is anchor-only: the value model trains on its own
batch whose token_mask is the anchor mask (the same mechanism the
privileged critic already uses). Since process_global_batch derives
global_valid_toks from token_mask, MseValueLossFn becomes an
equal-weighted mean over turns for free — today the longest decile of
turns owns ~43% of the critic loss. This costs no information: at
lambda_value=1 every target inside a trajectory is the same number, R.
gamma / lambda_value / lambda_policy are required config keys with no
silent defaults, so a sweep cannot appear to have run when it did not.
Verification: 25 CPU unit tests (span construction incl. consecutive and
empty assistant messages, ragged batches, all fail-loud paths, turn GAE
vs a hand-computed reference, the gamma=lambda=1 => G_k == R identity,
gather/scatter round trips, per-turn KL aggregation, and that the
estimator's returns are covered by the critic batch's mask); plus an
offline replay of the production estimator over stored SWE value dumps
(8 held-out groups, 13,330 turns, 3.24M response tokens) checking anchor
placement, the MC-return identity, within-turn constancy of advantages
and the returns/mask bijection on every group.
The default path is unchanged: with name=gae or raw_reward this is inert.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
The Gym submodule declares boto3>=1.34 (ECS Fargate sandbox provider) and is a uv workspace member, so its dependencies resolve into the root uv.lock. The pin crossed that boundary at 11b1240 ("bump gym") and no bump since re-locked, leaving HEAD's lockfile without boto3/botocore/s3transfer. Every worker runs `uv run --locked --directory {git_root}` (nemo_rl/distributed/virtual_cluster.py), which asserts the lockfile is up to date and refuses to run otherwise -- so a fresh clone of this branch could not start training at all. Regenerating adds only the three missing packages under nemo-gym's sandbox/all extras; no manifest changed. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
These two files were not formatter-clean at HEAD: earlier commits hand-wrapped calls that `ruff format` collapses. Running the project's pinned ruff (0.9.9) over HEAD's versions reproduces exactly this diff, so it is mechanical output, not hand editing. Split out from the critic work that follows so those diffs contain only behavioural change. Verified semantics-free: the parsed AST of each file is identical before and after. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
Newer mcore (NCCL 2.30.4 image) rejects more than one deprecated SM-count knob:
ValueError: Conflicting deprecated SM-count knobs
{'moe_deepep_num_sms': 20, 'moe_hybridep_num_sms': 32};
set a single moe_flex_dispatcher_num_sms instead.
mcore supplies moe_deepep_num_sms itself, so NeMo-RL setting the hybridep knob
is what tips it over. Feature-detect the replacement attribute and route the
configured value there, falling back to the deprecated knob on older builds.
This keeps moe_hybridep_num_sms: 32 effective in ppo_nano_v3_5_swe_cmh.yaml and
the grpo/ultra configs rather than discarding it. The implementation is
byte-identical to the fix already on main (ec46ab3), which this branch predates,
so the hunk will resolve to nothing when the branch meets main.
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
The Bayes-optimal value splits into a between-task and a within-task part, V*(s) = B(X) + C(s) with E[C | X] = 0. An absolute critic must fit both, and on this SWE workload it overwhelmingly fits the first and fits it badly: 88-93% of its output variance is between-task, yet EV ~0.157 -- below a two-value "all-fail vs rest" lookup (0.330) and far below the free leave-one-out group baseline (0.553). Substituting the critic for that baseline RAISES advantage variance 1.67-2.09x. So hand B to the rollout group and let the critic learn only C. ResidualBaselineEstimator wraps either value-based estimator (gae or turn_gae), keeps the batch in residual space end to end, and reconstructs V~ = B_LOO + C only for the inner GAE call -- which is what keeps the PPO value clip self-consistent. gamma == 1 is required and enforced: the baseline only cancels from nonterminal TD errors there. num_generations_per_prompt >= 2 is also enforced, since a single sibling makes B_LOO == R and would silently drop the reward from the policy gradient entirely. adv_estimator.residual_baseline has no default, so a sweep cannot appear to have run when it did not. It is required for gae/turn_gae and rejected for raw_reward. Diagnostics are the point of the change as much as the mechanism: the wrapper is applied even when residualization is OFF so an absolute run reports critic/ev_res on the same axis as a residual run (there 1 - ev_res is exactly the measured advantage-variance ratio). MseValueLossFn gains the per-space return statistics this needs, and critic/ev_res_mixed_group isolates EV where the target is actually nonzero -- homogeneous groups are ~58% of this pool and contribute Y = 0, so any prediction there is a pure penalty to the whole-batch number. value_loss_fn.homogeneous_group_weight can down-weight them; 1.0 keeps today's behaviour. The default path is unchanged: with residual_baseline: false the targets stay absolute and only the extra diagnostics appear. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
Extends the existing privileged-critic machinery with a SWE-specific reference block: the critic sees the accepted fix and the grading tests, the policy never does. Audited over all 9262 instances -- golden patch and test patch are present for 100% of them, and none appear anywhere in the agent's context (checked against a real 408k-char rollout), so this is genuine privileged information rather than a cheaper route to something already visible. The block is prefixed BEFORE the first assistant token. That is mandatory, not a detail: the value head is causal and ~150 assistant turns are spread across the context, so appending would be a silent no-op. It is byte-identical for all 16 siblings of a group, so it cannot introduce a within-task confound. Orthogonal to both other critic axes -- it composes with residual or absolute targets, and with token-level or turn_gae. For turn_gae the anchors are remapped into the augmented layout (build_turn_value_batch_augmented), which is why this variant lifts the turn_gae-plus-privileged-critic restriction that still applies to the math privileged critic. value.max_total_sequence_length and the packing/dynamic-batching token budgets are raised automatically by the reference block's cap. The budget interpolations are resolved before setup() runs, so without raising them too the packer fails with "Sequence length N exceeds bin capacity" hours into a run rather than at startup. max_total_tokens is a single TOTAL budget spent in priority order (fail_to_pass, golden_patch, test_patch, pass_to_pass); anything cut is marked inline. Measured over 400 instances the untruncated block is median 5467 / p90 17683 / p99 77911 tokens, so the 32768 default truncates ~4.8%. Watch privilege/frac_truncated. Off by default (value.swe_privileged_critic.enabled: false). Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
Six conflict regions in ppo.py, all in nemo_rl/algorithms/ppo.py:
* _pooled_explained_var (theirs) landed immediately before _calibration_ece,
whose signature this side had reformatted -- kept both.
* _build_ppo_rollout_dump_payload gained turn_spans (ours) and adv_raw_metrics
(theirs) independently -- kept both, at the definition and both call sites.
* The value_train_batch selection was hoisted out of the PPO epoch loop by
1fa70a8 so the extra critic-only passes can reuse it. This side had added
_prepare_value_train_batch immediately after that selection, inside the loop.
Resolved by taking the hoisted block and moving the _prepare_value_train_batch
call up with it: applied once, before the extra passes, so every critic pass
in a step trains on the same homogeneous-group weighting. Leaving it in the
loop would have dropped the weighting from those extra passes silently --
MseValueLossFn reads the offsets with .get(), so there would be no error.
multimodal_utils.py merged cleanly: this side's duplicate lazy-decord fix was
dropped before merging, since eb028ce already carries an identical one.
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
…r mask e7ac37b made critic/explained_var pre-update by overwriting it driver-side from the rollout-time values. Merged with the residual critic that is wrong in two ways, neither of which git could flag -- the two changes are in different functions and merged cleanly. 1. SPACE. ResidualBaselineEstimator keeps the batch in whichever space the run uses: with residual_baseline=true train_data holds values = C and returns = Y, so pooling it as-is yields 1 - Var(Y - C)/Var(Y), i.e. the RESIDUAL EV, stored under critic/explained_var. The control arm stores the absolute EV under the same key. The headline metric of the A/B would mean two different things in the two arms, and critic/ev_res -- which is derived from the value loss, hence a different critic snapshot -- would not be comparable to it either. _pooled_explained_var now returns both spaces from one numerator (the prediction error is offset-invariant) using the estimator's returns_to_abs / returns_to_res offsets, and both critic/explained_var and critic/ev_res are set from it. Both are pre-update, both on a fixed space. Without a residual estimator both offsets are absent and both collapse to today's number. 2. MASK. It pooled over train_data["token_mask"], the full response mask. In turn_gae mode returns are structurally zero off-anchor (scatter_turns_to_anchors) while the critic emits a dense value, so the EV would be dominated by V^2 at positions the critic is never supervised at. Use _value_metric_mask, which already exists for exactly this and is what the positional diagnostics use. critic/explained_var_post_update is untouched: it is still the loss-derived number from the post-update forward, which is the point of that key. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
ppo_nano_v3_5_swe_hsg.yaml arrived on this branch (d8ec189) in parallel with the residual-critic work, so it never got the keys that work added to the other four configs. adv_estimator.residual_baseline is the one that matters: _create_advantage_estimator requires it explicitly for gae/turn_gae with no default, so every hsg run would have died at setup -- after cluster spin-up and gym venv creation, i.e. minutes of allocation per attempt. The other two have safe fallbacks (homogeneous_group_weight defaults on the BaseModel, swe_privileged_critic is read with .get()) but are added for parity with the cmh config, so the two clusters can run the same arms. Values are identical to ppo_nano_v3_5_swe_cmh.yaml; nothing else in the file changed (verified by comparing the parsed YAML before and after). Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
42452ef changed _pooled_explained_var to return (ev_abs, ev_res) but left test_pooled_explained_var_pre_update -- which arrived with e7ac37b -- asserting against a scalar. Comparing a tuple to pytest.approx(...) does not raise; approx swallows the TypeError and returns False, so all four assertions simply failed. Update them to assert both spaces, and add test_pooled_explained_var_residual_offsets for the returns_to_abs/returns_to_res arms, which had no coverage at all. It pins the properties the merge fix relies on: a per-sample offset that is constant across samples moves neither EV (variance is shift-invariant); an offset applied only to the absolute space raises ev_abs and leaves ev_res untouched; and offsets never enter the shared numerator, so a perfect predictor stays at 1.0 in both spaces. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
value.train_micro_batch_size is a required field of ValueConfig
(nemo_rl/models/value/config.py:47, not NotRequired), so
value_config.get("train_micro_batch_size", 1) is a call-site default of the kind
the config conventions forbid: the value is supposed to come from the YAML and
nowhere else.
It is not harmless here. The fallback feeds the SWE privileged critic's
packing-budget bump, so if the key were ever absent or null the bins would be
sized to _needed * 1 instead of _needed * mbs -- under-sized exactly where that
block exists to prevent it, and surfacing as "Sequence length N exceeds bin
capacity" hours into a run rather than at startup. Fail loudly at config load
instead.
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
_create_advantage_estimator has raised "PPO only supports 'gae', 'turn_gae' or 'raw_reward'." since 2f5ff50 added the turn-level estimator, but test_create_advantage_estimator_rejects_unsupported_name still matched the older "only supports 'gae' or 'raw_reward'". The test has been red on this branch since that commit -- it fails identically on backup/ppo-dev-pre-merge, so neither the residual/privileged work nor the merge introduced it. Worth fixing rather than leaving: pyproject's addopts carries -x, so this single stale assertion aborted the whole tests/unit/algorithms run after 144 of 177 items, hiding everything ordered after it. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
TurnLevelGeneralizedAdvantageEstimator assigned last_metrics wholesale from turn_level_metrics(), so the pre-whitening advantage stats that GeneralizedAdvantageEstimator reports never appeared under turn_gae. The two features landed in parallel -- adv_raw/* on the remote side, turn_gae here -- so the gap only became visible once they were merged. It is not only a missing diagnostic. b5eb09e reconstructs unwhitened advantages in the rollout dump as `advantages * adv_raw_std + adv_raw_mean`, and `if adv_raw_metrics:` sees only turn/* keys, so the dump silently degraded to "absent" for exactly the runs the turn estimator exists for. advantage/turn_std_prenorm is not a substitute: it is computed over the [B, K] turn tensor, while the whitening is a token-level masked statistic, so it is not the affine constant the reconstruction needs. Hoist the helper to a module-level raw_advantage_metrics() -- it only ever used self.normalize_advantages -- and call it from both estimators on the TOKEN-level advantages under the same mask the whitening uses. In turn mode the per-turn advantage has already been broadcast to every token of its turn (scatter_turns_to_tokens), so the two arms are directly comparable and whiten_gain describes the rescale actually applied. Merge into last_metrics instead of assigning, so turn_level_metrics can no longer drop it; the two key families are disjoint (adv_raw/* vs advantage/*, critic/*, turn/*). Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Overview: Bring PPO to feature parity with GRPO — add non-colocated generation, asynchronous training, and NeMo-Gym
rollout support, plus the train/inference-mismatch diagnostics and resume/robustness fixes that go with them.
with an NCCL-collective weight refit, mirroring GRPO (vLLM-only; SGLang/Megatron-generation rejected). Ports GRPO's train/inference mismatch diagnostics to PPO: gen_kl_error, token_mult_prob_error, sequence-level mismatch masking, and the W&B mismatch plot.
V-shape); drop the survivorship-biased incomplete frontier target on resume (configurable, default off); and cycle
the async collector's dataloader across epochs (fixes a silent buffer-stall hang on resume once a finite dataset is
exhausted).
boundaries, replay-buffer restore, collector epoch-cycling) and functional smoke tests (ppo_non_colocated.sh,
ppo_async.sh, ppo_nemo_gym.sh, ppo_async_gym.sh). Sync + async NeMo-Gym validated end-to-end on GB200.
Issues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information