nemo_rl.models.generation.trtllm.trtllm_generation#

GenerationInterface implementation backed by TRT-LLM.

Non-colocated: separate train / inference GPU sets, NCCL broadcast for weight sync. Colocated: shares GPUs with the policy and uses sleep/wakeup to time-multiplex GPU memory between training and inference phases.

Module Contents#

Classes#

TrtllmGeneration

TRT-LLM generation backend (requires trtllm_cfg.async_engine=true).

API#

class nemo_rl.models.generation.trtllm.trtllm_generation.TrtllmGeneration(
cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster,
config: nemo_rl.models.generation.trtllm.config.TrtllmConfig,
name_prefix: str = 'trtllm_policy',
workers_per_node: Optional[Union[int, list[int]]] = None,
)#

Bases: nemo_rl.models.generation.interfaces.GenerationInterface

TRT-LLM generation backend (requires trtllm_cfg.async_engine=true).

Initialization

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

Pre-initialize placement groups matching TRT-LLM’s topology.

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

Calculate bundle indices for tensor-parallel worker groups.

Handles both unified placement groups (cross-node model parallelism) and per-node placement groups (node-local model parallelism). For unified PGs, bundles are reordered by physical node before slicing so each TP group stays as node-local as possible.

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

Collect HTTP server base URLs from each DP-rank-0 worker.

_report_device_id() → list[list[str]]#
init_collective(
ip: str,
port: int,
world_size: int,
*,
train_world_size: int,
) → list[ray.ObjectRef]#
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]#
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]#

Yield a single-sample generation result.

Called by run_async_multi_turn_rollout, which dispatches one sample at a time per coroutine. The async worker’s max_concurrency lets multiple in-flight Ray calls share the same AsyncLLM, which batches them internally via asyncio.gather.

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

Wake inference workers up. No-op for non-colocated.

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

Sleep workers (colocated) or reset prefix cache (non-colocated).

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

Grpo profiling protocol: start nsys capture on the GPU workers.

stop_gpu_profiling() → None#

Grpo profiling protocol: stop nsys capture on the GPU workers.

update_weights_from_collective() → list[ray.ObjectRef]#
update_weights_via_ipc_zmq() → list[ray.ObjectRef]#

Receive weights via CUDA-IPC + ZMQ (colocated mode).

invalidate_kv_cache() → bool#

No-op for TRT-LLM: KV-cache invalidation happens inside the refit path.

For async RL correctness, KV/prefix-cache invalidation must happen in the same engine step boundary as the weight update — otherwise in-flight requests forward several decode steps with new weights × old KV, opening a race window.

TRT-LLM avoids this by performing the invalidation inside the refit function itself, under the same control_action context:

  • NcclExtension.update_weights_from_collective (NCCL path)

  • NcclExtension.update_weights_via_ipc_zmq (IPC-ZMQ path)

clear_logger_metrics() → None#

No-op: TRT-LLM rollout telemetry is not yet wired up.

vLLM overrides this (vllm_generation.py) to reset AsyncLLM iteration stats (rollout throughput, KV-cache utilization, in-flight batch counters). TRT-LLM has no equivalent plumbing yet, so this stays a no-op rather than inheriting silently — making the gap explicit.

TODO: fetch AsyncLLM iteration stats from the TRT-LLM workers and clear them here, mirroring clear_vllm_logger_metrics.

get_logger_metrics() → dict[str, Any]#

Return empty metrics: TRT-LLM rollout telemetry is not yet wired up.

See :meth:clear_logger_metrics. Returning {} explicitly documents that rollout observability metrics (present for vLLM) are absent for TRT-LLM runs, rather than letting them silently disappear via the inherited no-op default.

TODO: surface AsyncLLM iteration stats here, mirroring get_vllm_logger_metrics.

shutdown() → bool#
__del__() → None#