`nemo_rl.modelopt.models.generation.vllm_quant_backend`#

Module Contents#

Classes#

`VllmQuantInternalWorkerExtension`
`VllmQuantInternalWorkerExtensionWithCheckpointEngine`	ModelOpt worker extension with checkpoint-engine refit support.

Functions#

`_match_fused_modelopt_moe_weight`
`_w13_num_shards_from_state_dict_info`	Validate complete fused-MoE families and resolve their W13 layout.
`_batch_fused_modelopt_moe_weights`	Map fused ModelOpt payloads to vLLM per-projection checkpoint names.
`_detach_pending_layerwise_weights`	Own deferred weights before a transport buffer may be reused.
`_iter_modelopt_quant_modules`	Return modules whose runtime layout is owned by vLLM ModelOpt methods.
`_modelopt_layerwise_reload_roots`	Select disjoint roots that require vLLM’s native reload lifecycle.
`_require_complete_modelopt_layerwise_reload`	Reject ModelOpt layers that vLLM would otherwise finalize partially.

Data#

_FUSED_MODELOPT_MOE_SUFFIXES

API#

nemo_rl.modelopt.models.generation.vllm_quant_backend._FUSED_MODELOPT_MOE_SUFFIXES#: None

nemo_rl.modelopt.models.generation.vllm_quant_backend._match_fused_modelopt_moe_weight(name: str) → tuple[str, str] | None#

nemo_rl.modelopt.models.generation.vllm_quant_backend._w13_num_shards_from_state_dict_info( state_dict_info: dict[str, Any], *, require_input_scales: bool = False, ) → dict[str, int]#: Validate complete fused-MoE families and resolve their W13 layout.

nemo_rl.modelopt.models.generation.vllm_quant_backend._batch_fused_modelopt_moe_weights( weights: list[tuple[str, torch.Tensor]], *, w13_num_shards_by_prefix: dict[str, int], ) → list[tuple[str, torch.Tensor]]#

Map fused ModelOpt payloads to vLLM per-projection checkpoint names.

w2 weights and block scales stay batched so vLLM can tensor-parallel-shard the full [E, ...] tensor at once. Its scalar loader still requires an expert id, so only the tiny per-expert global scales are exposed as scalar views.

Gated w13 payloads are the exception on vLLM >= 0.25: they are emitted as per-expert 2-D shards instead, because RoutedExperts.load_weights’ fused-3D branch mis-transposes packed NVFP4. See the comment at the emission site below.

nemo_rl.modelopt.models.generation.vllm_quant_backend._detach_pending_layerwise_weights( reload_roots: tuple[torch.nn.Module, ...], source_storage_ptrs: set[int], ) → None#

Own deferred weights before a transport buffer may be reused.

Completed layers have already released their buffered arguments, so this clones only tensors from a layer split across transport batches. Only the cached layerwise-reload subgraphs are inspected.

nemo_rl.modelopt.models.generation.vllm_quant_backend._iter_modelopt_quant_modules( model: torch.nn.Module, ) → list[tuple[str, torch.nn.Module]]#: Return modules whose runtime layout is owned by vLLM ModelOpt methods.

nemo_rl.modelopt.models.generation.vllm_quant_backend._modelopt_layerwise_reload_roots( model: torch.nn.Module, *, include_fp8_kv_cache: bool, ) → list[torch.nn.Module]#

Select disjoint roots that require vLLM’s native reload lifecycle.

Ordinary parameters are already updated in place by vLLM’s checkpoint loaders. Restricting layerwise reconstruction to ModelOpt runtime layouts and attention scale owners avoids materializing unrelated non-persistent buffers. In vLLM 0.20, whole-model reconstruction can otherwise break a derived buffer that aliases a child parameter (for example Nemotron-H’s conv_weights view of conv1d.weight).

nemo_rl.modelopt.models.generation.vllm_quant_backend._require_complete_modelopt_layerwise_reload( model: torch.nn.Module, ) → None#: Reject ModelOpt layers that vLLM would otherwise finalize partially.

class nemo_rl.modelopt.models.generation.vllm_quant_backend.VllmQuantInternalWorkerExtension#

Bases: nemo_rl.models.generation.vllm.vllm_backend.VllmInternalWorkerExtension

_nrl_w13_num_shards_by_prefix: dict[str, int]#: None

_nrl_modelopt_reload_roots: tuple[torch.nn.Module, ...] | None#: None

maybe_init_zmq() → None#: Use a longer timeout only for ModelOpt real-quant refits.

_is_real_quant_model() → bool#

_get_modelopt_reload_roots() → tuple[torch.nn.Module, ...]#: Return the invariant ModelOpt layerwise-reload subgraphs.

_weight_update_lifecycle( transport: nemo_rl.models.generation.vllm.vllm_backend.WeightUpdateTransport, ) → collections.abc.Iterator[nemo_rl.models.generation.vllm.vllm_backend.WeightUpdateFinalizer]#: Use vLLM’s native layerwise reload lifecycle for real quantization.

_weight_update_errors_are_fatal() → bool#

_synchronize_before_ipc_data_ack() → None#: Fence all accelerator streams used by ModelOpt post-load methods.

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

_patch_named_parameters_to_include_buffers(model)#

Temporarily patches model.named_parameters() to also yield input_quantizer buffers.

Weights arrive pre-folded from the Megatron side, so only input_quantizer amax buffers need to be loaded. Weight quantizer buffers are skipped.

_attach_input_quantizer_amax_loaders(model)#

Eagerly attach weight_loaders to input_quantizer amax buffers.

vLLM >= 0.25 loads refit weights through per-module load_weights (e.g. LinearBase.load_weights), which resolves targets via getattr and calls param.weight_loader(param, loaded_weight, shard_id) directly — it never iterates model.named_parameters(), so the lazy attach in _patch_named_parameters_to_include_buffers no longer fires and quantizer amax buffers arrive without a loader (AttributeError: ‘Tensor’ object has no attribute ‘weight_loader’).

_load_weights(weights)#

Load pre-folded weights and input_quantizer amax buffers.

Weights arrive already folded from the Megatron side (weight_quantizer applied during export), so no fold_weight step is needed here.

get_weight_snapshot(name: str) → torch.Tensor#: Return a CPU copy of a named parameter for before/after comparison.

get_quantizer_stats() → dict#

Return summary statistics for all TensorQuantizer modules.

Matches the interface of MegatronQuantPolicyWorker.get_quantizer_stats().

class nemo_rl.modelopt.models.generation.vllm_quant_backend.VllmQuantInternalWorkerExtensionWithCheckpointEngine#

Bases: nemo_rl.models.generation.vllm.checkpoint_engine.VllmCheckpointEngineMixin, nemo_rl.modelopt.models.generation.vllm_quant_backend.VllmQuantInternalWorkerExtension

ModelOpt worker extension with checkpoint-engine refit support.

nemo_rl.modelopt.models.generation.vllm_quant_backend#