nemo_rl.models.generation.vllm.vllm_worker_async#

Module Contents#

Classes#

VllmAsyncGenerationWorkerImpl

VllmAsyncGenerationWorker

Functions#

_replace_prefix_tokens

This is a subroutine used inside the vLLM Chat Completion server.

Data#

LOGGER

API#

nemo_rl.models.generation.vllm.vllm_worker_async.LOGGER#

‘getLogger(…)’

nemo_rl.models.generation.vllm.vllm_worker_async._replace_prefix_tokens(
tokenizer,
model_prefix_token_ids: list[int],
template_prefix_token_ids: list[int],
template_token_ids: list[int],
) → list[int]#

This is a subroutine used inside the vLLM Chat Completion server.

This function is for fixing up the chat template-tokenized messages history to match the model output tokenization up to the last assistant turn, in order to preserve the monotonic tokens property for optimized multi-turn training.

Some environments (namely NeMo-Gym) require an OpenAI compatible server endpoint rather than an inference engine handle. This is fine for the most part, but it may cause issues when the environment is used as a part of training.

RL training frameworks train models on token IDs, but the OpenAI compatible server communicates in what is basically de-tokenized text. When multiple model calls are made to the OpenAI compatible server in a single trajectory, model generations in previous model calls may be re-tokenized to something that is different than what was generated. This is not too big of an issue (that we know of) at inference time, but the log probs the model produces are different enough for the differently re-tokenized generation result that it causes the training to be off policy. Off policy isn’t necessarily a bad thing in isolation, but this source of off-policyness may cause unexpected issues if not properly accounted for. It also mis-aligns the token ID sequences across model calls, which feels very strange during training.

There are real cases where the model output string does not match the chat template tokenization of the parsed model output. A concrete example is inconsistent whitespace tokens around tool call special tokens.

TODO When NeMo RL supports training image generation models, we want to revisit and possibly update this function. This issue occurs when the model generates tokens that are de-tokenized into text or images, and then re-tokenized into tokens. So if there is a situation like that with images and image tokenization is non-unique, then we will need to uppdate this function.

Example (turn-by-turn, concise; eos_token_id = 2): Turn 1: - prefill_T1 (template prefill) = [11,12,13,40,41] - model output = [220,17,2] # decodes to “ 4” + EOS - model_prefix_token_ids = prefill_T1 + model output => [11,12,13,40,41,220,17,2]

Turn 2 (template retokenizes prior assistant text differently):
    - template_prefix_token_ids = [11,12,13,40,41,1001,2]  # 1001 decodes to " 4"
    - template_token_ids = [11,12,13,40,41,1001,2,21,22,40,41]

_replace_prefix_tokens keeps the exact prior model tokens up to EOS and
resumes from the template after that EOS:
    output => [11,12,13,40,41,220,17,2,21,22,40,41]
class nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorkerImpl(
config,
bundle_indices=None,
fraction_of_gpus: float = 1.0,
seed=None,
extra_env_vars: Optional[list[str]] = None,
defer_model_load: bool = False,
)#

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

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

_return_routed_experts_enabled() → bool#
_reserve_port() → None#

Bind and listen on a TCP socket to reserve a free port from the OS.

The socket is held open in LISTENING state and later passed directly to uvicorn via the sockets= parameter in server.serve(). The socket is never closed and re-opened, so there is zero gap where another process could steal the port.

load_model() → None#

Load the vLLM model and create the engine.

Called after a deferred init to perform the heavy model loading.

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

Serialise sends on AsyncMPClient.input_socket across OS threads to prevent race conditions that block the vLLM engine (e.g. during in flight weight updates in async grpo).

_start_vllm_metrics_logger() → None#

Start a background thread that periodically collects vLLM logger metrics.

Controlled by vllm_metrics_logger_interval (default: 0.5) in vllm_cfg. Runs only on the model-owner actor.

get_vllm_logger_metrics() → dict[str, Any]#
clear_vllm_logger_metrics() → None#
async post_init_async()#
async get_reserved_url() → Optional[str]#

Return the URL from the reserved socket, available before model loading.

async report_dp_openai_server_base_url() → Optional[str]#
_setup_vllm_openai_api_server(app: fastapi.FastAPI) → fastapi.FastAPI#
_setup_vllm_server() → tuple[threading.Thread, str, uvicorn.Server]#
async init_collective_async(
rank_prefix: int,
ip: str,
port: int,
world_size: int,
train_world_size: int,
) → None#
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 a batch of data using vLLM’s AsyncLLMEngine, yielding results as they are ready.

Parameters:

  • data – BatchedDataDict with input_ids and input_lengths

  • greedy – Whether to use greedy decoding instead of sampling

Yields:

Tuple of (original_index, BatchedDataDict conforming to GenerationOutputSpec for the single sequence)

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 report_device_id_async() → list[str]#

Async version of report_device_id.

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

Async version of prepare_refit_info.

async update_weights_via_ipc_zmq_async() → bool#

Async version of update_weights_via_ipc_zmq.

async update_weights_from_collective_async() → bool#

Async version of update_weights_from_collective.

async reset_prefix_cache_async()#

Async version of reset_prefix_cache.

async sleep_async()#

Async version of sleep.

async wake_up_async(**kwargs)#

Async version of wake_up.

async shutdown() → bool#

Clean up vLLM resources.

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

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