feat: fp32 LM head in vLLM and MCore + fix max-length handling - #3936
feat: fp32 LM head in vLLM and MCore + fix max-length handling#3936HeyyyyyyG wants to merge 9 commits into
Conversation
bf16 rounding of the LM-head GEMM output is the dominant contributor to train/token_mult_prob_error. Logits are ~15-30 in magnitude, where the bf16 ulp is 0.125-0.25, so rounding perturbs each logprob by ~0.02-0.03. Both engines round to the same grid, which is why fixing only one side makes the multiplicative error worse (1.029 -> 1.035) and both must be changed together. Megatron (megatron_cfg.fp32_lm_head: true | "tf32"): wrap output_layer.forward to upcast input and weight, inside the autograd graph so gradients still reach the bf16 weight. "tf32" runs that fp32 GEMM on TF32 tensor cores, which is numerically identical here (inputs are already exact bf16 values, and accumulation stays fp32) but avoids the 2.1x cost of full fp32: 19.75s vs 40.1s vs a 19.5s baseline on a fixed logprob pass. Also plumbs megatron_cfg.fp32_residual_connection. vLLM (NRL_VLLM_FP32_LM_HEAD=1): source-patch NemotronH.compute_logits to use an fp32 copy of lm_head. This must be a source patch, not a monkeypatch: the model executes in EngineCore subprocesses that never see a parent-process patch. Keeping that copy correct under RL is the subtle part. Engines start on load_format=dummy and receive trained weights by refit, so a copy cached on first use holds dummy weights; measured that way, every sequence exceeded the masking threshold (mean seq error 6.2e11, loss and reward 0). The copy is therefore invalidated in _load_weights, the leaf every transport funnels through, including the checkpoint engine, which bypasses the weight-update lifecycle. Refit logs show it is load-bearing per step, not just at startup: the live head drifts ~0.02 from the cache after each update. Building the copy is skipped under CUDA graph capture, since an allocation there lives in the graph's private memory pool and is not valid for later eager replays, and it is stored with object.__setattr__ so it does not enter named_parameters(), which the refit weight mapping is built from. The MTP/Eagle3 drafter is covered too: it is a separate module with its own head and its own weight stream. Measured on Nemotron 3.5 Super, 2-step SWE GRPO, iter10500: train/token_mult_prob_error 1.03522 -> 1.01543 and train/gen_kl_error 0.004042 -> 0.000797, with rewards unchanged and no measurable speed cost. Holds under MTP speculative decoding (1.03966 -> 1.01468). All three flags default off, so this is inert unless requested. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com> Declares both keys on MegatronConfig as NotRequired with documentation, and records the defaults in the exemplar configs plus the matching tests/unit/reference_configs snapshots, per the config conventions: no call-site fallbacks, defaults live in the YAML for TypedDict schemas. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
The Gym vllm_model proxy already handles a prompt that overflows the context window: it converts the failure into an empty completion with finish_reason="length" and lets the rollout end cleanly (responses_api_models/vllm_model/app.py). That handler fires only when the response is HTTP 400 *and* the body contains "context length". The NeMo-RL chat endpoint satisfied neither condition on two of its three overflow paths. create_chat_completion mapped only VLLMValidationError to 400, but the online renderer and _clamp_max_tokens raise a plain ValueError for the same condition, which escaped as a 500. The code was already half aware of this: preprocess_chat catches (ValueError, VLLMValidationError), logs "Prompt exceeds max_model_len", then re-raises into a handler that only catches one of them. Separately, the _clamp_max_tokens message said "max_model_len", which matches neither substring the proxy looks for, so even a 400 would not have been classified. The result was that an overflowing rollout burned its retry budget on 500s and then failed, instead of being ended gracefully and masked from the gradient. Seen on both agent harnesses: rarely with openhands (5 occurrences in a 2-step run, which survived) and often with opencode, where it stalled a step at 7 of 8 trajectories. Catch context-length ValueError and return 400, letting any other ValueError still surface as a 500, and reword the clamp error to say "maximum context length" so the proxy's substring check matches. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
|
Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
HeyyyyyyG
changed the title
HeyyyyyyG
changed the title
…aqiz/super_v35_fp32_lmhead Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com> # Conflicts: # nemo_rl/models/generation/vllm/patches.py
HeyyyyyyG
changed the title
HeyyyyyyG
changed the title
An overflowing prompt is recovered by Gym's vllm_model proxy: it returns an empty completion with finish_reason="length" instead of failing the rollout. That only happens when this endpoint answers HTTP 400 *and* the body contains "context length" (responses_api_models/vllm_model/app.py). Both halves broke silently before — the rollout just burned its retry budget on 500s and failed — and neither is visible from either repo alone, so pin the contract here. Extracts the marker, the message builder and the detection predicate to module scope so they are reachable from a test; create_chat_completion and _clamp_max_tokens now share them instead of repeating the literal. Covers: all three overflow shapes vLLM produces (clamp, online renderer, and vLLM's own serving message) are detected and satisfy Gym's classifier; unrelated ValueErrors are not, so they keep surfacing as 500s; and the endpoint still catches plain ValueError and answers 400 — asserted against the source, since that path is nested in a Ray actor method and needs vLLM to import. Both checks were mutation-tested: restoring the old "max_model_len" wording fails the detection tests, and narrowing the except clause back to VLLMValidationError alone fails the endpoint test. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
The first version restated Gym's classifier inside the test, which only pinned one side: Gym could change its substrings or status check and the test would stay green while production broke. Lift the expressions assigned to `is_out_of_context_length` out of responses_api_models/vllm_model/app.py and evaluate them against the response NeMo-RL actually sends, so the test passes only while both sides agree. Verified by mutating each side in turn: returning 500, or rewording the message back to "max_model_len", fails; and on the Gym side, renaming the substring to "context_window" or switching the check to HTTP 422 fails, while renaming the predicate variable raises with a message pointing at the moved handler. Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
| marked.append(f"{label}:none-yet") | ||
| # Logged once per worker: _load_weights runs per refit batch. | ||
| if not getattr(self, "_nrl_fp32_dirty_logged", False): | ||
| self._nrl_fp32_dirty_logged = True |
There was a problem hiding this comment.
2 action items.
TL;DR — pre-commit fails on this line, and separately the refresh log below it prints once per rank per training step.
AI-1 — pyrefly is red (CI-blocking)
ERROR Attribute `_nrl_fp32_dirty_logged` is implicitly defined by assignment in
method `_mark_fp32_lm_head_dirty`, which is not a constructor
--> nemo_rl/models/generation/vllm/vllm_backend.py:627:18
VllmInternalWorkerExtension has no __init__, which is why its siblings are declared at class level — _mtp_drafter_from_disk: bool = False and _nrl_named_parameters — precisely so pyrefly accepts their later in-method assignment.
Action: add _nrl_fp32_dirty_logged: bool = False beside them at line 187. (No suggestion block — that line is outside this PR's diff hunks.) Verified: with that one line added, pyrefly check reports errors shown: 0. Everything else in pre-commit passes.
AI-2 — the refresh log is ungated
_sync_fp32_lm_head runs from the _weight_update_lifecycle finalizer, i.e. every refit, on every rank, per model — so the print at :700 is unbounded. Note the sibling on this line is already gated behind _nrl_fp32_dirty_logged, and :1276 is rank-0 gated; those are the two house idioms. Demoting to logger.info would not help — the call site frequency is the problem, not the sink. See the separate comment on :686.
Minor, same file: megatron_policy_worker.py:554 re-imports nemo_rl.models.megatron.setup, which is already imported at module top (:78).
| output_layer = getattr(module, "output_layer", None) | ||
| if output_layer is None: | ||
| continue # not the last pipeline stage |
There was a problem hiding this comment.
1 action item.
TL;DR — apply_fp32_lm_head never checks that it wrapped anything, so on Nemotron VL/Omni it silently wraps nothing while the vLLM side still runs fp32 — the exact one-sided config this PR's own YAML comment says is worse than leaving the feature off.
The walk is while hasattr(module, "module"), then getattr(module, "output_layer", None). NemotronVLModel and NemotronOmniModel nest the LM under .llava_model.language_model / .language_model and define neither self.module nor __getattr__ delegation (modeling_nemotron_vl.py:54,70-74) — grep -c output_layer on that file is 0. So the loop stops immediately, output_layer is None, and line 505 takes continue # not the last pipeline stage on the last pipeline stage. The print at :529 sits after the assignment, so it never fires either: no error, no warning, no log.
The general problem is that zero matches is indistinguishable from success. The VLM shape is just the case that is reachable today — an upstream rename of output_layer or a new wrapper class would do the same.
Action: fail loud when a chunk that owns the output layer yields no wrap. post_process is the right predicate — it is what Megatron uses to mark that chunk, and it is present on GPTModel, on HybridModel:149 (so this fires for NemotronH too, not just VLMs), and on NemotronVLModel:
| output_layer = getattr(module, "output_layer", None) | |
| if output_layer is None: | |
| continue # not the last pipeline stage | |
| output_layer = getattr(module, "output_layer", None) | |
| if output_layer is None: | |
| if getattr(module, "post_process", False): | |
| raise ValueError( | |
| f"fp32_lm_head is enabled but no output_layer was found on a " | |
| f"post_process model chunk of type {type(module).__name__}. The " | |
| f"trainer would run bf16 while generation runs fp32, which is " | |
| f"worse than disabling both." | |
| ) | |
| continue # not the last pipeline stage |
Separately, the unwrap itself needs the VLM chain to actually work on those models — freeze_moe_router at setup.py:1789-1800 already has it (.thinker -> .llava_model -> .language_model), with a comment naming these exact classes.
| for label, model in ( | ||
| ("policy", self.model_runner.model), | ||
| ("drafter", self._get_drafter_model()), | ||
| ): |
There was a problem hiding this comment.
1 action item.
TL;DR — the drafter branch allocates a ~512 MiB/rank fp32 head that nothing ever reads, because the source patch does not cover the drafter's model file.
The docstring says this "covers the MTP/Eagle3 drafter too", but _patch_vllm_nemotron_h_fp32_lm_head rewrites only model_executor/models/nemotron_h.py. The NemotronH drafter is NemotronHMTP in a different file, with its own full-vocab ParallelLMHead and its own unpatched compute_logits. So the fp32 copy built for "drafter" is never consulted.
Cost on the shipped spec-decode recipe (num_speculative_tokens: 5, tensor_parallel_size: 4, vocab 131072 / hidden 4096): 131072/4 x 4096 x 4 B = 512 MiB per rank, and it is allocated at the first refit — after memory profiling and after the KV cache is sized, so it is outside vLLM's budget. Correctness is unaffected (the drafter only proposes; the target verifies).
Action: drop the ("drafter", self._get_drafter_model()) entry from both _sync_fp32_lm_head and _mark_fp32_lm_head_dirty, and correct the drafter-coverage sentence in this docstring. Also note the comment at :663 ("Drafters commonly tie their head to the policy's") is wrong for this drafter — NemotronHMTP builds its own head unconditionally, despite its own assert message claiming otherwise. If drafter coverage is genuinely wanted, nemotron_h_mtp.py needs its own patch.
…ides Production runs of the Super blend28 (VL-architecture) checkpoint ran on a random fp32 LM head: vLLM starts on dummy weights, the source patch built the fp32 cache lazily during warmup, and the post-refit sync looked for `lm_head` on `model_runner.model`, which is a CUDAGraphWrapper around a NemotronH_Nano_VL_V2 that nests the text model as `language_model`. The sync logged "has no lm_head; skipping" at info level and the cache was never refreshed, so rollouts collapsed to garbage generations (8192-rollout batches went from ~14 min to >90 min). Controlled probes did not reproduce it because they load real weights (load_format=auto) and never refit. The trainer side had the same shape: `apply_fp32_lm_head` walked `.module` and looked for `output_layer`, which NemotronVLModel/NemotronOmniModel nest under `.llava_model.language_model` / `.language_model` / `.thinker`. Zero matches looked like "not the last pipeline stage", so the trainer stayed bf16 while generation ran fp32. vLLM backend: - `_resolve_lm_head_owner` sheds `unwrap()` wrappers, then descends into `language_model`, then scans `named_modules`; sync and dirty-marking both use it so they act on the module the patched `compute_logits` reads from. - Sync raises if the policy has no resolvable head (drafters may still tie theirs) and asserts the fp32 copy equals `lm_head` after refresh. Megatron setup: - `_resolve_output_layer_owner` walks `.module`, `.thinker`, `.llava_model`, `.language_model` until `output_layer` appears (same chain as `freeze_moe_router`). - A `post_process` chunk with no `output_layer` now raises instead of silently continuing. Tests cover wrapper + VL resolution, build-then-refresh across an in-place refit, fail-loud on a headless policy, drafter tolerance, and the three nested Megatron shapes. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Conflicts in vllm_worker_async.py: both sides converted context-length overflow errors to HTTP 400. Kept this branch's version, which keys on the "context length" marker that Gym's vllm_model proxy actually checks and is pinned by tests/unit/models/generation/test_vllm_context_length_overflow.py. Dropped the now-unreachable plain-ValueError handler from #3878 and replaced its test parameter that lacked the marker with the two real vLLM messages. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
- Declare `_nrl_fp32_dirty_logged` (and the new `_nrl_fp32_refresh_logged`) at class level; VllmInternalWorkerExtension has no __init__, so pyrefly rejects in-method first assignment. - Gate the per-refit refresh log: print the first refresh per worker (the dummy->real transition whose drift is the useful signal), debug after. - Drop the drafter from the fp32 head targets. NemotronHMTP has its own unpatched compute_logits in nemotron_h_mtp.py, so its fp32 copy (~512 MiB/rank, allocated after KV-cache sizing) was never read. - Use the module-level `apply_fp32_lm_head` import in the Megatron worker. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…2 head `validate_fp32_lm_head_config` rejects megatron_cfg.fp32_lm_head without NRL_VLLM_FP32_LM_HEAD=1 in generation.vllm_cfg.env_vars (and vice versa), and rejects fp32_lm_head combined with use_fused_linear_logprobs, whose fused linear+CE kernel bypasses output_layer. Both leave one engine in bf16 while the other runs fp32, which the config comment already documents as worse than disabling the feature. Skipped when generation is not vLLM. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
2 participants
What does this PR do ?
Using FP32 for LM head for both vLLM and MCore
Handle the situation where model hits max length correctly. There's a handler in gym, but currently nemorl doesn't trigger it.
Usage