Skip to content

[WIP] Trtllm agentic training - SWE#3

Closed
hchings wants to merge 223 commits into
shuyi+rayenfrom
trtllm-agentic-swe
Closed

[WIP] Trtllm agentic training - SWE#3
hchings wants to merge 223 commits into
shuyi+rayenfrom
trtllm-agentic-swe

Conversation

@hchings

@hchings hchings commented Jun 30, 2026

Copy link
Copy Markdown
Owner

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

  • ...

terrykong and others added 30 commits May 15, 2026 11:35
@terrykong
fix(rm): raise clear error when CP is used with DTensor RM training (N…
f87cb1b 
…VIDIA-NeMo#2483)

Signed-off-by: Terry Kong <terryk@nvidia.com>
@seonjinn @terrykong @guyueh1
feat: Add HybridEP support for MoE expert parallelism (NVIDIA-NeMo#1942)
30108c3 
Signed-off-by: sna <sna@nvidia.com>
Signed-off-by: Seonjin <sna@nvidia.com>
Signed-off-by: Seonjin Na <sna@nvidia.com>
Signed-off-by: Guyue Huang <140554423+guyueh1@users.noreply.github.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>
Co-authored-by: Guyue Huang <140554423+guyueh1@users.noreply.github.com>
@zpqiu @claude
fix: enable TE FusedAdam for Qwen3.5 MoE & GLM-4.7-Flash automodel re…
1a97e0f 
…cipes (NVIDIA-NeMo#2320)

Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@seonjinn @terrykong
perf: enable MoE GroupedGEMM for MoE models (NVIDIA-NeMo#2278)
9053154 
Signed-off-by: sna <sna@nvidia.com>
Signed-off-by: Seonjin <sna@nvidia.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>
@kajalj22 @claude
fix: remove mlm workspace (NVIDIA-NeMo#2136)
30afecc 
Signed-off-by: Kajal Jain <kajalj@nvidia.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
@kajalj22 @claude
ci: skip sglang build and tests for external contributors (NVIDIA-NeM…
79a6290 
…o#2532)

Signed-off-by: Kajal Jain <kajalj@nvidia.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
@chtruong814
ci: Skip sglang prefetch for Lfast CI label (NVIDIA-NeMo#2534)
37afd80 
Signed-off-by: Charlie Truong <chtruong@nvidia.com>
@vinhngx
feat: Auto research skill (NVIDIA-NeMo#2419)
012bf17 
Signed-off-by: Vinh Nguyen <vinhn@nvidia.com>
@xxman-google @yuki-97
feat: add AIME-2026 benchmark. (NVIDIA-NeMo#2469)
a7fb764 
Signed-off-by: Xuehan Xiong <xxman@google.com>
Co-authored-by: Yuki Huang <yukih@nvidia.com>
@achartier
docs(fp8): fix YAML key names for pow2 scaling factors (NVIDIA-NeMo#2536
ad36496 
)

Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com>
@jinglinglingling
fix: reset dataloader index when dataset changes on resume (NVIDIA-Ne…
6afb047 
…Mo#2519)

Signed-off-by: Linglin Jing <linglinj@nvidia.com>
@guyueh1 Guyue Huang
feat: Discard weight when finish generation in the main loop (NVIDIA-…
2a2e7c1 
…NeMo#2495)

Signed-off-by: Guyue Huang <guyueh@login-lyris01.lyris.clusters.nvidia.com>
Signed-off-by: Guyue Huang <guyueh@nvidia.com>
Co-authored-by: Guyue Huang <guyueh@login-lyris01.lyris.clusters.nvidia.com>
@yfw
fix: set PackedSeqParams.total_tokens for mamba seq packing (NVIDIA-N…
247f2e3 
…eMo#2529)

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
@ashors1
fix: run dynamic sampling on unshaped rewards (NVIDIA-NeMo#2478)
fc60573 
Signed-off-by: ashors1 <ashors@nvidia.com>
Signed-off-by: Anna Shors <ashors@nvidia.com>
@qiaochuz-nv
fix(data): add 'aime2026' to AIMEEvalDataConfig literal (NVIDIA-NeMo#…
aa4107e 
…2469 follow-up) (NVIDIA-NeMo#2541)

Signed-off-by: Qiaochu Zhu <qiaochuz@nvidia.com>
Signed-off-by: qiaochuz <qiaochuz@nvidia.com>
@seonjinn @terrykong
perf: selective activation checkpointing feature support (NVIDIA-NeMo…
b603568 
…#2280)

Signed-off-by: sna <sna@nvidia.com>
Signed-off-by: Seonjin <sna@nvidia.com>
Signed-off-by: seonjinn <sna@nvidia.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>
@yfw @claude
fix: aggregate rollout metrics by semantic type (min/max/sum/mean) (N…
501cd12 
…VIDIA-NeMo#2175)

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@kajalj22 @claude
chore: Upgrade vLLM from 0.17.1 to 0.20.0 (NVIDIA-NeMo#2384)
241ece5 
Signed-off-by: Kajal Jain <kajalj@nvidia.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
@ZhiyuLi-Nvidia @claude @cursoragent
feat: data plane transfer queue integration (NVIDIA-NeMo#2439)
4454fa6 
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
@terrykong
chore: move contributor-facing skills to .agents/contributor-skills/ (N…
f3e9e81 
…VIDIA-NeMo#2561)

Signed-off-by: Terry Kong <terryk@nvidia.com>
@guyueh1
fix: Fix the performance page (NVIDIA-NeMo#2563)
66059ab 
Signed-off-by: Guyue Huang <guyueh@nvidia.com>
@guyueh1
perf: Performance script tuning (NVIDIA-NeMo#2473)
ed685b6 
Signed-off-by: Guyue Huang <guyueh@nvidia.com>
Signed-off-by: Guyue Huang <140554423+guyueh1@users.noreply.github.com>
@saumishr
feat: [NemoRL] Introduce WeightSynchronizer ABC with IPC/HTTP/NCCL tr…
bb7e845 
…ansports (NVIDIA-NeMo#2466)

Signed-off-by: Saurabh Mishra <sauramishra@nvidia.com>
@chtruong814
ci: Remove GHA workflow to update skills symlink (NVIDIA-NeMo#2575)
bff7286 
Signed-off-by: Charlie Truong <chtruong@nvidia.com>
@terrykong
fix: move session-memory back to skills/ with auto-research group (NV…
fad476e 
…IDIA-NeMo#2577)

Signed-off-by: Terry Kong <terryk@nvidia.com>
@yuki-97 @Ai-chan-0411 @claude
feat(sft): make only_unmask_final configurable in SFTConfig (NVIDIA-N…
f46ecf7 
…eMo#2516)

Signed-off-by: Yuki Huang <yukih@nvidia.com>
Co-authored-by: Ai-chan-0411 (藍) <aoikabu12@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
@yfw @HeyyyyyyG
@claude Anish Mahishi
fix: use prompt token length for advantage group extraction and fix t…
c3bb1d7 
…oken mask (NVIDIA-NeMo#2176)

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Signed-off-by: Anish Mahishi <amahishi@cw-dfw-cs-001-vscode-02.cm.cluster>
Co-authored-by: Jiaqi Zeng <jiaqiz@nvidia.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Anish Mahishi <amahishi@cw-dfw-cs-001-vscode-02.cm.cluster>
@ko3n1g
ci: validate release branch-rules (NVIDIA-NeMo#2546)
354b8d3 
Signed-off-by: oliver könig <okoenig@nvidia.com>
@yuki-97
fix: lower mooncake memory in functional test (NVIDIA-NeMo#2584)
52d7c93 
Signed-off-by: Yuki Huang <yukih@nvidia.com>
@kaloyan-inherent @yuki-97
fix(vllm): serialise AsyncMPClient input_socket sends to prevent zmq …
1ded27a 
…race (NVIDIA-NeMo#2513)

Signed-off-by: Kaloyan <253267049+kaloyan-inherent@users.noreply.github.com>
Co-authored-by: Yuki Huang <yukih@nvidia.com>
@yuki-97 @claude
chore: make ppo default to dtensor to keep consistent with other algos (
30b9133 
NVIDIA-NeMo#2974)

Signed-off-by: Yuki Huang <yukih@nvidia.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@hchings hchings self-assigned this Jun 30, 2026
avenkateshha and others added 5 commits July 1, 2026 12:51
@avenkateshha @claude @RayenTian
feat(xtoken): multi-teacher support for cross-tokenizer off-policy di…
ef5dd9a 
…stillation (NVIDIA-NeMo#2797)

Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>
Signed-off-by: ruit <ruit@nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Rayen <ruit@nvidia.com>
@qiaochuz-nv
docs(mopd): fix dataset override key data.val -> data.validation (NVI…
3f3eb44 
…DIA-NeMo#2991)

Signed-off-by: Qiaochu Zhu <qiaochuz@nvidia.com>
@zyzhou5 @yuki-97
feat: R3 gym notq router replay (NVIDIA-NeMo#2915)
af207e8 
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Co-authored-by: Yuki Huang <yukih@nvidia.com>
@zpqiu
test: disable flaky grpo-qwen3.5-35ba3b-2n8g-automodel-ep16 nightly (N…
26effe2 
…VIDIA-NeMo#3020)

Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
@ricklamers @shuyixiong
feat(trtllm): add TRT-LLM generation backend
195280a 
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>
hchings
hchings commented Jul 1, 2026
View reviewed changes
Comment thread nemo_rl/distributed/virtual_cluster.py
# NOTE: Ray's runtime_env can only SET env vars, not UNSET them; setting PMI
# vars to "" is NOT equivalent to unset — OpenPMIx still finds the SLURM
# PMIx server socket in /tmp and fails. Only truly UNSET vars work.
_this_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

@hchings hchings Jul 1, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TODO: need a better way to deal w/ nested actor env.

hchings
hchings commented Jul 2, 2026
View reviewed changes
Comment thread .gitmodules
url = https://github.com/NVIDIA-NeMo/Gym.git
branch = main
url = https://github.com/hchings/Gym.git
branch = trtllm-agentic-swe

@hchings hchings Jul 2, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tentative chagnes.

shuyixiong and others added 16 commits July 2, 2026 09:49
@shuyixiong
Custom build wheel
021a8ea 
Signed-off-by: shuyix <219646547+shuyixiong@users.noreply.github.com>
@shuyixiong
build(trtllm): migrate from UV_FIND_LINKS stub wheel to PEP 517 path …
21c15a5 
…source

Signed-off-by: shuyix <219646547+shuyixiong@users.noreply.github.com>
@shuyixiong
fix(trtllm): correct workspace deps and scope cutlass override to vllm
e6c93d6 
Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com>
@shuyixiong
Pin nvidia-cutlass-dsl version in vllm
0b8a437 
Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com>
@shuyixiong
Drop custom modelopt build
45eaad6 
Signed-off-by: shuyix <219646547+shuyixiong@users.noreply.github.com>
@shuyixiong
Fix eos_id issue
1d1b2b3 
Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com>
@shuyixiong
Rename trtllm workspace
9c2297c 
Signed-off-by: shuyix <219646547+shuyixiong@users.noreply.github.com>
@shuyixiong
Clean code
ac36770 
Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
@shikicloud
allow TRT-LLM backend for async GRPO
4c466f5 
Signed-off-by: Shiki Wu <shikiw@nvidia.com>
@hchings
add trtllm agentic scripts
c7d069e
@hchings
4b smoke e2e for trtllm (agentic): prompt_token_ids + max_input_len f…
61aeb77 
…ixes, agentic Dockerfile, 4b smoke script
@hchings
30B 16 nodes trtllm swe
3af32de
@hchings
fix multiturn, fix env
12e1aef
@hchings
fix empty gen token case
de7ed78
@hchings
refactor
b9622b9
@hchings
rebase fix
c6b1fd9
@hchings
hchings force-pushed the trtllm-agentic-swe branch from bc92d0a to c6b1fd9 Compare July 10, 2026 07:25
@github-actions github-actions Bot added documentation Improvements or additions to documentation CI labels Jul 10, 2026
@hchings

hchings commented Jul 10, 2026

Copy link
Copy Markdown
Owner Author

Closing for NVIDIA-NeMo#3130

@hchings hchings closed this Jul 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

@hchings hchings

Labels

CI documentation Improvements or additions to documentation

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.