Skip to content

feat(trtllm): add nsys GPU profiling for TRT-LLM generation backend#1

Draft
Superjomn wants to merge 10 commits into
shuyix/trtllmfrom
chunweiy/nemorl-trtllm-cleanup
Draft

feat(trtllm): add nsys GPU profiling for TRT-LLM generation backend#1
Superjomn wants to merge 10 commits into
shuyix/trtllmfrom
chunweiy/nemorl-trtllm-cleanup

Conversation

@Superjomn

Copy link
Copy Markdown
Collaborator

Wire the TRT-LLM generation backend into NeMo-RL's existing nsys worker-pattern interface (NRL_NSYS_WORKER_PATTERNS / NRL_NSYS_PROFILE_STEP_RANGE).

  • The sync/async generation actors resolve nsys options via a TRT-LLM adapter and forward them into TRT-LLM's inner Ray executor through ray_worker_nsight_options (setdefault, so user kwargs win), so the ranks that actually run the engine are profiled.
  • The adapter layers TRT-LLM-only options (python-gil,osrt trace, cuda-memory-usage, capture-range-end=repeat-shutdown:1, kill=none, %h in the output name) on top of the shared base config without mutating it, keeping other backends' defaults unchanged.
  • The TRT-LLM engine self-manages its capture window, so the trainer's start/stop_gpu_profiling hooks are no-ops for this backend.
  • Add focused contract tests and docs.

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

Issues

List issues that this PR closes (syntax):

Usage

  • You can potentially add a usage example below
# Add a code snippet demonstrating how to use this

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

  • ...

ricklamers and others added 10 commits July 1, 2026 08:47
Add the TRT-LLM PyTorch backend integration needed to validate GRPO with draft-target speculative decoding, Mongo rollout logging, and weight-sync timing on AIHub.

chore(trtllm): drop speculative decoding and MongoDB logging

Remove the experimental speculative-decoding hooks (config, worker
plumbing, spec_token_origins propagation, specdec exemplar config)
and the MongoDB logger backend that were bundled with the TRT-LLM
backend introduction. The TRT-LLM PyTorch backend itself is unchanged.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>

chore(trtllm): WIP — wire trtllm backend on top of CUDA 13

TMP commit (will amend or squash later). Layers TRT-LLM-specific
wiring on top of the CUDA 13 squash from NVIDIA-NeMo#2332:

  - pyproject.toml: pin trtllm extra to tensorrt_llm==1.3.0rc13
  - virtual_cluster.py: add PY_EXECUTABLES.TRTLLM
  - ray_actor_environment_registry.py: add TRTLLM_EXECUTABLE
    (with NEMO_RL_PY_EXECUTABLES_SYSTEM=1 fallback)
  - docker/Dockerfile: add SKIP_TRTLLM_BUILD arg + HPCX env exposure
    (cuda-compat removed — cu13.2 base image already provides it)
  - docs/docker.md: SKIP_TRTLLM_BUILD usage docs

Known TODOs before this can be promoted:
  - `uv lock` likely fails: tensorrt-llm 1.3.0rc13 has no cp313 wheel.
    Either move trtllm to its own workspace (cp312) or relax pin.
  - Build + smoke test on a GB200 node not yet validated.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>

Downgrade to Python 3.12

Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>

Successfully run nemorl+trtllm and verify convergence

Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>

Add async support

Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>

Support colocated case

Cleanup

Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>

drop cp313 flash-attn wheel pin to unblock cp312 lock

Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>

drop branch-local nemo_gym changes

Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>

Refactor

Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>

Refactor code

Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>

Change flash attn wheel

Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>

Verify perf

Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>

Discard weights

Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>

add py3.13 build infrastructure with torch 2.10

Signed-off-by: Shiki Wu <shikiw@ipp1-3010.ipp1a1.colossus.nvidia.com>

Initial Async support

Signed-off-by: Shiki Wu <shikiw@ipp1-3010.ipp1a1.colossus.nvidia.com>

add async support for trtllm

Signed-off-by: Shiki Wu <shikiw@nvidia.com>
Signed-off-by: shuyix <219646547+shuyixiong@users.noreply.github.com>
…source

Signed-off-by: shuyix <219646547+shuyixiong@users.noreply.github.com>
Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com>
Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com>
Signed-off-by: shuyix <219646547+shuyixiong@users.noreply.github.com>
Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com>
Signed-off-by: shuyix <219646547+shuyixiong@users.noreply.github.com>
Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
Wire the TRT-LLM generation backend into NeMo-RL's existing nsys
worker-pattern interface (NRL_NSYS_WORKER_PATTERNS /
NRL_NSYS_PROFILE_STEP_RANGE).

- The sync/async generation actors resolve nsys options via a TRT-LLM
  adapter and forward them into TRT-LLM's inner Ray executor through
  ray_worker_nsight_options (setdefault, so user kwargs win), so the
  ranks that actually run the engine are profiled.
- The adapter layers TRT-LLM-only options (python-gil,osrt trace,
  cuda-memory-usage, capture-range-end=repeat-shutdown:1, kill=none, %h
  in the output name) on top of the shared base config without mutating
  it, keeping other backends' defaults unchanged.
- The TRT-LLM engine self-manages its capture window, so the trainer's
  start/stop_gpu_profiling hooks are no-ops for this backend.
- Add focused contract tests and docs.

Signed-off-by: Superjomn <yanchunwei@outlook.com>
@Superjomn
Superjomn marked this pull request as draft July 7, 2026 02:24
@github-actions github-actions Bot added the documentation Improvements or additions to documentation label Jul 7, 2026
@shuyixiong
shuyixiong force-pushed the shuyix/trtllm branch 3 times, most recently from 82f4407 to afa959f Compare July 21, 2026 03:56
@joyang-nv
joyang-nv force-pushed the shuyix/trtllm branch 2 times, most recently from 6099141 to 340c8b5 Compare July 23, 2026 00:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants