From be2ceb5cc42d0fe92237469320b4512d4db97d93 Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Thu, 10 Sep 2026 17:41:34 -0700 Subject: [PATCH 01/18] feat: add fp32 LM head toggle Signed-off-by: Guyue Huang --- examples/configs/distillation_math.yaml | 1 + examples/configs/evals/eval.yaml | 1 + examples/configs/grpo_math_1B.yaml | 1 + examples/configs/ppo_math_1B.yaml | 1 + nemo_rl/models/generation/vllm/config.py | 12 + nemo_rl/models/generation/vllm/patches.py | 230 ++++++++++++++++ nemo_rl/models/generation/vllm/vllm_worker.py | 11 + nemo_rl/models/megatron/setup.py | 134 +++++++++ nemo_rl/models/policy/__init__.py | 14 + .../policy/workers/megatron_policy_worker.py | 8 + .../models/generation/test_vllm_patches.py | 258 ++++++++++++++++++ .../models/megatron/test_megatron_setup.py | 150 ++++++++++ .../reference_configs/distillation_math.yaml | 1 + tests/unit/reference_configs/eval.yaml | 1 + .../unit/reference_configs/grpo_math_1B.yaml | 1 + .../ppo_math_1B_megatron.yaml | 1 + 16 files changed, 825 insertions(+) diff --git a/examples/configs/distillation_math.yaml b/examples/configs/distillation_math.yaml index c288b50999c..000937735bd 100644 --- a/examples/configs/distillation_math.yaml +++ b/examples/configs/distillation_math.yaml @@ -228,6 +228,7 @@ policy: &POLICY_BASE precision: ${...precision} kv_cache_dtype: "auto" logprobs_mode: processed_logprobs + fp32_lm_head: false # Compute Nemotron-H logits with an fp32 LM head. Pair with policy.megatron_cfg.fp32_lm_head. # false: use NeMo-RL's legacy refit loader; true: opt into vLLM reload_weights. refit_with_reload_api: false tensor_parallel_size: 1 diff --git a/examples/configs/evals/eval.yaml b/examples/configs/evals/eval.yaml index 8f4da66a50e..c9393174ac3 100644 --- a/examples/configs/evals/eval.yaml +++ b/examples/configs/evals/eval.yaml @@ -22,6 +22,7 @@ generation: vllm_cfg: async_engine: false precision: "bfloat16" + fp32_lm_head: false # Compute Nemotron-H logits with an fp32 LM head. tensor_parallel_size: 1 pipeline_parallel_size: 1 expert_parallel_size: 1 diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index ba1b704f2b3..62163fcfe4e 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -426,6 +426,7 @@ policy: precision: ${policy.precision} kv_cache_dtype: "auto" logprobs_mode: processed_logprobs + fp32_lm_head: false # Compute Nemotron-H logits with an fp32 LM head. Pair with policy.megatron_cfg.fp32_lm_head. # false: use NeMo-RL's legacy refit loader; true: opt into vLLM reload_weights. refit_with_reload_api: false tensor_parallel_size: 1 diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index c125d1d7930..f7bd4ac20e2 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -284,6 +284,7 @@ policy: precision: ${policy.precision} kv_cache_dtype: "auto" logprobs_mode: processed_logprobs + fp32_lm_head: false # Compute Nemotron-H logits with an fp32 LM head. Pair with policy.megatron_cfg.fp32_lm_head. # false: use NeMo-RL's legacy refit loader; true: opt into vLLM reload_weights. refit_with_reload_api: false tensor_parallel_size: 1 diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py index 0a7b5c03f5b..e1d235a4923 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -28,6 +28,7 @@ VllmRefitTransportName = Literal["s3", "zmq"] VllmRefitSelector = Literal["vllm_s3_sparse", "vllm_zmq_sparse", "nixl", "nccl_reshard"] VLLM_SPARSE_REFIT_TRANSPORTS = frozenset({"vllm_s3_sparse", "vllm_zmq_sparse"}) +VLLM_FP32_LM_HEAD_ENV_VAR = "NRL_VLLM_FP32_LM_HEAD" # TODO(rohitrango): Move model-specific video fields behind ProcessorInterface. @@ -60,6 +61,9 @@ class VllmSpecificArgs(TypedDict): # with generation-time processors should request ``raw_logprobs`` when # comparing generation and policy logprobs. logprobs_mode: NotRequired[Literal["processed_logprobs", "raw_logprobs"]] + # Compute Nemotron-H logits with an fp32 LM head in vLLM. Pair this with + # policy.megatron_cfg.fp32_lm_head when using a Megatron trainer. + fp32_lm_head: NotRequired[bool] # Cap each request's generated tokens so the training prompt plus response # fits within max_model_len. This is needed when multimodal processing makes # the training prompt longer than its text-only representation. @@ -106,6 +110,14 @@ class VllmSpecificArgs(TypedDict): reasoning_parser_plugin: NotRequired[str] +def vllm_fp32_lm_head_enabled(vllm_cfg: VllmSpecificArgs | dict[str, Any]) -> bool: + """Return whether vLLM should run Nemotron-H logits with an fp32 head.""" + if vllm_cfg.get("fp32_lm_head"): + return True + env_vars = vllm_cfg.get("env_vars") + return env_vars is not None and str(env_vars.get(VLLM_FP32_LM_HEAD_ENV_VAR)) == "1" + + class VllmDeltaCompressionConfig(BaseModel, extra="allow"): encoding: Literal["xor", "overwrite"] = "xor" sparse_bucket_size_bytes: PositiveInt = 512 * 1024**2 diff --git a/nemo_rl/models/generation/vllm/patches.py b/nemo_rl/models/generation/vllm/patches.py index f306eadc715..9b1edbaea29 100644 --- a/nemo_rl/models/generation/vllm/patches.py +++ b/nemo_rl/models/generation/vllm/patches.py @@ -16,6 +16,8 @@ from contextlib import contextmanager from importlib.util import find_spec +from nemo_rl.models.generation.vllm.config import VLLM_FP32_LM_HEAD_ENV_VAR + def _get_vllm_file(relative_path: str) -> str: """Return absolute path to a vLLM file or raise if it cannot be found. @@ -622,6 +624,221 @@ def _patch_vllm_glm_decoder_sequence_parallel_moe(logger) -> None: logger.info("Successfully disabled decoder-level SP-MoE for GLM DSA models.") +def _patch_vllm_nemotron_h_fp32_lm_head(logger) -> bool: + """Compute NemotronH logits with an fp32 LM head (MiniMax-M1-style). + + bf16 rounding of the logits GEMM output is the dominant contributor to + generation/training logprob mismatch (train/token_mult_prob_error). With + this patch the sampled-token logprobs come from fp32 logits, matching a + trainer that enables megatron_cfg.fp32_lm_head. + + This must be a source patch (not a monkeypatch): the model executes in + vLLM's EngineCore worker subprocesses, which import vllm independently of + this process. The patched code is opt-in at runtime via + NRL_VLLM_FP32_LM_HEAD=1 (set by policy.generation.vllm_cfg.fp32_lm_head). + When enabled, the live ParallelLMHead keeps its original parameter dtype + and quantization config; only the projection path casts hidden states, + weights, and optional bias to fp32 at runtime. + """ + try: + file_to_patch = _get_vllm_file("model_executor/models/nemotron_h.py") + except RuntimeError: + logger.warning("Could not locate nemotron_h.py for the fp32 LM head patch.") + return False + + old_import_snippet = """import torch +from torch import nn""" + old_fp32_import_snippet = """import os + +import torch +from torch import nn""" + functional_import_snippet = """import os + +import torch +import torch.nn.functional as F +from torch import nn""" + new_import_snippet = old_fp32_import_snippet + old_lm_head_snippet = """ self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + )""" + previous_lm_head_snippet = f""" self._nrl_fp32_lm_head = ( + os.environ.get("{VLLM_FP32_LM_HEAD_ENV_VAR}", "0") == "1" + ) + if self._nrl_fp32_lm_head: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + params_dtype=torch.float32, + quant_config=None, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + )""" + new_lm_head_snippet = f""" self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self._nrl_fp32_lm_head = ( + os.environ.get("{VLLM_FP32_LM_HEAD_ENV_VAR}", "0") == "1" + )""" + old_logits_processor_snippet = ( + " self.logits_processor = LogitsProcessor(config.vocab_size)" + ) + new_logits_processor_snippet = """ self.logits_processor = LogitsProcessor(config.vocab_size) + if self._nrl_fp32_lm_head: + + def _nrl_fp32_lm_head_forward( + input_, embedding_bias=None, _lm_head=self.lm_head + ): + logits = torch.matmul( + input_.to(dtype=torch.float32), + _lm_head.weight.to(dtype=torch.float32).t(), + ) + if embedding_bias is not None: + logits = logits + embedding_bias.to(dtype=torch.float32) + return logits + + self.lm_head.forward = _nrl_fp32_lm_head_forward + _orig_quant_apply = self.lm_head.quant_method.apply + + def _nrl_fp32_lm_head_apply( + layer, + input_, + bias=None, + _lm_head=self.lm_head, + _orig_apply=_orig_quant_apply, + **kwargs, + ): + if layer is _lm_head: + return _lm_head(input_, bias) + return _orig_apply(layer, input_, bias=bias, **kwargs) + + self.lm_head.quant_method.apply = _nrl_fp32_lm_head_apply""" + old_snippet = """ logits = self.logits_processor(self.lm_head, hidden_states) + return logits""" + previous_compute_logits_snippet = """ if self._nrl_fp32_lm_head: + hidden_states = hidden_states.to(dtype=torch.float32) + logits = self.logits_processor(self.lm_head, hidden_states) + return logits""" + # Worker environments can persist across launches. Migrate an installed + # source file that still contains the previous lazy-deepcopy patch. + legacy_snippet = f""" import os as _os + + if _os.environ.get("{VLLM_FP32_LM_HEAD_ENV_VAR}", "0") == "1": + # NeMo-RL patch: fp32 LM head (MiniMax-M1-style). bf16 rounding of + # the logits is the dominant gen/train logprob mismatch source. + _fp32_head = getattr(self, "_nrl_lm_head_fp32", None) + if _fp32_head is None and not torch.cuda.is_current_stream_capturing(): + # Skipped under graph capture: an allocation there lives in the + # graph's memory pool and is not valid for later eager replays. + # Capture output is discarded anyway, so bf16 is fine for it. + import copy as _copy + + _fp32_head = _copy.deepcopy(self.lm_head).float() + # object.__setattr__ bypasses nn.Module.__setattr__: registering + # this as a submodule would add a vocab-sized parameter to + # named_parameters(), which the refit weight mapping is built from. + object.__setattr__(self, "_nrl_lm_head_fp32", _fp32_head) + self._nrl_lm_head_fp32_dirty = False + print( + "[fp32_lm_head] built fp32 head in forward shape=%s" + % (tuple(_fp32_head.weight.shape),), + flush=True, + ) + elif _fp32_head is not None and getattr( + self, "_nrl_lm_head_fp32_dirty", False + ): + # Refreshed in place: replacing the module would leave any + # captured CUDA graph pointing at the old storage. + _fp32_head.weight.data.copy_(self.lm_head.weight) + if getattr(_fp32_head, "bias", None) is not None: + _fp32_head.bias.data.copy_(self.lm_head.bias) + self._nrl_lm_head_fp32_dirty = False + print("[fp32_lm_head] refreshed cached head in forward", flush=True) + if _fp32_head is not None: + return self.logits_processor(_fp32_head, hidden_states.float()) + logits = self.logits_processor(self.lm_head, hidden_states) + return logits""" + + with _locked_file_patch(file_to_patch) as (content, write_back): + if ( + new_import_snippet in content + and new_lm_head_snippet in content + and new_logits_processor_snippet in content + and previous_compute_logits_snippet not in content + ): + logger.info("NemotronH fp32 LM head patch already present.") + return True + + if legacy_snippet in content: + content = content.replace(legacy_snippet, old_snippet, 1) + if previous_compute_logits_snippet in content: + content = content.replace(previous_compute_logits_snippet, old_snippet, 1) + if previous_lm_head_snippet in content: + content = content.replace(previous_lm_head_snippet, old_lm_head_snippet, 1) + if functional_import_snippet in content: + content = content.replace(functional_import_snippet, new_import_snippet, 1) + + if new_import_snippet not in content: + if old_fp32_import_snippet in content: + content = content.replace( + old_fp32_import_snippet, new_import_snippet, 1 + ) + elif content.count(old_import_snippet) == 1: + content = content.replace(old_import_snippet, new_import_snippet, 1) + else: + logger.warning( + "NemotronH fp32 LM head import anchor not found exactly once " + "in %s; patch not applied.", + file_to_patch, + ) + return False + + if new_lm_head_snippet not in content: + if content.count(old_lm_head_snippet) != 1: + logger.warning( + "NemotronH fp32 LM head constructor anchor not found exactly " + "once in %s; patch not applied.", + file_to_patch, + ) + return False + content = content.replace(old_lm_head_snippet, new_lm_head_snippet, 1) + + if new_logits_processor_snippet not in content: + if content.count(old_logits_processor_snippet) != 1: + logger.warning( + "NemotronH fp32 logits_processor anchor not found exactly once " + "in %s; patch not applied.", + file_to_patch, + ) + return False + content = content.replace( + old_logits_processor_snippet, new_logits_processor_snippet, 1 + ) + + if content.count(old_snippet) != 1: + logger.warning( + "NemotronH fp32 compute_logits anchor not found exactly once " + "in %s; patch not applied.", + file_to_patch, + ) + return False + write_back(content) + + logger.info("Applied NemotronH fp32 LM head source patch.") + return True + + def ensure_vllm_source_compat() -> None: """Apply interpreter-independent vLLM source-compat patches. @@ -643,12 +860,19 @@ def _apply_vllm_patches( py_executable: str, *, extra_env_vars: list[str] | None = None, + fp32_lm_head: bool | None = None, ) -> None: # Import lazily so importing the worker module does not import vLLM. import vllm.envs as envs from vllm.logger import init_logger patch_logger = init_logger("vllm_patch") + fp32_lm_head_enabled = ( + fp32_lm_head or os.environ.get(VLLM_FP32_LM_HEAD_ENV_VAR) == "1" + ) + if fp32_lm_head_enabled: + os.environ[VLLM_FP32_LM_HEAD_ENV_VAR] = "1" + extra_env_vars = [*(extra_env_vars or []), VLLM_FP32_LM_HEAD_ENV_VAR] # Whether the v1 patch matters at all depends on which executor vLLM will # select. 0.25 defaults this to "1" (RayExecutorV2), which has no @@ -692,3 +916,9 @@ def _apply_vllm_patches( _patch_vllm_shm_broadcast_bind_retry(patch_logger) _patch_vllm_radio_layerscale_loader(patch_logger) _patch_vllm_glm_decoder_sequence_parallel_moe(patch_logger) + if fp32_lm_head_enabled and not _patch_vllm_nemotron_h_fp32_lm_head(patch_logger): + raise RuntimeError( + "vllm_cfg.fp32_lm_head is enabled, but the Nemotron-H fp32 LM head " + "source patch could not be applied. Disable the flag or update the " + "patch anchors for this vLLM version." + ) diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index f503a89c3b5..32b10f737ab 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -46,6 +46,7 @@ VLLM_SPARSE_REFIT_TRANSPORTS, VllmConfig, resolve_vllm_video_config, + vllm_fp32_lm_head_enabled, ) from nemo_rl.models.generation.vllm.patches import _apply_vllm_patches from nemo_rl.models.generation.vllm.utils import ( @@ -437,9 +438,19 @@ def _init_config( # Store the Python executable being used by this worker self.py_executable = sys.executable + vllm_cfg = self.cfg["vllm_cfg"] + configured_env_vars = vllm_cfg.get("env_vars") + if configured_env_vars is not None: + extra_env_vars = [ + *(extra_env_vars or []), + *(str(k) for k in configured_env_vars), + ] + self._extra_env_vars = extra_env_vars + _apply_vllm_patches( self.py_executable, extra_env_vars=extra_env_vars, + fp32_lm_head=vllm_fp32_lm_head_enabled(vllm_cfg), ) # Skip model loading if we're not the model owner diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index afc14e8dcf0..b97ef64a83e 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -76,6 +76,14 @@ _HF_CONFIG_PATCHED = False _NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT = "expanded_sequence_v1" +_VLLM_FP32_LM_HEAD_ENV_VAR = "NRL_VLLM_FP32_LM_HEAD" + + +def _vllm_fp32_lm_head_enabled(vllm_cfg: Mapping[str, Any]) -> bool: + if vllm_cfg.get("fp32_lm_head"): + return True + env_vars = vllm_cfg.get("env_vars") + return env_vars is not None and str(env_vars.get(_VLLM_FP32_LM_HEAD_ENV_VAR)) == "1" def _patch_hf_config_double_instantiation(): @@ -526,6 +534,132 @@ def _resolve_iter_dir_from_root(path: str, not_found_msg: str) -> str: return os.path.join(path, iter_subdirs[-1]) +def validate_fp32_lm_head_config(config: PolicyConfig) -> None: + """Reject an fp32 LM head that is enabled on only one engine. + + ``megatron_cfg.fp32_lm_head`` and ``generation.vllm_cfg.fp32_lm_head`` must + agree: with bf16 heads on both sides the logits round to the same grid, so + enabling fp32 on one side alone makes train/token_mult_prob_error worse than + leaving the feature off. The fused linear+CE path bypasses ``output_layer`` + entirely, so the trainer head would silently stay bf16 there too. + + Only checked when generation uses the vLLM backend; SFT/DPO have no + generation engine to disagree with. + """ + megatron_cfg = config.get("megatron_cfg") or {} + trainer_fp32 = bool(megatron_cfg.get("fp32_lm_head")) + generation = config.get("generation") or {} + if generation.get("backend") != "vllm": + return + vllm_cfg = generation.get("vllm_cfg") or {} + env_vars = vllm_cfg.get("env_vars") + vllm_env_value = ( + None if env_vars is None else env_vars.get(_VLLM_FP32_LM_HEAD_ENV_VAR) + ) + vllm_fp32 = _vllm_fp32_lm_head_enabled(vllm_cfg) + if trainer_fp32 != vllm_fp32: + raise ValueError( + "fp32 LM head must be enabled on both engines or neither: " + f"megatron_cfg.fp32_lm_head={megatron_cfg.get('fp32_lm_head')!r} but " + f"generation.vllm_cfg.fp32_lm_head={vllm_cfg.get('fp32_lm_head')!r} " + f"({_VLLM_FP32_LM_HEAD_ENV_VAR}={vllm_env_value!r}). " + "A one-sided fp32 head " + "increases the generation/training logprob mismatch instead of " + "reducing it." + ) + if trainer_fp32 and megatron_cfg.get("use_fused_linear_logprobs"): + raise ValueError( + "megatron_cfg.fp32_lm_head has no effect with " + "use_fused_linear_logprobs=true (the fused linear+CE kernel bypasses " + "output_layer), which would leave the trainer in bf16 while vLLM " + "runs fp32. Disable one of them." + ) + + +def _resolve_output_layer_owner(chunk: Any) -> Any: + """Return the module that owns ``output_layer`` for a Megatron model chunk.""" + module = chunk + while hasattr(module, "module"): + module = module.module + for _ in range(4): + if getattr(module, "output_layer", None) is not None: + break + for attr in ("thinker", "llava_model", "language_model"): + inner = getattr(module, attr, None) + if inner is not None: + module = inner + break + else: + break + return module + + +def apply_fp32_lm_head(model_chunks: list, use_tf32: bool = False) -> None: + """Run the LM output-layer GEMM in fp32 (MiniMax-M1-style, arXiv:2506.13585). + + bf16 rounding of the logits (magnitude ~15-30, bf16 ulp 0.125-0.25) is the + dominant contributor to generation/training logprob mismatch + (train/token_mult_prob_error). Upcasting the head input and weight to fp32 + removes that rounding. The casts are part of the autograd graph, so + training gradients flow to the bf16 weight through the fp32 cast. + + Note: has no effect on the fused linear+CE path + (megatron_cfg.use_fused_linear_logprobs), which bypasses output_layer's + standalone forward. + + With ``use_tf32`` (megatron_cfg.fp32_lm_head: "tf32"), the fp32 head GEMM + runs with TF32 tensor cores. The inputs are exact bf16 values, so TF32's + 10-bit input rounding loses nothing; accumulation and output stay fp32. + Numerically equivalent to full fp32 here, at near-bf16 tensor-core + throughput. + """ + if not isinstance(model_chunks, (list, tuple)): + model_chunks = [model_chunks] + for chunk in model_chunks: + module = _resolve_output_layer_owner(chunk) + output_layer = getattr(module, "output_layer", None) + if output_layer is None: + # A post-process chunk should own the LM head; silently wrapping + # nothing leaves trainer/generator precision mismatched. + if getattr(module, "post_process", False) or getattr( + chunk, "post_process", False + ): + raise ValueError( + "fp32_lm_head is enabled but no output_layer was found on a " + f"post_process model chunk of type {type(module).__name__} " + f"(chunk type {type(chunk).__name__}). The trainer would run " + "the LM head in bf16 while generation runs fp32, which is " + "worse than disabling both." + ) + continue + original_forward = output_layer.forward + + def _fp32_forward( + input_, + *args, + weight=None, + _orig_forward=original_forward, + _layer=output_layer, + _tf32=use_tf32, + **kwargs, + ): + w = weight if weight is not None else _layer.weight + if not _tf32: + return _orig_forward(input_.float(), *args, weight=w.float(), **kwargs) + prev = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = True + try: + return _orig_forward(input_.float(), *args, weight=w.float(), **kwargs) + finally: + torch.backends.cuda.matmul.allow_tf32 = prev + + output_layer.forward = _fp32_forward + print( + "[fp32_lm_head] output layer will compute logits in fp32" + + (" (tf32 tensor cores)" if use_tf32 else "") + ) + + def validate_model_paths(config: PolicyConfig) -> tuple[str, str, bool]: """Validate and setup model paths. diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 300995b7697..bb9b1fd8d20 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -467,6 +467,20 @@ class MegatronConfig(TypedDict): # Number of tokens per chunk when computing fused linear logprobs. # Smaller values reduce peak memory further but may decrease throughput. fused_linear_logprobs_chunk_size: NotRequired[int] + # Compute the LM output-layer GEMM in fp32 instead of bf16. bf16 rounding of + # the logits is the dominant source of generation/training logprob mismatch + # (train/token_mult_prob_error). Set the matching generation.vllm_cfg + # fp32_lm_head flag for vLLM generation: applying this to only one engine + # makes the multiplicative error worse, since both otherwise round to the + # same grid. + # False - bf16 head (default) + # True - full fp32 head (~2x cost on the logprob pass) + # "tf32" - fp32 head on TF32 tensor cores; numerically identical here + # because the inputs are already exact bf16 values, at ~baseline + # speed. Prefer this. + # No effect when use_fused_linear_logprobs is set, which bypasses the + # output layer's standalone forward. + fp32_lm_head: NotRequired[bool | Literal["tf32"]] # When mtp_num_layers=0, Multi-Token Prediction is disabled. mtp_num_layers: NotRequired[int] # MTP loss weight added to the main next-token loss (0.0 disables the MTP loss contribution). diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 3b1dae699fe..ab3d30afadb 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -82,12 +82,14 @@ router_replay_enabled, ) from nemo_rl.models.megatron.setup import ( + apply_fp32_lm_head, build_inference_model, finalize_megatron_setup, handle_model_import, setup_distributed, setup_model_and_optimizer, setup_reference_model_state, + validate_fp32_lm_head_config, validate_and_set_config, validate_model_paths, ) @@ -584,6 +586,12 @@ def __init__( self.mcore_state = model_and_optimizer_state.state self.model = model_and_optimizer_state.model + validate_fp32_lm_head_config(self.cfg) + if self.cfg["megatron_cfg"].get("fp32_lm_head"): + apply_fp32_lm_head( + self.model, + use_tf32=self.cfg["megatron_cfg"]["fp32_lm_head"] == "tf32", + ) self.optimizer = model_and_optimizer_state.optimizer self.scheduler = model_and_optimizer_state.scheduler self.checkpointing_context = model_and_optimizer_state.checkpointing_context diff --git a/tests/unit/models/generation/test_vllm_patches.py b/tests/unit/models/generation/test_vllm_patches.py index a26a3d8ef11..362f3acde3e 100644 --- a/tests/unit/models/generation/test_vllm_patches.py +++ b/tests/unit/models/generation/test_vllm_patches.py @@ -33,8 +33,11 @@ import ast import logging import os +import sys +import types import pytest +import torch from nemo_rl.models.generation.vllm import patches from tests.unit.models.generation.vllm_patch_source_utils import ( @@ -50,6 +53,82 @@ _GLM_DSA_SOURCE = "model_executor/models/deepseek_v2.py" _GLM_DSA_PATCH_FN = "_patch_vllm_glm_decoder_sequence_parallel_moe" _GLM_DSA_MARKER = 'getattr(config, "model_type", None) != "glm_moe_dsa"' +_NEMOTRON_H_SOURCE = """import torch +from torch import nn + + +def maybe_prefix(prefix, name): + return f"{prefix}.{name}" + + +class LogitsProcessor: + def __init__(self, vocab_size): + self.vocab_size = vocab_size + + def __call__(self, lm_head, hidden_states): + return lm_head.quant_method.apply(lm_head, hidden_states) + + +class QuantMethod: + def __init__(self): + self.seen_dtypes = [] + + def apply(self, lm_head, hidden_states, bias=None): + self.seen_dtypes.append( + ( + hidden_states.dtype, + lm_head.weight.dtype, + None if bias is None else bias.dtype, + ) + ) + logits = hidden_states @ lm_head.weight.t() + if bias is not None: + logits = logits + bias + return logits + + +class ParallelLMHead(nn.Module): + def __init__( + self, + vocab_size, + hidden_size, + params_dtype=None, + quant_config=None, + prefix="", + ): + super().__init__() + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.params_dtype = params_dtype + self.quant_config = quant_config + self.prefix = prefix + self.weight = nn.Parameter( + torch.ones(vocab_size, hidden_size, dtype=torch.bfloat16), + requires_grad=False, + ) + self.bias = None + self.quant_method = QuantMethod() + + def forward(self, input_): + del input_ + raise RuntimeError("LMHead's weights should be used in the sampler.") + + +class NemotronHForCausalLM: + def __init__(self, config, prefix): + self.quant_config = object() + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + + def compute_logits(self, hidden_states): + logits = self.logits_processor(self.lm_head, hidden_states) + return logits +""" @pytest.fixture @@ -81,6 +160,15 @@ def patched_glm_dsa_source(tmp_path, monkeypatch): return copied +@pytest.fixture +def patched_nemotron_h_source(tmp_path, monkeypatch): + source = tmp_path / "nemotron_h.py" + source.write_text(_NEMOTRON_H_SOURCE) + monkeypatch.setattr(patches, "_get_vllm_file", lambda _relative: str(source)) + patches._patch_vllm_nemotron_h_fp32_lm_head(logging.getLogger(__name__)) + return source + + @pytest.mark.vllm def test_namespace_tool_patch_anchor_still_matches_installed_vllm( patched_tool_parser_source, @@ -212,6 +300,176 @@ def test_glm_decoder_sp_moe_patch_warns_on_unknown_source( assert "vLLM 0.25.1 source shape was not found" in caplog.text +@pytest.mark.parametrize("env_value", [None, "0", "1"]) +def test_nemotron_h_fp32_lm_head_patch_is_env_gated( + patched_nemotron_h_source, monkeypatch, env_value +): + if env_value is None: + monkeypatch.delenv("NRL_VLLM_FP32_LM_HEAD", raising=False) + else: + monkeypatch.setenv("NRL_VLLM_FP32_LM_HEAD", env_value) + + namespace = {} + source = patched_nemotron_h_source.read_text() + exec(compile(source, str(patched_nemotron_h_source), "exec"), namespace) + config = types.SimpleNamespace(vocab_size=16, hidden_size=8) + model = namespace["NemotronHForCausalLM"](config, "model") + hidden_states = torch.ones(2, 8, dtype=torch.bfloat16) + + logits = model.compute_logits(hidden_states) + + if env_value == "1": + assert model._nrl_fp32_lm_head is True + assert model.lm_head.params_dtype is None + assert model.lm_head.quant_config is model.quant_config + assert model.lm_head.weight.dtype is torch.bfloat16 + assert logits.dtype is torch.float32 + assert model.lm_head(hidden_states).dtype is torch.float32 + assert model.lm_head.quant_method.seen_dtypes == [] + else: + assert model._nrl_fp32_lm_head is False + assert model.lm_head.params_dtype is None + assert model.lm_head.quant_config is model.quant_config + assert model.lm_head.weight.dtype is torch.bfloat16 + assert logits.dtype is torch.bfloat16 + assert model.lm_head.quant_method.seen_dtypes == [ + (torch.bfloat16, torch.bfloat16, None) + ] + + assert "deepcopy" not in source + assert "params_dtype=torch.float32" not in source + assert "torch.matmul(" in source + ast.parse(source) + + +def test_nemotron_h_fp32_lm_head_patch_is_idempotent( + patched_nemotron_h_source, monkeypatch +): + before = patched_nemotron_h_source.read_text() + monkeypatch.setattr( + patches, "_get_vllm_file", lambda _relative: str(patched_nemotron_h_source) + ) + + patches._patch_vllm_nemotron_h_fp32_lm_head(logging.getLogger(__name__)) + + assert patched_nemotron_h_source.read_text() == before + + +def _install_fake_vllm_modules(monkeypatch): + vllm_module = types.ModuleType("vllm") + envs_module = types.ModuleType("vllm.envs") + envs_module.VLLM_USE_RAY_V2_EXECUTOR_BACKEND = True + logger_module = types.ModuleType("vllm.logger") + logger_module.init_logger = lambda name: logging.getLogger(name) + vllm_module.envs = envs_module + vllm_module.logger = logger_module + monkeypatch.setitem(sys.modules, "vllm", vllm_module) + monkeypatch.setitem(sys.modules, "vllm.envs", envs_module) + monkeypatch.setitem(sys.modules, "vllm.logger", logger_module) + + +def _stub_non_fp32_vllm_patches(monkeypatch, captured_extra_env_vars): + monkeypatch.setattr( + patches, + "_patch_vllm_init_workers_ray", + lambda _py, extra: captured_extra_env_vars.append(extra) or False, + ) + for patch_name in ( + "_patch_vllm_llama_eagle3_own_lm_head", + "_patch_vllm_tool_parser_namespace_tool", + "_patch_vllm_ray_executor_v2_tcpstore_port", + "_patch_vllm_shm_broadcast_bind_retry", + "_patch_vllm_radio_layerscale_loader", + "_patch_vllm_glm_decoder_sequence_parallel_moe", + ): + monkeypatch.setattr(patches, patch_name, lambda _logger: None) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_apply_vllm_patches_gates_nemotron_h_fp32_lm_head(monkeypatch, enabled): + _install_fake_vllm_modules(monkeypatch) + monkeypatch.delenv(patches.VLLM_FP32_LM_HEAD_ENV_VAR, raising=False) + captured_extra_env_vars = [] + fp32_patch_calls = [] + _stub_non_fp32_vllm_patches(monkeypatch, captured_extra_env_vars) + monkeypatch.setattr( + patches, + "_patch_vllm_nemotron_h_fp32_lm_head", + lambda _logger: fp32_patch_calls.append(True) or True, + ) + + patches._apply_vllm_patches("py", extra_env_vars=["USER_VAR"], fp32_lm_head=enabled) + + assert bool(fp32_patch_calls) is enabled + if enabled: + assert os.environ[patches.VLLM_FP32_LM_HEAD_ENV_VAR] == "1" + assert captured_extra_env_vars == [ + ["USER_VAR", patches.VLLM_FP32_LM_HEAD_ENV_VAR] + ] + else: + assert patches.VLLM_FP32_LM_HEAD_ENV_VAR not in os.environ + assert captured_extra_env_vars == [["USER_VAR"]] + + +def test_apply_vllm_patches_accepts_legacy_fp32_lm_head_env_toggle(monkeypatch): + _install_fake_vllm_modules(monkeypatch) + monkeypatch.setenv(patches.VLLM_FP32_LM_HEAD_ENV_VAR, "1") + captured_extra_env_vars = [] + fp32_patch_calls = [] + _stub_non_fp32_vllm_patches(monkeypatch, captured_extra_env_vars) + monkeypatch.setattr( + patches, + "_patch_vllm_nemotron_h_fp32_lm_head", + lambda _logger: fp32_patch_calls.append(True) or True, + ) + + patches._apply_vllm_patches("py") + + assert fp32_patch_calls == [True] + assert captured_extra_env_vars == [[patches.VLLM_FP32_LM_HEAD_ENV_VAR]] + + +def test_vllm_worker_threads_fp32_lm_head_cfg_into_source_patches(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_worker + + patch_calls = [] + monkeypatch.setattr( + vllm_worker, + "_apply_vllm_patches", + lambda py, *, extra_env_vars, fp32_lm_head: patch_calls.append( + { + "py": py, + "extra_env_vars": extra_env_vars, + "fp32_lm_head": fp32_lm_head, + } + ), + ) + + vllm_worker.BaseVllmGenerationWorker( + { + "model_name": "model", + "vllm_cfg": { + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "expert_parallel_size": 1, + "gpu_memory_utilization": 0.6, + "precision": "bfloat16", + "env_vars": {"USER_VAR": "value"}, + "fp32_lm_head": True, + }, + }, + extra_env_vars=["EXPLICIT_VAR"], + ) + + assert patch_calls == [ + { + "py": sys.executable, + "extra_env_vars": ["EXPLICIT_VAR", "USER_VAR"], + "fp32_lm_head": True, + } + ] + + @pytest.mark.parametrize( "existing,extra,expected", [ diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 22994ccd10d..daf74d99e75 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -3755,6 +3755,156 @@ def attach_fresh_draft(): ) +# --------------------------------------------------------------------------- +# apply_fp32_lm_head: output_layer resolution through multimodal wrappers +# --------------------------------------------------------------------------- + + +class _FakeOutputLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter( + torch.ones(4, 2, dtype=torch.bfloat16), requires_grad=False + ) + self.seen_dtypes: list[tuple[torch.dtype, torch.dtype]] = [] + + def forward(self, input_, *args, weight=None, **kwargs): + w = weight if weight is not None else self.weight + self.seen_dtypes.append((input_.dtype, w.dtype)) + return input_ @ w.t() + + +def _assert_fp32_wrapped(output_layer: _FakeOutputLayer) -> None: + out = output_layer.forward(torch.ones(3, 2, dtype=torch.bfloat16)) + assert out.dtype == torch.float32 + assert output_layer.seen_dtypes == [(torch.float32, torch.float32)] + + +def test_apply_fp32_lm_head_wraps_plain_last_stage_chunk(): + from nemo_rl.models.megatron.setup import apply_fp32_lm_head + + # Float16Module/DDP-style `.module` nesting around a GPTModel-like chunk. + layer = _FakeOutputLayer() + chunk = SimpleNamespace( + module=SimpleNamespace(output_layer=layer, post_process=True) + ) + apply_fp32_lm_head([chunk]) + _assert_fp32_wrapped(layer) + + +@pytest.mark.parametrize( + "build", + [ + # NemotronVLModel with an LLaVA wrapper: .llava_model.language_model + lambda layer: SimpleNamespace( + post_process=True, + llava_model=SimpleNamespace( + language_model=SimpleNamespace(output_layer=layer, post_process=True) + ), + ), + # NemotronVLModel without LLaVA: .language_model + lambda layer: SimpleNamespace( + post_process=True, + llava_model=None, + language_model=SimpleNamespace(output_layer=layer, post_process=True), + ), + # NemotronOmniModel: .thinker.language_model + lambda layer: SimpleNamespace( + post_process=True, + thinker=SimpleNamespace( + language_model=SimpleNamespace(output_layer=layer, post_process=True) + ), + ), + ], + ids=["vl_llava", "vl_language_model", "omni_thinker"], +) +def test_apply_fp32_lm_head_resolves_nested_language_model(build): + from nemo_rl.models.megatron.setup import apply_fp32_lm_head + + layer = _FakeOutputLayer() + chunk = SimpleNamespace(module=build(layer)) + apply_fp32_lm_head([chunk]) + _assert_fp32_wrapped(layer) + + +def test_apply_fp32_lm_head_raises_when_post_process_chunk_has_no_output_layer(): + from nemo_rl.models.megatron.setup import apply_fp32_lm_head + + # A post_process chunk with no reachable output_layer must not be mistaken + # for a non-last pipeline stage: that leaves the trainer in bf16 while + # generation runs fp32. + chunk = SimpleNamespace(module=SimpleNamespace(post_process=True)) + with pytest.raises(ValueError, match="no output_layer was found"): + apply_fp32_lm_head([chunk]) + + +def test_apply_fp32_lm_head_skips_non_last_pipeline_stage(): + from nemo_rl.models.megatron.setup import apply_fp32_lm_head + + chunk = SimpleNamespace(module=SimpleNamespace(post_process=False)) + apply_fp32_lm_head([chunk]) # no output_layer, not post_process: no-op + + +def _fp32_policy_cfg(trainer, vllm_cfg, fused=False, backend="vllm"): + megatron_cfg = {"fp32_lm_head": trainer} + if fused: + megatron_cfg["use_fused_linear_logprobs"] = True + return { + "megatron_cfg": megatron_cfg, + "generation": { + "backend": backend, + "vllm_cfg": vllm_cfg, + }, + } + + +@pytest.mark.parametrize( + "trainer, vllm_cfg", + [ + (False, {}), + ("tf32", {"fp32_lm_head": True}), + (True, {"fp32_lm_head": True}), + (True, {"env_vars": {"NRL_VLLM_FP32_LM_HEAD": "1"}}), + ], +) +def test_validate_fp32_lm_head_config_accepts_matched_engines(trainer, vllm_cfg): + from nemo_rl.models.megatron.setup import validate_fp32_lm_head_config + + validate_fp32_lm_head_config(_fp32_policy_cfg(trainer, vllm_cfg)) + + +@pytest.mark.parametrize( + "trainer, vllm_cfg", + [ + ("tf32", {}), # trainer fp32, vLLM bf16: the production misconfiguration + (False, {"fp32_lm_head": True}), # vLLM fp32, trainer bf16 + (False, {"env_vars": {"NRL_VLLM_FP32_LM_HEAD": "1"}}), + ], +) +def test_validate_fp32_lm_head_config_rejects_one_sided(trainer, vllm_cfg): + from nemo_rl.models.megatron.setup import validate_fp32_lm_head_config + + with pytest.raises(ValueError, match="both engines or neither"): + validate_fp32_lm_head_config(_fp32_policy_cfg(trainer, vllm_cfg)) + + +def test_validate_fp32_lm_head_config_rejects_fused_logprobs(): + from nemo_rl.models.megatron.setup import validate_fp32_lm_head_config + + with pytest.raises(ValueError, match="use_fused_linear_logprobs"): + validate_fp32_lm_head_config( + _fp32_policy_cfg("tf32", {"fp32_lm_head": True}, fused=True) + ) + + +def test_validate_fp32_lm_head_config_ignores_non_vllm_generation(): + from nemo_rl.models.megatron.setup import validate_fp32_lm_head_config + + # SFT/DPO-style configs have no vLLM engine to disagree with. + validate_fp32_lm_head_config({"megatron_cfg": {"fp32_lm_head": "tf32"}}) + validate_fp32_lm_head_config(_fp32_policy_cfg("tf32", {}, backend="megatron")) + + @pytest.mark.mcore class TestForceSyncOptimizerFp32FromModel: """Tests for _force_sync_optimizer_fp32_from_model. diff --git a/tests/unit/reference_configs/distillation_math.yaml b/tests/unit/reference_configs/distillation_math.yaml index 73455bb41ea..0bb8e50f9e3 100644 --- a/tests/unit/reference_configs/distillation_math.yaml +++ b/tests/unit/reference_configs/distillation_math.yaml @@ -207,6 +207,7 @@ policy: &POLICY_BASE precision: ${...precision} kv_cache_dtype: "auto" logprobs_mode: processed_logprobs + fp32_lm_head: false # false: use NeMo-RL's legacy refit loader; true: opt into vLLM reload_weights. refit_with_reload_api: false tensor_parallel_size: 1 diff --git a/tests/unit/reference_configs/eval.yaml b/tests/unit/reference_configs/eval.yaml index 966336428c2..bd0f3a516fc 100644 --- a/tests/unit/reference_configs/eval.yaml +++ b/tests/unit/reference_configs/eval.yaml @@ -22,6 +22,7 @@ generation: vllm_cfg: async_engine: false precision: "bfloat16" + fp32_lm_head: false tensor_parallel_size: 1 pipeline_parallel_size: 1 expert_parallel_size: 1 diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 3c378dc38f7..5ddf1bf66e8 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -408,6 +408,7 @@ policy: precision: ${policy.precision} kv_cache_dtype: "auto" logprobs_mode: processed_logprobs + fp32_lm_head: false # false: use NeMo-RL's legacy refit loader; true: opt into vLLM reload_weights. refit_with_reload_api: false tensor_parallel_size: 1 diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 9bb901085a5..06df5198eb9 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -261,6 +261,7 @@ policy: precision: ${policy.precision} kv_cache_dtype: "auto" logprobs_mode: processed_logprobs + fp32_lm_head: false # false: use NeMo-RL's legacy refit loader; true: opt into vLLM reload_weights. refit_with_reload_api: false tensor_parallel_size: 1 From a4d4cc9297eaf616ee49ada1e9a608a116683238 Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 09:55:54 -0700 Subject: [PATCH 02/18] fix: address fp32 lm head review comments Signed-off-by: Guyue Huang --- nemo_rl/models/megatron/setup.py | 26 ++--- nemo_rl/models/policy/__init__.py | 8 +- .../policy/workers/megatron_policy_worker.py | 2 +- .../models/generation/test_vllm_patches.py | 102 ++++++++++++++++++ .../models/megatron/test_megatron_setup.py | 26 +++++ 5 files changed, 144 insertions(+), 20 deletions(-) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index dec436880a3..4722cac4461 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -75,18 +75,14 @@ from transformers import PreTrainedTokenizerBase from nemo_rl.distributed.model_utils import patch_gpt_model_forward_for_linear_ce_fusion +from nemo_rl.models.generation.vllm.config import ( + VLLM_FP32_LM_HEAD_ENV_VAR, + vllm_fp32_lm_head_enabled, +) _HF_CONFIG_PATCHED = False _NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT = "expanded_sequence_v1" -_VLLM_FP32_LM_HEAD_ENV_VAR = "NRL_VLLM_FP32_LM_HEAD" - - -def _vllm_fp32_lm_head_enabled(vllm_cfg: Mapping[str, Any]) -> bool: - if vllm_cfg.get("fp32_lm_head"): - return True - env_vars = vllm_cfg.get("env_vars") - return env_vars is not None and str(env_vars.get(_VLLM_FP32_LM_HEAD_ENV_VAR)) == "1" def _patch_hf_config_double_instantiation(): @@ -626,15 +622,15 @@ def validate_fp32_lm_head_config(config: PolicyConfig) -> None: vllm_cfg = generation.get("vllm_cfg") or {} env_vars = vllm_cfg.get("env_vars") vllm_env_value = ( - None if env_vars is None else env_vars.get(_VLLM_FP32_LM_HEAD_ENV_VAR) + None if env_vars is None else env_vars.get(VLLM_FP32_LM_HEAD_ENV_VAR) ) - vllm_fp32 = _vllm_fp32_lm_head_enabled(vllm_cfg) + vllm_fp32 = vllm_fp32_lm_head_enabled(vllm_cfg) if trainer_fp32 != vllm_fp32: raise ValueError( "fp32 LM head must be enabled on both engines or neither: " f"megatron_cfg.fp32_lm_head={megatron_cfg.get('fp32_lm_head')!r} but " f"generation.vllm_cfg.fp32_lm_head={vllm_cfg.get('fp32_lm_head')!r} " - f"({_VLLM_FP32_LM_HEAD_ENV_VAR}={vllm_env_value!r}). " + f"({VLLM_FP32_LM_HEAD_ENV_VAR}={vllm_env_value!r}). " "A one-sided fp32 head " "increases the generation/training logprob mismatch instead of " "reducing it." @@ -680,10 +676,10 @@ def apply_fp32_lm_head(model_chunks: list, use_tf32: bool = False) -> None: standalone forward. With ``use_tf32`` (megatron_cfg.fp32_lm_head: "tf32"), the fp32 head GEMM - runs with TF32 tensor cores. The inputs are exact bf16 values, so TF32's - 10-bit input rounding loses nothing; accumulation and output stay fp32. - Numerically equivalent to full fp32 here, at near-bf16 tensor-core - throughput. + allows CUDA matmul to use TF32 tensor cores where available. The operands + originate as bf16 values, so TF32 preserves the input values while keeping + fp32 accumulation/output, but exact throughput and tolerance should be + validated on the target workload. """ if not isinstance(model_chunks, (list, tuple)): model_chunks = [model_chunks] diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index a03f7e70ae7..24610f8ce03 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -507,10 +507,10 @@ class MegatronConfig(TypedDict): # makes the multiplicative error worse, since both otherwise round to the # same grid. # False - bf16 head (default) - # True - full fp32 head (~2x cost on the logprob pass) - # "tf32" - fp32 head on TF32 tensor cores; numerically identical here - # because the inputs are already exact bf16 values, at ~baseline - # speed. Prefer this. + # True - full fp32 head with regular fp32 matmul + # "tf32" - fp32 head with CUDA TF32 matmul enabled where available; + # generally faster than regular fp32 on supported GPUs, but + # validate throughput and tolerance for the target workload # No effect when use_fused_linear_logprobs is set, which bypasses the # output layer's standalone forward. fp32_lm_head: NotRequired[bool | Literal["tf32"]] diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 06775122391..88c0894b4da 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -92,8 +92,8 @@ setup_distributed, setup_model_and_optimizer, setup_reference_model_state, - validate_fp32_lm_head_config, validate_and_set_config, + validate_fp32_lm_head_config, validate_megatron_config, validate_model_paths, ) diff --git a/tests/unit/models/generation/test_vllm_patches.py b/tests/unit/models/generation/test_vllm_patches.py index 362f3acde3e..74723ab6c8b 100644 --- a/tests/unit/models/generation/test_vllm_patches.py +++ b/tests/unit/models/generation/test_vllm_patches.py @@ -40,6 +40,10 @@ import torch from nemo_rl.models.generation.vllm import patches +from nemo_rl.models.generation.vllm.config import ( + VLLM_FP32_LM_HEAD_ENV_VAR, + vllm_fp32_lm_head_enabled, +) from tests.unit.models.generation.vllm_patch_source_utils import ( write_unpatched_copy, ) @@ -129,6 +133,55 @@ def compute_logits(self, hidden_states): logits = self.logits_processor(self.lm_head, hidden_states) return logits """ +_NEMOTRON_H_LEGACY_FP32_HEAD_COMPUTE = f""" def compute_logits(self, hidden_states): + import os as _os + + if _os.environ.get("{VLLM_FP32_LM_HEAD_ENV_VAR}", "0") == "1": + # NeMo-RL patch: fp32 LM head (MiniMax-M1-style). bf16 rounding of + # the logits is the dominant gen/train logprob mismatch source. + _fp32_head = getattr(self, "_nrl_lm_head_fp32", None) + if _fp32_head is None and not torch.cuda.is_current_stream_capturing(): + # Skipped under graph capture: an allocation there lives in the + # graph's memory pool and is not valid for later eager replays. + # Capture output is discarded anyway, so bf16 is fine for it. + import copy as _copy + + _fp32_head = _copy.deepcopy(self.lm_head).float() + # object.__setattr__ bypasses nn.Module.__setattr__: registering + # this as a submodule would add a vocab-sized parameter to + # named_parameters(), which the refit weight mapping is built from. + object.__setattr__(self, "_nrl_lm_head_fp32", _fp32_head) + self._nrl_lm_head_fp32_dirty = False + print( + "[fp32_lm_head] built fp32 head in forward shape=%s" + % (tuple(_fp32_head.weight.shape),), + flush=True, + ) + elif _fp32_head is not None and getattr( + self, "_nrl_lm_head_fp32_dirty", False + ): + # Refreshed in place: replacing the module would leave any + # captured CUDA graph pointing at the old storage. + _fp32_head.weight.data.copy_(self.lm_head.weight) + if getattr(_fp32_head, "bias", None) is not None: + _fp32_head.bias.data.copy_(self.lm_head.bias) + self._nrl_lm_head_fp32_dirty = False + print("[fp32_lm_head] refreshed cached head in forward", flush=True) + if _fp32_head is not None: + return self.logits_processor(_fp32_head, hidden_states.float()) + logits = self.logits_processor(self.lm_head, hidden_states) + return logits +""" +_NEMOTRON_H_LEGACY_FP32_HEAD_SOURCE = _NEMOTRON_H_SOURCE.replace( + "import torch\nfrom torch import nn", + "import os\n\nimport torch\nfrom torch import nn", +).replace( + """ def compute_logits(self, hidden_states): + logits = self.logits_processor(self.lm_head, hidden_states) + return logits +""", + _NEMOTRON_H_LEGACY_FP32_HEAD_COMPUTE, +) @pytest.fixture @@ -169,6 +222,15 @@ def patched_nemotron_h_source(tmp_path, monkeypatch): return source +@pytest.fixture +def patched_legacy_nemotron_h_source(tmp_path, monkeypatch): + source = tmp_path / "nemotron_h.py" + source.write_text(_NEMOTRON_H_LEGACY_FP32_HEAD_SOURCE) + monkeypatch.setattr(patches, "_get_vllm_file", lambda _relative: str(source)) + patches._patch_vllm_nemotron_h_fp32_lm_head(logging.getLogger(__name__)) + return source + + @pytest.mark.vllm def test_namespace_tool_patch_anchor_still_matches_installed_vllm( patched_tool_parser_source, @@ -300,6 +362,20 @@ def test_glm_decoder_sp_moe_patch_warns_on_unknown_source( assert "vLLM 0.25.1 source shape was not found" in caplog.text +@pytest.mark.parametrize( + ("vllm_cfg", "expected"), + [ + ({}, False), + ({"fp32_lm_head": False}, False), + ({"fp32_lm_head": True}, True), + ({"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}, True), + ({"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "0"}}, False), + ], +) +def test_vllm_fp32_lm_head_enabled(vllm_cfg, expected): + assert vllm_fp32_lm_head_enabled(vllm_cfg) is expected + + @pytest.mark.parametrize("env_value", [None, "0", "1"]) def test_nemotron_h_fp32_lm_head_patch_is_env_gated( patched_nemotron_h_source, monkeypatch, env_value @@ -342,6 +418,32 @@ def test_nemotron_h_fp32_lm_head_patch_is_env_gated( ast.parse(source) +def test_nemotron_h_fp32_lm_head_patch_migrates_legacy_cached_head_source( + patched_legacy_nemotron_h_source, monkeypatch +): + monkeypatch.setenv(VLLM_FP32_LM_HEAD_ENV_VAR, "1") + + source = patched_legacy_nemotron_h_source.read_text() + namespace = {} + exec(compile(source, str(patched_legacy_nemotron_h_source), "exec"), namespace) + config = types.SimpleNamespace(vocab_size=16, hidden_size=8) + model = namespace["NemotronHForCausalLM"](config, "model") + hidden_states = torch.ones(2, 8, dtype=torch.bfloat16) + + logits = model.compute_logits(hidden_states) + + assert model._nrl_fp32_lm_head is True + assert model.lm_head.params_dtype is None + assert model.lm_head.quant_config is model.quant_config + assert model.lm_head.weight.dtype is torch.bfloat16 + assert logits.dtype is torch.float32 + assert "_nrl_lm_head_fp32" not in source + assert "deepcopy" not in source + assert "params_dtype=torch.float32" not in source + assert "torch.matmul(" in source + ast.parse(source) + + def test_nemotron_h_fp32_lm_head_patch_is_idempotent( patched_nemotron_h_source, monkeypatch ): diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 73624260df8..e5b582ea8b0 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -4009,6 +4009,17 @@ def test_apply_fp32_lm_head_wraps_plain_last_stage_chunk(): _assert_fp32_wrapped(layer) +def test_apply_fp32_lm_head_tf32_path_produces_fp32_output(): + from nemo_rl.models.megatron.setup import apply_fp32_lm_head + + layer = _FakeOutputLayer() + chunk = SimpleNamespace( + module=SimpleNamespace(output_layer=layer, post_process=True) + ) + apply_fp32_lm_head([chunk], use_tf32=True) + _assert_fp32_wrapped(layer) + + @pytest.mark.parametrize( "build", [ @@ -4122,6 +4133,21 @@ def test_validate_fp32_lm_head_config_ignores_non_vllm_generation(): validate_fp32_lm_head_config(_fp32_policy_cfg("tf32", {}, backend="megatron")) +def test_validate_fp32_lm_head_config_handles_missing_megatron_cfg(): + from nemo_rl.models.megatron.setup import validate_fp32_lm_head_config + + validate_fp32_lm_head_config({"generation": {"backend": "vllm", "vllm_cfg": {}}}) + with pytest.raises(ValueError, match="both engines or neither"): + validate_fp32_lm_head_config( + { + "generation": { + "backend": "vllm", + "vllm_cfg": {"fp32_lm_head": True}, + }, + } + ) + + @pytest.mark.mcore class TestForceSyncOptimizerFp32FromModel: """Tests for _force_sync_optimizer_fp32_from_model. From 20779d8da649495f0d158d9747a788fef07f9687 Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 10:22:41 -0700 Subject: [PATCH 03/18] fix: narrow vllm fp32 env propagation Signed-off-by: Guyue Huang --- nemo_rl/models/generation/vllm/vllm_worker.py | 8 -------- .../unit/models/generation/test_vllm_patches.py | 16 ++++++++++++---- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index a37616cce64..a35ce53ec99 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -440,14 +440,6 @@ def _init_config( self.py_executable = sys.executable vllm_cfg = self.cfg["vllm_cfg"] - configured_env_vars = vllm_cfg.get("env_vars") - if configured_env_vars is not None: - extra_env_vars = [ - *(extra_env_vars or []), - *(str(k) for k in configured_env_vars), - ] - self._extra_env_vars = extra_env_vars - _apply_vllm_patches( self.py_executable, extra_env_vars=extra_env_vars, diff --git a/tests/unit/models/generation/test_vllm_patches.py b/tests/unit/models/generation/test_vllm_patches.py index 74723ab6c8b..33daab2630e 100644 --- a/tests/unit/models/generation/test_vllm_patches.py +++ b/tests/unit/models/generation/test_vllm_patches.py @@ -531,7 +531,16 @@ def test_apply_vllm_patches_accepts_legacy_fp32_lm_head_env_toggle(monkeypatch): assert captured_extra_env_vars == [[patches.VLLM_FP32_LM_HEAD_ENV_VAR]] -def test_vllm_worker_threads_fp32_lm_head_cfg_into_source_patches(monkeypatch): +@pytest.mark.parametrize( + "vllm_cfg_overrides", + [ + {"env_vars": {"USER_VAR": "value"}, "fp32_lm_head": True}, + {"env_vars": {"USER_VAR": "value", VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}, + ], +) +def test_vllm_worker_threads_fp32_lm_head_cfg_into_source_patches( + monkeypatch, vllm_cfg_overrides +): from nemo_rl.models.generation.vllm import vllm_worker patch_calls = [] @@ -556,8 +565,7 @@ def test_vllm_worker_threads_fp32_lm_head_cfg_into_source_patches(monkeypatch): "expert_parallel_size": 1, "gpu_memory_utilization": 0.6, "precision": "bfloat16", - "env_vars": {"USER_VAR": "value"}, - "fp32_lm_head": True, + **vllm_cfg_overrides, }, }, extra_env_vars=["EXPLICIT_VAR"], @@ -566,7 +574,7 @@ def test_vllm_worker_threads_fp32_lm_head_cfg_into_source_patches(monkeypatch): assert patch_calls == [ { "py": sys.executable, - "extra_env_vars": ["EXPLICIT_VAR", "USER_VAR"], + "extra_env_vars": ["EXPLICIT_VAR"], "fp32_lm_head": True, } ] From 3153f1e27271103b67ce675f877b0d0f0f9d6f9d Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 10:27:02 -0700 Subject: [PATCH 04/18] fix: validate vllm fp32 lm head config earlier Signed-off-by: Guyue Huang --- .../models/generation/vllm/vllm_generation.py | 42 +++++++++++ nemo_rl/models/megatron/setup.py | 46 ------------ .../policy/workers/megatron_policy_worker.py | 2 - .../models/generation/test_vllm_generation.py | 66 ++++++++++++++++ .../models/megatron/test_megatron_setup.py | 75 ------------------- 5 files changed, 108 insertions(+), 123 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 9b089a2b960..3796ba36b1f 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -46,7 +46,9 @@ ) from nemo_rl.models.generation.vllm.config import ( REFITTABLE_FP8_KV_CACHE_DTYPES, + VLLM_FP32_LM_HEAD_ENV_VAR, VllmConfig, + vllm_fp32_lm_head_enabled, ) from nemo_rl.models.generation.vllm.utils import ( aggregate_spec_decode_counters, @@ -74,6 +76,45 @@ logger = logging.getLogger(__name__) +def _validate_fp32_lm_head_config(master_config: "MasterConfig") -> None: + """Reject vLLM fp32 LM-head configs that the trainer cannot match.""" + policy_config = master_config.policy + generation_config = cast(VllmConfig, policy_config["generation"]) + vllm_cfg = generation_config.get("vllm_cfg") + env_vars = None if vllm_cfg is None else vllm_cfg.get("env_vars") + vllm_env_value = ( + None if env_vars is None else env_vars.get(VLLM_FP32_LM_HEAD_ENV_VAR) + ) + vllm_fp32 = vllm_cfg is not None and vllm_fp32_lm_head_enabled(vllm_cfg) + + megatron_cfg = policy_config.get("megatron_cfg") + megatron_fp32_value = ( + None if megatron_cfg is None else megatron_cfg.get("fp32_lm_head") + ) + trainer_fp32 = bool(megatron_fp32_value) + if trainer_fp32 != vllm_fp32: + raise ValueError( + "fp32 LM head must be enabled on both engines or neither: " + f"policy.megatron_cfg.fp32_lm_head={megatron_fp32_value!r} but " + f"policy.generation.vllm_cfg.fp32_lm_head=" + f"{None if vllm_cfg is None else vllm_cfg.get('fp32_lm_head')!r} " + f"({VLLM_FP32_LM_HEAD_ENV_VAR}={vllm_env_value!r}). " + "A one-sided fp32 head increases the generation/training logprob " + "mismatch instead of reducing it." + ) + if ( + trainer_fp32 + and megatron_cfg is not None + and megatron_cfg.get("use_fused_linear_logprobs") + ): + raise ValueError( + "policy.megatron_cfg.fp32_lm_head has no effect with " + "use_fused_linear_logprobs=true (the fused linear+CE kernel bypasses " + "output_layer), which would leave the trainer in bf16 while vLLM " + "runs fp32. Disable one of them." + ) + + def _record_vllm_generation_metrics( model_name: str | None, data: BatchedDataDict, @@ -114,6 +155,7 @@ def validate_settings(cls, master_config: "MasterConfig") -> None: """Reject pure-config vLLM settings the SC entrypoint cannot honor.""" generation_config = cast(VllmConfig, master_config.policy["generation"]) assert_reload_refit_config_supported(generation_config) + _validate_fp32_lm_head_config(master_config) @staticmethod def init_cluster_placement_groups( diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 4722cac4461..638a3daf21a 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -75,10 +75,6 @@ from transformers import PreTrainedTokenizerBase from nemo_rl.distributed.model_utils import patch_gpt_model_forward_for_linear_ce_fusion -from nemo_rl.models.generation.vllm.config import ( - VLLM_FP32_LM_HEAD_ENV_VAR, - vllm_fp32_lm_head_enabled, -) _HF_CONFIG_PATCHED = False @@ -602,48 +598,6 @@ def _resolve_iter_dir_from_root( return os.path.join(path, iter_subdirs[-1]) -def validate_fp32_lm_head_config(config: PolicyConfig) -> None: - """Reject an fp32 LM head that is enabled on only one engine. - - ``megatron_cfg.fp32_lm_head`` and ``generation.vllm_cfg.fp32_lm_head`` must - agree: with bf16 heads on both sides the logits round to the same grid, so - enabling fp32 on one side alone makes train/token_mult_prob_error worse than - leaving the feature off. The fused linear+CE path bypasses ``output_layer`` - entirely, so the trainer head would silently stay bf16 there too. - - Only checked when generation uses the vLLM backend; SFT/DPO have no - generation engine to disagree with. - """ - megatron_cfg = config.get("megatron_cfg") or {} - trainer_fp32 = bool(megatron_cfg.get("fp32_lm_head")) - generation = config.get("generation") or {} - if generation.get("backend") != "vllm": - return - vllm_cfg = generation.get("vllm_cfg") or {} - env_vars = vllm_cfg.get("env_vars") - vllm_env_value = ( - None if env_vars is None else env_vars.get(VLLM_FP32_LM_HEAD_ENV_VAR) - ) - vllm_fp32 = vllm_fp32_lm_head_enabled(vllm_cfg) - if trainer_fp32 != vllm_fp32: - raise ValueError( - "fp32 LM head must be enabled on both engines or neither: " - f"megatron_cfg.fp32_lm_head={megatron_cfg.get('fp32_lm_head')!r} but " - f"generation.vllm_cfg.fp32_lm_head={vllm_cfg.get('fp32_lm_head')!r} " - f"({VLLM_FP32_LM_HEAD_ENV_VAR}={vllm_env_value!r}). " - "A one-sided fp32 head " - "increases the generation/training logprob mismatch instead of " - "reducing it." - ) - if trainer_fp32 and megatron_cfg.get("use_fused_linear_logprobs"): - raise ValueError( - "megatron_cfg.fp32_lm_head has no effect with " - "use_fused_linear_logprobs=true (the fused linear+CE kernel bypasses " - "output_layer), which would leave the trainer in bf16 while vLLM " - "runs fp32. Disable one of them." - ) - - def _resolve_output_layer_owner(chunk: Any) -> Any: """Return the module that owns ``output_layer`` for a Megatron model chunk.""" module = chunk diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 88c0894b4da..7f773a10d9c 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -93,7 +93,6 @@ setup_model_and_optimizer, setup_reference_model_state, validate_and_set_config, - validate_fp32_lm_head_config, validate_megatron_config, validate_model_paths, ) @@ -718,7 +717,6 @@ def __init__( self.mcore_state = model_and_optimizer_state.state self.model = model_and_optimizer_state.model - validate_fp32_lm_head_config(self.cfg) if self.cfg["megatron_cfg"].get("fp32_lm_head"): apply_fp32_lm_head( self.model, diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 1b2ba597aac..d3b6b228f56 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -39,6 +39,7 @@ ) from nemo_rl.models.generation.openai_server_utils import replace_prefix_tokens from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration +from nemo_rl.models.generation.vllm.config import VLLM_FP32_LM_HEAD_ENV_VAR from nemo_rl.models.generation.vllm.vllm_worker import ( VllmGenerationWorkerImpl, _context_capped_max_new_tokens, @@ -1480,6 +1481,71 @@ def test_vllm_validate_settings_accepts_missing_vllm_cfg_as_refit_disabled(): VllmGeneration.validate_settings(master_config) +def _fp32_master_config( + trainer, + vllm_cfg, + *, + fused=False, + include_megatron_cfg=True, +): + generation = deepcopy(basic_vllm_test_config) + generation["vllm_cfg"].update(vllm_cfg) + policy = {"generation": generation} + if include_megatron_cfg: + megatron_cfg = {"fp32_lm_head": trainer} + if fused: + megatron_cfg["use_fused_linear_logprobs"] = True + policy["megatron_cfg"] = megatron_cfg + return types.SimpleNamespace(policy=policy) + + +@pytest.mark.parametrize( + "trainer, vllm_cfg", + [ + (False, {}), + ("tf32", {"fp32_lm_head": True}), + (True, {"fp32_lm_head": True}), + (True, {"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}), + ], +) +def test_vllm_validate_settings_accepts_matched_fp32_lm_head(trainer, vllm_cfg): + VllmGeneration.validate_settings(_fp32_master_config(trainer, vllm_cfg)) + + +@pytest.mark.parametrize( + "trainer, vllm_cfg", + [ + ("tf32", {}), # trainer fp32, vLLM bf16: the production misconfiguration + (False, {"fp32_lm_head": True}), # vLLM fp32, trainer bf16 + (False, {"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}), + ], +) +def test_vllm_validate_settings_rejects_one_sided_fp32_lm_head(trainer, vllm_cfg): + with pytest.raises(ValueError, match="both engines or neither"): + VllmGeneration.validate_settings(_fp32_master_config(trainer, vllm_cfg)) + + +def test_vllm_validate_settings_rejects_fp32_lm_head_with_fused_logprobs(): + with pytest.raises(ValueError, match="use_fused_linear_logprobs"): + VllmGeneration.validate_settings( + _fp32_master_config("tf32", {"fp32_lm_head": True}, fused=True) + ) + + +def test_vllm_validate_settings_handles_missing_megatron_cfg_for_fp32_lm_head(): + VllmGeneration.validate_settings( + _fp32_master_config(False, {}, include_megatron_cfg=False) + ) + with pytest.raises(ValueError, match="both engines or neither"): + VllmGeneration.validate_settings( + _fp32_master_config( + False, + {"fp32_lm_head": True}, + include_megatron_cfg=False, + ) + ) + + def test_vllm_policy_generation(policy, test_input_data, tokenizer): """Test vLLM policy generation capabilities.""" # Test generation diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index e5b582ea8b0..98445827385 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -4073,81 +4073,6 @@ def test_apply_fp32_lm_head_skips_non_last_pipeline_stage(): apply_fp32_lm_head([chunk]) # no output_layer, not post_process: no-op -def _fp32_policy_cfg(trainer, vllm_cfg, fused=False, backend="vllm"): - megatron_cfg = {"fp32_lm_head": trainer} - if fused: - megatron_cfg["use_fused_linear_logprobs"] = True - return { - "megatron_cfg": megatron_cfg, - "generation": { - "backend": backend, - "vllm_cfg": vllm_cfg, - }, - } - - -@pytest.mark.parametrize( - "trainer, vllm_cfg", - [ - (False, {}), - ("tf32", {"fp32_lm_head": True}), - (True, {"fp32_lm_head": True}), - (True, {"env_vars": {"NRL_VLLM_FP32_LM_HEAD": "1"}}), - ], -) -def test_validate_fp32_lm_head_config_accepts_matched_engines(trainer, vllm_cfg): - from nemo_rl.models.megatron.setup import validate_fp32_lm_head_config - - validate_fp32_lm_head_config(_fp32_policy_cfg(trainer, vllm_cfg)) - - -@pytest.mark.parametrize( - "trainer, vllm_cfg", - [ - ("tf32", {}), # trainer fp32, vLLM bf16: the production misconfiguration - (False, {"fp32_lm_head": True}), # vLLM fp32, trainer bf16 - (False, {"env_vars": {"NRL_VLLM_FP32_LM_HEAD": "1"}}), - ], -) -def test_validate_fp32_lm_head_config_rejects_one_sided(trainer, vllm_cfg): - from nemo_rl.models.megatron.setup import validate_fp32_lm_head_config - - with pytest.raises(ValueError, match="both engines or neither"): - validate_fp32_lm_head_config(_fp32_policy_cfg(trainer, vllm_cfg)) - - -def test_validate_fp32_lm_head_config_rejects_fused_logprobs(): - from nemo_rl.models.megatron.setup import validate_fp32_lm_head_config - - with pytest.raises(ValueError, match="use_fused_linear_logprobs"): - validate_fp32_lm_head_config( - _fp32_policy_cfg("tf32", {"fp32_lm_head": True}, fused=True) - ) - - -def test_validate_fp32_lm_head_config_ignores_non_vllm_generation(): - from nemo_rl.models.megatron.setup import validate_fp32_lm_head_config - - # SFT/DPO-style configs have no vLLM engine to disagree with. - validate_fp32_lm_head_config({"megatron_cfg": {"fp32_lm_head": "tf32"}}) - validate_fp32_lm_head_config(_fp32_policy_cfg("tf32", {}, backend="megatron")) - - -def test_validate_fp32_lm_head_config_handles_missing_megatron_cfg(): - from nemo_rl.models.megatron.setup import validate_fp32_lm_head_config - - validate_fp32_lm_head_config({"generation": {"backend": "vllm", "vllm_cfg": {}}}) - with pytest.raises(ValueError, match="both engines or neither"): - validate_fp32_lm_head_config( - { - "generation": { - "backend": "vllm", - "vllm_cfg": {"fp32_lm_head": True}, - }, - } - ) - - @pytest.mark.mcore class TestForceSyncOptimizerFp32FromModel: """Tests for _force_sync_optimizer_fp32_from_model. From 75f312f43dee99ab7fa6c2b14f172aaeaade0da8 Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 10:29:07 -0700 Subject: [PATCH 05/18] fix: make megatron fp32 lm head patch idempotent Signed-off-by: Guyue Huang --- nemo_rl/models/megatron/setup.py | 6 ++++++ .../models/megatron/test_megatron_setup.py | 19 +++++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 638a3daf21a..67133384a46 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -79,6 +79,8 @@ _HF_CONFIG_PATCHED = False _NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT = "expanded_sequence_v1" +_FP32_LM_HEAD_PATCHED_ATTR = "_nrl_fp32_lm_head_patched" +_FP32_LM_HEAD_USE_TF32_ATTR = "_nrl_fp32_lm_head_use_tf32" def _patch_hf_config_double_instantiation(): @@ -654,6 +656,8 @@ def apply_fp32_lm_head(model_chunks: list, use_tf32: bool = False) -> None: "worse than disabling both." ) continue + if getattr(output_layer.forward, _FP32_LM_HEAD_PATCHED_ATTR, False): + continue original_forward = output_layer.forward def _fp32_forward( @@ -675,6 +679,8 @@ def _fp32_forward( finally: torch.backends.cuda.matmul.allow_tf32 = prev + setattr(_fp32_forward, _FP32_LM_HEAD_PATCHED_ATTR, True) + setattr(_fp32_forward, _FP32_LM_HEAD_USE_TF32_ATTR, use_tf32) output_layer.forward = _fp32_forward print( "[fp32_lm_head] output layer will compute logits in fp32" diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 98445827385..d1368fc27e1 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -3992,6 +3992,7 @@ def forward(self, input_, *args, weight=None, **kwargs): def _assert_fp32_wrapped(output_layer: _FakeOutputLayer) -> None: + assert getattr(output_layer.forward, "_nrl_fp32_lm_head_patched") is True out = output_layer.forward(torch.ones(3, 2, dtype=torch.bfloat16)) assert out.dtype == torch.float32 assert output_layer.seen_dtypes == [(torch.float32, torch.float32)] @@ -4017,6 +4018,24 @@ def test_apply_fp32_lm_head_tf32_path_produces_fp32_output(): module=SimpleNamespace(output_layer=layer, post_process=True) ) apply_fp32_lm_head([chunk], use_tf32=True) + assert getattr(layer.forward, "_nrl_fp32_lm_head_use_tf32") is True + _assert_fp32_wrapped(layer) + + +def test_apply_fp32_lm_head_is_idempotent(): + from nemo_rl.models.megatron.setup import apply_fp32_lm_head + + layer = _FakeOutputLayer() + chunk = SimpleNamespace( + module=SimpleNamespace(output_layer=layer, post_process=True) + ) + apply_fp32_lm_head([chunk]) + first_forward = layer.forward + + apply_fp32_lm_head([chunk]) + + assert layer.forward is first_forward + assert getattr(layer.forward, "_nrl_fp32_lm_head_use_tf32") is False _assert_fp32_wrapped(layer) From 094406d76efdf467b2675ebf2dec7ce3f9a1c10c Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 10:44:14 -0700 Subject: [PATCH 06/18] fix: use config-only vllm fp32 lm head toggle Signed-off-by: Guyue Huang --- nemo_rl/models/generation/vllm/config.py | 5 +---- nemo_rl/models/generation/vllm/patches.py | 11 +++++----- .../models/generation/vllm/vllm_generation.py | 14 +++++++----- .../models/generation/test_vllm_generation.py | 9 ++++++-- .../models/generation/test_vllm_patches.py | 22 +++++++++++-------- 5 files changed, 35 insertions(+), 26 deletions(-) diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py index fbce1f9fef2..20171ea5daf 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -118,10 +118,7 @@ class VllmSpecificArgs(TypedDict): def vllm_fp32_lm_head_enabled(vllm_cfg: VllmSpecificArgs | dict[str, Any]) -> bool: """Return whether vLLM should run Nemotron-H logits with an fp32 head.""" - if vllm_cfg.get("fp32_lm_head"): - return True - env_vars = vllm_cfg.get("env_vars") - return env_vars is not None and str(env_vars.get(VLLM_FP32_LM_HEAD_ENV_VAR)) == "1" + return bool(vllm_cfg.get("fp32_lm_head")) class VllmDeltaCompressionConfig(BaseModel, extra="allow"): diff --git a/nemo_rl/models/generation/vllm/patches.py b/nemo_rl/models/generation/vllm/patches.py index 9b1edbaea29..6ba9934d41b 100644 --- a/nemo_rl/models/generation/vllm/patches.py +++ b/nemo_rl/models/generation/vllm/patches.py @@ -634,8 +634,9 @@ def _patch_vllm_nemotron_h_fp32_lm_head(logger) -> bool: This must be a source patch (not a monkeypatch): the model executes in vLLM's EngineCore worker subprocesses, which import vllm independently of - this process. The patched code is opt-in at runtime via - NRL_VLLM_FP32_LM_HEAD=1 (set by policy.generation.vllm_cfg.fp32_lm_head). + this process. The patched code is opt-in at runtime via an internal + NRL_VLLM_FP32_LM_HEAD=1 environment variable set from + policy.generation.vllm_cfg.fp32_lm_head. When enabled, the live ParallelLMHead keeps its original parameter dtype and quantization config; only the projection path casts hidden states, weights, and optional bias to fp32 at runtime. @@ -867,12 +868,12 @@ def _apply_vllm_patches( from vllm.logger import init_logger patch_logger = init_logger("vllm_patch") - fp32_lm_head_enabled = ( - fp32_lm_head or os.environ.get(VLLM_FP32_LM_HEAD_ENV_VAR) == "1" - ) + fp32_lm_head_enabled = bool(fp32_lm_head) if fp32_lm_head_enabled: os.environ[VLLM_FP32_LM_HEAD_ENV_VAR] = "1" extra_env_vars = [*(extra_env_vars or []), VLLM_FP32_LM_HEAD_ENV_VAR] + else: + os.environ.pop(VLLM_FP32_LM_HEAD_ENV_VAR, None) # Whether the v1 patch matters at all depends on which executor vLLM will # select. 0.25 defaults this to "1" (RayExecutorV2), which has no diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 3796ba36b1f..f69f200ac44 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -81,10 +81,13 @@ def _validate_fp32_lm_head_config(master_config: "MasterConfig") -> None: policy_config = master_config.policy generation_config = cast(VllmConfig, policy_config["generation"]) vllm_cfg = generation_config.get("vllm_cfg") - env_vars = None if vllm_cfg is None else vllm_cfg.get("env_vars") - vllm_env_value = ( - None if env_vars is None else env_vars.get(VLLM_FP32_LM_HEAD_ENV_VAR) - ) + env_vars = {} if vllm_cfg is None else vllm_cfg.get("env_vars") or {} + if VLLM_FP32_LM_HEAD_ENV_VAR in env_vars: + raise ValueError( + f"{VLLM_FP32_LM_HEAD_ENV_VAR} is reserved for NeMo-RL internal " + "vLLM patch plumbing; configure fp32 LM head with " + "policy.generation.vllm_cfg.fp32_lm_head instead." + ) vllm_fp32 = vllm_cfg is not None and vllm_fp32_lm_head_enabled(vllm_cfg) megatron_cfg = policy_config.get("megatron_cfg") @@ -97,8 +100,7 @@ def _validate_fp32_lm_head_config(master_config: "MasterConfig") -> None: "fp32 LM head must be enabled on both engines or neither: " f"policy.megatron_cfg.fp32_lm_head={megatron_fp32_value!r} but " f"policy.generation.vllm_cfg.fp32_lm_head=" - f"{None if vllm_cfg is None else vllm_cfg.get('fp32_lm_head')!r} " - f"({VLLM_FP32_LM_HEAD_ENV_VAR}={vllm_env_value!r}). " + f"{None if vllm_cfg is None else vllm_cfg.get('fp32_lm_head')!r}. " "A one-sided fp32 head increases the generation/training logprob " "mismatch instead of reducing it." ) diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index d3b6b228f56..1f54001446e 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -1505,7 +1505,6 @@ def _fp32_master_config( (False, {}), ("tf32", {"fp32_lm_head": True}), (True, {"fp32_lm_head": True}), - (True, {"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}), ], ) def test_vllm_validate_settings_accepts_matched_fp32_lm_head(trainer, vllm_cfg): @@ -1517,7 +1516,6 @@ def test_vllm_validate_settings_accepts_matched_fp32_lm_head(trainer, vllm_cfg): [ ("tf32", {}), # trainer fp32, vLLM bf16: the production misconfiguration (False, {"fp32_lm_head": True}), # vLLM fp32, trainer bf16 - (False, {"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}), ], ) def test_vllm_validate_settings_rejects_one_sided_fp32_lm_head(trainer, vllm_cfg): @@ -1525,6 +1523,13 @@ def test_vllm_validate_settings_rejects_one_sided_fp32_lm_head(trainer, vllm_cfg VllmGeneration.validate_settings(_fp32_master_config(trainer, vllm_cfg)) +def test_vllm_validate_settings_rejects_fp32_lm_head_env_var_toggle(): + with pytest.raises(ValueError, match="policy.generation.vllm_cfg.fp32_lm_head"): + VllmGeneration.validate_settings( + _fp32_master_config(False, {"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}) + ) + + def test_vllm_validate_settings_rejects_fp32_lm_head_with_fused_logprobs(): with pytest.raises(ValueError, match="use_fused_linear_logprobs"): VllmGeneration.validate_settings( diff --git a/tests/unit/models/generation/test_vllm_patches.py b/tests/unit/models/generation/test_vllm_patches.py index 33daab2630e..20bf274aeea 100644 --- a/tests/unit/models/generation/test_vllm_patches.py +++ b/tests/unit/models/generation/test_vllm_patches.py @@ -368,7 +368,7 @@ def test_glm_decoder_sp_moe_patch_warns_on_unknown_source( ({}, False), ({"fp32_lm_head": False}, False), ({"fp32_lm_head": True}, True), - ({"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}, True), + ({"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}, False), ({"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "0"}}, False), ], ) @@ -513,7 +513,7 @@ def test_apply_vllm_patches_gates_nemotron_h_fp32_lm_head(monkeypatch, enabled): assert captured_extra_env_vars == [["USER_VAR"]] -def test_apply_vllm_patches_accepts_legacy_fp32_lm_head_env_toggle(monkeypatch): +def test_apply_vllm_patches_ignores_ambient_fp32_lm_head_env_toggle(monkeypatch): _install_fake_vllm_modules(monkeypatch) monkeypatch.setenv(patches.VLLM_FP32_LM_HEAD_ENV_VAR, "1") captured_extra_env_vars = [] @@ -527,19 +527,23 @@ def test_apply_vllm_patches_accepts_legacy_fp32_lm_head_env_toggle(monkeypatch): patches._apply_vllm_patches("py") - assert fp32_patch_calls == [True] - assert captured_extra_env_vars == [[patches.VLLM_FP32_LM_HEAD_ENV_VAR]] + assert fp32_patch_calls == [] + assert patches.VLLM_FP32_LM_HEAD_ENV_VAR not in os.environ + assert captured_extra_env_vars == [None] @pytest.mark.parametrize( - "vllm_cfg_overrides", + ("vllm_cfg_overrides", "expected_fp32_lm_head"), [ - {"env_vars": {"USER_VAR": "value"}, "fp32_lm_head": True}, - {"env_vars": {"USER_VAR": "value", VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}, + ({"env_vars": {"USER_VAR": "value"}, "fp32_lm_head": True}, True), + ( + {"env_vars": {"USER_VAR": "value", VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}, + False, + ), ], ) def test_vllm_worker_threads_fp32_lm_head_cfg_into_source_patches( - monkeypatch, vllm_cfg_overrides + monkeypatch, vllm_cfg_overrides, expected_fp32_lm_head ): from nemo_rl.models.generation.vllm import vllm_worker @@ -575,7 +579,7 @@ def test_vllm_worker_threads_fp32_lm_head_cfg_into_source_patches( { "py": sys.executable, "extra_env_vars": ["EXPLICIT_VAR"], - "fp32_lm_head": True, + "fp32_lm_head": expected_fp32_lm_head, } ] From 4a47caf875a51ccad8bbaae79926d943c23128ab Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 10:53:08 -0700 Subject: [PATCH 07/18] fix: validate fp32 lm head in policy init Signed-off-by: Guyue Huang --- .../models/generation/vllm/vllm_generation.py | 44 ----- nemo_rl/models/policy/lm_policy.py | 69 ++++++++ .../models/generation/test_vllm_generation.py | 71 -------- .../models/policy/test_policy_validation.py | 163 ++++++++++++++++++ 4 files changed, 232 insertions(+), 115 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index f69f200ac44..9b089a2b960 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -46,9 +46,7 @@ ) from nemo_rl.models.generation.vllm.config import ( REFITTABLE_FP8_KV_CACHE_DTYPES, - VLLM_FP32_LM_HEAD_ENV_VAR, VllmConfig, - vllm_fp32_lm_head_enabled, ) from nemo_rl.models.generation.vllm.utils import ( aggregate_spec_decode_counters, @@ -76,47 +74,6 @@ logger = logging.getLogger(__name__) -def _validate_fp32_lm_head_config(master_config: "MasterConfig") -> None: - """Reject vLLM fp32 LM-head configs that the trainer cannot match.""" - policy_config = master_config.policy - generation_config = cast(VllmConfig, policy_config["generation"]) - vllm_cfg = generation_config.get("vllm_cfg") - env_vars = {} if vllm_cfg is None else vllm_cfg.get("env_vars") or {} - if VLLM_FP32_LM_HEAD_ENV_VAR in env_vars: - raise ValueError( - f"{VLLM_FP32_LM_HEAD_ENV_VAR} is reserved for NeMo-RL internal " - "vLLM patch plumbing; configure fp32 LM head with " - "policy.generation.vllm_cfg.fp32_lm_head instead." - ) - vllm_fp32 = vllm_cfg is not None and vllm_fp32_lm_head_enabled(vllm_cfg) - - megatron_cfg = policy_config.get("megatron_cfg") - megatron_fp32_value = ( - None if megatron_cfg is None else megatron_cfg.get("fp32_lm_head") - ) - trainer_fp32 = bool(megatron_fp32_value) - if trainer_fp32 != vllm_fp32: - raise ValueError( - "fp32 LM head must be enabled on both engines or neither: " - f"policy.megatron_cfg.fp32_lm_head={megatron_fp32_value!r} but " - f"policy.generation.vllm_cfg.fp32_lm_head=" - f"{None if vllm_cfg is None else vllm_cfg.get('fp32_lm_head')!r}. " - "A one-sided fp32 head increases the generation/training logprob " - "mismatch instead of reducing it." - ) - if ( - trainer_fp32 - and megatron_cfg is not None - and megatron_cfg.get("use_fused_linear_logprobs") - ): - raise ValueError( - "policy.megatron_cfg.fp32_lm_head has no effect with " - "use_fused_linear_logprobs=true (the fused linear+CE kernel bypasses " - "output_layer), which would leave the trainer in bf16 while vLLM " - "runs fp32. Disable one of them." - ) - - def _record_vllm_generation_metrics( model_name: str | None, data: BatchedDataDict, @@ -157,7 +114,6 @@ def validate_settings(cls, master_config: "MasterConfig") -> None: """Reject pure-config vLLM settings the SC entrypoint cannot honor.""" generation_config = cast(VllmConfig, master_config.policy["generation"]) assert_reload_refit_config_supported(generation_config) - _validate_fp32_lm_head_config(master_config) @staticmethod def init_cluster_placement_groups( diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index a9e02f94f58..6462ea71590 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -40,6 +40,10 @@ GenerationOutputSpec, RefitPayloadMode, ) +from nemo_rl.models.generation.vllm.config import ( + VLLM_FP32_LM_HEAD_ENV_VAR, + vllm_fp32_lm_head_enabled, +) from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.interfaces import ( ColocatablePolicyInterface, @@ -89,6 +93,68 @@ def _aggregate_megatron_flops_metrics( return aggregated +def _validate_fp32_lm_head_config( + config: PolicyConfig, *, megatron_enabled: bool, dtensor_enabled: bool +) -> None: + """Reject fp32 LM-head settings that the selected backends cannot match.""" + generation_config = config.get("generation") + if generation_config is None: + return + + generation_backend = generation_config["backend"] + megatron_cfg = config.get("megatron_cfg") + megatron_fp32_value = ( + megatron_cfg.get("fp32_lm_head") + if megatron_enabled and megatron_cfg is not None + else None + ) + megatron_fp32 = bool(megatron_fp32_value) + + if ( + megatron_fp32 + and megatron_cfg is not None + and megatron_cfg.get("use_fused_linear_logprobs") + ): + raise ValueError( + "policy.megatron_cfg.fp32_lm_head has no effect with " + "use_fused_linear_logprobs=true (the fused linear+CE kernel bypasses " + "output_layer). Disable one of them." + ) + + if generation_backend != "vllm": + return + + vllm_cfg = generation_config.get("vllm_cfg") + if vllm_cfg is None: + return + + env_vars = vllm_cfg.get("env_vars") or {} + if VLLM_FP32_LM_HEAD_ENV_VAR in env_vars: + raise ValueError( + f"{VLLM_FP32_LM_HEAD_ENV_VAR} is reserved for NeMo-RL internal " + "vLLM patch plumbing; configure fp32 LM head with " + "policy.generation.vllm_cfg.fp32_lm_head instead." + ) + + vllm_fp32 = vllm_fp32_lm_head_enabled(vllm_cfg) + if dtensor_enabled and vllm_fp32: + raise ValueError( + "policy.generation.vllm_cfg.fp32_lm_head=true is only supported " + "with the Megatron trainer because DTensor has no matching " + "policy.dtensor_cfg fp32 LM-head implementation." + ) + if megatron_enabled and megatron_fp32 != vllm_fp32: + raise ValueError( + "fp32 LM head must be enabled on both Megatron training and vLLM " + "generation or neither: " + f"policy.megatron_cfg.fp32_lm_head={megatron_fp32_value!r} but " + f"policy.generation.vllm_cfg.fp32_lm_head=" + f"{vllm_cfg.get('fp32_lm_head')!r}. " + "A one-sided fp32 head increases the generation/training logprob " + "mismatch instead of reducing it." + ) + + class Policy(ColocatablePolicyInterface, GenerationInterface): def __init__( self, @@ -156,6 +222,9 @@ def __init__( "Configure either Megatron (policy.megatron_cfg.enabled=true) or " "DTensor (policy.dtensor_cfg.enabled=true), not both." ) + _validate_fp32_lm_head_config( + config, megatron_enabled=megatron_enable, dtensor_enabled=dtensor_enable + ) if reserved_http_server_ports is not None and not megatron_enable: raise ValueError( "reserved_http_server_ports is only supported by the Megatron " diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 1f54001446e..1b2ba597aac 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -39,7 +39,6 @@ ) from nemo_rl.models.generation.openai_server_utils import replace_prefix_tokens from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration -from nemo_rl.models.generation.vllm.config import VLLM_FP32_LM_HEAD_ENV_VAR from nemo_rl.models.generation.vllm.vllm_worker import ( VllmGenerationWorkerImpl, _context_capped_max_new_tokens, @@ -1481,76 +1480,6 @@ def test_vllm_validate_settings_accepts_missing_vllm_cfg_as_refit_disabled(): VllmGeneration.validate_settings(master_config) -def _fp32_master_config( - trainer, - vllm_cfg, - *, - fused=False, - include_megatron_cfg=True, -): - generation = deepcopy(basic_vllm_test_config) - generation["vllm_cfg"].update(vllm_cfg) - policy = {"generation": generation} - if include_megatron_cfg: - megatron_cfg = {"fp32_lm_head": trainer} - if fused: - megatron_cfg["use_fused_linear_logprobs"] = True - policy["megatron_cfg"] = megatron_cfg - return types.SimpleNamespace(policy=policy) - - -@pytest.mark.parametrize( - "trainer, vllm_cfg", - [ - (False, {}), - ("tf32", {"fp32_lm_head": True}), - (True, {"fp32_lm_head": True}), - ], -) -def test_vllm_validate_settings_accepts_matched_fp32_lm_head(trainer, vllm_cfg): - VllmGeneration.validate_settings(_fp32_master_config(trainer, vllm_cfg)) - - -@pytest.mark.parametrize( - "trainer, vllm_cfg", - [ - ("tf32", {}), # trainer fp32, vLLM bf16: the production misconfiguration - (False, {"fp32_lm_head": True}), # vLLM fp32, trainer bf16 - ], -) -def test_vllm_validate_settings_rejects_one_sided_fp32_lm_head(trainer, vllm_cfg): - with pytest.raises(ValueError, match="both engines or neither"): - VllmGeneration.validate_settings(_fp32_master_config(trainer, vllm_cfg)) - - -def test_vllm_validate_settings_rejects_fp32_lm_head_env_var_toggle(): - with pytest.raises(ValueError, match="policy.generation.vllm_cfg.fp32_lm_head"): - VllmGeneration.validate_settings( - _fp32_master_config(False, {"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}) - ) - - -def test_vllm_validate_settings_rejects_fp32_lm_head_with_fused_logprobs(): - with pytest.raises(ValueError, match="use_fused_linear_logprobs"): - VllmGeneration.validate_settings( - _fp32_master_config("tf32", {"fp32_lm_head": True}, fused=True) - ) - - -def test_vllm_validate_settings_handles_missing_megatron_cfg_for_fp32_lm_head(): - VllmGeneration.validate_settings( - _fp32_master_config(False, {}, include_megatron_cfg=False) - ) - with pytest.raises(ValueError, match="both engines or neither"): - VllmGeneration.validate_settings( - _fp32_master_config( - False, - {"fp32_lm_head": True}, - include_megatron_cfg=False, - ) - ) - - def test_vllm_policy_generation(policy, test_input_data, tokenizer): """Test vLLM policy generation capabilities.""" # Test generation diff --git a/tests/unit/models/policy/test_policy_validation.py b/tests/unit/models/policy/test_policy_validation.py index d8fc112367e..83195c5c6f6 100644 --- a/tests/unit/models/policy/test_policy_validation.py +++ b/tests/unit/models/policy/test_policy_validation.py @@ -24,6 +24,7 @@ import pytest +from nemo_rl.models.generation.vllm.config import VLLM_FP32_LM_HEAD_ENV_VAR from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.lm_policy import Policy @@ -171,6 +172,168 @@ def create_megatron_config( } +def set_vllm_generation( + config: PolicyConfig, vllm_cfg_overrides: dict[str, object] +) -> PolicyConfig: + config["generation"] = { + "backend": "vllm", + "temperature": 1.0, + "top_p": 1.0, + "top_k": None, + "max_new_tokens": 16, + "stop_token_ids": None, + "stop_strings": None, + "colocated": { + "enabled": False, + "resources": { + "gpus_per_node": 1, + "num_nodes": 1, + }, + }, + "vllm_cfg": { + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "expert_parallel_size": 1, + "gpu_memory_utilization": 0.6, + "max_model_len": 128, + "skip_tokenizer_init": True, + "async_engine": False, + "kv_cache_dtype": "auto", + **vllm_cfg_overrides, + }, + } + return config + + +def construct_policy_with_mocks(config: PolicyConfig) -> Policy: + model_config = MagicMock() + with ( + patch.dict("os.environ", {"TORCH_CUDA_ARCH_LIST": "9.0"}), + patch("nemo_rl.models.policy.lm_policy.RayQueue"), + patch("nemo_rl.models.policy.lm_policy.RayWorkerBuilder"), + patch("nemo_rl.models.policy.lm_policy.RayWorkerGroup"), + patch( + "nemo_rl.models.policy.lm_policy.get_hf_config", return_value=model_config + ), + patch("nemo_rl.models.policy.lm_policy.FLOPTracker.from_config"), + ): + return Policy( + cluster=create_mock_cluster(world_size=1), + config=config, + tokenizer=create_mock_tokenizer(), + ) + + +@pytest.mark.parametrize("trainer_fp32", [True, "tf32"]) +def test_policy_accepts_matched_vllm_and_megatron_fp32_lm_head(trainer_fp32): + config = create_megatron_config("test-model", tp=1) + config["megatron_cfg"]["fp32_lm_head"] = trainer_fp32 + set_vllm_generation(config, {"fp32_lm_head": True}) + + policy = construct_policy_with_mocks(config) + + assert policy.worker_group is not None + + +@pytest.mark.parametrize( + ("trainer_fp32", "vllm_fp32"), + [ + ("tf32", False), + (False, True), + ], +) +def test_policy_rejects_mismatched_vllm_and_megatron_fp32_lm_head( + trainer_fp32, vllm_fp32 +): + config = create_megatron_config("test-model", tp=1) + config["megatron_cfg"]["fp32_lm_head"] = trainer_fp32 + set_vllm_generation(config, {"fp32_lm_head": vllm_fp32}) + + with ( + patch("nemo_rl.models.policy.lm_policy.RayWorkerGroup") as worker_group, + pytest.raises(ValueError, match="both Megatron training and vLLM generation"), + ): + Policy( + cluster=create_mock_cluster(world_size=1), + config=config, + tokenizer=create_mock_tokenizer(), + ) + + worker_group.assert_not_called() + + +def test_policy_rejects_vllm_fp32_lm_head_with_dtensor_trainer(): + config = create_dtensor_config("test-model", tp=1) + set_vllm_generation(config, {"fp32_lm_head": True}) + + with ( + patch("nemo_rl.models.policy.lm_policy.RayWorkerGroup") as worker_group, + pytest.raises(ValueError, match="DTensor has no matching"), + ): + Policy( + cluster=create_mock_cluster(world_size=1), + config=config, + tokenizer=create_mock_tokenizer(), + ) + + worker_group.assert_not_called() + + +def test_policy_accepts_vllm_fp32_lm_head_disabled_with_dtensor_trainer(): + config = create_dtensor_config("test-model", tp=1) + set_vllm_generation(config, {}) + + policy = construct_policy_with_mocks(config) + + assert policy.worker_group is not None + + +def test_policy_rejects_fp32_lm_head_env_var_toggle(): + config = create_dtensor_config("test-model", tp=1) + set_vllm_generation(config, {"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}) + + with ( + patch("nemo_rl.models.policy.lm_policy.RayWorkerGroup") as worker_group, + pytest.raises(ValueError, match="policy.generation.vllm_cfg.fp32_lm_head"), + ): + Policy( + cluster=create_mock_cluster(world_size=1), + config=config, + tokenizer=create_mock_tokenizer(), + ) + + worker_group.assert_not_called() + + +def test_policy_rejects_megatron_fp32_lm_head_with_fused_logprobs(): + config = create_megatron_config("test-model", tp=1) + config["megatron_cfg"]["fp32_lm_head"] = "tf32" + config["megatron_cfg"]["use_fused_linear_logprobs"] = True + set_vllm_generation(config, {"fp32_lm_head": True}) + + with ( + patch("nemo_rl.models.policy.lm_policy.RayWorkerGroup") as worker_group, + pytest.raises(ValueError, match="use_fused_linear_logprobs"), + ): + Policy( + cluster=create_mock_cluster(world_size=1), + config=config, + tokenizer=create_mock_tokenizer(), + ) + + worker_group.assert_not_called() + + +def test_policy_accepts_megatron_fp32_lm_head_with_megatron_generation(): + config = create_megatron_config("test-model", tp=1) + config["megatron_cfg"]["fp32_lm_head"] = "tf32" + config["generation"]["backend"] = "megatron" + + policy = construct_policy_with_mocks(config) + + assert policy.worker_group is not None + + def test_policy_flops_tracker_uses_hf_config_overrides() -> None: cluster = create_mock_cluster(world_size=1) tokenizer = create_mock_tokenizer() From 87cea5ef751fe519d2fc7a1f5d3930a15f89ae9e Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 11:02:06 -0700 Subject: [PATCH 08/18] refactor: move fp32 lm head policy validation Signed-off-by: Guyue Huang --- nemo_rl/models/policy/lm_policy.py | 69 +---------------------------- nemo_rl/models/policy/utils.py | 71 +++++++++++++++++++++++++++++- 2 files changed, 72 insertions(+), 68 deletions(-) diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index 6462ea71590..c07e054b6b6 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -40,10 +40,6 @@ GenerationOutputSpec, RefitPayloadMode, ) -from nemo_rl.models.generation.vllm.config import ( - VLLM_FP32_LM_HEAD_ENV_VAR, - vllm_fp32_lm_head_enabled, -) from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.interfaces import ( ColocatablePolicyInterface, @@ -55,6 +51,7 @@ from nemo_rl.models.policy.utils import ( aggregate_per_sample_handles, resolve_policy_worker_cls, + validate_fp32_lm_head_config, ) from nemo_rl.utils.checkpoint import CheckpointingConfig from nemo_rl.utils.flops_tracker import ( @@ -93,68 +90,6 @@ def _aggregate_megatron_flops_metrics( return aggregated -def _validate_fp32_lm_head_config( - config: PolicyConfig, *, megatron_enabled: bool, dtensor_enabled: bool -) -> None: - """Reject fp32 LM-head settings that the selected backends cannot match.""" - generation_config = config.get("generation") - if generation_config is None: - return - - generation_backend = generation_config["backend"] - megatron_cfg = config.get("megatron_cfg") - megatron_fp32_value = ( - megatron_cfg.get("fp32_lm_head") - if megatron_enabled and megatron_cfg is not None - else None - ) - megatron_fp32 = bool(megatron_fp32_value) - - if ( - megatron_fp32 - and megatron_cfg is not None - and megatron_cfg.get("use_fused_linear_logprobs") - ): - raise ValueError( - "policy.megatron_cfg.fp32_lm_head has no effect with " - "use_fused_linear_logprobs=true (the fused linear+CE kernel bypasses " - "output_layer). Disable one of them." - ) - - if generation_backend != "vllm": - return - - vllm_cfg = generation_config.get("vllm_cfg") - if vllm_cfg is None: - return - - env_vars = vllm_cfg.get("env_vars") or {} - if VLLM_FP32_LM_HEAD_ENV_VAR in env_vars: - raise ValueError( - f"{VLLM_FP32_LM_HEAD_ENV_VAR} is reserved for NeMo-RL internal " - "vLLM patch plumbing; configure fp32 LM head with " - "policy.generation.vllm_cfg.fp32_lm_head instead." - ) - - vllm_fp32 = vllm_fp32_lm_head_enabled(vllm_cfg) - if dtensor_enabled and vllm_fp32: - raise ValueError( - "policy.generation.vllm_cfg.fp32_lm_head=true is only supported " - "with the Megatron trainer because DTensor has no matching " - "policy.dtensor_cfg fp32 LM-head implementation." - ) - if megatron_enabled and megatron_fp32 != vllm_fp32: - raise ValueError( - "fp32 LM head must be enabled on both Megatron training and vLLM " - "generation or neither: " - f"policy.megatron_cfg.fp32_lm_head={megatron_fp32_value!r} but " - f"policy.generation.vllm_cfg.fp32_lm_head=" - f"{vllm_cfg.get('fp32_lm_head')!r}. " - "A one-sided fp32 head increases the generation/training logprob " - "mismatch instead of reducing it." - ) - - class Policy(ColocatablePolicyInterface, GenerationInterface): def __init__( self, @@ -222,7 +157,7 @@ def __init__( "Configure either Megatron (policy.megatron_cfg.enabled=true) or " "DTensor (policy.dtensor_cfg.enabled=true), not both." ) - _validate_fp32_lm_head_config( + validate_fp32_lm_head_config( config, megatron_enabled=megatron_enable, dtensor_enabled=dtensor_enable ) if reserved_http_server_ports is not None and not megatron_enable: diff --git a/nemo_rl/models/policy/utils.py b/nemo_rl/models/policy/utils.py index 96ddd6ebd43..8817c460c75 100644 --- a/nemo_rl/models/policy/utils.py +++ b/nemo_rl/models/policy/utils.py @@ -18,7 +18,7 @@ import warnings from datetime import timedelta from enum import Enum -from typing import Any, Dict, Iterable, Optional, cast +from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional, cast import torch import torch.distributed as dist @@ -55,6 +55,13 @@ NEMO_AUTOMODEL_AVAILABLE = False from nemo_rl.distributed.worker_group_utils import get_nsight_config_if_pattern_matches +from nemo_rl.models.generation.vllm.config import ( + VLLM_FP32_LM_HEAD_ENV_VAR, + vllm_fp32_lm_head_enabled, +) + +if TYPE_CHECKING: + from nemo_rl.models.policy import PolicyConfig # Plain Hugging Face classes remain separate from the NeMo AutoModel wrappers so # callers that manage distribution can request them when NeMo AutoModel is installed. @@ -133,6 +140,68 @@ def resolve_policy_worker_cls(default_cls: str, config: dict) -> str: return POLICY_WORKER_OVERRIDES.get(default_cls, default_cls) +def validate_fp32_lm_head_config( + config: "PolicyConfig", *, megatron_enabled: bool, dtensor_enabled: bool +) -> None: + """Reject fp32 LM-head settings that the selected backends cannot match.""" + generation_config = config.get("generation") + if generation_config is None: + return + + generation_backend = generation_config["backend"] + megatron_cfg = config.get("megatron_cfg") + megatron_fp32_value = ( + megatron_cfg.get("fp32_lm_head") + if megatron_enabled and megatron_cfg is not None + else None + ) + megatron_fp32 = bool(megatron_fp32_value) + + if ( + megatron_fp32 + and megatron_cfg is not None + and megatron_cfg.get("use_fused_linear_logprobs") + ): + raise ValueError( + "policy.megatron_cfg.fp32_lm_head has no effect with " + "use_fused_linear_logprobs=true (the fused linear+CE kernel bypasses " + "output_layer). Disable one of them." + ) + + if generation_backend != "vllm": + return + + vllm_cfg = generation_config.get("vllm_cfg") + if vllm_cfg is None: + return + + env_vars = vllm_cfg.get("env_vars") or {} + if VLLM_FP32_LM_HEAD_ENV_VAR in env_vars: + raise ValueError( + f"{VLLM_FP32_LM_HEAD_ENV_VAR} is reserved for NeMo-RL internal " + "vLLM patch plumbing; configure fp32 LM head with " + "policy.generation.vllm_cfg.fp32_lm_head instead." + ) + + vllm_fp32 = vllm_fp32_lm_head_enabled(vllm_cfg) + if dtensor_enabled and vllm_fp32: + raise ValueError( + "policy.generation.vllm_cfg.fp32_lm_head=true is only supported " + "with the Megatron trainer because DTensor has no matching " + "policy.dtensor_cfg fp32 LM-head implementation." + ) + if megatron_enabled and megatron_fp32 != vllm_fp32: + raise ValueError( + "fp32 LM head must be enabled on both Megatron training and vLLM " + "generation or neither: " + f"policy.megatron_cfg.fp32_lm_head={megatron_fp32_value!r} but " + f"policy.generation.vllm_cfg.fp32_lm_head=" + f"{vllm_cfg.get('fp32_lm_head')!r}. " + "A one-sided fp32 head increases the generation/training logprob " + "mismatch instead of reducing it." + ) + + def resolve_model_class( model_name: str, *, From 02ddbf43f9897cde4ad4e2897d2e9928e22065c3 Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 11:07:19 -0700 Subject: [PATCH 09/18] test: add fp32 lm head functional coverage Signed-off-by: Guyue Huang --- .../L1_Functional_Tests_Megatron_1.sh | 1 + .../grpo_megatron_vllm_fp32_lm_head.sh | 54 +++++++++++++++++++ 2 files changed, 55 insertions(+) create mode 100755 tests/functional/grpo_megatron_vllm_fp32_lm_head.sh diff --git a/tests/functional/L1_Functional_Tests_Megatron_1.sh b/tests/functional/L1_Functional_Tests_Megatron_1.sh index 000867eda15..c625341c703 100644 --- a/tests/functional/L1_Functional_Tests_Megatron_1.sh +++ b/tests/functional/L1_Functional_Tests_Megatron_1.sh @@ -36,6 +36,7 @@ run_test() { run_test fast uv run --no-sync bash ./tests/functional/audio_grpo_megatron.sh run_test uv run --no-sync bash ./tests/functional/grpo_megatron.sh +run_test uv run --no-sync bash ./tests/functional/grpo_megatron_vllm_fp32_lm_head.sh run_test uv run --no-sync bash ./tests/functional/grpo_megatron_mbridge_restore.sh run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_eagle3_online.sh diff --git a/tests/functional/grpo_megatron_vllm_fp32_lm_head.sh b/tests/functional/grpo_megatron_vllm_fp32_lm_head.sh new file mode 100755 index 00000000000..bd9f70fc1e1 --- /dev/null +++ b/tests/functional/grpo_megatron_vllm_fp32_lm_head.sh @@ -0,0 +1,54 @@ +#!/bin/bash + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) +# Mark the current repo as safe, since wandb fetches metadata about the repo +git config --global --add safe.directory $PROJECT_ROOT + +set -eou pipefail + +EXP_NAME=$(basename $0 .sh) +EXP_DIR=$SCRIPT_DIR/$EXP_NAME +LOG_DIR=$EXP_DIR/logs +JSON_METRICS=$EXP_DIR/metrics.json +RUN_LOG=$EXP_DIR/run.log +export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} + +rm -rf $EXP_DIR $LOG_DIR +mkdir -p $EXP_DIR $LOG_DIR + +# Qwen2.5 keeps this L1 aligned with the existing Megatron+vLLM GRPO smoke +# while still exercising the fp32_lm_head config and launch path. +cd $PROJECT_ROOT +uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ + $PROJECT_ROOT/examples/run_grpo.py \ + --config $PROJECT_ROOT/examples/configs/grpo_math_1B_megatron.yaml \ + policy.model_name=Qwen/Qwen2.5-0.5B \ + grpo.num_prompts_per_step=2 \ + grpo.num_generations_per_prompt=4 \ + policy.train_global_batch_size=4 \ + policy.logprob_batch_size=4 \ + policy.train_micro_batch_size=1 \ + policy.megatron_cfg.fp32_lm_head=tf32 \ + policy.generation.vllm_cfg.fp32_lm_head=true \ + cluster.gpus_per_node=2 \ + grpo.max_num_steps=2 \ + logger.tensorboard_enabled=true \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=false \ + logger.monitor_gpus=true \ + checkpointing.enabled=false \ + $@ \ + 2>&1 | tee $RUN_LOG + +grep -q "\[fp32_lm_head\] output layer will compute logits in fp32" $RUN_LOG +grep -Eq "Applied NemotronH fp32 LM head source patch|NemotronH fp32 LM head patch already present" $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/token_mult_prob_error"]) < 1.05' \ + 'min(data["train/probs_ratio_clamped_min"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_min"]) < 1.21' \ + 'min(data["train/probs_ratio_clamped_max"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_max"]) < 1.21' From ed47246592a2277a7c76c4e1a13c0c4392fff1a1 Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 11:25:35 -0700 Subject: [PATCH 10/18] fix: clarify nemotron h fp32 lm head scope Signed-off-by: Guyue Huang --- nemo_rl/models/generation/vllm/config.py | 15 +++- nemo_rl/models/generation/vllm/patches.py | 37 +++++---- nemo_rl/models/generation/vllm/vllm_worker.py | 4 +- nemo_rl/models/policy/lm_policy.py | 26 ++++++- nemo_rl/models/policy/utils.py | 78 +++++++++++++++++-- .../L1_Functional_Tests_Megatron_1.sh | 1 - .../grpo_megatron_vllm_fp32_lm_head.sh | 54 ------------- .../models/generation/test_vllm_patches.py | 55 +++++++------ .../models/policy/test_policy_validation.py | 34 +++++++- 9 files changed, 190 insertions(+), 114 deletions(-) delete mode 100755 tests/functional/grpo_megatron_vllm_fp32_lm_head.sh diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py index 20171ea5daf..35c3d05568a 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -28,7 +28,11 @@ VllmRefitTransportName = Literal["s3", "zmq"] VllmRefitSelector = Literal["vllm_s3_sparse", "vllm_zmq_sparse", "nixl", "nccl_reshard"] VLLM_SPARSE_REFIT_TRANSPORTS = frozenset({"vllm_s3_sparse", "vllm_zmq_sparse"}) -VLLM_FP32_LM_HEAD_ENV_VAR = "NRL_VLLM_FP32_LM_HEAD" +VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR = "NRL_VLLM_FP32_LM_HEAD" +# Backward-compatible alias for already-patched workers and older imports. The +# value stays unchanged because it is internal patch plumbing, not a user-facing +# configuration surface. +VLLM_FP32_LM_HEAD_ENV_VAR = VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR REFITTABLE_FP8_KV_CACHE_DTYPES = frozenset({"fp8", "fp8_e4m3"}) @@ -62,8 +66,9 @@ class VllmSpecificArgs(TypedDict): # with generation-time processors should request ``raw_logprobs`` when # comparing generation and policy logprobs. logprobs_mode: NotRequired[Literal["processed_logprobs", "raw_logprobs"]] - # Compute Nemotron-H logits with an fp32 LM head in vLLM. Pair this with - # policy.megatron_cfg.fp32_lm_head when using a Megatron trainer. + # Nemotron-H only: compute vLLM Nemotron-H logits with an fp32 LM head. + # Pair this with policy.megatron_cfg.fp32_lm_head when using a Megatron + # trainer. fp32_lm_head: NotRequired[bool] # Cap each request's generated tokens so the training prompt plus response # fits within max_model_len. This is needed when multimodal processing makes @@ -116,7 +121,9 @@ class VllmSpecificArgs(TypedDict): reasoning_parser_plugin: NotRequired[str] -def vllm_fp32_lm_head_enabled(vllm_cfg: VllmSpecificArgs | dict[str, Any]) -> bool: +def vllm_nemotron_h_fp32_lm_head_enabled( + vllm_cfg: VllmSpecificArgs | dict[str, Any], +) -> bool: """Return whether vLLM should run Nemotron-H logits with an fp32 head.""" return bool(vllm_cfg.get("fp32_lm_head")) diff --git a/nemo_rl/models/generation/vllm/patches.py b/nemo_rl/models/generation/vllm/patches.py index 6ba9934d41b..98125bca06c 100644 --- a/nemo_rl/models/generation/vllm/patches.py +++ b/nemo_rl/models/generation/vllm/patches.py @@ -16,7 +16,9 @@ from contextlib import contextmanager from importlib.util import find_spec -from nemo_rl.models.generation.vllm.config import VLLM_FP32_LM_HEAD_ENV_VAR +from nemo_rl.models.generation.vllm.config import ( + VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, +) def _get_vllm_file(relative_path: str) -> str: @@ -666,7 +668,7 @@ def _patch_vllm_nemotron_h_fp32_lm_head(logger) -> bool: prefix=maybe_prefix(prefix, "lm_head"), )""" previous_lm_head_snippet = f""" self._nrl_fp32_lm_head = ( - os.environ.get("{VLLM_FP32_LM_HEAD_ENV_VAR}", "0") == "1" + os.environ.get("{VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR}", "0") == "1" ) if self._nrl_fp32_lm_head: self.lm_head = ParallelLMHead( @@ -690,7 +692,7 @@ def _patch_vllm_nemotron_h_fp32_lm_head(logger) -> bool: prefix=maybe_prefix(prefix, "lm_head"), ) self._nrl_fp32_lm_head = ( - os.environ.get("{VLLM_FP32_LM_HEAD_ENV_VAR}", "0") == "1" + os.environ.get("{VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR}", "0") == "1" )""" old_logits_processor_snippet = ( " self.logits_processor = LogitsProcessor(config.vocab_size)" @@ -735,7 +737,7 @@ def _nrl_fp32_lm_head_apply( # source file that still contains the previous lazy-deepcopy patch. legacy_snippet = f""" import os as _os - if _os.environ.get("{VLLM_FP32_LM_HEAD_ENV_VAR}", "0") == "1": + if _os.environ.get("{VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR}", "0") == "1": # NeMo-RL patch: fp32 LM head (MiniMax-M1-style). bf16 rounding of # the logits is the dominant gen/train logprob mismatch source. _fp32_head = getattr(self, "_nrl_lm_head_fp32", None) @@ -861,19 +863,22 @@ def _apply_vllm_patches( py_executable: str, *, extra_env_vars: list[str] | None = None, - fp32_lm_head: bool | None = None, + nemotron_h_fp32_lm_head: bool | None = None, ) -> None: # Import lazily so importing the worker module does not import vLLM. import vllm.envs as envs from vllm.logger import init_logger patch_logger = init_logger("vllm_patch") - fp32_lm_head_enabled = bool(fp32_lm_head) - if fp32_lm_head_enabled: - os.environ[VLLM_FP32_LM_HEAD_ENV_VAR] = "1" - extra_env_vars = [*(extra_env_vars or []), VLLM_FP32_LM_HEAD_ENV_VAR] + nemotron_h_fp32_lm_head_enabled = bool(nemotron_h_fp32_lm_head) + if nemotron_h_fp32_lm_head_enabled: + os.environ[VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR] = "1" + extra_env_vars = [ + *(extra_env_vars or []), + VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, + ] else: - os.environ.pop(VLLM_FP32_LM_HEAD_ENV_VAR, None) + os.environ.pop(VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, None) # Whether the v1 patch matters at all depends on which executor vLLM will # select. 0.25 defaults this to "1" (RayExecutorV2), which has no @@ -917,9 +922,13 @@ def _apply_vllm_patches( _patch_vllm_shm_broadcast_bind_retry(patch_logger) _patch_vllm_radio_layerscale_loader(patch_logger) _patch_vllm_glm_decoder_sequence_parallel_moe(patch_logger) - if fp32_lm_head_enabled and not _patch_vllm_nemotron_h_fp32_lm_head(patch_logger): + if ( + nemotron_h_fp32_lm_head_enabled + and not _patch_vllm_nemotron_h_fp32_lm_head(patch_logger) + ): raise RuntimeError( - "vllm_cfg.fp32_lm_head is enabled, but the Nemotron-H fp32 LM head " - "source patch could not be applied. Disable the flag or update the " - "patch anchors for this vLLM version." + "vllm_cfg.fp32_lm_head is enabled, but that flag currently maps to " + "the Nemotron-H-only vLLM fp32 LM head source patch, and the patch " + "could not be applied. Disable the flag or update the patch anchors " + "for this vLLM version." ) diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index a35ce53ec99..d08b1e429f1 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -46,7 +46,7 @@ VLLM_SPARSE_REFIT_TRANSPORTS, VllmConfig, resolve_vllm_video_config, - vllm_fp32_lm_head_enabled, + vllm_nemotron_h_fp32_lm_head_enabled, ) from nemo_rl.models.generation.vllm.patches import _apply_vllm_patches from nemo_rl.models.generation.vllm.utils import ( @@ -443,7 +443,7 @@ def _init_config( _apply_vllm_patches( self.py_executable, extra_env_vars=extra_env_vars, - fp32_lm_head=vllm_fp32_lm_head_enabled(vllm_cfg), + nemotron_h_fp32_lm_head=vllm_nemotron_h_fp32_lm_head_enabled(vllm_cfg), ) # Skip model loading if we're not the model owner diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index c07e054b6b6..a19b2fea291 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -160,6 +160,22 @@ def __init__( validate_fp32_lm_head_config( config, megatron_enabled=megatron_enable, dtensor_enabled=dtensor_enable ) + hf_config = None + hf_config_overrides = config.get("hf_config_overrides") or {} + generation_config = config.get("generation") + if generation_config is not None and generation_config["backend"] == "vllm": + vllm_cfg = generation_config.get("vllm_cfg") + if vllm_cfg is not None and vllm_cfg.get("fp32_lm_head"): + hf_config = get_hf_config( + config["model_name"], + **hf_config_overrides, + ) + validate_fp32_lm_head_config( + config, + megatron_enabled=megatron_enable, + dtensor_enabled=dtensor_enable, + model_config=hf_config, + ) if reserved_http_server_ports is not None and not megatron_enable: raise ValueError( "reserved_http_server_ports is only supported by the Megatron " @@ -416,12 +432,14 @@ def __init__( # initialize FLOPs tracker try: + if hf_config is None: + hf_config = get_hf_config( + config["model_name"], + **hf_config_overrides, + ) self.flops_tracker = FLOPTracker.from_config( config["model_name"], - get_hf_config( - config["model_name"], - **(config.get("hf_config_overrides") or {}), - ), + hf_config, ) except ValueError as e: self.flops_tracker = None diff --git a/nemo_rl/models/policy/utils.py b/nemo_rl/models/policy/utils.py index 8817c460c75..40568b1bb1b 100644 --- a/nemo_rl/models/policy/utils.py +++ b/nemo_rl/models/policy/utils.py @@ -56,8 +56,8 @@ from nemo_rl.distributed.worker_group_utils import get_nsight_config_if_pattern_matches from nemo_rl.models.generation.vllm.config import ( - VLLM_FP32_LM_HEAD_ENV_VAR, - vllm_fp32_lm_head_enabled, + VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, + vllm_nemotron_h_fp32_lm_head_enabled, ) if TYPE_CHECKING: @@ -127,6 +127,9 @@ class IPCProtocol(Enum): "nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2": "nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker_v2.DTensorQuantPolicyWorkerV2", } +_NEMOTRON_H_MODEL_TYPES = frozenset({"nemotron_h"}) +_NEMOTRON_H_ARCHITECTURES = frozenset({"NemotronHForCausalLM"}) + def resolve_policy_worker_cls(default_cls: str, config: dict) -> str: """Return the quantized policy worker FQN if ``quant_cfg`` is set, else ``default_cls``. @@ -140,8 +143,54 @@ def resolve_policy_worker_cls(default_cls: str, config: dict) -> str: return POLICY_WORKER_OVERRIDES.get(default_cls, default_cls) +def _normalize_model_type(model_type: object) -> str: + return str(model_type).lower().replace("-", "_") + + +def _get_config_model_type(model_config: object) -> object | None: + return getattr(model_config, "model_type", None) or getattr( + model_config.__class__, "model_type", None + ) + + +def _get_config_architectures(model_config: object) -> list[str]: + architectures = getattr(model_config, "architectures", None) or [] + if isinstance(architectures, str): + return [architectures] + try: + return [str(architecture) for architecture in architectures] + except TypeError: + return [] + + +def _is_nemotron_h_model_config(model_config: object) -> bool: + model_type = _get_config_model_type(model_config) + if ( + model_type is not None + and _normalize_model_type(model_type) in _NEMOTRON_H_MODEL_TYPES + ): + return True + + return any( + architecture in _NEMOTRON_H_ARCHITECTURES + for architecture in _get_config_architectures(model_config) + ) + + +def _describe_model_config(model_config: object) -> str: + model_type = _get_config_model_type(model_config) + architectures = _get_config_architectures(model_config) + if architectures: + return f"architectures={architectures!r}, model_type={model_type!r}" + return f"model_type={model_type!r}" + + def validate_fp32_lm_head_config( - config: "PolicyConfig", *, megatron_enabled: bool, dtensor_enabled: bool + config: "PolicyConfig", + *, + megatron_enabled: bool, + dtensor_enabled: bool, + model_config: object | None = None, ) -> None: """Reject fp32 LM-head settings that the selected backends cannot match.""" generation_config = config.get("generation") @@ -176,14 +225,15 @@ def validate_fp32_lm_head_config( return env_vars = vllm_cfg.get("env_vars") or {} - if VLLM_FP32_LM_HEAD_ENV_VAR in env_vars: + if VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR in env_vars: raise ValueError( - f"{VLLM_FP32_LM_HEAD_ENV_VAR} is reserved for NeMo-RL internal " - "vLLM patch plumbing; configure fp32 LM head with " + f"{VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR} is reserved for " + "NeMo-RL internal vLLM Nemotron-H patch plumbing; configure fp32 " + "LM head with " "policy.generation.vllm_cfg.fp32_lm_head instead." ) - vllm_fp32 = vllm_fp32_lm_head_enabled(vllm_cfg) + vllm_fp32 = vllm_nemotron_h_fp32_lm_head_enabled(vllm_cfg) if dtensor_enabled and vllm_fp32: raise ValueError( "policy.generation.vllm_cfg.fp32_lm_head=true is only supported " @@ -200,6 +250,20 @@ def validate_fp32_lm_head_config( "A one-sided fp32 head increases the generation/training logprob " "mismatch instead of reducing it." ) + if ( + vllm_fp32 + and model_config is not None + and not _is_nemotron_h_model_config(model_config) + ): + warnings.warn( + "policy.generation.vllm_cfg.fp32_lm_head=true currently only " + "patches vLLM's Nemotron-H model implementation " + "(NemotronHForCausalLM). The configured policy model does not " + f"look like Nemotron-H ({_describe_model_config(model_config)}), " + "so vLLM generation will not execute an fp32 LM-head path.", + UserWarning, + stacklevel=2, + ) def resolve_model_class( diff --git a/tests/functional/L1_Functional_Tests_Megatron_1.sh b/tests/functional/L1_Functional_Tests_Megatron_1.sh index c625341c703..000867eda15 100644 --- a/tests/functional/L1_Functional_Tests_Megatron_1.sh +++ b/tests/functional/L1_Functional_Tests_Megatron_1.sh @@ -36,7 +36,6 @@ run_test() { run_test fast uv run --no-sync bash ./tests/functional/audio_grpo_megatron.sh run_test uv run --no-sync bash ./tests/functional/grpo_megatron.sh -run_test uv run --no-sync bash ./tests/functional/grpo_megatron_vllm_fp32_lm_head.sh run_test uv run --no-sync bash ./tests/functional/grpo_megatron_mbridge_restore.sh run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_eagle3_online.sh diff --git a/tests/functional/grpo_megatron_vllm_fp32_lm_head.sh b/tests/functional/grpo_megatron_vllm_fp32_lm_head.sh deleted file mode 100755 index bd9f70fc1e1..00000000000 --- a/tests/functional/grpo_megatron_vllm_fp32_lm_head.sh +++ /dev/null @@ -1,54 +0,0 @@ -#!/bin/bash - -SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) -PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) -# Mark the current repo as safe, since wandb fetches metadata about the repo -git config --global --add safe.directory $PROJECT_ROOT - -set -eou pipefail - -EXP_NAME=$(basename $0 .sh) -EXP_DIR=$SCRIPT_DIR/$EXP_NAME -LOG_DIR=$EXP_DIR/logs -JSON_METRICS=$EXP_DIR/metrics.json -RUN_LOG=$EXP_DIR/run.log -export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} - -rm -rf $EXP_DIR $LOG_DIR -mkdir -p $EXP_DIR $LOG_DIR - -# Qwen2.5 keeps this L1 aligned with the existing Megatron+vLLM GRPO smoke -# while still exercising the fp32_lm_head config and launch path. -cd $PROJECT_ROOT -uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ - $PROJECT_ROOT/examples/run_grpo.py \ - --config $PROJECT_ROOT/examples/configs/grpo_math_1B_megatron.yaml \ - policy.model_name=Qwen/Qwen2.5-0.5B \ - grpo.num_prompts_per_step=2 \ - grpo.num_generations_per_prompt=4 \ - policy.train_global_batch_size=4 \ - policy.logprob_batch_size=4 \ - policy.train_micro_batch_size=1 \ - policy.megatron_cfg.fp32_lm_head=tf32 \ - policy.generation.vllm_cfg.fp32_lm_head=true \ - cluster.gpus_per_node=2 \ - grpo.max_num_steps=2 \ - logger.tensorboard_enabled=true \ - logger.log_dir=$LOG_DIR \ - logger.wandb_enabled=false \ - logger.monitor_gpus=true \ - checkpointing.enabled=false \ - $@ \ - 2>&1 | tee $RUN_LOG - -grep -q "\[fp32_lm_head\] output layer will compute logits in fp32" $RUN_LOG -grep -Eq "Applied NemotronH fp32 LM head source patch|NemotronH fp32 LM head patch already present" $RUN_LOG - -uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS - -uv run tests/check_metrics.py $JSON_METRICS \ - 'max(data["train/token_mult_prob_error"]) < 1.05' \ - 'min(data["train/probs_ratio_clamped_min"]) > 0.79' \ - 'max(data["train/probs_ratio_clamped_min"]) < 1.21' \ - 'min(data["train/probs_ratio_clamped_max"]) > 0.79' \ - 'max(data["train/probs_ratio_clamped_max"]) < 1.21' diff --git a/tests/unit/models/generation/test_vllm_patches.py b/tests/unit/models/generation/test_vllm_patches.py index 20bf274aeea..d070a7792cb 100644 --- a/tests/unit/models/generation/test_vllm_patches.py +++ b/tests/unit/models/generation/test_vllm_patches.py @@ -41,8 +41,8 @@ from nemo_rl.models.generation.vllm import patches from nemo_rl.models.generation.vllm.config import ( - VLLM_FP32_LM_HEAD_ENV_VAR, - vllm_fp32_lm_head_enabled, + VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, + vllm_nemotron_h_fp32_lm_head_enabled, ) from tests.unit.models.generation.vllm_patch_source_utils import ( write_unpatched_copy, @@ -136,7 +136,7 @@ def compute_logits(self, hidden_states): _NEMOTRON_H_LEGACY_FP32_HEAD_COMPUTE = f""" def compute_logits(self, hidden_states): import os as _os - if _os.environ.get("{VLLM_FP32_LM_HEAD_ENV_VAR}", "0") == "1": + if _os.environ.get("{VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR}", "0") == "1": # NeMo-RL patch: fp32 LM head (MiniMax-M1-style). bf16 rounding of # the logits is the dominant gen/train logprob mismatch source. _fp32_head = getattr(self, "_nrl_lm_head_fp32", None) @@ -368,12 +368,12 @@ def test_glm_decoder_sp_moe_patch_warns_on_unknown_source( ({}, False), ({"fp32_lm_head": False}, False), ({"fp32_lm_head": True}, True), - ({"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}, False), - ({"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "0"}}, False), + ({"env_vars": {VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR: "1"}}, False), + ({"env_vars": {VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR: "0"}}, False), ], ) -def test_vllm_fp32_lm_head_enabled(vllm_cfg, expected): - assert vllm_fp32_lm_head_enabled(vllm_cfg) is expected +def test_vllm_nemotron_h_fp32_lm_head_enabled(vllm_cfg, expected): + assert vllm_nemotron_h_fp32_lm_head_enabled(vllm_cfg) is expected @pytest.mark.parametrize("env_value", [None, "0", "1"]) @@ -381,9 +381,9 @@ def test_nemotron_h_fp32_lm_head_patch_is_env_gated( patched_nemotron_h_source, monkeypatch, env_value ): if env_value is None: - monkeypatch.delenv("NRL_VLLM_FP32_LM_HEAD", raising=False) + monkeypatch.delenv(VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, raising=False) else: - monkeypatch.setenv("NRL_VLLM_FP32_LM_HEAD", env_value) + monkeypatch.setenv(VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, env_value) namespace = {} source = patched_nemotron_h_source.read_text() @@ -421,7 +421,7 @@ def test_nemotron_h_fp32_lm_head_patch_is_env_gated( def test_nemotron_h_fp32_lm_head_patch_migrates_legacy_cached_head_source( patched_legacy_nemotron_h_source, monkeypatch ): - monkeypatch.setenv(VLLM_FP32_LM_HEAD_ENV_VAR, "1") + monkeypatch.setenv(VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, "1") source = patched_legacy_nemotron_h_source.read_text() namespace = {} @@ -490,7 +490,7 @@ def _stub_non_fp32_vllm_patches(monkeypatch, captured_extra_env_vars): @pytest.mark.parametrize("enabled", [False, True]) def test_apply_vllm_patches_gates_nemotron_h_fp32_lm_head(monkeypatch, enabled): _install_fake_vllm_modules(monkeypatch) - monkeypatch.delenv(patches.VLLM_FP32_LM_HEAD_ENV_VAR, raising=False) + monkeypatch.delenv(patches.VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, raising=False) captured_extra_env_vars = [] fp32_patch_calls = [] _stub_non_fp32_vllm_patches(monkeypatch, captured_extra_env_vars) @@ -500,22 +500,24 @@ def test_apply_vllm_patches_gates_nemotron_h_fp32_lm_head(monkeypatch, enabled): lambda _logger: fp32_patch_calls.append(True) or True, ) - patches._apply_vllm_patches("py", extra_env_vars=["USER_VAR"], fp32_lm_head=enabled) + patches._apply_vllm_patches( + "py", extra_env_vars=["USER_VAR"], nemotron_h_fp32_lm_head=enabled + ) assert bool(fp32_patch_calls) is enabled if enabled: - assert os.environ[patches.VLLM_FP32_LM_HEAD_ENV_VAR] == "1" + assert os.environ[patches.VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR] == "1" assert captured_extra_env_vars == [ - ["USER_VAR", patches.VLLM_FP32_LM_HEAD_ENV_VAR] + ["USER_VAR", patches.VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR] ] else: - assert patches.VLLM_FP32_LM_HEAD_ENV_VAR not in os.environ + assert patches.VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR not in os.environ assert captured_extra_env_vars == [["USER_VAR"]] def test_apply_vllm_patches_ignores_ambient_fp32_lm_head_env_toggle(monkeypatch): _install_fake_vllm_modules(monkeypatch) - monkeypatch.setenv(patches.VLLM_FP32_LM_HEAD_ENV_VAR, "1") + monkeypatch.setenv(patches.VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, "1") captured_extra_env_vars = [] fp32_patch_calls = [] _stub_non_fp32_vllm_patches(monkeypatch, captured_extra_env_vars) @@ -528,22 +530,27 @@ def test_apply_vllm_patches_ignores_ambient_fp32_lm_head_env_toggle(monkeypatch) patches._apply_vllm_patches("py") assert fp32_patch_calls == [] - assert patches.VLLM_FP32_LM_HEAD_ENV_VAR not in os.environ + assert patches.VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR not in os.environ assert captured_extra_env_vars == [None] @pytest.mark.parametrize( - ("vllm_cfg_overrides", "expected_fp32_lm_head"), + ("vllm_cfg_overrides", "expected_nemotron_h_fp32_lm_head"), [ ({"env_vars": {"USER_VAR": "value"}, "fp32_lm_head": True}, True), ( - {"env_vars": {"USER_VAR": "value", VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}, + { + "env_vars": { + "USER_VAR": "value", + VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR: "1", + } + }, False, ), ], ) -def test_vllm_worker_threads_fp32_lm_head_cfg_into_source_patches( - monkeypatch, vllm_cfg_overrides, expected_fp32_lm_head +def test_vllm_worker_threads_nemotron_h_fp32_lm_head_cfg_into_source_patches( + monkeypatch, vllm_cfg_overrides, expected_nemotron_h_fp32_lm_head ): from nemo_rl.models.generation.vllm import vllm_worker @@ -551,11 +558,11 @@ def test_vllm_worker_threads_fp32_lm_head_cfg_into_source_patches( monkeypatch.setattr( vllm_worker, "_apply_vllm_patches", - lambda py, *, extra_env_vars, fp32_lm_head: patch_calls.append( + lambda py, *, extra_env_vars, nemotron_h_fp32_lm_head: patch_calls.append( { "py": py, "extra_env_vars": extra_env_vars, - "fp32_lm_head": fp32_lm_head, + "nemotron_h_fp32_lm_head": nemotron_h_fp32_lm_head, } ), ) @@ -579,7 +586,7 @@ def test_vllm_worker_threads_fp32_lm_head_cfg_into_source_patches( { "py": sys.executable, "extra_env_vars": ["EXPLICIT_VAR"], - "fp32_lm_head": expected_fp32_lm_head, + "nemotron_h_fp32_lm_head": expected_nemotron_h_fp32_lm_head, } ] diff --git a/tests/unit/models/policy/test_policy_validation.py b/tests/unit/models/policy/test_policy_validation.py index 83195c5c6f6..e15f0b96dea 100644 --- a/tests/unit/models/policy/test_policy_validation.py +++ b/tests/unit/models/policy/test_policy_validation.py @@ -20,11 +20,14 @@ when the cluster size is insufficient for the specified parallelism configuration. """ +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest -from nemo_rl.models.generation.vllm.config import VLLM_FP32_LM_HEAD_ENV_VAR +from nemo_rl.models.generation.vllm.config import ( + VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, +) from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.lm_policy import Policy @@ -205,8 +208,13 @@ def set_vllm_generation( return config -def construct_policy_with_mocks(config: PolicyConfig) -> Policy: - model_config = MagicMock() +def construct_policy_with_mocks( + config: PolicyConfig, model_config: object | None = None +) -> Policy: + if model_config is None: + model_config = SimpleNamespace( + architectures=["NemotronHForCausalLM"], model_type="nemotron_h" + ) with ( patch.dict("os.environ", {"TORCH_CUDA_ARCH_LIST": "9.0"}), patch("nemo_rl.models.policy.lm_policy.RayQueue"), @@ -235,6 +243,22 @@ def test_policy_accepts_matched_vllm_and_megatron_fp32_lm_head(trainer_fp32): assert policy.worker_group is not None +def test_policy_warns_when_vllm_fp32_lm_head_model_is_not_nemotron_h(): + config = create_megatron_config("test-model", tp=1) + config["megatron_cfg"]["fp32_lm_head"] = True + set_vllm_generation(config, {"fp32_lm_head": True}) + + with pytest.warns(UserWarning, match="Nemotron-H"): + policy = construct_policy_with_mocks( + config, + model_config=SimpleNamespace( + architectures=["Qwen2ForCausalLM"], model_type="qwen2" + ), + ) + + assert policy.worker_group is not None + + @pytest.mark.parametrize( ("trainer_fp32", "vllm_fp32"), [ @@ -290,7 +314,9 @@ def test_policy_accepts_vllm_fp32_lm_head_disabled_with_dtensor_trainer(): def test_policy_rejects_fp32_lm_head_env_var_toggle(): config = create_dtensor_config("test-model", tp=1) - set_vllm_generation(config, {"env_vars": {VLLM_FP32_LM_HEAD_ENV_VAR: "1"}}) + set_vllm_generation( + config, {"env_vars": {VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR: "1"}} + ) with ( patch("nemo_rl.models.policy.lm_policy.RayWorkerGroup") as worker_group, From c9ab6062aa8f65f9c17f1080b896f9f8c7270ce1 Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 11:32:31 -0700 Subject: [PATCH 11/18] test: add nemotron h fp32 lm head l1 smoke Signed-off-by: Guyue Huang --- nemo_rl/models/generation/vllm/patches.py | 7 ++ .../functional/L1_Functional_Tests_Other_1.sh | 1 + .../vllm_nemotron_h_fp32_lm_head.py | 93 +++++++++++++++++++ .../vllm_nemotron_h_fp32_lm_head.sh | 35 +++++++ .../models/generation/test_vllm_patches.py | 1 + 5 files changed, 137 insertions(+) create mode 100644 tests/functional/vllm_nemotron_h_fp32_lm_head.py create mode 100755 tests/functional/vllm_nemotron_h_fp32_lm_head.sh diff --git a/nemo_rl/models/generation/vllm/patches.py b/nemo_rl/models/generation/vllm/patches.py index 98125bca06c..5bcca6b7011 100644 --- a/nemo_rl/models/generation/vllm/patches.py +++ b/nemo_rl/models/generation/vllm/patches.py @@ -703,6 +703,13 @@ def _patch_vllm_nemotron_h_fp32_lm_head(logger) -> bool: def _nrl_fp32_lm_head_forward( input_, embedding_bias=None, _lm_head=self.lm_head ): + if not getattr(_lm_head, "_nrl_fp32_lm_head_forward_logged", False): + print( + "[fp32_lm_head] NemotronH vLLM lm_head.forward casts " + "input and weight to fp32", + flush=True, + ) + _lm_head._nrl_fp32_lm_head_forward_logged = True logits = torch.matmul( input_.to(dtype=torch.float32), _lm_head.weight.to(dtype=torch.float32).t(), diff --git a/tests/functional/L1_Functional_Tests_Other_1.sh b/tests/functional/L1_Functional_Tests_Other_1.sh index 8ebcb08326e..74ecbd732c1 100644 --- a/tests/functional/L1_Functional_Tests_Other_1.sh +++ b/tests/functional/L1_Functional_Tests_Other_1.sh @@ -39,6 +39,7 @@ run_test bash ./tests/functional/test_frozen_env.sh run_test fast uv run --no-sync bash ./tests/functional/test_converters.sh run_test uv run --no-sync bash ./tests/functional/test_decode_vs_prefill.sh +run_test uv run --no-sync bash ./tests/functional/vllm_nemotron_h_fp32_lm_head.sh run_test uv run --no-sync bash ./tests/functional/test_mcore_extra_installed_correctly.sh # Research functional tests (self-discovery) diff --git a/tests/functional/vllm_nemotron_h_fp32_lm_head.py b/tests/functional/vllm_nemotron_h_fp32_lm_head.py new file mode 100644 index 00000000000..2e5cbeea6a6 --- /dev/null +++ b/tests/functional/vllm_nemotron_h_fp32_lm_head.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from nemo_rl.algorithms.utils import get_tokenizer +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.virtual_cluster import RayVirtualCluster +from nemo_rl.models.generation import configure_generation_config +from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration + +MODEL_NAME = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" + + +def main() -> None: + config: VllmConfig = { + "backend": "vllm", + "model_name": MODEL_NAME, + "tokenizer": {"name": MODEL_NAME}, + "max_new_tokens": 4, + "temperature": 1.0, + "top_p": 1.0, + "top_k": None, + "val_temperature": 1.0, + "val_top_p": 1.0, + "val_top_k": None, + "stop_token_ids": None, + "stop_strings": None, + "vllm_cfg": { + "precision": "bfloat16", + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "expert_parallel_size": 1, + "gpu_memory_utilization": 0.8, + "max_model_len": 256, + "async_engine": False, + "skip_tokenizer_init": False, + "load_format": "auto", + "enforce_eager": True, + "kv_cache_dtype": "auto", + "fp32_lm_head": True, + "use_tqdm": False, + }, + "vllm_kwargs": { + "mamba_ssm_cache_dtype": "float32", + "compilation_config": {"backend": "eager"}, + }, + "colocated": { + "enabled": True, + "resources": { + "gpus_per_node": None, + "num_nodes": None, + }, + }, + } + + tokenizer = get_tokenizer(config["tokenizer"]) + config = configure_generation_config(config, tokenizer, is_eval=True) + cluster = RayVirtualCluster( + bundle_ct_per_node_list=[1], + use_gpus=True, + max_colocated_worker_groups=1, + num_gpus_per_node=1, + name="vllm-nemotron-h-fp32-lm-head-functional", + ) + vllm_generation = None + try: + vllm_generation = VllmGeneration(cluster, config) + output = vllm_generation.generate_text( + BatchedDataDict({"prompts": ["The capital of France is"]}), + greedy=True, + ) + texts = output["texts"] + assert len(texts) == 1 + assert texts[0], "Nemotron-H vLLM generation returned an empty string" + print(f"[PASS] Nemotron-H fp32 lm_head generated text: {texts[0]!r}") + finally: + if vllm_generation is not None: + vllm_generation.shutdown() + cluster.shutdown() + + +if __name__ == "__main__": + main() diff --git a/tests/functional/vllm_nemotron_h_fp32_lm_head.sh b/tests/functional/vllm_nemotron_h_fp32_lm_head.sh new file mode 100755 index 00000000000..436a5b2dbf3 --- /dev/null +++ b/tests/functional/vllm_nemotron_h_fp32_lm_head.sh @@ -0,0 +1,35 @@ +#!/bin/bash + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath "$SCRIPT_DIR/../..") +# Mark the current repo as safe, since wandb fetches metadata about the repo +git config --global --add safe.directory "$PROJECT_ROOT" + +set -eou pipefail + +EXP_NAME=$(basename "$0" .sh) +EXP_DIR=$SCRIPT_DIR/$EXP_NAME +RUN_LOG=$EXP_DIR/run.log +export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} + +rm -rf "$EXP_DIR" +mkdir -p "$EXP_DIR" + +assert_grep() { + local pattern=$1 + local file=$2 + grep -Eq "$pattern" "$file" || { + echo "[FAIL] expected '$pattern' in $file" + exit 1 + } +} + +cd "$PROJECT_ROOT" +uv run --extra vllm coverage run -a --data-file="$PROJECT_ROOT/tests/.coverage" --source="$PROJECT_ROOT/nemo_rl" \ + "$PROJECT_ROOT/tests/functional/vllm_nemotron_h_fp32_lm_head.py" \ + "$@" \ + 2>&1 | tee "$RUN_LOG" + +assert_grep "Applied NemotronH fp32 LM head source patch|NemotronH fp32 LM head patch already present" "$RUN_LOG" +assert_grep "\\[fp32_lm_head\\] NemotronH vLLM lm_head.forward casts input and weight to fp32" "$RUN_LOG" +assert_grep "\\[PASS\\] Nemotron-H fp32 lm_head generated text" "$RUN_LOG" diff --git a/tests/unit/models/generation/test_vllm_patches.py b/tests/unit/models/generation/test_vllm_patches.py index d070a7792cb..442f309f26e 100644 --- a/tests/unit/models/generation/test_vllm_patches.py +++ b/tests/unit/models/generation/test_vllm_patches.py @@ -414,6 +414,7 @@ def test_nemotron_h_fp32_lm_head_patch_is_env_gated( assert "deepcopy" not in source assert "params_dtype=torch.float32" not in source + assert "NemotronH vLLM lm_head.forward casts input and weight to fp32" in source assert "torch.matmul(" in source ast.parse(source) From b1f3afb3608de58893b58ce0b1107877d412120f Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 12:07:08 -0700 Subject: [PATCH 12/18] fix: initialize ray in nemotron h fp32 lm head l1 Signed-off-by: Guyue Huang --- tests/functional/vllm_nemotron_h_fp32_lm_head.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/functional/vllm_nemotron_h_fp32_lm_head.py b/tests/functional/vllm_nemotron_h_fp32_lm_head.py index 2e5cbeea6a6..0102846f0a7 100644 --- a/tests/functional/vllm_nemotron_h_fp32_lm_head.py +++ b/tests/functional/vllm_nemotron_h_fp32_lm_head.py @@ -12,9 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os + +import ray + from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.distributed.batched_data_dict import BatchedDataDict -from nemo_rl.distributed.virtual_cluster import RayVirtualCluster +from nemo_rl.distributed.virtual_cluster import RayVirtualCluster, init_ray from nemo_rl.models.generation import configure_generation_config from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration @@ -65,6 +69,7 @@ def main() -> None: tokenizer = get_tokenizer(config["tokenizer"]) config = configure_generation_config(config, tokenizer, is_eval=True) + init_ray(log_dir=os.environ.get("RAY_TMPDIR")) cluster = RayVirtualCluster( bundle_ct_per_node_list=[1], use_gpus=True, @@ -87,6 +92,7 @@ def main() -> None: if vllm_generation is not None: vllm_generation.shutdown() cluster.shutdown() + ray.shutdown() if __name__ == "__main__": From 35fb4dc41bb68e7d4bb4553119087571c1c500a4 Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 12:13:39 -0700 Subject: [PATCH 13/18] fix: use short ray temp dir in nemotron h l1 Signed-off-by: Guyue Huang --- .../vllm_nemotron_h_fp32_lm_head.py | 49 ++++++++++--------- 1 file changed, 25 insertions(+), 24 deletions(-) diff --git a/tests/functional/vllm_nemotron_h_fp32_lm_head.py b/tests/functional/vllm_nemotron_h_fp32_lm_head.py index 0102846f0a7..c3f4a7ee034 100644 --- a/tests/functional/vllm_nemotron_h_fp32_lm_head.py +++ b/tests/functional/vllm_nemotron_h_fp32_lm_head.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os +import tempfile import ray @@ -69,30 +69,31 @@ def main() -> None: tokenizer = get_tokenizer(config["tokenizer"]) config = configure_generation_config(config, tokenizer, is_eval=True) - init_ray(log_dir=os.environ.get("RAY_TMPDIR")) - cluster = RayVirtualCluster( - bundle_ct_per_node_list=[1], - use_gpus=True, - max_colocated_worker_groups=1, - num_gpus_per_node=1, - name="vllm-nemotron-h-fp32-lm-head-functional", - ) - vllm_generation = None - try: - vllm_generation = VllmGeneration(cluster, config) - output = vllm_generation.generate_text( - BatchedDataDict({"prompts": ["The capital of France is"]}), - greedy=True, + with tempfile.TemporaryDirectory(prefix="nrl-ray-") as ray_log_dir: + init_ray(log_dir=ray_log_dir) + cluster = RayVirtualCluster( + bundle_ct_per_node_list=[1], + use_gpus=True, + max_colocated_worker_groups=1, + num_gpus_per_node=1, + name="vllm-nemotron-h-fp32-lm-head-functional", ) - texts = output["texts"] - assert len(texts) == 1 - assert texts[0], "Nemotron-H vLLM generation returned an empty string" - print(f"[PASS] Nemotron-H fp32 lm_head generated text: {texts[0]!r}") - finally: - if vllm_generation is not None: - vllm_generation.shutdown() - cluster.shutdown() - ray.shutdown() + vllm_generation = None + try: + vllm_generation = VllmGeneration(cluster, config) + output = vllm_generation.generate_text( + BatchedDataDict({"prompts": ["The capital of France is"]}), + greedy=True, + ) + texts = output["texts"] + assert len(texts) == 1 + assert texts[0], "Nemotron-H vLLM generation returned an empty string" + print(f"[PASS] Nemotron-H fp32 lm_head generated text: {texts[0]!r}") + finally: + if vllm_generation is not None: + vllm_generation.shutdown() + cluster.shutdown() + ray.shutdown() if __name__ == "__main__": From 04f7ebd54315cefd020cd52b317f17f17d9d265c Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 12:19:48 -0700 Subject: [PATCH 14/18] fix: keep nemotron h l1 ray sockets under tmp Signed-off-by: Guyue Huang --- tests/functional/vllm_nemotron_h_fp32_lm_head.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/functional/vllm_nemotron_h_fp32_lm_head.py b/tests/functional/vllm_nemotron_h_fp32_lm_head.py index c3f4a7ee034..7bc8faeb8aa 100644 --- a/tests/functional/vllm_nemotron_h_fp32_lm_head.py +++ b/tests/functional/vllm_nemotron_h_fp32_lm_head.py @@ -69,7 +69,7 @@ def main() -> None: tokenizer = get_tokenizer(config["tokenizer"]) config = configure_generation_config(config, tokenizer, is_eval=True) - with tempfile.TemporaryDirectory(prefix="nrl-ray-") as ray_log_dir: + with tempfile.TemporaryDirectory(prefix="nrl-ray-", dir="/tmp") as ray_log_dir: init_ray(log_dir=ray_log_dir) cluster = RayVirtualCluster( bundle_ct_per_node_list=[1], From 18b9e24ef76d816b40d31e7520a5def69f36c290 Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 12:29:11 -0700 Subject: [PATCH 15/18] fix: assert nemotron h runtime markers in l1 Signed-off-by: Guyue Huang --- tests/functional/vllm_nemotron_h_fp32_lm_head.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/functional/vllm_nemotron_h_fp32_lm_head.sh b/tests/functional/vllm_nemotron_h_fp32_lm_head.sh index 436a5b2dbf3..dd607611f04 100755 --- a/tests/functional/vllm_nemotron_h_fp32_lm_head.sh +++ b/tests/functional/vllm_nemotron_h_fp32_lm_head.sh @@ -30,6 +30,6 @@ uv run --extra vllm coverage run -a --data-file="$PROJECT_ROOT/tests/.coverage" "$@" \ 2>&1 | tee "$RUN_LOG" -assert_grep "Applied NemotronH fp32 LM head source patch|NemotronH fp32 LM head patch already present" "$RUN_LOG" +assert_grep "Resolved architecture: NemotronHForCausalLM" "$RUN_LOG" assert_grep "\\[fp32_lm_head\\] NemotronH vLLM lm_head.forward casts input and weight to fp32" "$RUN_LOG" assert_grep "\\[PASS\\] Nemotron-H fp32 lm_head generated text" "$RUN_LOG" From a59da2e489903217baa2b3a86c9c97efcc8bcff5 Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 13:19:48 -0700 Subject: [PATCH 16/18] fix: satisfy fp32 lm head lint checks Signed-off-by: Guyue Huang --- nemo_rl/models/generation/vllm/patches.py | 5 ++--- nemo_rl/models/policy/utils.py | 2 ++ 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/nemo_rl/models/generation/vllm/patches.py b/nemo_rl/models/generation/vllm/patches.py index 5bcca6b7011..376e3106b12 100644 --- a/nemo_rl/models/generation/vllm/patches.py +++ b/nemo_rl/models/generation/vllm/patches.py @@ -929,9 +929,8 @@ def _apply_vllm_patches( _patch_vllm_shm_broadcast_bind_retry(patch_logger) _patch_vllm_radio_layerscale_loader(patch_logger) _patch_vllm_glm_decoder_sequence_parallel_moe(patch_logger) - if ( - nemotron_h_fp32_lm_head_enabled - and not _patch_vllm_nemotron_h_fp32_lm_head(patch_logger) + if nemotron_h_fp32_lm_head_enabled and not _patch_vllm_nemotron_h_fp32_lm_head( + patch_logger ): raise RuntimeError( "vllm_cfg.fp32_lm_head is enabled, but that flag currently maps to " diff --git a/nemo_rl/models/policy/utils.py b/nemo_rl/models/policy/utils.py index 40568b1bb1b..492d525a930 100644 --- a/nemo_rl/models/policy/utils.py +++ b/nemo_rl/models/policy/utils.py @@ -57,6 +57,7 @@ from nemo_rl.distributed.worker_group_utils import get_nsight_config_if_pattern_matches from nemo_rl.models.generation.vllm.config import ( VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, + VllmSpecificArgs, vllm_nemotron_h_fp32_lm_head_enabled, ) @@ -223,6 +224,7 @@ def validate_fp32_lm_head_config( vllm_cfg = generation_config.get("vllm_cfg") if vllm_cfg is None: return + vllm_cfg = cast(VllmSpecificArgs | dict[str, Any], vllm_cfg) env_vars = vllm_cfg.get("env_vars") or {} if VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR in env_vars: From 3143fb90b77a141977286ab5c53d170b01022a8e Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 14:47:03 -0700 Subject: [PATCH 17/18] fix: mark fp32 lm head tests as mcore Signed-off-by: Guyue Huang --- tests/unit/models/megatron/test_megatron_setup.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index d1368fc27e1..26e873162d7 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -3998,6 +3998,7 @@ def _assert_fp32_wrapped(output_layer: _FakeOutputLayer) -> None: assert output_layer.seen_dtypes == [(torch.float32, torch.float32)] +@pytest.mark.mcore def test_apply_fp32_lm_head_wraps_plain_last_stage_chunk(): from nemo_rl.models.megatron.setup import apply_fp32_lm_head @@ -4010,6 +4011,7 @@ def test_apply_fp32_lm_head_wraps_plain_last_stage_chunk(): _assert_fp32_wrapped(layer) +@pytest.mark.mcore def test_apply_fp32_lm_head_tf32_path_produces_fp32_output(): from nemo_rl.models.megatron.setup import apply_fp32_lm_head @@ -4022,6 +4024,7 @@ def test_apply_fp32_lm_head_tf32_path_produces_fp32_output(): _assert_fp32_wrapped(layer) +@pytest.mark.mcore def test_apply_fp32_lm_head_is_idempotent(): from nemo_rl.models.megatron.setup import apply_fp32_lm_head @@ -4065,6 +4068,7 @@ def test_apply_fp32_lm_head_is_idempotent(): ], ids=["vl_llava", "vl_language_model", "omni_thinker"], ) +@pytest.mark.mcore def test_apply_fp32_lm_head_resolves_nested_language_model(build): from nemo_rl.models.megatron.setup import apply_fp32_lm_head @@ -4074,6 +4078,7 @@ def test_apply_fp32_lm_head_resolves_nested_language_model(build): _assert_fp32_wrapped(layer) +@pytest.mark.mcore def test_apply_fp32_lm_head_raises_when_post_process_chunk_has_no_output_layer(): from nemo_rl.models.megatron.setup import apply_fp32_lm_head @@ -4085,6 +4090,7 @@ def test_apply_fp32_lm_head_raises_when_post_process_chunk_has_no_output_layer() apply_fp32_lm_head([chunk]) +@pytest.mark.mcore def test_apply_fp32_lm_head_skips_non_last_pipeline_stage(): from nemo_rl.models.megatron.setup import apply_fp32_lm_head From 723bc95fdfa26763330606feda97039ae8f3d7d8 Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Wed, 16 Sep 2026 18:41:29 -0700 Subject: [PATCH 18/18] test: fix fp32 lm head patch source assertion Signed-off-by: Guyue Huang --- tests/unit/models/generation/test_vllm_patches.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/models/generation/test_vllm_patches.py b/tests/unit/models/generation/test_vllm_patches.py index 442f309f26e..7bc8e8a0209 100644 --- a/tests/unit/models/generation/test_vllm_patches.py +++ b/tests/unit/models/generation/test_vllm_patches.py @@ -414,7 +414,8 @@ def test_nemotron_h_fp32_lm_head_patch_is_env_gated( assert "deepcopy" not in source assert "params_dtype=torch.float32" not in source - assert "NemotronH vLLM lm_head.forward casts input and weight to fp32" in source + assert "NemotronH vLLM lm_head.forward casts " in source + assert "input and weight to fp32" in source assert "torch.matmul(" in source ast.parse(source)