`nemo_rl.models.generation.vllm.vllm_worker`#

Module Contents#

Classes#

`BaseVllmGenerationWorker`
`VllmGenerationWorkerImpl`
`VllmGenerationWorker`

Functions#

`_context_capped_max_new_tokens`	Cap generation so the training prompt and response fit the context.
`_maybe_enable_vllm_native_tracing`	Optionally enable vLLM’s native OpenTelemetry tracing on the engine.
`_resolve_enable_prefix_caching`
`_merge_fp8_kwargs`	Merge fp8 init kwargs into `vllm_kwargs` in place, preserving user overrides.
`_validate_worker_extension_cls`	Reserve `worker_extension_cls` for NeMo-RL’s refit protocol.
`_log_effective_quantization_ignore_patterns`
`_configure_nvfp4_pertoken_engine_kwargs`	Configure per-token NVFP4 on the standard vLLM worker when enabled.

Data#

logger

API#

nemo_rl.models.generation.vllm.vllm_worker.logger#: ‘getLogger(…)’

nemo_rl.models.generation.vllm.vllm_worker._context_capped_max_new_tokens( *, configured_max_new_tokens: int, input_length: int, max_model_len: int, ) → int#: Cap generation so the training prompt and response fit the context.

nemo_rl.models.generation.vllm.vllm_worker._maybe_enable_vllm_native_tracing( llm_kwargs: dict[str, Any], ) → None#

Optionally enable vLLM’s native OpenTelemetry tracing on the engine.

Requires both telemetry.enabled and telemetry.vllm_native_tracing (plus an OTLP endpoint). vLLM builds its own span exporter and picks the protocol from OTEL_EXPORTER_OTLP_TRACES_PROTOCOL alone, defaulting to gRPC; it does not consult the generic OTEL_EXPORTER_OTLP_PROTOCOL. So an http/protobuf endpoint inherited from lens needs that traces-specific var set as well, or vLLM will speak gRPC at an HTTP port.

Note this only enables tracing in the engine: vLLM’s worker processes gate on OTEL_EXPORTER_OTLP_TRACES_ENDPOINT being present in their environment and never read the engine args, so they trace whenever that var reaches them – which the Ray runtime_env does for the whole cluster. This switch therefore governs engine spans, not every vLLM span.

Degrades to a no-op if the installed vLLM lacks these engine args.

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. Required FP8 ignore entries are combined with user-provided ignore entries so generated safety exclusions cannot be removed. This regression was reintroduced once already; see #1413/#2904.

nemo_rl.models.generation.vllm.vllm_worker._validate_worker_extension_cls( vllm_kwargs: dict[str, Any], required_worker_extension_cls: str, ) → None#: Reserve worker_extension_cls for NeMo-RL’s refit protocol.

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

nemo_rl.models.generation.vllm.vllm_worker._configure_nvfp4_pertoken_engine_kwargs( cfg: nemo_rl.models.generation.vllm.config.VllmConfig, llm_kwargs: dict[str, Any], ) → None#: Configure per-token NVFP4 on the standard vLLM worker when enabled.

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, num_gpus_per_node: Optional[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)
num_gpus_per_node – Number of GPUs per node in the cluster. Used to map a bundle id to a node-local engine slot when deriving VLLM_PORT. When None, the original per-engine index is used unchanged.

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.

_refit_with_reload_api_enabled() → bool#

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

_tokenize_prompt_with_bos( prompt: str | dict[str, Any], ) → str | dict[str, Any]#: Tokenize an explicit BOS once, before vLLM expands media placeholders.

_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 _spec_decode_max_tokens( base_max_tokens: int, input_len: int, max_model_len: int, spec_lookahead: int, ) → int#

Clamp max_tokens so speculative decoding never reads past max_model_len.

The drafter looks spec_lookahead tokens ahead, so generation must stop at least spec_lookahead + 1 tokens before the boundary.

classmethod _request_max_new_tokens( *, configured_max_new_tokens: int, input_length: int, max_model_len: int, cap_to_context: bool, spec_lookahead: int, ) → int#: Apply context and speculative-decoding limits to one request.

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 the vLLM engine’s Prometheus counters.

Returns the engine’s whole snapshot – the spec-decode counters plus the token, sequence-length and request-outcome series – so one RPC per snapshot serves every metric family the driver reports on.

Two shapes in the snapshot do not survive a plain {name: value} read and so are encoded into the key instead (see utils.encode_counter_key): a labelled series arrives as several objects sharing one name, so reading by name alone would keep only whichever came last, and a histogram carries count/sum/buckets with no scalar at all, so reading by name alone would drop it entirely.

Returns:: Dictionary mapping metric names to their values. Values may be floats or lists of floats (for per-position metrics). Empty before the engine is initialized.

report_refit_server_base_url() → str | None#

start_zmq_sparse_refit_relay() → str#

configure_zmq_sparse_refit_relay(relay_addresses: list[str]) → None#

stop_zmq_sparse_refit_relay() → None#

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.checkpoint_engine.VllmCheckpointEngineRpcMixin, nemo_rl.models.generation.vllm.vllm_worker.BaseVllmGenerationWorker

_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], ) → Optional[list[str]]#

Prepare the info for refit.

Returns the parameter names the engine wants pre-quantized on the trainer (vllm_cfg.refit_prequantize), or None.

update_weights_via_ipc_zmq() → bool#: Update weights from IPC handles via ZMQ socket.

update_weights_from_collective( refit_timeout_s: Optional[float] = None, ) → bool#: Update the model weights from collective communication.

init_nccl_reshard_comm_group( rank_prefix: int, pp_ips: list[str], pp_ports: list[int], pp_size: int, train_ranks_per_stage: int, sub_world_size: int, ) → None#: Forward nccl_reshard bulk-path comm group init to vLLM backend workers.

prepare_nccl_reshard_refit_info(refit_info: dict) → None#: Forward refit info to vLLM backend workers.

nccl_reshard_refit( refit_timeout_s: Optional[float] = None, ) → bool#: Receive weights from training workers via nccl_reshard (xferdtensor).

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

nemo_rl.models.generation.vllm.vllm_worker#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.models.generation.vllm.vllm_worker`#