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

Module Contents#

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

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

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

update_weights_via_ipc_zmq() → bool#: Receive and update weights through CUDA IPC.

update_weights_from_collective() → bool#: Receive and update weights through collective communication.

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().