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

Module Contents#

Classes#

`_IPCWeightManifest`	Validate an IPC stream against its prepared state-dict manifest.
`_ReloadWeightPreparer`	Turn transport batches into reload-safe checkpoint tensors.
`NixlVllmWorker`	vLLM worker that establishes NIXL/UCX before vLLM initialization.
`VllmReshardParamSpecCollection`	Stateless factories for vLLM’s nccl_reshard receive hooks.
`VllmInternalWorkerExtension`
`VllmInternalWorkerExtensionWithCheckpointEngine`	vLLM worker extension with checkpoint-engine refit support.

Functions#

`_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.
`_unquantized_flashinfer_trtllm_modules`	Return modules that realized the unquantized TRTLLM MoE backend.
`_reload_target_module_ids`	Return every module covered by the native reload targets.
`_process_mxfp8_modules_after_native_reload`	Rebuild MXFP8 layouts outside the partial native reload.
`_model_uses_unquantized_flashinfer_trtllm`	Return whether a model realized the unquantized TRTLLM MoE backend.
`_local_shard_slices`	Return this destination rank’s slices in an HF-global tensor.
`fix_gemma3_vision_weight_name`	Re-insert the `vision_model` segment into Gemma3 vision-tower weights.
`_is_gemma4_unified_text_only`	Return whether a vLLM model is a text-only Gemma 4 Unified model.
`_filter_gemma4_unified_multimodal_weights`	Drop frozen Gemma 4 Unified vision and audio embedder weights.
`_tied_embedding_aliases`	Return vLLM’s `{alias qualname: canonical qualname}` for tied embeddings.
`_drop_tied_embedding_aliases`	Drop checkpoint weights that vLLM would skip as tied-embedding aliases.
`_read_mtp_layer_weights_from_checkpoint`	Read only the MTP draft layer weights from a sharded HF safetensors checkpoint.

Data#

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

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

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._unquantized_flashinfer_trtllm_modules( model: torch.nn.Module, ) → list[torch.nn.Module]#: Return modules that realized the unquantized TRTLLM MoE backend.

nemo_rl.models.generation.vllm.vllm_backend._reload_target_module_ids( reload_targets: collections.abc.Sequence[torch.nn.Module], ) → set[int]#: Return every module covered by the native reload targets.

nemo_rl.models.generation.vllm.vllm_backend._process_mxfp8_modules_after_native_reload( model: torch.nn.Module, reloaded_module_ids: set[int], ) → None#: Rebuild MXFP8 layouts outside the partial native 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.

nemo_rl.models.generation.vllm.vllm_backend._local_shard_slices( param_info: dict[str, Any], rank: int, ) → tuple[slice, ...]#: Return this destination rank’s slices in an HF-global tensor.

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._ReloadWeightPreparer#

Bases: typing.Protocol

Turn transport batches into reload-safe checkpoint tensors.

Tensors returned by process must not reference the source IPC buffer, because it is released as soon as the batch is acknowledged.

reset() → None#

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

finish() → 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._GEMMA4_UNIFIED_MULTIMODAL_WEIGHT_MARKERS#: (‘model.embed_vision.’, ‘model.embed_audio.’)

nemo_rl.models.generation.vllm.vllm_backend._is_gemma4_unified_text_only(model_config: Any) → bool#: Return whether a vLLM model is a text-only Gemma 4 Unified model.

nemo_rl.models.generation.vllm.vllm_backend._filter_gemma4_unified_multimodal_weights( weights: collections.abc.Iterable[tuple[str, torch.Tensor]], ) → collections.abc.Iterator[tuple[str, torch.Tensor]]#: Drop frozen Gemma 4 Unified vision and audio embedder weights.

nemo_rl.models.generation.vllm.vllm_backend._tied_embedding_aliases( model: torch.nn.Module | None, ) → dict[str, str]#

Return vLLM’s {alias qualname: canonical qualname} for tied embeddings.

Uses vLLM’s own detector so this set is exactly what AutoWeightsLoader skips. Empty on a vLLM without the helper, which also has no alias check, so the filter below becomes a no-op there. Also empty for a stand-in that is not a module tree (a bare load_weights callable), which has no tied embeddings for the detector to walk.

nemo_rl.models.generation.vllm.vllm_backend._drop_tied_embedding_aliases( weights: collections.abc.Iterable[tuple[str, torch.Tensor]], aliases: collections.abc.Mapping[str, str], mapper: Any | None, ) → collections.abc.Iterator[tuple[str, torch.Tensor]]#

Drop checkpoint weights that vLLM would skip as tied-embedding aliases.

vLLM 0.29’s AutoWeightsLoader skips e.g. lm_head.weight when it is tied to the input embedding and then asserts that the canonical embedding weight was loaded in the same load_weights call (vllm-project/vllm#51665). Refit streams weights in transport-sized batches, so the two routinely arrive in different calls and every refit of a tied-embedding model died with 'lm_head.weight' was skipped because it is tied to 'model.embed_tokens.weight' ... was not found in the checkpoint. The alias never loads anything, so dropping it up front is lossless. mapper is the model’s hf_to_vllm_mapper (if any): the alias set is keyed by vLLM parameter names while refit sends checkpoint names, and the loader applies the same mapper before its check.

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.VllmReshardParamSpecCollection#

Stateless factories for vLLM’s nccl_reshard receive hooks.

static _direct_param_spec( vllm_param: torch.Tensor, wire_dtype: torch.dtype, ) → nemo_rl.weight_sync.nccl_reshard_utils.LocalParamSpec#: Receive into the destination storage captured when the map is built.

static _default_staged_param_spec( vllm_param: torch.Tensor, merged_slice: tuple[slice, ...] | None, wire_dtype: torch.dtype, ) → nemo_rl.weight_sync.nccl_reshard_utils.LocalParamSpec#: Stage a canonical weight and copy it into the current destination storage.

static _trtllm_grouped_expert_spec( param_info: dict[str, Any], wire_dtype: torch.dtype, *, refit_comm_groups: dict[int, Any] | None, device: torch.device, registered_vllm_name: str | None, load_weights: collections.abc.Callable[[collections.abc.Iterable[tuple[str, torch.Tensor]]], set[str] | None], ) → nemo_rl.weight_sync.nccl_reshard_utils.LocalParamSpec#: Stage EP-local expert weights for the native vLLM weight loader.

static _mxfp8_receiver_quant_spec( value_param: torch.Tensor, scale_param: torch.Tensor, merged_slice: tuple[slice, ...] | None, wire_dtype: torch.dtype, ) → nemo_rl.weight_sync.nccl_reshard_utils.LocalParamSpec#: Quantize a received weight and commit its MXFP8 values and scales.

static _blockwise_fp8_receiver_quant_spec( value_param: torch.Tensor, scale_param: torch.Tensor, merged_slice: tuple[slice, ...] | None, block_size: tuple[int, int], wire_dtype: torch.dtype, ) → nemo_rl.weight_sync.nccl_reshard_utils.LocalParamSpec#: Quantize a received weight and commit its blockwise FP8 values and scales.

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

Bases: nemo_rl.weight_sync.nccl_reshard_utils.RefitBuilderInterface

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

_mtp_drafter_weights_from_refit: bool#: True

_DRAFTER_OWNER_ATTRS: tuple[str, ...]#: (‘drafter’, ‘speculator’)

_warned_missing_mtp_drafter: bool#: False

_logged_gemma4_unified_drop: bool#: False

_sparse_delta_applier: Any#: None

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

hf_to_local_param_map: nemo_rl.weight_sync.nccl_reshard_utils.HFToLocalParamMap#: 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( weights: collections.abc.Iterable[tuple[str, torch.Tensor]], ) → set[str] | None#

Load HF weights and detach any deferred reload tensors from transport storage.

Returns the set of weight names vLLM reported as loaded, or None when the model’s load_weights does not report one.

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

_without_tied_embedding_aliases( policy_weights: list[tuple[str, torch.Tensor]], ) → list[tuple[str, torch.Tensor]]#: Drop the tied-embedding aliases vLLM would skip (see module helper).

_prepare_reload_weight_iterator( weights: collections.abc.Iterable[tuple[str, torch.Tensor]], ) → collections.abc.Iterable[tuple[str, torch.Tensor]]#: Prepare checkpoint-format weights for vLLM’s native reload API.

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], serialized_fp8_config: Optional[dict[str, Any]] = None, ) → Optional[list[str]]#

Prepare state dict metadata for weight refitting and IPC streaming.

For co-trained speculative decoding (dspark/dflash/eagle3), the trainer-provided draft.* keys are validated here against the drafter’s own loadable layout. Combined with the per-refit manifests (the IPC path enforces every state_dict_info key exactly once, and the collective path iterates state_dict_info by construction), this proves on every transport that the full expected draft key set is delivered.

Parameters:: state_dict_info (dict) – A dictionary containing the info for refit. e.g. {tensor_name: (shape, dtype)}
Returns:: When MXFP8 trainer-side pre-quantization is enabled (vllm_cfg.refit_prequantize), the list of parameter names this worker will quantize at load time; the trainer quantizes exactly these and streams E4M3 data plus *_scale_from_checkpoint scales. None otherwise.
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 separately refittable FP8 KV 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.

_speculative_method() → Optional[str]#

_get_drafter_model() → Any#

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

The drafter holds the speculative-decoding draft model (Eagle3, MTP, or DSpark), which vLLM keeps as a module separate from the main model, under model_runner.drafter (legacy runner) or model_runner.speculator (v2 runner, used by DSpark); see _DRAFTER_OWNER_ATTRS. Typed Any because these are dynamic vLLM model classes whose load_weights / mtp_start_layer_idx members are not visible through nn.Module.

configure_mtp_drafter_weight_source(weights_from_refit: bool) → None#: Record whether the trainer owns and refreshes the MTP weights.

_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)#

Apply model-specific transforms and load policy and draft weights.

Checkpoint-format weights are normalized for the target vLLM model, then routed to the policy model and any supported speculative drafter.

_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.

_unquantized_flashinfer_trtllm_param_ids() → set[int]#: Return parameters owned by realized unquantized TRTLLM MoE modules.

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

_uses_deepseek_v4_fp8_refit() → bool#: Return whether the realized rollout model needs DSV4 FP8 reload hooks.

_validate_native_layerwise_refit( transport: nemo_rl.models.generation.vllm.vllm_backend.WeightUpdateTransport | None = None, ) → 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.

_get_reload_weight_preparer() → nemo_rl.models.generation.vllm.vllm_backend._ReloadWeightPreparer | None#: Return this worker’s transport-neutral checkpoint preparer.

_drain_ipc_reload_sender() → None#: Release a REQ sender after a native reload has failed.

_update_weights_via_ipc_zmq_with_reload( preparer: nemo_rl.models.generation.vllm.vllm_backend._ReloadWeightPreparer, ) → bool#: Receive IPC batches through vLLM’s native layerwise reload API.

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, refit_with_reload_api: bool = False, ) → 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.

Parameters:

refit_timeout_s – Refit watchdog timeout in seconds.
refit_with_reload_api – Whether to use vLLM’s native reload API.

_update_weights_from_collective(refit_with_reload_api: bool) → bool#

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

synchronize_sparse_refit_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:

ordinary storage: unmerged params matching the wire dtype receive in place; otherwise pre stages a buffer and post copies it back.
FP8 storage: post quantizes the received component and commits its value and scale together before the backend’s final layout conversion.
TRTLLM grouped experts: pre allocates canonical EP-local wire storage, and post sends each expert through vLLM’s native weight loader.

_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 during prepare or after PP communicator setup for TRTLLM experts) provides a canonical receive buffer in the advertised wire dtype. When needed, its post hook casts or quantizes into the live destination.

_nccl_reshard_refit_impl( finalize: nemo_rl.models.generation.vllm.vllm_backend.WeightUpdateFinalizer, ) → bool#

Receive weights from training workers via xferdtensor.

Each HF param’s LocalParamSpec (from hf_to_local_param_map, built once in prepare_nccl_reshard_refit_info) receives the canonical shard in the advertised wire dtype, staging and committing in post when needed. TRTLLM grouped experts use vLLM’s native weight loader.

_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#