nemo_rl.models.generation.vllm.vllm_generation#

Module Contents#

Classes#

VllmGeneration

Functions#

_record_vllm_generation_metrics

Record vLLM token-usage metrics to nemo-lens (no-op unless exporting).

Data#

logger

API#

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

‘getLogger(…)’

nemo_rl.models.generation.vllm.vllm_generation._record_vllm_generation_metrics(
model_name: str | None,
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict,
combined: nemo_rl.distributed.batched_data_dict.BatchedDataDict,
) → None#

Record vLLM token-usage metrics to nemo-lens (no-op unless exporting).

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.

set_refit_membership(
membership: nemo_rl.weight_sync.membership.RefitMembership,
) → None#

Record which shards take part in refits from now on.

Rebuilding the communicator is not enough on its own. Every refit dispatch – update_weights_from_collective, nccl_reshard_refit – goes through run_all_workers_*, which walks the whole worker group. Left alone they would keep calling the dead shard’s Ray actor after the rebuild and fail the refit with RayActorError, so the run would still die, just differently.

_refit_leader_workers() → list[Any]#

DP leaders that should receive refit calls, in rank order.

Falls back to every leader when no membership has been recorded, which is the state for the entire life of a run that never loses a shard.

rebuild_collective(
membership: nemo_rl.weight_sync.membership.RefitMembership,
ip: str,
port: int,
) → list[ray.ObjectRef]#

Re-init the collective over the surviving shards only.

Deliberately not init_collective with a filter. run_all_workers_multiple_data walks every worker in the group, so it would dispatch to the shard we are rebuilding because it is gone – and calling into a dead Ray actor is the hang this is meant to end. Here the surviving DP leaders are addressed directly.

Only leaders are called: each one collective_rpcs into its own TP/PP workers, so one Ray call per shard reaches every rank in it.

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)

attach_fleet_health(
monitor: nemo_rl.models.generation.fleet_health.GenerationFleetHealth,
selector: nemo_rl.models.generation.fleet_health.HealthyShardSelector,
) → None#

Route generation through fleet health from now on.

Parameters:

  • monitor – Owns shard eligibility and receives observed failures.

  • selector – Picks among the shards the monitor considers serving.

_next_dp_shard_idx() → int#

Return the data-parallel shard that should serve the next request.

async _generate_on_shard(
*,
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
method_name: str,
greedy: bool,
dp_shard_idx: int,
leader_worker_idx: int,
) → AsyncGenerator[tuple[int, nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]], None]#

Run one generation on a chosen shard, reporting its failures to the fleet.

A dead worker surfaces here as a Ray actor error. Reporting it is what lets the next request skip this shard instead of rediscovering the same corpse, and re-raising it as GenerationUnavailable is what tells the rollout retry policy the prompt is fine and worth re-dispatching.

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(
refit_timeout_s: Optional[float] = None,
) → list[ray.ObjectRef]#

Update weights of the policy using collective communication.

init_nccl_reshard_comm_group(
pp_ips: list[str],
pp_ports: list[int],
pp_size: int,
train_ranks_per_stage: int,
sub_world_size: int,
) → list[ray.ObjectRef]#

Initialize the nccl_reshard bulk-path comm group(s) on all gen workers.

One group per PP stage (non-PP = pp_size 1).

prepare_nccl_reshard_refit_info(refit_info: dict) → None#

Forward per-layer param metadata to vLLM workers for nccl_reshard refit.

rebuild_nccl_reshard_comm_group(
membership: nemo_rl.weight_sync.membership.RefitMembership,
pp_ips: list[str],
pp_ports: list[int],
pp_size: int,
train_ranks_per_stage: int,
sub_world_size: int,
) → list[ray.ObjectRef]#

Re-init the bulk-path comm groups over the surviving shards only.

The bulk groups are sized train_ranks_per_stage + inference_world_size, so losing a shard changes their world size as well as the shared model_update_group’s – both families have to be rebuilt together or the two disagree about who is present.

nccl_reshard_refit(
refit_timeout_s: Optional[float] = None,
) → list[ray.ObjectRef]#

Receive weights from training workers via nccl_reshard (xferdtensor).

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.

pause_generation_for_refit(*, clear_cache: bool) → bool#

Pause every async vLLM engine while preserving in-flight requests.

resume_generation_after_refit() → bool#

Resume every async vLLM engine paused for refit.

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.