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 a8f7ae9b131..464f2081e31 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -449,6 +449,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 8c107cb2402..a1496ca8cbd 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -296,6 +296,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 85579ac4c9b..35c3d05568a 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -28,6 +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_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"}) @@ -61,6 +66,10 @@ 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"]] + # 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 # the training prompt longer than its text-only representation. @@ -112,6 +121,13 @@ class VllmSpecificArgs(TypedDict): reasoning_parser_plugin: NotRequired[str] +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")) + + 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..376e3106b12 100644 --- a/nemo_rl/models/generation/vllm/patches.py +++ b/nemo_rl/models/generation/vllm/patches.py @@ -16,6 +16,10 @@ from contextlib import contextmanager from importlib.util import find_spec +from nemo_rl.models.generation.vllm.config import ( + VLLM_NEMOTRON_H_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 +626,229 @@ 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 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. + """ + 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_NEMOTRON_H_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_NEMOTRON_H_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 + ): + 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(), + ) + 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_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) + 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 +870,22 @@ def _apply_vllm_patches( py_executable: str, *, extra_env_vars: list[str] | 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") + 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_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 @@ -692,3 +929,12 @@ 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 + ): + raise RuntimeError( + "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 225b6cf5b6f..d08b1e429f1 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_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 ( @@ -438,9 +439,11 @@ def _init_config( # Store the Python executable being used by this worker self.py_executable = sys.executable + vllm_cfg = self.cfg["vllm_cfg"] _apply_vllm_patches( self.py_executable, extra_env_vars=extra_env_vars, + 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/megatron/setup.py b/nemo_rl/models/megatron/setup.py index ff5793a753c..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(): @@ -598,6 +600,94 @@ def _resolve_iter_dir_from_root( return os.path.join(path, iter_subdirs[-1]) +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 + 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] + 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 + if getattr(output_layer.forward, _FP32_LM_HEAD_PATCHED_ATTR, False): + 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 + + 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" + + (" (tf32 tensor cores)" if use_tf32 else "") + ) + + def _resolve_peft_restore_dir(restore_from: str) -> str: """Resolve a ``megatron_cfg.peft.restore_from`` path to an iteration directory. diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 3325edfbfb4..24610f8ce03 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -500,6 +500,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 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"]] # 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/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index a9e02f94f58..a19b2fea291 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -51,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 ( @@ -156,6 +157,25 @@ 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 + ) + 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 " @@ -412,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 96ddd6ebd43..492d525a930 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,14 @@ 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_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, + VllmSpecificArgs, + vllm_nemotron_h_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. @@ -120,6 +128,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``. @@ -133,6 +144,130 @@ 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, + model_config: object | None = None, +) -> 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 + 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: + raise ValueError( + 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_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 " + "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." + ) + 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( model_name: str, *, diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 30ee2043850..7f773a10d9c 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -84,6 +84,7 @@ router_replay_enabled, ) from nemo_rl.models.megatron.setup import ( + apply_fp32_lm_head, build_inference_model, finalize_megatron_setup, handle_model_import, @@ -716,6 +717,11 @@ def __init__( self.mcore_state = model_and_optimizer_state.state self.model = model_and_optimizer_state.model + 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/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..7bc8faeb8aa --- /dev/null +++ b/tests/functional/vllm_nemotron_h_fp32_lm_head.py @@ -0,0 +1,100 @@ +# 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. + +import tempfile + +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, init_ray +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) + 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], + 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() + ray.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..dd607611f04 --- /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 "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" diff --git a/tests/unit/models/generation/test_vllm_patches.py b/tests/unit/models/generation/test_vllm_patches.py index a26a3d8ef11..7bc8e8a0209 100644 --- a/tests/unit/models/generation/test_vllm_patches.py +++ b/tests/unit/models/generation/test_vllm_patches.py @@ -33,10 +33,17 @@ import ast import logging import os +import sys +import types import pytest +import torch from nemo_rl.models.generation.vllm import patches +from nemo_rl.models.generation.vllm.config import ( + 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, ) @@ -50,6 +57,131 @@ _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 +""" +_NEMOTRON_H_LEGACY_FP32_HEAD_COMPUTE = f""" def compute_logits(self, hidden_states): + import os as _os + + 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) + 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 @@ -81,6 +213,24 @@ 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.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, @@ -212,6 +362,237 @@ 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_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR: "1"}}, False), + ({"env_vars": {VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR: "0"}}, False), + ], +) +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"]) +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(VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, raising=False) + else: + monkeypatch.setenv(VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, 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 "NemotronH vLLM lm_head.forward casts " in source + assert "input and weight to fp32" in source + assert "torch.matmul(" in source + 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_NEMOTRON_H_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 +): + 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_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) + 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"], nemotron_h_fp32_lm_head=enabled + ) + + assert bool(fp32_patch_calls) is enabled + if enabled: + assert os.environ[patches.VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR] == "1" + assert captured_extra_env_vars == [ + ["USER_VAR", patches.VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR] + ] + else: + 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_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) + 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 == [] + 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_nemotron_h_fp32_lm_head"), + [ + ({"env_vars": {"USER_VAR": "value"}, "fp32_lm_head": True}, True), + ( + { + "env_vars": { + "USER_VAR": "value", + VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR: "1", + } + }, + False, + ), + ], +) +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 + + patch_calls = [] + monkeypatch.setattr( + vllm_worker, + "_apply_vllm_patches", + lambda py, *, extra_env_vars, nemotron_h_fp32_lm_head: patch_calls.append( + { + "py": py, + "extra_env_vars": extra_env_vars, + "nemotron_h_fp32_lm_head": nemotron_h_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", + **vllm_cfg_overrides, + }, + }, + extra_env_vars=["EXPLICIT_VAR"], + ) + + assert patch_calls == [ + { + "py": sys.executable, + "extra_env_vars": ["EXPLICIT_VAR"], + "nemotron_h_fp32_lm_head": expected_nemotron_h_fp32_lm_head, + } + ] + + @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 eb26b886d03..26e873162d7 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -3972,6 +3972,132 @@ 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: + 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)] + + +@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 + + # 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.mcore +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 getattr(layer.forward, "_nrl_fp32_lm_head_use_tf32") is True + _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 + + 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) + + +@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"], +) +@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 + + layer = _FakeOutputLayer() + chunk = SimpleNamespace(module=build(layer)) + apply_fp32_lm_head([chunk]) + _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 + + # 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]) + + +@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 + + chunk = SimpleNamespace(module=SimpleNamespace(post_process=False)) + apply_fp32_lm_head([chunk]) # no output_layer, not post_process: no-op + + @pytest.mark.mcore class TestForceSyncOptimizerFp32FromModel: """Tests for _force_sync_optimizer_fp32_from_model. diff --git a/tests/unit/models/policy/test_policy_validation.py b/tests/unit/models/policy/test_policy_validation.py index d8fc112367e..e15f0b96dea 100644 --- a/tests/unit/models/policy/test_policy_validation.py +++ b/tests/unit/models/policy/test_policy_validation.py @@ -20,10 +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_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, +) from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.lm_policy import Policy @@ -171,6 +175,191 @@ 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, 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"), + 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 + + +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"), + [ + ("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_NEMOTRON_H_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() 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 d4d92ad603c..958bdbfbf8e 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -428,6 +428,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 168b36d0b65..c138044e4de 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -273,6 +273,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