fix: Qwen3.5-397B training stability and async-GRPO robustness - #3364
Merged
Superjomn merged 5 commits intoJul 29, 2026
Merged
Conversation
- NCCL process-group timeout configurable via NRL_NCCL_TIMEOUT_MINUTES, applied to mcore sub-group creation as well - check_for_nan_in_grad and overlap_grad_reduce plumbed through distributed_data_parallel_config (mcore defaults restored) - megatron_cfg.attention_backend knob for NVTE backend selection (flash required for head_dim=256 on sm10.3) - expandable_segments allocator config for megatron workers - NCCL flight-recorder env plumbing for deadlock forensics Squashed from mlperf-training-qwen35-main (4a910e9, d9c9d61, 94d2083, 78d82ef, df2242b, 2113d9a, f941d27, f61f6f5). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
- split stacked fused-expert exports into per-expert HF tensors so megatron->vLLM refit works with triton moe backend - fix AttributeError on self.cfg in HTTP chat preprocessing by binding worker_self (500s on every chat/tokenize request) - enable prefix caching made refit-safe: invalidate vLLM prefix cache on every refit path - FULL_AND_PIECEWISE CUDA graphs for decode; compiled engines re-enabled Squashed from mlperf-training-qwen35-main (897ad87, f281097, 48e3526, 3a643e7, 41352f7, 315d238, 2edce20, a8def3b). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
- start trajectory collection only after the initial weight refit (engines with load_format=dummy produced random-weight rollouts) - drain in-flight rollouts before async in-loop validation - close pause races that launched train batches into validation windows - optional env.nemo_gym.backfill_failed_rollouts: back-fill failed/NaN/ empty rollouts as zero-reward trajectories instead of aborting the batch (absent/false keeps the fail-fast default; long multi-node benchmark runs enable it) - optional grpo.val_start_at protocol key delaying periodic validation (benchmark recipes set it; absent means no delay) and training-dynamics alignment with the reference branch - validation-only sampling overrides (grpo.validation_generation) with request tagging so the generation server exempts validation requests from the strict training sampling checks - grouped validation pass@k metrics (grpo.num_val_generations_per_prompt) - TrainTelemetryHooks protocol + optional mlperf_logger hooks in the train loops (implementations injected by benchmark launch stacks) - opt-in per-sequence logprob-mask debug line (NRL_LP_MASK_DEBUG=1) Squashed from mlperf-training-qwen35-main (ea2bad2, fd90386, 9402b7d, 983a1f3, f7013bc, 96252d0, 20f5c8a). Co-authored-by: Al Rigazzi <arigazzi@nvidia.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
Add grpo.mask_env_flagged_samples (absent/true keeps the upstream behavior: env-driven mask_sample flags ride in async NeMo-Gym rollout batches and zero loss_multiplier via _apply_mask_sample_filter). Setting it false keeps the key out of the batch entirely. Motivation: a 6-run A/B on GB200+GB300 (2026-07-16, Qwen3.5-397B SWE, sequence packing on) showed the mere presence of the [B]-bool key collapsing training — per-sequence train-vs-gen logprob-error masking ramped from <10/256 to 252/256 by step ~9-14 with rewards and val dropping to zero — even though the tensor was provably all-False and the filter measurably masked 0 samples (2/2 branches with the key collapsed; 4/4 without stayed clean, including a pair differing only in this line). Suspected interaction with the sequence-packing/microbatch path; until that is root-caused, affected runs need a supported way to opt out. num_mask_sample_filtered stays 0 in collapsing runs, so it is not a sufficient health signal on its own. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
michal2409
marked this pull request as draft
July 27, 2026 20:28
michal2409
changed the title
michal2409
marked this pull request as ready for review
July 28, 2026 12:26
michal2409
force-pushed
the
mlperf-training-qwen35-main-v2
branch
2 times, most recently
from
July 28, 2026 14:23
6c484d7 to
c147ae1
Compare
Removes everything the GB300 bring-up needed that the finished recipe does not, restoring upstream defaults: - NCCL watchdog override (NRL_NCCL_TIMEOUT_MINUTES, the timeout= on init_process_group, mcore distributed_timeout_minutes, flight_recorder_dump_on_timeout): the 10-minute firings it papered over were early kills of two real deadlocks, both fixed structurally. Worst measured collective-bearing phase on converged GB300 runs is 227 s against the 600 s default. - first_last_layers_bf16 / num_layers_at_start_in_bf16 / num_layers_at_end_in_bf16 forwarding: mcore consumes these only in fp8/fp4 code paths; inert under pure bf16 training. - _normalize_optimizer_dtypes: no config on the Megatron path sets string optimizer dtype names anymore. - truncate_prompt_tokens config key and its async-server request shim: never set by any config; replaced an env var that was itself dead. - NRL_MEM_DEBUG pre-fwd/bwd memory snapshot: OOM forensics, recipe is memory-stable now. - One-shot [refit-debug] expert-shape sender log: its receiver-side counterpart no longer exists, and refit load failures raise loudly since the async update-weights fix. - Comment trim on the check_for_nan_in_grad plumbing to match its bare siblings. Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
michal2409
force-pushed
the
mlperf-training-qwen35-main-v2
branch
from
July 28, 2026 14:33
c147ae1 to
f82b9ba
Compare
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
2 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?
Merges the Qwen3.5-397B reference training changes into
opt/dev, rebased onto currentmain(branch is 0 commits behindopt/dev).Included commits:
Convergence data
Setup
64×GB300 (4 GPUs/node): 16 training nodes (Megatron TP4·PP2·EP32, sequence
packing @ 65 536 tokens, bf16) + 48 generation nodes (24 vLLM engines, TP8/EP8, async
engine). GBS 256 = 16 prompts × 16 generations/step, async GRPO lag-1, NeMo-Gym
OpenHands agents on R2E-Gym easy (curriculum-v2 train set), validation = 256 instances
× 4 generations @ temp 0.1 / top-p 0.95, seed 3978.
Training reward (fraction of 256 rollouts solved, per step):