nemo_rl.models.generation.vllm.vllm_worker#

Module Contents#

Classes#

BaseVllmGenerationWorker

VllmGenerationWorkerImpl

VllmGenerationWorker

Functions#

_resolve_enable_prefix_caching

_merge_fp8_kwargs

Merge fp8 init kwargs into vllm_kwargs in place, preserving user overrides.

Data#

logger

API#

nemo_rl.models.generation.vllm.vllm_worker.logger#

‘getLogger(…)’

nemo_rl.models.generation.vllm.vllm_worker._resolve_enable_prefix_caching(
vllm_cfg: dict[str, Any],
) → bool#
nemo_rl.models.generation.vllm.vllm_worker._merge_fp8_kwargs(
vllm_kwargs: dict[str, Any],
fp8_kwargs: dict[str, Any],
) → None#

Merge fp8 init kwargs into vllm_kwargs in place, preserving user overrides.

init_fp8 returns a nested hf_overrides (holding quantization_config), so a blanket vllm_kwargs.update(fp8_kwargs) would wholesale-replace any user-supplied hf_overrides. We pop hf_overrides before the shallow update and merge it separately so that fp8’s quantization_config is the base while user overrides (e.g. max_position_embeddings) survive and take precedence. This regression was reintroduced once already; see #1413/#2904.

class nemo_rl.models.generation.vllm.vllm_worker.BaseVllmGenerationWorker(
config: nemo_rl.models.generation.vllm.config.VllmConfig,
bundle_indices: Optional[list[int]] = None,
fraction_of_gpus: float = 1.0,
seed: Optional[int] = None,
extra_env_vars: Optional[list[str]] = None,
defer_model_load: bool = False,
)#

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.

__repr__() → str#

Customizes the actor’s prefix in the Ray logs.

This makes it easier to identify which worker is producing specific log messages.

static configure_worker(
num_gpus: int | float,
bundle_indices: Optional[tuple[int, list[int]]] = None,
) → tuple[dict[str, Any], dict[str, str], dict[str, Any], dict[str, Any]]#

Provides complete worker configuration for vLLM tensor and pipeline parallelism.

This method configures the worker based on its role in tensor and pipeline parallelism, which is determined directly from the bundle_indices parameter.

Parameters:

  • num_gpus – Original GPU allocation for this worker based on the placement group

  • bundle_indices – Tuple of (node_idx, local_bundle_indices) for parallelism (if applicable)

Returns:

  • ‘resources’: Resource allocation (e.g., num_gpus)

  • ’env_vars’: Environment variables for this worker

  • ’init_kwargs’: Parameters to pass to init of the worker

  • ’runtime_env’: Additional runtime_env options (e.g., nsight config)

Return type:

tuple with complete worker configuration

_init_config(
config: nemo_rl.models.generation.vllm.config.VllmConfig,
bundle_indices: Optional[list[int]],
fraction_of_gpus: float,
seed: Optional[int],
extra_env_vars: Optional[list[str]],
)#

Lightweight config setup. No model loading, no heavy imports.

_load_model(bundle_indices, seed)#

Perform the heavy model loading and engine creation.

Split out from init so it can be deferred (defer_model_load=True) and run after ports are reserved, overlapping with NeMo Gym init.

llm()#
is_alive()#

Check if the worker is alive.

_merge_stop_strings(batch_stop_strings)#
_build_sampling_params(
*,
greedy: bool,
stop_strings,
max_new_tokens: Optional[int] = None,
)#
start_gpu_profiling() → None#

Start GPU profiling.

stop_gpu_profiling() → None#

Stop GPU profiling.

static _patch_vllm_nsight_config() → None#

Override vLLM’s nsight config for internal TP workers to use deferred capture.

vLLM’s default _configure_ray_workers_use_nsight applies an always-on nsight config (no capture-range, cuda-graph-trace=node) which causes significant overhead and hangs Ray’s compiled DAG. This patch replaces it with NeMo RL’s lighter config that uses capture-range=cudaProfilerApi, deferring actual tracing until start_gpu_profiling() triggers cudaProfilerStart() on each internal worker via collective_rpc.

_get_raw_spec_counters() → dict[str, float | list[float]]#

Get speculative decoding metrics from the vLLM engine.

Collects spec decode counters including number of drafts, draft tokens, and accepted tokens for monitoring acceptance rates.

Returns:

Dictionary mapping metric names to their values. Values may be floats or lists of floats (for per-position metrics).

Raises:

AssertionError – If called before vLLM engine is initialized.

class nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorkerImpl(
config: nemo_rl.models.generation.vllm.config.VllmConfig,
bundle_indices: Optional[list[int]] = None,
fraction_of_gpus: float = 1.0,
seed: Optional[int] = None,
extra_env_vars: Optional[list[str]] = None,
defer_model_load: bool = False,
)#

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

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#
post_init()#
init_collective(
rank_prefix: int,
ip: str,
port: int,
world_size: int,
train_world_size: int,
) → None#
generate(
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
greedy: bool = False,
) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]#

Generate a batch of data using vLLM generation.

Parameters:

  • data – BatchedDataDict containing input_ids and input_lengths tensors

  • greedy – Whether to use greedy decoding instead of sampling

Returns:

  • output_ids: input + generated token IDs with proper padding

  • logprobs: Log probabilities for tokens

  • generation_lengths: Lengths of each response

  • unpadded_sequence_lengths: Lengths of each input + generated sequence

Return type:

BatchedDataDict conforming to GenerationOutputSpec

generate_text(
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
greedy: bool = False,
) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]#

Generate text responses using vLLM generation.

Parameters:

  • data – BatchedDataDict containing prompts with text strings

  • greedy – Whether to use greedy decoding instead of sampling

Returns:

  • texts: List of generated text responses

Return type:

BatchedDataDict containing

report_device_id() → list[str]#

Report device ID from the vLLM worker.

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

Prepare the info for refit.

update_weights_via_ipc_zmq() → bool#

Update weights from IPC handles via ZMQ socket.

update_weights_from_collective() → bool#

Update the model weights from collective communication.

reset_prefix_cache()#

Reset the prefix cache of vLLM engine.

sleep()#

Put the vLLM engine to sleep.

wake_up(**kwargs)#

Wake up the vLLM engine.

shutdown() → bool#

Clean up vLLM resources.

class nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker(
config: nemo_rl.models.generation.vllm.config.VllmConfig,
bundle_indices: Optional[list[int]] = None,
fraction_of_gpus: float = 1.0,
seed: Optional[int] = None,
extra_env_vars: Optional[list[str]] = None,
defer_model_load: bool = False,
)#

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.