feat(generation): add TensorRT-LLM rollout backend

[shuyixiong]

What does this PR do?

Adds TensorRT-LLM (TRT-LLM) as a new generation backend for RL rollouts. This lets GRPO training drive high-throughput TRT-LLM inference with in-flight weight updates (refit), in both colocated and non-colocated deployment modes.

Key changes

Generation backend (nemo_rl/models/generation/trtllm/)

• trtllm_generation.py — TrtllmGeneration implementing GenerationInterface. Manages Ray worker groups with tensor parallelism, dispatches samples round-robin across data-parallel shards, exposes generate() / generate_async() / prepare_for_generation() / finish_generation(), and handles colocated sleep/wake for GPU memory sharing.
• trtllm_worker_async.py — TrtllmAsyncGenerationWorker Ray actor wrapping AsyncLLM; concurrent per-sample generation and collective_rpc hooks for weight updates.
• trtllm_backend.py — NcclExtension WorkerExtension injected into TRT-LLM's RayGPUWorker. NCCL collective broadcasts for non-colocated refit; CUDA IPC + ZMQ weight streaming for colocated refit; wraps reload with stream sync + prefix-cache invalidation.
• config.py — TrtllmConfig (extends GenerationConfig) andTrtllmSpecificArgs (TP, MoE expert parallelism, async engine mode, in-flight weight updates, memory/engine knobs).

Algorithm / rollout wiring

• grpo.py: trtllm backend case in worker init, async-rollout detection, and refit path (shared ZMQ IPC with vLLM). NeMo-gym raises NotImplementedError for trtllm (no HTTP server yet — deferred to a follow-up agent-RL PR).
• async_utils/trajectory_collector.py and experience/rollouts.py: trtllm backend handling (async engine required).
• distributed/virtual_cluster.py + ray_actor_environment_registry.py: new PY_EXECUTABLES.TRTLLM executable (uv run --locked --extra trtllm) and worker registration.
• utils/packed_tensor.py: added stream.synchronize() after NCCL broadcasts.

Build & dependencies

• pyproject.toml: new trtllm extra pinning tensorrt-llm==1.3.0rc21, path source into 3rdparty/TensorRT-LLM-workspace, no-build-isolation config, and conflict declarations between trtllm and vllm/sglang/automodel/mcore.
• 3rdparty/TensorRT-LLM-workspace/: custom PEP 517 build backend (_backend.py) with content-addressed wheel caching + tools/build-custom-trtllm.sh that clones and compiles TRT-LLM (x86_64 and aarch64).
• docker/Dockerfile: ccache, CMake 4.0.3, TRT-LLM build args, and a SKIP_TRTLLM_BUILD flag to conditionally skip TRTLLM build and dependency.

Recipes

Adds a TRT-LLM generation exemplar plus three recipes that inherit from it:

• examples/configs/grpo_math_1B_trtllm.yaml — TRT-LLM generation exemplar (Qwen 1B math, TP1, async_engine: true); the base all recipes below inherit from.
• recipes/llm/grpo-qwen2.5-0.5b-1n4g-megatron-trtllm-noncolocated-async.yaml — Qwen2.5-0.5B, Megatron training + TRT-LLM generation + Async + Noncolocated, 1n4g.
  wandb link: https://wandb.ai/nvidia/nemo-rl/runs/b5r05hnd?nw=nwusershikiwu

• recipes/llm/grpo-qwen3-1.7b-2n4g-fsdp2-trtllm.yaml — Qwen3-1.7B, FSDP2 training + TRT-LLM generation, 2n4g.
  wandb link: https://wandb.ai/nvidia/nemo-rl/runs/upc7lm0j

• recipes/llm/grpo-dapomath17k-nanov3-30BA3B-8n4g-megatron-trtllm.yaml — Nemotron-3-Nano-30B-A3B-Base (MoE) on DAPO-Math-17K, Megatron training + colocated TRT-LLM generation, 8n4g.
  wandb link: https://wandb.ai/nvidia/nemorl-trtllm-run/groups/trtllm-nanov3-dapo

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yuki-97]

thanks for adding this and overall looks clean! left some comments.
you may also need to fix the copyright-check and check lint by run pre-commit run --all-files locally.

Scope: skipped pyproject.toml, docker/Dockerfile, tools/build-custom-trtllm.sh, 3rdparty/TensorRT-LLM-workspace/*, and uv.lock per author instruction — those are still under discussion.

Focus areas covered: the new nemo_rl/models/generation/trtllm/*.py module (worker, generation interface, NCCL/IPC backend), the algorithm-layer integration in grpo.py and friends, and the PY_EXECUTABLES.TRTLLM registration.

[shuyixiong]

/ok to test 7f3a748

[terrykong]

/ok to test 7f3a748

[shuyixiong]

/ok to test 68f56cb

[shuyixiong]

/ok to test a90b543

[shuyixiong]

/ok to test a5c8345

[terrykong]

Team review (6 agents + adversarial verification pass; all findings verified at head 68f56cb). The backend integration is solid — it closely mirrors the vLLM sibling, and we verified a number of scary-looking things are correct (NCCL rank math, IPC-ZMQ protocol/ACK pairing, logprob extraction, top_k=0 sentinel — vLLM's -1 would raise in TRT-LLM, the packed_tensor stream syncs, mcore×trtllm uv-conflict + per-actor venvs, and the full version-pin chain pyproject → workspace → _backend.py → build script → uv.lock → ccache tags, all agreeing on 1.3.0rc21 @ bf2ef86). The 14 inline comments are ranked; the ray.sub topology one is the only high-severity item.

Open items from earlier review threads (no new inline comments to avoid duplication):

• Rollout logger metrics: TrtllmGeneration still inherits the no-op get_logger_metrics/clear_logger_metrics, so throughput/KV-util metrics silently vanish for TRT-LLM runs.
• EP nightly coverage: still unanswered; generation-side moe_expert_parallel_size is unset in every recipe.
• Recipe matrix + curves: see the partial-fix reply — both mcore recipes are colocated+sync rather than the agreed non-colocated+async-grpo, and convergence curves are still missing from the description.

PR description nits: the Recipes section still describes the old 4-config set ("Qwen3 1.7B / 8B…"); the actual set is Qwen2.5-0.5B + Qwen3-1.7B-fsdp2 + nanov3-30BA3B plus the grpo_math_1B_trtllm.yaml exemplar.

Non-blocking notes: a mid-refit exception can leave the IPC-ZMQ REP socket mid-state (trtllm_backend.py:190-273) — self-limiting since refit failure is already fatal, but worth a try/finally if refit ever becomes retryable · CMake 3.31→4.0.3 is image-global; this PR's own image builds are the evidence other CUDA extensions still compile · the Python-side init_ray SLURM/MPI strip was checked and is safe (topology comes from Ray custom resources registered before Python starts).

FYI — notable uv.lock changes (for readers skimming the 2.4k-line lock diff):

• The trtllm extra adds a fourth concurrent transformers version (5.5.4, alongside the existing 5.5.0/5.6.0/5.8.1 per-extra forks), plus per-extra forks of torchao (0.15.0), nvidia-cutlass-dsl (4.5.0/4.5.1), and flashinfer-python (0.6.12) — expected given the declared extra conflicts, but each backend venv now runs different versions of these.
• Two global bumps that affect all extras, not just trtllm: zipp 3.23.1 → 4.1.0 (major) and librt 0.11.0 → 0.13.0.
• flash-attn-4 enters at a beta pin (4.0.0b15) via tensorrt-llm.
• ~45 new transitive packages ride in with tensorrt-llm, including some heavyweight baggage worth knowing about: boto3/botocore (AWS SDK), an audio stack (librosa, audioread, audioop-lts), hydra-core, etcd-sdk-python, modelscope-hub, mpi4py, and the tensorrt 10.16.1.11 suite (cu13 split packages pinned via hand-written [[tool.uv.dependency-metadata]]).

Generated by Claude Code

[joyang-nv]

/ok to test 340c8b5

[terrykong]

Build sense-check — TRT-LLM caching (joyang-nv's changes)

Focused review of the TRT-LLM build/caching work only (docker/Dockerfile, tools/build-custom-trtllm.sh, .github/workflows/_build_container.yml, 3rdparty/TensorRT-LLM-workspace/_backend.py). Static design review — the ~60-min compile was not run locally (no GPU).

Verdict: the design is sound and I found no correctness bugs. In particular:

• The layer ordering genuinely achieves the TE-reuse goal: the TE-building syncs run in one layer before the TRT-LLM inputs are COPY'd, and a [tool.trtllm].ref bump changes neither uv.lock nor the root pyproject, so it only busts the last layer.
• The wheel cache key in _backend.py correctly folds url+ref+version+real platform tag+arch+torch/CUDA toolchain.
• CCACHE_COMPILERCHECK=content and the incremental-vs-compact registry layering with a nonce-forced export RUN are correct, and the assert_patch_target fail-loud guards are the right instinct.

The inline comments are the few things worth eyeballing before trusting it in prod — none are blockers. The lead item (unconditional compile) is the one that actually matters for a build sense-check.

Generated by Claude Code

[terrykong]

/ok to test e927101

[joyang-nv]

/ok to test b479c39

[joyang-nv]

/ok to test c9dd7ec