nemo_rl.modelopt.models.generation.vllm_quant_worker#

Module Contents#

Classes#

VllmQuantGenerationWorker

VllmQuantAsyncGenerationWorker

Functions#

_configure_quant_engine_kwargs

Data#

_EXTRA_ENV_VARS

API#

nemo_rl.modelopt.models.generation.vllm_quant_worker._EXTRA_ENV_VARS#

(‘VLLM_QUANT_CFG’, ‘VLLM_MODELOPT_REAL_QUANT’)

nemo_rl.modelopt.models.generation.vllm_quant_worker._configure_quant_engine_kwargs(
cfg: nemo_rl.models.generation.vllm.config.VllmConfig,
llm_kwargs: dict[str, Any],
) → None#
class nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantGenerationWorker(*args, **kwargs)#

Bases: nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorkerImpl

Initialization

Initialize a vLLM worker for distributed inference.

Parameters:

  • config – Configuration dictionary for the policy

  • bundle_indices – List of local bundle indices within a node for parallelism. Only needed for the first worker in each tied worker group.

  • fraction_of_gpus – Fraction of GPUs to use for this worker

  • seed – Random seed for initialization

  • extra_env_vars – Additional environment variable names to forward into the vLLM worker subprocess (e.g. for quantization configs).

  • defer_model_load – If True, skip model loading during init. Call _load_model() later to perform the heavy model loading. This enables overlapping vLLM model loading with NeMo Gym init.

_create_engine(llm_kwargs: dict[str, Any]) → None#
get_quantizer_stats() → dict[str, Any]#

Return quantizer statistics. Mirrors MegatronQuantPolicyWorker.get_quantizer_stats().

get_weight_snapshot(name: str) → Any#

Return a CPU copy of a named parameter for before/after comparison.

class nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantAsyncGenerationWorker(*args, **kwargs)#

Bases: nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorkerImpl

Initialization

Initialize an async vLLM worker.

When defer_model_load=True, only stores config and reserves a port for the HTTP server (if expose_http_server is enabled). Call load_model() later to perform the heavy model loading. This enables overlapping vLLM model loading with NeMo Gym init.

Parameters:

  • config – Configuration dictionary for the policy

  • bundle_indices – List of local bundle indices within a node for parallelism.

  • fraction_of_gpus – Fraction of GPUs to use for this worker

  • seed – Random seed for initialization

  • extra_env_vars – Additional environment variable names to forward into the vLLM worker subprocess.

  • defer_model_load – If True, skip model loading and only reserve port

_create_engine(llm_kwargs: dict[str, Any]) → None#
async get_quantizer_stats() → dict[str, Any]#

Return quantizer statistics. Mirrors MegatronQuantPolicyWorker.get_quantizer_stats().

async get_weight_snapshot(name: str) → Any#

Return a CPU copy of a named parameter for before/after comparison.