`nemo_rl.models.generation.trtllm.trtllm_backend`#

TRT-LLM WorkerExtension for NCCL / IPC weight synchronisation.

Injected into TRT-LLM’s RayGPUWorker via ray_worker_extension_cls.

update_weights_from_collective — NCCL broadcast via packed_broadcast_consumer, used in non-colocated mode.
update_weights_via_ipc_zmq — CUDA IPC handles streamed over a per-GPU ZMQ socket, used in colocated mode (NCCL can’t form a group when train and inference processes share the same physical GPU).

Module Contents#

NCCL-based weight update extension for TRT-LLM Ray workers.

Call a refit lifecycle hook when supported by the installed TRT-LLM.

nemo_rl.models.generation.trtllm.trtllm_backend._call_model_loader_hook_if_available( model_loader: Any, hook_name: str, ) → bool#: Call a refit lifecycle hook when supported by the installed TRT-LLM.

class nemo_rl.models.generation.trtllm.trtllm_backend.NcclExtension#

Bases: tensorrt_llm.llmapi.rlhf_utils.WorkerExtension

NCCL-based weight update extension for TRT-LLM Ray workers.

Attributes set by TRT-LLM’s mixin injection (from RayGPUWorker): self.engine – PyExecutor instance self.device_id – int GPU ordinal

init_collective( rank_prefix: int, ip: str, port: int, world_size: int, train_world_size: int, ) → None#

start_gpu_profiling() → None#: Start CUDA profiler on this GPU worker (nsys capture-range trigger).

_finalize_weight_update() → None#: Finalize refit using TRT-LLM’s CUDA-graph-safe path when available.

update_weights_from_collective( *, drain: bool = True, recompute_kv: bool = False, ) → bool#

Receive weights via NCCL broadcast and update model parameters.

Parameters:

drain – If True (default), wait for all in-flight requests to drain before applying weights — exclusive engine access. If False, the swap happens at a scheduler step boundary with in-flight requests still in the engine (in-flight weight update).
recompute_kv – Only meaningful with drain=False. If True, preempt in-flight requests so they re-prefill under the new weights. Otherwise, they keep decoding with their current KV cache. The reusable prefix cache is cleared after every weight update.

update_weights_via_ipc_zmq() → bool#

Receive weights via CUDA-IPC + ZMQ, reload model.

Trainer sends (ipc_handle, list_keys, used_bytes) chunks; end of refit is signalled by IPCProtocol.COMPLETE.

cleanup_zmq() → None#: Close ZMQ socket if open — called from worker shutdown.