`nemo_rl.models.generation.vllm.vllm_backend`#

Module Contents#

Classes#

`_IPCWeightManifest`	Validate an IPC stream against its prepared state-dict manifest.
`NixlVllmWorker`	vLLM worker that establishes NIXL/UCX before vLLM initialization.
`VllmInternalWorkerExtension`
`VllmInternalWorkerExtensionWithCheckpointEngine`	vLLM worker extension with checkpoint-engine refit support.

Functions#

`_format_refit_key_error`	Format a bounded refit-key diagnostic.
`_detach_pending_layerwise_weights`	Clone deferred reload weights that still alias a transport buffer.
`_refresh_hpc_modules_after_layerwise_reload`	Rebuild kernel-specific state omitted by vLLM’s layerwise finalizer.
`_model_uses_unquantized_flashinfer_trtllm`	Return whether a model realized the unquantized TRTLLM MoE backend.
`fix_gemma3_vision_weight_name`	Re-insert the `vision_model` segment into Gemma3 vision-tower weights.
`_read_mtp_layer_weights_from_checkpoint`	Read only the MTP draft layer weights from a sharded HF safetensors checkpoint.

Data#

`logger`
`WeightUpdateTransport`
`UnsupportedNativeRefitTransport`
`WeightUpdateFinalizer`

API#

nemo_rl.models.generation.vllm.vllm_backend.logger#: ‘getLogger(…)’

nemo_rl.models.generation.vllm.vllm_backend.WeightUpdateTransport#: None

nemo_rl.models.generation.vllm.vllm_backend.UnsupportedNativeRefitTransport#: None

nemo_rl.models.generation.vllm.vllm_backend.WeightUpdateFinalizer#: None

nemo_rl.models.generation.vllm.vllm_backend._format_refit_key_error(label: str, keys: set[str]) → str#: Format a bounded refit-key diagnostic.

exception nemo_rl.models.generation.vllm.vllm_backend.IPCWeightManifestError#

Bases: RuntimeError

An IPC transfer did not match the prepared state-dict manifest.

Initialization

Initialize self. See help(type(self)) for accurate signature.

nemo_rl.models.generation.vllm.vllm_backend._detach_pending_layerwise_weights( model: torch.nn.Module, source_storage_ptrs: set[int], ) → None#

Clone deferred reload weights that still alias a transport buffer.

vLLM 0.25.1 replays deferred weight-loader arguments during layerwise finalization, after NeMo-RL may have reused the source transport buffer.

nemo_rl.models.generation.vllm.vllm_backend._refresh_hpc_modules_after_layerwise_reload( model: torch.nn.Module, ) → None#

Rebuild kernel-specific state omitted by vLLM’s layerwise finalizer.

HpcModule implementations derive runtime state from loaded weights via process_weights_after_loading. vLLM 0.25.1 runs that model-wide pass on normal loads, but not after a layerwise reload.

nemo_rl.models.generation.vllm.vllm_backend._model_uses_unquantized_flashinfer_trtllm( model: torch.nn.Module, ) → bool#: Return whether a model realized the unquantized TRTLLM MoE backend.

class nemo_rl.models.generation.vllm.vllm_backend._IPCWeightManifest(expected_keys: collections.abc.Iterable[str])#

Validate an IPC stream against its prepared state-dict manifest.

Initialization

validate_batch(keys: collections.abc.Sequence[str]) → set[str] | None#

record_loaded(keys: set[str]) → None#

record_load_failure(error: Exception) → None#

require_complete() → None#

class nemo_rl.models.generation.vllm.vllm_backend.NixlVllmWorker#

Bases: vllm.v1.worker.gpu_worker.Worker

vLLM worker that establishes NIXL/UCX before vLLM initialization.

__new__(

vllm_config: Any,

*args: Any,

**kwargs: Any,

) → nemo_rl.models.generation.vllm.vllm_backend.NixlVllmWorker#

nemo_rl.models.generation.vllm.vllm_backend.fix_gemma3_vision_weight_name(key: str) → str#

Re-insert the vision_model segment into Gemma3 vision-tower weights.

When performing refit, the vision-tower weight paths are flattened. This unflattens them.

nemo_rl.models.generation.vllm.vllm_backend._read_mtp_layer_weights_from_checkpoint( model_path: str, mtp_layer_indices: set[int], ) → list[tuple[str, torch.Tensor]]#

Read only the MTP draft layer weights from a sharded HF safetensors checkpoint.

Uses the checkpoint’s model.safetensors.index.json to open only the shards that contain the requested transformer layer indices, so the multi-terabyte base-model weights are never read from disk.

Parameters:

model_path – Path to the HF checkpoint directory.
mtp_layer_indices – Transformer layer indices belonging to the MTP module(s).

Returns:

A list of (weight_name, tensor) pairs for the requested layers, with tensors on CPU.

class nemo_rl.models.generation.vllm.vllm_backend.VllmInternalWorkerExtension#

pp_comm_groups: Optional[dict[int, Any]]#: None

_mtp_drafter_from_disk: bool#: False

_sparse_delta_applier: Any#: None

_nrl_named_parameters: dict[str, torch.nn.Parameter]#: None

_nrl_layerwise_reload_active: bool#: False

_nrl_layerwise_reload_failure: Exception | None#: None

model_update_group: Any#: None

_get_named_parameters() → dict[str, torch.nn.Parameter]#

_load_full_hf_weights( policy_weights: list[tuple[str, torch.Tensor]], ) → None#: Load HF weights and detach any deferred reload tensors from transport storage.

_load_hf_weights( policy_weights: list[tuple[str, torch.Tensor]], ) → None#

bind_numa() → bool#

Pin this TP worker to its GPU’s NUMA-local CPUs/memory.

Invoked via collective_rpc on each vLLM TP worker once the engine (and CUDA) is up, so the worker’s physical GPU id is resolved from its local device index (see resolve_visible_gpu_id).

init_collective( rank_prefix: int, ip: str, port: int, world_size: int, train_world_size: int, ) → None#: Initialize the collective communication.

init_nccl_reshard_comm_group( rank_prefix: int, pp_ips: list[str], pp_ports: list[int], pp_size: int, train_ranks_per_stage: int, sub_world_size: int, ) → None#

Bootstrap this gen worker’s nccl_reshard bulk-path comm group(s).

One comm group per PP stage; gen workers join ALL pp_size groups (they need every stage’s layers), created in stage order so the train ranks (each in only their own stage) unblock deterministically. Non-PP is simply pp_size == 1 that contains all the gen ranks.

report_device_id() → str#: Retrieve the UUID of the current CUDA device.

report_node_hostname() → str#: Return the host shared by worker processes on this node.

get_zmq_address()#: Get the ZMQ address for the current device.

maybe_init_zmq()#: Initialize the ZMQ socket if it doesn’t exist.

prepare_refit_info(state_dict_info: dict[str, Any]) → None#

Prepare state dict metadata for weight refitting and IPC streaming.

Parameters:: state_dict_info (dict) – A dictionary containing the info for refit. e.g. {tensor_name: (shape, dtype)}
Raises:: RuntimeError – If the model realizes the unquantized FlashInfer TRTLLM MoE backend while a co-trained MTP drafter is enabled (unsupported by the native layerwise refit lifecycle).

prepare_sparse_delta_refit_info( state_dict_info: dict[str, tuple[tuple[int, ...], torch.dtype]], ) → list[str]#: Reserve scratch space and report weights that require overwrite.

_uses_fp8_kv_cache() → bool#: Return whether this worker owns an FP8 KV cache.

_maybe_process_fp8_kv_cache() → None#: Process weights after loading for FP8 KV cache (static scales).

static _split_policy_and_draft_weights( weights: list[tuple[str, torch.Tensor]], ) → tuple[list[tuple[str, torch.Tensor]], list[tuple[str, torch.Tensor]]]#

Split trainer-owned draft weights from policy weights.

This path is only used for the Eagle3 online-training flow, where the trainer exports draft parameters under a draft. prefix before sending them to vLLM. MTP parameters do not use the draft. prefix; they remain in the policy stream and are forwarded separately by _maybe_refit_mtp_drafter. The “draft.” prefix is added here https://github.com/isomap/RL/blob/d3a5e1396d00f82fb888d9ec6800687a23bb4017/nemo_rl/models/policy/workers/megatron_policy_worker.py#L967-L997

static _trim_vocab_padding( draft_model: torch.nn.Module, draft_weights: list[tuple[str, torch.Tensor]], ) → list[tuple[str, torch.Tensor]]#

Trim padded vocab dimensions from draft weights.

Megatron pads vocab to a multiple, but vLLM’s autoloader strictly asserts loaded_weight.shape[0] == org_vocab_size on VocabParallelEmbedding layers. Each such layer may have a different org_vocab_size (e.g. embed_tokens uses vocab_size while lm_head uses draft_vocab_size), so we match each weight to its target module by name.

_get_drafter_model() → Any#

Return the vLLM drafter’s underlying model, or None if absent.

The drafter holds the speculative-decoding draft model (Eagle3 or MTP), which vLLM keeps as a module separate from the main model. Typed Any because these are dynamic vLLM model classes whose load_weights / mtp_start_layer_idx members are not visible through nn.Module.

_load_draft_weights( draft_weights: list[tuple[str, torch.Tensor]], ) → None#

_mtp_drafter_refit_enabled() → bool#

Whether MTP drafter weights should be refreshed from the refit stream.

For MTP speculative decoding where the trainer co-trains the MTP layer (mtp_num_layers > 0), the MTP weights are exported as part of the policy weight stream during refit (without the draft. prefix used by Eagle3), so the drafter must be fed those weights on every refit.

Returns False when the MTP weights were instead loaded once from disk (see load_mtp_weights_from_disk) — the path used when the trainer does not co-train the MTP layer — to avoid clobbering and re-processing those static weights.

_maybe_refit_mtp_drafter( weights: list[tuple[str, torch.Tensor]], ) → None#

Load refit weights into an MTP drafter co-trained with the policy.

The drafter’s load_weights selects the MTP-specific parameters (and shared embed_tokens / lm_head) it needs from the full policy weight stream. Megatron pads the vocab dimension, so weights are trimmed to the drafter’s expected vocab size first, matching _load_draft_weights.

_maybe_process_mtp_drafter_after_loading() → None#

Finalize MTP drafter weights after a refit (e.g. MoE grouped-GEMM layout).

Mirrors the main-model post-processing so the freshly refit MTP layers are converted to their runtime layout. Skipped for the disk-load path, which already processes its weights once at startup.

load_mtp_weights_from_disk(model_path: str) → bool#

Load only the MTP (multi-token-prediction) draft weights from disk.

Used when an MTP speculative-decoding policy runs with load_format="dummy": the main model receives real weights via refit, but the MTP draft layer is not covered by refit (the trainer runs with mtp_num_layers=0), so its weights must come from the checkpoint. Only the MTP layer(s) are read, avoiding a full base-model load (~1.3 TB for DeepSeek-V3) on every inference replica.

Parameters:: model_path – Path to the HF checkpoint directory.
Returns:: True if MTP weights were loaded.
Return type:: bool

_load_weights(weights)#

Load weights with Gemma3 vision-tower weight name fix, FP8, and draft-weight support.

Applies Gemma3 vision-tower weight name fix if needed, splits policy/draft weights, dispatches policy weights through the configured refit loader, and loads draft weights into the drafter model.

_get_sparse_delta_applier() → Any#

_supports_unquantized_flashinfer_trtllm_refit() → bool#: Whether this worker supports native unquantized TRTLLM refits.

_uses_unquantized_flashinfer_trtllm() → bool#: Detect a realized unquantized FlashInfer TRTLLM MoE backend.

_uses_native_layerwise_refit( transport: nemo_rl.models.generation.vllm.vllm_backend.WeightUpdateTransport, ) → bool#: Return whether this transport needs vLLM’s layerwise lifecycle.

_validate_native_layerwise_refit() → None#: Reject unsupported features on the native layerwise reload path.

_reject_unsupported_native_refit( transport: nemo_rl.models.generation.vllm.vllm_backend.UnsupportedNativeRefitTransport, ) → None#: Reject transports that cannot run the native layerwise lifecycle.

_weight_update_lifecycle( transport: nemo_rl.models.generation.vllm.vllm_backend.WeightUpdateTransport, ) → collections.abc.Iterator[nemo_rl.models.generation.vllm.vllm_backend.WeightUpdateFinalizer]#

Provide setup/finalization around a transport-owned weight update.

Native reload initialization invalidates the old runtime layout. Any subsequent exception therefore marks this worker permanently unusable.

_weight_update_errors_are_fatal() → bool#: Whether transport errors should propagate instead of returning False.

_synchronize_before_ipc_data_ack() → None#: Fence work consuming one IPC data batch before its acknowledgment.

update_weights_via_ipc_zmq() → bool#

Receive and update model weights via ZMQ IPC socket.

Returns:: True if weights were successfully updated.
Return type:: bool

update_weights_from_collective( refit_timeout_s: float | None = None, ) → bool#

Update the model weights from collective communication.

Guarded for the same reason as the producer side: if a peer rank dies mid-refit this blocks in NCCL forever. Note the buffers hold PARTIAL weights once aborted, so the caller must not serve from this engine until a later refit completes.

_update_weights_from_collective() → bool#

update_weights_from_decoded_sparse_payload( *payloads: bytes | str, ) → dict[str, Any]#

synchronize_device() → None#

finish_sparse_delta_refit() → dict[str, Any]#

prepare_nccl_reshard_refit_info(refit_info: dict) → None#

Restore per-layer param metadata and build the HF→vLLM mapping.

Done once ahead of refit; the cached mapping is reused by every nccl_reshard_refit call.

build_hf_to_local_param_map( refit_info: dict, ) → nemo_rl.weight_sync.nccl_reshard_utils.HFToLocalParamMap#

Build the vLLM-backend hf_to_local_param_map (HFToLocalParamMap).

Wraps the (vllm_param, merged_slice) resolution from _build_hf_to_gen_backend_mapping into LocalParamSpecs:

direct (slice None): base is the live vLLM param; receive in place.
merged (dense gate_up_proj / grouped-expert w13): pre allocs a recv buffer for this component’s region slice, post copies it back (region recomputed each refit to track live storage).

_build_hf_to_gen_backend_mapping(refit_info)#

Map each FFN HF param name to its gen-backend param and slice.

Only gate_proj / up_proj / down_proj .weight (dense MLP and MoE experts) reach here. Returns hf_name -> (vllm_param, merged_param_slice or None); the slice (None for a 1:1 direct map) is the local region of a fused vLLM param this HF piece occupies, applied by the LocalParamSpec pre/post hooks. The three shapes:

grouped MoE experts: gate/up -> w13_weight halves (dim 1), down -> w2_weight (direct).
dense MLP gate/up -> gate_up_proj halves (dim 0).
dense MLP down -> down_proj (direct 1:1).

nccl_reshard_refit(refit_timeout_s: float | None = None) → bool#

Receive weights from training workers via xferdtensor, under a deadline.

Guarded like the collective receive: a peer dying mid-refit blocks this in NCCL forever, and the buffers hold PARTIAL weights once aborted, so the caller must not serve from this engine until a later refit completes.

Both communicator families are handed to the watchdog – the per-PP-stage bulk groups and the shared model_update_group – because the transfer uses them in sequence and a hang can be in either.

Each HF param’s LocalParamSpec (from hf_to_local_param_map, built once in prepare_nccl_reshard_refit_info) provides the dst buffer: for a direct param xferdtensor receives straight into the live vLLM param (no hooks); for a merged param (dense gate_up_proj, grouped w13) pre allocates a temp recv buffer and post copies the TP-local slice back into the live merged param.

_nccl_reshard_refit() → bool#

_receive_and_load_misc_params() → None#: Receive misc params via packed_broadcast and load via vLLM.

cleanup() → None#: Shutdown and cleanup resources.

start_gpu_profiling() → None#: Start GPU profiling.

stop_gpu_profiling() → None#: Stop GPU profiling.

class nemo_rl.models.generation.vllm.vllm_backend.VllmInternalWorkerExtensionWithCheckpointEngine#

Bases: nemo_rl.models.generation.vllm.checkpoint_engine.VllmCheckpointEngineMixin, nemo_rl.models.generation.vllm.vllm_backend.VllmInternalWorkerExtension

vLLM worker extension with checkpoint-engine refit support.

_validate_checkpoint_engine_weight_update() → None#

nemo_rl.models.generation.vllm.vllm_backend#