`nemo_rl.models.generation.vllm.patches`#

Module Contents#

Functions#

`_get_vllm_file`	Return absolute path to a vLLM file or raise if it cannot be found.
`_locked_file_patch`	Yield (content, writer) under an exclusive file lock.
`_patch_vllm_init_workers_ray`	Patch vLLM’s Ray executor env propagation and worker runtime_env.
`_patch_vllm_llama_eagle3_own_lm_head`	Patch LlamaEagle3 to keep truncated draft lm_head ownership.
`_patch_vllm_tool_parser_namespace_tool`	Guard vLLM’s NamespaceTool import for openai < 2.25.
`_patch_vllm_ray_executor_v2_tcpstore_port`	Keep RayExecutorV2’s TCPStore port out of the MessageQueue’s scan range.
`_patch_vllm_shm_broadcast_bind_retry`	Keep MessageQueue’s remote socket in the reserved band with bind retries.
`_patch_vllm_radio_layerscale_loader`	Load explicit RADIO LayerScale weights and initialize folded weights.
`_patch_vllm_radio_native_loader`	Accept transformers-native RADIO weight names in vLLM’s RADIO loader.
`_patch_vllm_radio_final_layernorm`	Port checkpoint-backed Super Omni final vision LayerNorm to vLLM 0.29.
`_patch_vllm_glm_decoder_sequence_parallel_moe`	Restore the vLLM 0.24 decoder boundary for GLM DSA models.
`_patch_vllm_moe_routed_experts_capture`	Fire the routed-experts capture hook on the monolithic fused-MoE path.
`_patch_vllm_routed_experts_capture_router_fallback`	Let monolithic MoE kernels without in-kernel capture use the router hook.
`_patch_vllm_nemotron_h_fp32_lm_head`	Compute NemotronH logits with an fp32 LM head (MiniMax-M1-style).
`modelopt_moe_amax_aliases`	Temporarily expose nested ModelOpt MoE amax buffers to vLLM’s loader.
`ensure_vllm_source_compat`	Apply interpreter-independent vLLM source-compat patches.
`_apply_vllm_patches`

Data#

VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR

API#

nemo_rl.models.generation.vllm.patches.VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR#: ‘NRL_VLLM_FP32_LM_HEAD’

nemo_rl.models.generation.vllm.patches._get_vllm_file(relative_path: str) → str#

Return absolute path to a vLLM file or raise if it cannot be found.

The relative_path should be a POSIX-style path under the vllm package root, e.g. “v1/executor/ray_executor.py” or “attention/layer.py”.

nemo_rl.models.generation.vllm.patches._locked_file_patch(file_path: str)#: Yield (content, writer) under an exclusive file lock.

nemo_rl.models.generation.vllm.patches._patch_vllm_init_workers_ray( py_executable: str, extra_env_vars: list[str] | None, ) → bool#

Patch vLLM’s Ray executor env propagation and worker runtime_env.

Pass custom runtime_env in _init_workers_ray call (file patch).
- This allows passing custom py_executable to worker initialization.
Forward extra env vars to the Ray workers via vLLM’s additive VLLM_RAY_EXTRA_ENV_VARS_TO_COPY hook (vLLM >= 0.25). NCCL_, HF_, and HUGGING_FACE_* vars are already copied by vLLM’s default prefix list (this includes the NCCL_CUMEM_ENABLE/NCCL_NVLS_ENABLE workaround from https://github.com/NVIDIA-NeMo/RL/pull/898).

.. note:: Step 1 patches the v1 Ray executor, which vLLM 0.25 no longer selects by default: VLLM_USE_RAY_V2_EXECUTOR_BACKEND flipped from "0" (0.20) to "1" (0.25), so Executor.get_class returns RayExecutorV2 for ray-backed engines. RayExecutorV2 has no _init_workers_ray at all – it creates workers inline, and its _build_runtime_env never sets py_executable.

The patch is kept because it is still load-bearing when
``VLLM_USE_RAY_V2_EXECUTOR_BACKEND=0`` selects the v1 executor. Under
the 0.25 default it is inert, and workers get the right interpreter
from Ray's per-field ``runtime_env`` inheritance instead: the parent
NeMo-RL actor sets ``py_executable``, and a child created with a
``runtime_env`` that omits it inherits the parent's value.

So a ``True`` return means "the anchor is in place", not "this is what
put the workers on the right interpreter". The caller logs
accordingly.

Returns:: Whether the v1 runtime_env source patch is in place. The env-var merge in step 2 cannot fail, but step 1 is anchored on a call-site string; if that moves upstream the py_executable injection silently stops happening, so the caller must not report success unconditionally.

nemo_rl.models.generation.vllm.patches._patch_vllm_llama_eagle3_own_lm_head(logger) → None#: Patch LlamaEagle3 to keep truncated draft lm_head ownership.

nemo_rl.models.generation.vllm.patches._patch_vllm_tool_parser_namespace_tool(logger) → None#

Guard vLLM’s NamespaceTool import for openai < 2.25.

vLLM 0.25 imports openai.types.responses.NamespaceTool (added in openai 2.25.0) at the top of tool_parsers/utils.py, but nemo-gym pins openai<=2.7.2 and its child server venvs must match the parent’s openai version exactly. NamespaceTool is only used in isinstance checks for Responses-API namespace tools, which cannot be constructed by an openai client that predates the feature, so a never-matching stub is a faithful fallback.

nemo_rl.models.generation.vllm.patches._patch_vllm_ray_executor_v2_tcpstore_port(logger) → None#

Keep RayExecutorV2’s TCPStore port out of the MessageQueue’s scan range.

vLLM 0.25’s RayExecutorV2._init_executor picks the torch.distributed TCPStore port with a bind-probe (Step 3) but only binds it much later, in the rank-0 worker’s init_process_group. In between, Step 4 builds the broadcast MessageQueue; when the engine spans nodes that queue needs a real TCP socket, so it calls get_open_port() and binds and holds the result (shm_broadcast.py: remote_subscribe_port = get_open_port() then remote_socket.bind(...)). Both searches start at VLLM_PORT, so the queue deterministically takes the very port the probe just released and engine startup dies with EADDRINUSE (DeepSeek-V3 generation TP=32, observed on port 7000). Engines that fit on one node use a shm/ipc socket instead and never allocate a TCP port here, which is why only node-spanning engines are affected.

Offsetting the TCPStore search past the queue’s scan range removes the collision while keeping both ports inside the engine’s 100-port window, and therefore below the OS ephemeral floor. That band is deliberate: leaving VLLM_PORT unset would send vLLM to kernel-assigned ephemeral ports and reintroduce the TOCTOU contention this layout exists to prevent (#2380, #3103).

The offset must be applied before the local_dp_rank is None test, not inside it. vLLM’s own disjoint-window branch below reads as if it only applies to DP engines, but ParallelConfig.__post_init__ takes the “offline SPMD” path for every engine NeMo-RL builds and assigns data_parallel_rank_local = envs.VLLM_DP_RANK_LOCAL (0 by default) and data_parallel_master_port = envs.VLLM_DP_MASTER_PORT (0 by default). So a plain non-DP engine arrives here with local_dp_rank=0, not None: the None branch is dead, and the DP branch searches from 0 + 100 + 0 * 32 = 100, fails all 32 attempts on the privileged range, and falls through to get_open_port() — straight back to VLLM_PORT. That is exactly the port the MessageQueue takes. See RL-1104.

vLLM 0.29 fixes the race upstream (vllm-project/vllm#53666, #50969): the rank-0 actor now binds the TCPStore itself, on a kernel-assigned port, and holds that socket (self._dist_init_store = store) until init_process_group reuses it, so there is no probe/bind window for the MessageQueue to land in. That is not the TOCTOU pattern the reserved band guards against (the port is never released between selection and use), and _select_tcpstore_port no longer exists to patch. When that upstream marker is present this function logs at info level and leaves the file alone.

Returns without raising when neither form is found, but logs at warning level so a silent no-op is visible in worker logs.

nemo_rl.models.generation.vllm.patches._patch_vllm_shm_broadcast_bind_retry(logger) → None#

Keep MessageQueue’s remote socket in the reserved band with bind retries.

vLLM 0.28 binds port zero directly, avoiding the old probe/bind race but ignoring VLLM_PORT. Restore reserved-band selection with retries so engine sockets do not consume the ephemeral ports used by other services. A probe alone releases its socket before ZMQ binds, leaving a TOCTOU race.

On vLLM 0.25 that race is lost reliably, not occasionally. Every RayWorkerProc on a non-driver node takes n_local_reader=0 (ray_executor_v2.py::_init_message_queues), so every one of them needs a real TCP port, and they all scan from the same VLLM_PORT – 7000 for a node-spanning engine. _init_message_queues runs immediately after init_device(), whose process-group setup is a collective barrier, so all workers on the node arrive at the probe within microseconds of each other, all see the same port free, and all but one die with::

zmq.error.ZMQError: Address already in use (addr='tcp://10.65.1.9:7000')

Workers on the driver node take n_local_reader=1 and use an ipc:// socket instead, which is why only node-spanning engines are affected – and why no nightly test catches it (none runs an engine whose tensor_parallel_size * pipeline_parallel_size exceeds cluster.gpus_per_node). See RL-1111.

Fix the race at the bind rather than the probe: retry, advancing past the port that was lost. This is safe and terminating because a port a peer already holds with ZMQ is visible to the next _get_open_port probe (a plain bind(("", port)) on it fails with EADDRINUSE), so each retry makes forward progress.

Deliberately keeps the search anchored at VLLM_PORT instead of letting vLLM fall back to bind(("", 0)): kernel-assigned ephemeral ports are exactly the TOCTOU contention the reserved sub-ephemeral band exists to prevent (#2380, #3103).

Patching the bind (rather than handing each worker a private start port) also covers every other MessageQueue with a remote reader – notably the executor’s own rpc_broadcast_mq – instead of the one call site that happens to be failing today.

Returns without raising when the snippet is missing, but logs at warning level so a silent no-op is visible in worker logs.

nemo_rl.models.generation.vllm.patches._patch_vllm_radio_layerscale_loader(logger) → None#

Load explicit RADIO LayerScale weights and initialize folded weights.

RadioVisionEncoderLayer scales each residual branch by ls1/ls2 (... + self.attn(...) * self.ls1), but RadioModel.load_weights discards those entries for legacy radio_model.* checkpoints, and folded checkpoints ship no LayerScale tensors at all. NeMo-RL forces load_format=dummy because weights arrive by refit, so anything the loader does not populate keeps its random initialization: the vision tower then multiplies every residual branch by noise. The failure is silent – unchanged step time, a collapsed reward, and a large train/token_mult_prob_error.

Applied as two INDEPENDENT edits rather than one block replacement, each anchored on text that is stable across loader variants. Both anchors are verified against stock vLLM 0.25.1 and 0.29.0. Custom vLLM builds that reshape the surrounding loader body (for example by gating the legacy branch or adding a native-name mapping branch) still take the required edit, where matching the whole body as one snippet would silently no-op.

Refit calls load_weights once per size-limited batch, so names loaded by any earlier call are remembered on the module and never re-filled.

nemo_rl.models.generation.vllm.patches._patch_vllm_radio_native_loader(logger) → None#

Accept transformers-native RADIO weight names in vLLM’s RADIO loader.

vLLM’s RadioModel.load_weights only understands the legacy timm-style layout (radio_model.model.blocks.N.attn.qkv with a fused qkv). Nemotron 3.5 Super VL ships transformers-native remote code, so the Automodel policy holds encoder.layer.N.attention.attention.{query,key,value} and layer_scale{1,2}.lambda1 parameters. Its per-tensor refit path cannot fuse the three attention tensors back into one attn.qkv, so those names reach vLLM unchanged (Automodel main) or under the legacy block prefix (Automodel r0.6.0: radio_model.model.blocks.N.attention.attention.query). Stock vLLM silently drops both forms and keeps generating with stale vision weights once the vision tower is trained.

Port the native branch of the super_vl_rl_v0.25.1 fork loader (TomerBN-Nvidia/vllm commit 436aefb8): map native embedding and encoder names onto vLLM’s keys and write q/k/v as shards of the fused QKVParallelLinear. Legacy names keep the stock code path, so the LayerScale loader patch and its anchors are unaffected.

nemo_rl.models.generation.vllm.patches._patch_vllm_radio_final_layernorm(logger) → None#

Port checkpoint-backed Super Omni final vision LayerNorm to vLLM 0.29.

The layer runs in FP32 before pixel shuffle/projection. Both affine parameters must load before activation, including dummy startup + refit. Supports stock vLLM 0.29 source or this patch’s complete output. Custom forks, including the already-fixed serving fork below, are not recognized; unexpected source layouts fail before writing rather than being skipped. References: RL #4036 and TomerBN-Nvidia/vllm commit 10908b9f.

nemo_rl.models.generation.vllm.patches._patch_vllm_glm_decoder_sequence_parallel_moe(logger) → None#

Restore the vLLM 0.24 decoder boundary for GLM DSA models.

vLLM 0.25.1 keeps hidden states sequence-parallel across attention and MoE decoder layers when TP, DP, and EP are all enabled. GLM-5.1/5.2 decode diverges on that new path: the first generated token is correct, while subsequent decode-token logprobs collapse. Keep vLLM’s existing MoE-local sequence parallelism, but disable the new decoder-level optimization for glm_moe_dsa so the MoE gathers its output as it did in vLLM 0.24.

The upstream bug and proposed fix are tracked at https://github.com/vllm-project/vllm/issues/50154 and https://github.com/vllm-project/vllm/pull/50155. Remove this patch after upgrading to a vLLM release containing the fix and validating iterative GLM-5.1/5.2 decode with TP, DP, and EP all enabled.

nemo_rl.models.generation.vllm.patches._patch_vllm_moe_routed_experts_capture( logger, *, required: bool = False, ) → bool#

Fire the routed-experts capture hook on the monolithic fused-MoE path.

RoutedExpertsCapturer (used by router replay / R3) is driven by the capture_fn that only fires inside BaseRouter._select_experts. But MoERunner._apply_quant_method calls select_experts only on the modular kernel branch; monolithic kernels (e.g. the FlashInfer TRT-LLM NVFP4-per-token fused MoE) compute top-k routing internally via forward_monolithic and never call it. The capture buffer therefore stays zero, and the returned routed_experts are all-zero -> Megatron’s router replay sees duplicate expert ids and dies with “Split sizes doesn’t match total dim 0” in the MoE all_to_all during get_logprobs.

This inserts an explicit select_experts call on the monolithic branch, guarded by capture_fn is not None so it only runs during rollout when routing capture is active (no cost otherwise).

nemo_rl.models.generation.vllm.patches._patch_vllm_routed_experts_capture_router_fallback( logger, *, required: bool = False, ) → bool#

Let monolithic MoE kernels without in-kernel capture use the router hook.

vLLM 0.29 moved the routed-experts binding into routed_experts_capturer.bind_routed_experts_capturer. For a monolithic kernel it requires fused_experts.supports_routing_replay_capture() and binds the capture function to that experts object (the kernel then writes routing_replay_out itself); any other monolithic kernel is rejected with ValueError. Two things make the object binding unusable for NeMo-RL’s NVFP4 per-token method: the kernel is rebuilt on every refit, so the bound capture function is dropped after the first weight update and the returned routes go back to all-zero; and the per-token kernel is the one FlashInfer launch that has not been validated with a replay buffer attached. Kernels that report no in-kernel capture (see nvfp4_pertoken.host_captured_experts_cls) therefore fall back to router.set_capture_fn — the hook that _patch_vllm_moe_routed_experts_capture fires on the monolithic branch and the path vLLM 0.26 used for every monolithic kernel.

nemo_rl.models.generation.vllm.patches._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.

nemo_rl.models.generation.vllm.patches.modelopt_moe_amax_aliases( model: torch.nn.Module, ) → collections.abc.Iterator[None]#

Temporarily expose nested ModelOpt MoE amax buffers to vLLM’s loader.

Verified against vLLM 0.26.0: RoutedExperts.load_weights maps expert amax keys to names such as w13_input_quantizer._amax, then resolves them with one getattr. This refit path sends ModelOpt buffers through that loader, which offers no hook for resolving nested target names.

Nemotron-H reached this loader after vLLM removed the inner NemotronHModel.load_weights in 0.26.0. Its old flat parameter lookup accepted dotted keys exposed by our buffer-to-parameter adapter: https://github.com/vllm-project/vllm/commit/c233d90aa826df072872df47b201450059be8e71

Aliases reference the original tensors without registering additional buffers or state-dict entries. Existing attributes belong to the caller or an outer context and are preserved. Only aliases created here are removed, including when setup or weight loading raises.

Parameters:: model – ModelOpt model whose MoE quantizer buffers will be refitted.

nemo_rl.models.generation.vllm.patches.ensure_vllm_source_compat() → None#

Apply interpreter-independent vLLM source-compat patches.

Safe to call from any process that imports vLLM directly (e.g. the tools/model_diagnostics scripts, which construct vllm.LLM without going through a NeMo-RL generation worker). Must be called BEFORE the first import vllm submodule that pulls in vllm.tool_parsers. Worker processes get this via _apply_vllm_patches at init.

nemo_rl.models.generation.vllm.patches._apply_vllm_patches( py_executable: str, *, extra_env_vars: list[str] | None = None, nemotron_h_fp32_lm_head: bool | None = None, require_moe_routed_experts_capture: bool = False, ) → None#

nemo_rl.models.generation.vllm.patches#

Module Contents#

Functions#

Data#

API#

`nemo_rl.models.generation.vllm.patches`#