nemo_rl.models.generation.vllm.vllm_generation#

Module Contents#

Classes#

VllmGeneration

Data#

logger

API#

nemo_rl.models.generation.vllm.vllm_generation.logger#

‘getLogger(…)’

class nemo_rl.models.generation.vllm.vllm_generation.VllmGeneration(
cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster,
config: nemo_rl.models.generation.vllm.config.VllmConfig,
name_prefix: str = 'vllm_policy',
workers_per_node: Optional[Union[int, list[int]]] = None,
defer_model_load: bool = False,
)#

Bases: nemo_rl.models.generation.interfaces.GenerationInterface

static init_cluster_placement_groups(
cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster,
config: nemo_rl.models.generation.vllm.config.VllmConfig,
) → None#

Pre-initialize placement groups matching the strategy VllmGeneration expects.

Call this before constructing VllmGeneration when other components compete for the same Ray resources and you need deterministic ordering — topology-constrained inference PGs should be created before unconstrained ones so they claim domain-aligned nodes first.

VllmGeneration.__init__ calls _init_placement_groups internally, but that call early-returns when PGs already exist, so calling this method first is safe.

_get_tied_worker_bundle_indices(
cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster,
) → list[tuple[int, list[int]]]#

Calculate bundle indices for tensor and pipeline parallel workers.

Handles both unified placement groups (for cross-node model parallelism) and per-node placement groups (for node-local model parallelism).

_report_device_id() → list[list[str]]#

Report the device ID of vllm workers.

_report_dp_openai_server_base_urls() → list[Optional[str]]#

Report the data parallel OpenAI server base URLs of vLLM workers, only populated if it is async vLLM engine and the HTTP server is active.

_collect_reserved_urls() → list[Optional[str]]#

Collect reserved URLs from DP leaders before model loading.

Only called when defer_model_load=True. Workers have bound ports during init and can report their reserved URLs immediately.

load_and_start() → None#

Load models on all workers and start HTTP servers.

Called after a deferred init (defer_model_load=True) to perform the heavy model loading. Updates dp_openai_server_base_urls with the actual running server URLs and populates device_uuids.

_post_init()#
_get_raw_spec_counters() → dict[str | tuple[str, int], float]#

Collect raw spec decode counters from workers.

snapshot_step_metrics() → None#

Snapshot current spec decode counters to begin tracking a training step.

Call this before generation to establish a baseline for metrics delta.

Raises:

RuntimeWarning – If called twice without get_step_metrics() in between.

get_step_metrics() → dict[str, float]#

Get speculative decoding metrics delta since snapshot_step_metrics().

Returns:

Dictionary of delta metrics with ‘vllm/’ prefix. Returns empty dict if snapshot_step_metrics() was not called.

Raises:

RuntimeWarning – If called without snapshot_step_metrics() first.

init_collective(
ip: str,
port: int,
world_size: int,
*,
train_world_size: int,
) → list[ray.ObjectRef]#

Initialize the collective communication.

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.

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.

async _async_generate_base(
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
method_name: str,
data_validation_fn,
greedy: bool = False,
) → AsyncGenerator[tuple[int, nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]], None]#

Base async generation method that handles common worker management logic.

Parameters:

  • data – Input data for generation

  • method_name – Name of the worker method to call (‘generate_async’ or ‘generate_text_async’)

  • data_validation_fn – Function to validate input data

  • greedy – Whether to use greedy decoding

Yields:

Tuple of (original_index, BatchedDataDict containing generation result)

async generate_text_async(
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
greedy: bool = False,
) → AsyncGenerator[tuple[int, nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]], None]#

Generate text responses asynchronously, yielding results as they are ready.

Parameters:

  • data – BatchedDataDict containing prompts with text strings

  • greedy – Whether to use greedy decoding instead of sampling

Yields:

Tuple of (original_index, BatchedDataDict containing single text response)

async generate_async(
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
greedy: bool = False,
) → AsyncGenerator[tuple[int, nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]], None]#

Generate responses asynchronously, yielding individual samples as they complete.

This method provides per-sample streaming across all workers, yielding each sample result as soon as it’s ready, regardless of which worker processed it.

prepare_for_generation(
*args: Any,
**kwargs: Any,
) → bool#

Wake workers up for colocated inference.

finish_generation(*args: Any, **kwargs: Any) → bool#

Sleep workers and reset prefix cache.

shutdown() → bool#

Shut down all vLLM workers and clean up resources.

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

Prepare the info for refit.

update_weights_via_ipc_zmq() → list[ray.ObjectRef]#

Update weights of the policy using IPC handles via ZMQ socket.

update_weights_from_collective() → list[ray.ObjectRef]#

Update weights of the policy using collective communication.

start_gpu_profiling() → None#

Start GPU profiling.

stop_gpu_profiling() → None#

Stop GPU profiling.

get_vllm_logger_metrics() → dict[str, Any]#

Collect vLLM logger metrics from vLLM workers (model-owner actors only).

clear_vllm_logger_metrics() → None#
clear_logger_metrics() → None#

Clear logger metrics for performance reporting.

get_logger_metrics() → dict[str, Any]#

Get logger metrics for performance reporting.

__del__() → None#

Shuts down the worker groups when the object is deleted or is garbage collected.

This is an extra safety net in case the user forgets to call shutdown() and the pointer to the object is lost due to leaving a function scope. It’s always recommended that the user calls shutdown().

invalidate_kv_cache() → bool#

Invalidate reusable caches in vLLM (e.g., prefix/KV cache) after weight updates.

For async_engine, calls reset_prefix_cache_async on workers. For sync, calls reset_prefix_cache. Returns True if all workers report success.

property requires_kv_scale_sync: bool#

Check if KV cache scales should be synchronized during refit.

Returns True if kv_cache_dtype is fp8/fp8_e4m3.