Skip to content

fix: Qwen3.5-397B training stability and async-GRPO robustness #3382

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
michal2409 wants to merge 3 commits into
base: main
Choose a base branch
from
Open

fix: Qwen3.5-397B training stability and async-GRPO robustness #3382

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
26e3777
feat(megatron): make DDP check_for_nan_in_grad configurable
michal2409 Jul 29, 2026
248be16
fix(vllm): async-server robustness for multi-turn agentic rollouts
michal2409 Jul 29, 2026
83c082a
feat(grpo): async-GRPO robustness for long-horizon SWE rollouts
michal2409 Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions nemo_rl/algorithms/async_utils/replay_buffer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,10 +186,11 @@ def sample(
) -> Optional[dict[str, Any]]:
"""Sample per-prompt trajectory groups intended for the current training step.

Only returns trajectories with target_weight_version == current_weight_version.
If insufficient trajectories are available, returns None to stall training
until the remaining trajectories are generated. This ensures no trajectory
loses its last chance to be used for its intended training step.
Only returns trajectories with target_weight_version >= current_weight_version,
so trajectories staged for a future step may be consumed early instead of
stalling. If insufficient trajectories are available, returns None to stall
training until the remaining trajectories are generated. This ensures no
trajectory loses its last chance to be used for its intended training step.

Returns:
Dictionary with 'trajectories' and 'avg_trajectory_age' keys, or None if insufficient data
Expand Down Expand Up @@ -252,7 +253,7 @@ def sample(
intended_indices = [
i
for i in valid_indices
if self.target_weight_versions[i] == current_weight_version
if self.target_weight_versions[i] >= current_weight_version
]

print(
Expand Down Expand Up @@ -283,8 +284,10 @@ def sample(
f"✅ Selected counts by generation weight-version: {Counter(sampled_weights)}"
)
print(f"📊 Average trajectory age: {avg_trajectory_age:.2f} steps")
sampled_targets = Counter(self.target_weight_versions[i] for i in selected)
print(
f"🎯 All selected trajectories target step {current_weight_version} (100% target match)"
f"🎯 Selected trajectory targets for step {current_weight_version}: "
f"{dict(sampled_targets)}"
)

# Remove selected items in reverse order to maintain correct indices
Expand Down
86 changes: 52 additions & 34 deletions nemo_rl/algorithms/async_utils/trajectory_collector.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -281,43 +281,51 @@ def _collection_loop(self):
if not self.running:
break

# Check if manually paused and wait
if not self._manual_pause_cleared.is_set() and self.running:
self._manual_pause_cleared.wait()

# Check if refit is in progress and wait
if not self._refit_pause_cleared.is_set() and self.running:
print("⏸️ Pausing collection for refit...")
with self._efficiency_timer.time("idle/refit_event_wait"):
self._refit_pause_cleared.wait()
print("▶️ Refit completed, resuming collection")

# Check if generation limits require pausing collection
if self._should_pause_for_generation_limits() and self.running:
# Only log warning once per weight version
if self._last_limit_warning_version != self.current_weight_version:
async_cfg = self.master_config.grpo.get("async_grpo", {})
max_trajectory_age = async_cfg["max_trajectory_age_steps"]
target_weights = [
self.current_weight_version + i
for i in range(max_trajectory_age)
]

print(
f"⏸️ Pausing collection: all target weights {target_weights} for weight version {self.current_weight_version} "
f"already exist in buffer. Waiting for weight update..."
)
self._last_limit_warning_version = self.current_weight_version
# Wait until NO pause condition holds, re-checking every
# condition after each wake-up: sequential checks race with
# pause() at validation boundaries and let full batches launch
# into the validation window.
while self.running:
if not self._manual_pause_cleared.is_set():
self._manual_pause_cleared.wait()
continue # re-check all conditions after waking

if not self._refit_pause_cleared.is_set():
print("⏸️ Pausing collection for refit...")
with self._efficiency_timer.time("idle/refit_event_wait"):
self._refit_pause_cleared.wait()
print("▶️ Refit completed, resuming collection")
continue # re-check all conditions after waking

if self._should_pause_for_generation_limits():
# Only log warning once per weight version
if (
self._last_limit_warning_version
!= self.current_weight_version
):
async_cfg = self.master_config.grpo.get("async_grpo", {})
max_trajectory_age = async_cfg["max_trajectory_age_steps"]
target_weights = [
self.current_weight_version + i
for i in range(max_trajectory_age)
]

print(
f"⏸️ Pausing collection: all target weights {target_weights} for weight version {self.current_weight_version} "
f"already exist in buffer. Waiting for weight update..."
)
self._last_limit_warning_version = (
self.current_weight_version
)

self._generation_limit_cleared.clear() # Clear the event to pause
self._generation_limit_cleared.clear() # Clear the event to pause

# Efficiently wait for generation limits to be cleared (no polling!)
with self._efficiency_timer.time("idle/generation_limit_pause"):
self._generation_limit_cleared.wait()
# Efficiently wait for generation limits to be cleared (no polling!)
with self._efficiency_timer.time("idle/generation_limit_pause"):
self._generation_limit_cleared.wait()
continue # re-check all conditions after waking

# Double-check we're still running after being woken up
if not self.running:
break
break # nothing requires pausing; clear to launch

if not self.running:
break
Expand Down Expand Up @@ -416,6 +424,12 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None:

use_nemo_gym = _should_use_nemo_gym(self.master_config)

# Honor a manual pause (e.g. a validation boundary) before
# spawning, so the batch cannot launch into the val window.
if not self._manual_pause_cleared.is_set() and self.running:
print("⏸️ Manual pause before batch spawn: holding launch")
self._manual_pause_cleared.wait()

if not self._refit_pause_cleared.is_set() and self.running:
with self._threads_lock:
active_threads = len(self._inflight_threads)
Expand Down Expand Up @@ -550,6 +564,10 @@ def resume_after_refit(self) -> None:
# Invalidate&recompute vLLM caches after the in-flight weight updates if
# recompute_kv_cache_after_weight_updates is True (AREAL-style implementation).
# Otherwise, keep using the stale KV caches (Magistral-style implementation).
# NOTE: for drained (non-in-flight) refits with prefix caching enabled,
# cache invalidation is handled unconditionally in
# refit_policy_generation (grpo.py), which also covers the
# pre-validation refit path that never reaches this method.
async_cfg = self.master_config.grpo.get("async_grpo", {})
if async_cfg.get("in_flight_weight_updates", False) and async_cfg.get(
"recompute_kv_cache_after_weight_updates", False
Expand Down
118 changes: 107 additions & 11 deletions nemo_rl/algorithms/grpo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1688,6 +1688,25 @@ def _scale(reward_tensor: torch.Tensor) -> torch.Tensor:
return repeated_batch


def _stable_group_ids(prompt_ids_for_adv, num_generations_per_prompt):
"""Stable per-prompt grouping key for GRPO advantage computation.

GRPO groups samples by prompt (torch.unique) to compute the leave-one-out baseline. The default
key is the rendered prompt token-ids, but for agentic gym rollouts each generation's first-turn
prompt tokenizes slightly differently (observed on Qwen3-Instruct + hermes: every generation
becomes its own singleton group -> leave-one-out baseline == reward -> advantage == 0 -> zero
gradient; Exp 26). The training batch is laid out as contiguous num_gen blocks per prompt
(async: BatchedDataDict.from_batches of per-prompt groups; sync: repeat_interleave), so the
correct, model-agnostic group id is positional: index // num_gen. Falls back to the original
token-id grouping if the batch is not an exact multiple of num_gen (e.g. dynamic sampling).
"""
n = int(prompt_ids_for_adv.shape[0])
g = int(num_generations_per_prompt)
if g <= 0 or n % g != 0:
return prompt_ids_for_adv
return (torch.arange(n, device=prompt_ids_for_adv.device) // g).unsqueeze(1)


def extract_initial_prompt_messages(
message_logs: list,
original_prompt_lengths: torch.Tensor,
Expand Down Expand Up @@ -2218,7 +2237,9 @@ def refit_policy_generation(
"""
synchronizer = getattr(policy_generation, "weight_synchronizer", None)
if synchronizer is not None:
return synchronizer.sync_weights(timer=timer, kv_scales=kv_scales) or {}
sync_metrics = synchronizer.sync_weights(timer=timer, kv_scales=kv_scales) or {}
_invalidate_prefix_cache_after_refit(policy_generation)
return sync_metrics

# Megatron generation backend needs explicit suspend/resume around refits.
if isinstance(policy_generation, MegatronGeneration):
Expand Down Expand Up @@ -2321,12 +2342,36 @@ def refit_policy_generation(
if colocated_inference or isinstance(policy_generation, MegatronGeneration):
policy_generation.prepare_for_generation(tags=["kv_cache"])

_invalidate_prefix_cache_after_refit(policy_generation)

if isinstance(policy_generation, MegatronGeneration):
policy_generation.resume_after_refit()

return {}


def _invalidate_prefix_cache_after_refit(
policy_generation: GenerationInterface,
) -> None:
"""Drop reusable KV blocks after a weight update.

vLLM prefix-cache blocks are keyed only by token ids, so blocks
prefilled under the old weights would be silently reused after refit
(stale KV -> train/gen logprob divergence). Called on every refit path;
backends without reusable caches inherit the no-op interface default.
"""
generation_cfg = getattr(policy_generation, "cfg", None) or {}
vllm_cfg = generation_cfg.get("vllm_cfg") or {}
if vllm_cfg.get("enable_prefix_caching"):
if not policy_generation.invalidate_kv_cache():
raise RuntimeError(
"❌ Error: prefix caching is enabled but invalidating the "
"vLLM prefix/KV cache after refit failed; continuing would "
"sample rollouts against stale KV computed under the "
"pre-refit weights."
)


def _initial_policy_generation_stale(
policy_generation: GenerationInterface, completed_steps: int
) -> bool:
Expand Down Expand Up @@ -2450,6 +2495,30 @@ def compute_and_apply_seq_logprob_error_masking(
)
masked_correct_pct = masked_correct_count / num_masked_seqs

# [lp-mask-debug] one parseable line per step attributing train/gen
# logprob divergence; opt in via NRL_LP_MASK_DEBUG=1.
if os.environ.get("NRL_LP_MASK_DEBUG") == "1":
seq_lens = mask.sum(dim=-1)
masked_rows = [
(
int(seq_lens[i]),
float(seq_mult_prob_error[i]),
float(rewards.view(-1)[i]),
)
for i in torch.nonzero(diff_mask_bool).flatten().tolist()
]
kept_bool = seq_error_mask.bool() & valid_seq_mask
kept_lens = seq_lens[kept_bool]
print(
"[lp-mask-debug] masked(len,err,rew)="
+ ";".join(f"{l},{e:.2f},{r:.0f}" for l, e, r in masked_rows[:200])
+ f" | kept_len mean={float(kept_lens.float().mean()):.0f}"
f" p90={float(kept_lens.float().quantile(0.9)):.0f}"
f" max={int(kept_lens.max())}"
f" | masked_len mean={sum(r[0] for r in masked_rows) / len(masked_rows):.0f}",
flush=True,
)

# Compute after-mask metrics (only for sequences that passed the threshold)
kept_mask = seq_error_mask.bool() & valid_seq_mask
if kept_mask.sum() > 0:
Expand Down Expand Up @@ -3064,8 +3133,17 @@ def grpo_train(
sample_mask = train_data["sample_mask"]
mask = token_mask * sample_mask.unsqueeze(-1)

# Positional grouping is only needed for agentic gym rollouts,
# where per-generation prompt tokenization is non-deterministic;
# keep main's token-id grouping for every other GRPO user.
advantage_group_ids = prompt_ids_for_adv
if _should_use_nemo_gym(master_config):
advantage_group_ids = _stable_group_ids(
prompt_ids_for_adv,
master_config.grpo["num_generations_per_prompt"],
)
train_data["advantages"] = adv_estimator.compute_advantage(
prompt_ids=prompt_ids_for_adv,
prompt_ids=advantage_group_ids,
rewards=rewards,
mask=mask,
repeated_batch=repeated_batch,
Expand Down Expand Up @@ -3972,14 +4050,6 @@ def async_grpo_train(
next_nemo_gym_task_index=next_nemo_gym_task_index,
)

# Start trajectory collection in background
collection_task = trajectory_collector.start_collection.remote(dataloader)

# Ensure collector knows initial weight version
trajectory_collector.set_weight_version.remote(weight_version)

print("📦 Started continuous background trajectory collection")

print(
f"🚀 Starting async GRPO training with buffer_size={optimal_buffer_size}, max_age={max_trajectory_age_steps} steps"
)
Expand Down Expand Up @@ -4009,6 +4079,17 @@ def async_grpo_train(
traceback.print_exc()
return

# Start trajectory collection only after generation holds real weights.
# The engines come up with load_format=dummy (weights arrive via the refit
# above); collecting before the refit fills the buffer with garbage
# rollouts sampled from randomly initialized weights.
collection_task = trajectory_collector.start_collection.remote(dataloader) # noqa: F841

# Ensure collector knows initial weight version
trajectory_collector.set_weight_version.remote(weight_version)

print("📦 Started continuous background trajectory collection")

print("✅ Policy generation setup complete, proceeding to validation...")

# Run validation at start if configured
Expand Down Expand Up @@ -4448,8 +4529,17 @@ def async_grpo_train(
sample_mask = train_data["sample_mask"]
mask = token_mask * sample_mask.unsqueeze(-1)

# Positional grouping is only needed for agentic gym rollouts,
# where per-generation prompt tokenization is non-deterministic;
# keep main's token-id grouping for every other GRPO user.
advantage_group_ids = prompt_ids_for_adv
if _should_use_nemo_gym(master_config):
advantage_group_ids = _stable_group_ids(
prompt_ids_for_adv,
master_config.grpo["num_generations_per_prompt"],
)
train_data["advantages"] = adv_estimator.compute_advantage(
prompt_ids=prompt_ids_for_adv,
prompt_ids=advantage_group_ids,
rewards=rewards,
mask=mask,
repeated_batch=repeated_batch,
Expand Down Expand Up @@ -4558,6 +4648,12 @@ def async_grpo_train(
with timer.time("idle/validation"):
# Pause trajectory collection during validation to reduce memory pressure
trajectory_collector.pause.remote()
# Drain in-flight rollouts too: pause only stops new
# launches, and in-flight train rollouts sharing the
# engines push val agents into timeout.
ray.get(
trajectory_collector.wait_for_pending_generations.remote()
)

if NEED_REFIT and POLICY_GENERATION_STALE:
refit_metrics = refit_policy_generation(
Expand Down
Loading
Loading