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

[michal2409]

What does this PR do?

Core-algorithm remainder of the Qwen3.5-397B reference training bring-up, targeting main. Per review feedback the independent pieces were split into their own PRs and removed from this branch; what stays here are the changes that touch the core async-GRPO path. Same lineage as #3364 (which targets opt/dev per the agreed workflow).

Included commits:

• feat(megatron): make DDP check_for_nan_in_grad configurable — plumbed through distributed_data_parallel_config (absent keeps the always-on default).
• fix(vllm): async-server robustness for multi-turn agentic rollouts — normalize non-list/str message content before template rendering, and rebuild the isolated prefix render from vLLM's normalized message representation (Qwen3.5 re-renders assistant history with normalized reasoning content, so the raw deepcopy diverged token-wise). Includes replace_prefix_tokens unit tests.
• feat(grpo): async-GRPO robustness for long-horizon SWE rollouts — start trajectory collection only after the initial refit; drain in-flight rollouts before in-loop validation; close pause races at validation boundaries; optional env.nemo_gym.backfill_failed_rollouts; refit-safe prefix caching (invalidate vLLM prefix cache on every refit path); positional advantage grouping for NeMo-Gym rollouts (_stable_group_ids); replay buffer consumes trajectories staged for future steps instead of stalling; opt-in NRL_LP_MASK_DEBUG diagnostics.

Split out per review (separate PRs, removed from this branch):

• feat(grpo): add val_start_at to delay periodic validation #3400 — grpo.val_start_at (delayed periodic validation)
• feat(grpo): validation-only sampling params and grouped pass@k validation #3401 — validation-only sampling params + grouped pass@k validation accuracy
• feat(grpo): make env-flagged sample masking optional #3402 — configurable env-flagged sample masking (grpo.mask_env_flagged_samples)
• feat(grpo): optional early stop at a validation accuracy threshold #3404 — grpo.stop_at_validation_accuracy early stop

Also removed from this branch:

• MLPerf telemetry hooks (TrainTelemetryHooks / mlperf_logger) — MLPerf compliance logging now lives entirely in the benchmark launch stack as a Logger wrapper, so nemo-rl needs no hooks.
• The Qwen3.5 MoE refit expert-split (fix: make Qwen3.5 MoE refit work and fail loudly on weight-update errors #3403) — isolation experiments showed vLLM 0.20 loads the stacked fused-expert export natively (byte-identical to cold load) once moe_backend=triton is set; the historical refit failure was the FlashInfer TRT-LLM post-load repack, which no export format survives. See fix: make Qwen3.5 MoE refit work and fail loudly on weight-update errors #3403 for details.

Diff is now 11 files, +843/−116 (was 17 files, +1542/−204).

Convergence data

The commits here are a subset of the branch validated end-to-end on the benchmark (results below from that full stack; the removed pieces are config-gated features that were active in the run via their split-PR equivalents).

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):

step	1	2	3	4	5	6	7	8	9	10
reward	.406	.441	.406	.434	.520	.387	.410	.457	.629	.598

step	11	12	13	14	15	16	17	18	19	→ eval@20
reward	.602	.461	.617	.586	.637	.613	.648	.574	.605	0.789

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.