Zongfeij/dev - #4131
Open
zongfeijing wants to merge 45 commits into
Open
Zongfeij/dev#4131zongfeijing wants to merge 45 commits into
zongfeijing wants to merge 45 commits into
Conversation
Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com> Signed-off-by: Erin Ho <erinh@nvidia.com> Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Erin Ho <erinh@nvidia.com> Co-authored-by: Erin Ho <14718778+hchings@users.noreply.github.com>
NVIDIA-NeMo#3428) Signed-off-by: Superjomn <yanchunwei@outlook.com> Co-authored-by: Erin Ho <14718778+hchings@users.noreply.github.com>
NVIDIA-NeMo#3466) Signed-off-by: Superjomn <yanchunwei@outlook.com>
Brings up PD-disaggregated generation end-to-end on GB200: a replica's inference GPUs are split into context (prefill) and generation (decode) engines fronted by one OpenAI-compatible disagg server, which is the single URL NeMo Gym talks to. Engine plumbing. TrtllmGeneration plans engines per replica from trtllm_cfg.disaggregation (engine counts, per-role TP/EP overrides, routers, cache transceiver backend) and hands each worker its role. DisaggServerActor wraps TRT-LLM's OpenAIDisaggServer; trtllm_disagg_server.py adapts it to the NeMo Gym request shape. config.py gains the disaggregation schema, and build-custom-trtllm.sh plus the Dockerfile pick up UCX and NIXL so the cache transceiver is actually compiled in -- without them an engine aborts on the first KV transfer. Six bring-up fixes, each of which silently broke the path rather than failing loudly: - DisaggServerActor ran in the driver's environment and died with ModuleNotFoundError: tensorrt_llm. Give it the engine workers' interpreter, which RayWorkerGroup now exposes as py_executable. - OpenAIDisaggServer builds a prometheus MultiProcessCollector in register_routes(), which raises unless PROMETHEUS_MULTIPROC_DIR is set. TRT-LLM's own entrypoint calls set_prometheus_multiproc_dir() first; we construct the server directly, so call it too. - The middleware stripping NeMo Gym's vLLM-only request fields was a Starlette BaseHTTPMiddleware, which hands the downstream app its own captured receive channel, so reassigning request._receive never reached FastAPI's validation and every request 400'd on extra_forbidden. Rewritten as raw ASGI. - Aggregated serving lost its rollout fields: they moved off the message onto declared response fields for the disagg path, but with no disagg server to re-attach them NeMo Gym silently dropped every assistant turn. Attach them on the message when no disaggregation is in play. - Under disaggregation the unit that must stay inside one NVLink domain is the replica -- its context and generation engines exchange KV every turn -- not the engine. Sizing gpus_per_instance by the engine yielded nodes_per_instance=1 and skipped domain pinning entirely. - Per-node placement groups were consumed in creation order, so two adjacent pg_idx values could sit in different NVLink domains and split a replica across the fabric. Consume them in topology order instead. The engine HTTP server also moves from asyncio.to_thread(llm.generate) to llm.generate_async: the blocking API parks a worker thread per in-flight request, capping concurrency at the default executor size rather than at the engine's scheduler. Error fidelity. OpenAIDisaggServer._handle_exception only re-raises HTTPException, so a 4xx from an engine arrives as an aiohttp.ClientResponseError, falls into the catch-all, and reaches the caller as 500. The aggregated server returns it as a 4xx, and Gym accounts for the two classes differently -- which would make masking statistics incomparable between the aggregated and disaggregated paths, the exact comparison this work exists to support. Re-raise 400-499 with the original status; 5xx still goes to super(). Empty rollouts. A Gym rollout can return without a single assistant turn (the agent stalls before its first completion and Gym's wall-clock timeout kills it). That raised ValueError, taking the run down over one sample and losing the step's other 127 rollouts. NRL_SKIP_FAILING_EMPTY_ROLLOUT gates it: the default "0" keeps the raise, since the usual causes are misconfigurations worth surfacing; "1" stands the sample up as prompt-only and masks it out of the loss, reported as train/num_masked_seqs_by_empty_rollout. That count overlaps num_mask_sample_filtered by design and must not be summed with it -- num_valid_samples stays authoritative. Note there is no circuit breaker on a sustained rate. Profiling. Under disaggregation every engine runs the same worker class, so nsys reports differed only by %p pid and matching a trace to the context or generation side meant grepping the driver log. The -o filename now carries the role and an ordinal (_context0, _context1, _generation0), appended rather than prefixed so the report names documented in docs/nsys-profiling.md stay prefix-matchable. Also bumps TRT-LLM to 1.3.0rc24. Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
The HTTP server built its TrtSamplingParams from temperature and top_p only, so a recipe setting policy.generation.top_k was silently ignored whenever rollouts went through the server -- which is every NeMo Gym and PD-disaggregated run. The direct path (_build_sampling_params) has always applied it, so the two paths sampled from different distributions for the same config, and a run that switched between them was not comparable with itself. Pass top_k through, using the same convention as the direct path: TRT-LLM spells "no top-k restriction" as 0 while the generation config spells it as null. Add it to the request-validation loop too, so a request that disagrees with the server's config is rejected rather than quietly overridden -- the same treatment temperature and top_p already get. That loop now reads sampling_config with .get(): top_k is absent from configs written before this change, and a KeyError there would reject every request instead of the mismatched ones. Also set logprobs_simple_format=True. Without it TRT-LLM returns the verbose logprob structure, which the adapter has to walk per token; the simple format is what the direct path consumes. Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
…kens The Rubin TRT-LLM base image ships PIP_CONSTRAINT=/etc/pip/constraint.txt, pinning the versions vendored into the image's own Python. build_wheel.py's setup_venv() shells out to real pip, which honours it, and those versions are not published anywhere: the image pins cuda-python==13.4.0, which does not exist on PyPI (latest 13.x is 13.3.1, already installed by uv). So the otherwise-satisfied `cuda-python>=13` in requirements.txt became unresolvable and the build died before cmake ran. Unset the constraint -- this venv is uv-managed and owes the base image's site-packages nothing. This only appeared after moving to the NGC-derived Rubin base; cuda-dl-base does not set PIP_CONSTRAINT. Also redact the clone token from two places that printed it verbatim into the build log: the "TRT-LLM Git URL" echo, and _backend.py's CalledProcessError, which stringifies the whole argv (the expanded url included) on failure. And drop nvidia-modelopt from requirements.txt, preventively: build_wheel.py pip-installs into the live nemo-rl venv, so the ref's `~=0.39.0` pin would let pip downgrade the modelopt uv resolved from a git rev at 0.46.0.dev*. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
Adding 10.7 to TORCH_CUDA_ARCH_LIST / HYBRID_EP_CUDA_ARCH_LIST fails to build: the torch extensions compiled during `uv sync` (mamba-ssm, causal-conv1d, transformer-engine) and DeepEP cannot target sm_107 with this toolchain. Fall back to Blackwell, matching what the opt/dev-backup-rubin reference shipped. TRT-LLM itself still targets Rubin via _DEFAULT_ARCH / BUILD_CUSTOM_TRTLLM_ARCH =107-real; only the torch-side extensions are affected here. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
Extend the default SM arch list to 100-real;103-real;107-real (GB200/B200, GB300/B300, Rubin) in both places that must agree: _DEFAULT_ARCH in _backend.py, which is folded into the wheel cache key, and the ARCH fallback in build-custom-trtllm.sh. The nvshmem patch gets the same three archs as 100\;103\;107 -- bare, because nvshmem rejects the suffixed names CMake generates, which is the reason that patch exists. Its semicolons stay backslash-escaped: the string lands inside CMAKE_CACHE_ARGS of an ExternalProject_Add, where an unescaped ';' would split one cache entry into three arguments and silently build sm_100 only. That patch does not read BUILD_CUSTOM_TRTLLM_ARCH, so it and _DEFAULT_ARCH have to be edited together; both comments now say so. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
The pinned ref's requirements.txt carries tensorrt~=10.16.1, but it was absent from [project].dependencies, so uv never resolved it -- the wheel build merely pip-installed it transiently into the venv via requirements-dev.txt. Declare it so the version uv manages matches the ref and the base image's TRT 10.16.1.11. tensorrt on PyPI ships only an sdist whose wheel_stub downloads multi-GB binaries at build time, which would make `uv lock` fetch them. Supply static [[tool.uv.dependency-metadata]] for tensorrt / tensorrt-cu13 / -libs / -bindings so the graph resolves from metadata alone; the real binaries are fetched at `uv sync --extra trtllm` time inside Docker. Same approach, same placement and versions as origin/main -- our base d5fb8d0 predates it. Lock goes 553 -> 557 packages (the four tensorrt entries); `uv lock --check` passes with submodules at their committed pins. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
torch_dist checkpointing opens with a save-planning gather_object over the default process group: Megatron's dist_checkpointing/strategies/torch.py passes process_group=None into save_state_dict_async_plan, and DCP's _DistWrapper forwards that to dist.gather_object. On a NCCL default group that is a device collective, and NCCL allocates its buffers with its own cudaMalloc rather than through PyTorch's caching allocator, so memory the allocator is merely holding is unreachable to it. With the device full the cudaMalloc fails and NCCL surfaces it as the opaque "NCCL Error 1: unhandled cuda error" -- no torch OOM is raised anywhere, because torch never requested the memory. Job 2724162 hit this at its first save: every training rank sat at 277.5 GiB of 277.5 GiB and the run died in that gather. Reclaiming first measured 68 GiB of headroom per rank (device-free 9.8 -> 77.7 GiB) and the 740 GiB step_1 checkpoint then completed. nccl_reshard_refit and prepare_for_lp_inference already clear the cache ahead of their own large NCCL phases; save_checkpoint was the path that skipped it. The [CKPT_MEM] line is kept because the failure it guards against is silent -- NCCL reports only "unhandled cuda error", so a recurrence would otherwise leave nothing to reason from. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
Ports 4a910e9 and d9c9d61 (Michal Futrega, branch trtllm-agentic-swe-mlperf) onto opt/dev, which has neither. The default 10-minute NCCL watchdog SIGABRTs ranks that wait on peers stuck in transient first-hit stalls -- triton JIT of GDN kernels for new packed-sequence shapes being the prime suspect -- observed as coordinated "Terminating the process after attempting to dump debug info" during step-2 collectives in jobs 2274348/2280396. Step 1 always passed, so these are warm-up effects rather than real hangs; true hangs stay bounded by the SLURM walltime. Both call sites are needed. setup_distributed() raises the default group's timeout, which sub-groups normally inherit, but initialize_megatron creates the TP/PP/EP/DP sub-groups with an EXPLICIT dist.distributed_timeout_minutes (mcore default 10) that overrides that inheritance -- job 2282178 still died at Timeout(ms)=600000 with only the first fix. Both read the same NRL_NCCL_TIMEOUT_MINUTES env var, default 60. Applied by hand rather than cherry-picked: the original commits sit on a base whose init_process_group already passed device_id=cuda:{local_rank}, which opt/dev does not have. Only the timeout change is taken, so the device-binding difference between the branches is left alone. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
Moves the custom TensorRT-LLM wheel to user/zongfeij/rl @ 7035705878, which adds recompute_active_requests and the Qwen3.5 MXFP8 refit path. Still 1.3.0rc23, so [project].version and the root pyproject's trtllm extra pin are unchanged. [project].dependencies are left as-is: the ref's requirements.txt relaxes the tensorrt, flashinfer-python, nvidia-cutlass-dsl and apache-tvm-ffi pins and adds nvidia-cuda-nvrtc==13.4.46rc1 (all from the 0.8 drop container update), but re-syncing those is deferred to a separate change. Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com>
Launchers already export MOUNT_LOG_DIR_IN_CONTAINER=1 alongside BASE_LOG_DIR (qwen35_397b_grpo_trtllm/launch_ray_cluster.sh does), but nothing here read it, so /logs stayed container-local: a per-node tmpfs in the enroot data dir. Anything written under /logs on one node was therefore invisible to every other node. NeMo-Gym is the case that matters -- its server resolves agent trajectories and llm_completions under /logs while the Ray runners that produce them are spread across the allocation, so the server found nothing, every rollout came back with zero output items, and the trajectory collector failed the step. mlperf's run_and_time.sh and logger.log_dir hardcode /logs the same way. Mount LOG_DIR (already required to be on a shared filesystem -- the STARTED_RAY_HEAD / ray_worker_units / ENDED signalling depends on it) at /logs when the flag is set. Guarded on empty MOUNTS like the UV_CACHE_DIR_OVERRIDE block above, and appended before COMMON_SRUN_ARGS consumes MOUNTS. Side benefit: driver and per-agent apptainer logs now survive the job instead of dying with the container's tmpfs.
Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
Force fla-core (and its flash-linear-attention meta) to 0.5.2 via the [tool.uv] override-dependencies. 0.5.2 carries the two Blackwell GDN autotune restrictions that 0.5.1 lacks: - PR NVIDIA-NeMo#953 (commit 3eeef4e): restrict the fwd h kernel (chunk_gated_delta_rule_fwd_kernel_h_blockdim64) to num_warps=2 on Blackwell — Triton tl.dot recurrence race (fla issue NVIDIA-NeMo#945). - PR NVIDIA-NeMo#1000 (commit ac6c648): restrict the gated-delta bwd autotune (prepare_wy_repr_bwd_kernel) on Blackwell. megatron-core's dev extra pins flash-linear-attention==0.5.1, which is the racy release; upstream 0.5.2 fixes it. The recipes keep FLA_TILELANG=0 (triton backend) — the tilelang GDN backward fault is not addressed by 0.5.2.
configure_generation_config gave vllm_cfg["load_format"] the value "auto" if is_eval else "dummy" but left trtllm_cfg alone, so a TRT-LLM engine read the full checkpoint at startup even though refit_policy_generation overwrites every weight before the trajectory collector is allowed to issue its first request. On a 397B engine that read is minutes of startup the run throws away. Mirror the vLLM contract: set trtllm_cfg["load_format"] in configure_generation_config and pass it through to the AsyncLLM constructor. precision="fp8" stays "dummy" in both modes -- the engine is built with the quantized layout the BF16 refit populates, and configure_fp8_llm_kwargs already rejects anything else -- so this changes nothing for FP8. Evaluation keeps "auto": it has no refit to supply weights.
The collective (non-colocated) refit path called reset_prefix_cache(), which only drops the reusable prefix blocks. Requests already in flight keep the KV they computed under the old weights, so the tail of every in-flight trajectory is generated against a mixed weight state. recompute_active_requests() — added to PyExecutor by the tekit ref this branch pins — instead sends the active requests back through context so their KV is rebuilt with the weights that just landed.
The guard pinned the exact private class name `fastokens._compat._TokenizerShim`. fastokens 0.3.1 renamed it to `fastokens._ConfiguredTokenizerShim`, so `patch_transformers()` succeeded but the worker still refused to start: RuntimeError: NRL_USE_FASTOKENS=1, but the TRT-LLM HTTP tokenizer backend is 'fastokens._ConfiguredTokenizerShim'; expected 'fastokens._compat._TokenizerShim' What the check is for is whether the monkey-patch landed at all, so match on the `fastokens.` module prefix instead of an internal class name that moves with every release.
init_collective calls ncclCommInitRank to build the refit communicator across all train + inference ranks. Since tekit 97b62625 the executor loop runs its per-iteration _broadcast_request_count as an NCCL broadcast on the engine's TP group rather than over CPU/gloo, so the two now contend for the same device and deadlock: the broadcast waits on peer ranks whose main thread is inside ncclCommInitRank, which in turn waits for every rank to arrive. PG5 stalls at work 85 after 84 clean iterations and ProcessGroupNCCL aborts the engine 600 s later. Wrap init_collective in control_action so the loop parks at a step boundary first. This is the mechanism the refit path already uses (update_weights_from_collective, update_weights_via_ipc_zmq); only init_collective was missing it. Verified on a 16-node 397B MXFP8 run with the NCCL object collectives left enabled: step 7/30, six refits, the step-5 checkpoint save, zero watchdog timeouts. Before the change the same configuration died in init_collective every time.
get_logprobs returned BatchedDataDict(...).to("cpu"), leaving the
device-to-host transfer to BatchedDataDict.to(); that path hung at scale.
Allocate an explicit pinned host tensor and copy into it before returning,
so the transfer is direct and the returned logprobs are pinned.
Applied from the optimized repo's carried patch
qwen35_397b_grpo_trtllm/pytorch/patches/nemo-rl/0010-fix-megatron-copy-logprobs-to-pinned-CPU-memory.patch
(optimized commit c885684bc), which ports 6c8c46cbc from qwen35_397b_grpo.
Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com>
Rebases the RL tekit branch onto feat/rubin-bringup, which replaces the
feat_sm107 flashinfer source build with the 0.6.18 internal wheel and drops
the cutlass-dsl 4.5.0 pin (tekit ff4598531 / 559775ea8 / de0cf4d8a).
Two build-custom-trtllm.sh patches become no-ops on this ref and would have
aborted via assert_patch_target:
- requirements.txt already ships `setuptools>=80`
- cutlass_kernels/CMakeLists.txt dropped the `setup_library.py develop
--user` execute_process in favour of PYTHONPATH for generate_kernels.py
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Picks up tekit 780b94e65, which passes the three trailing arguments trtllm_paged_attention_decode gained in the base image's flashinfer 0.6.18+8c3bbc00. The image can now keep the base's flashinfer instead of force-installing the older +cf3c3a3e build. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
CI passes GITLAB_CLONE_ACCESS_TOKEN already as "user:token", so the prefix expanded to "oauth2:user:token". URL userinfo splits on the first colon, which made the password "user:token" and every CI build failed the tekit fetch with "HTTP Basic: Access denied. You must use a token instead of a password". A bare token hid the bug locally, where "oauth2:glpat-x" is well formed. The url now carries the token alone, matching how the optimized Dockerfile passes it elsewhere; callers holding a bare token add their own username. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…e skill The skill's NeMo-RL patch instruments the vLLM worker. TRT-LLM has no equivalent, so the per-request milestones its schema expects are adapted here from RequestPerfMetrics, giving a two-layer trace (engine + Gym/OpenHands) for the generation-only flow. - timer.py: emit Chrome-trace events when NRL_TRAINING_TIMELINE is set, applied from the skill's nemo-rl-timeline.patch. - trtllm_worker_async.py: dump the resolved AsyncLLM kwargs, so a run's actual engine config is verifiable from the log rather than inferred from the RUN_ID. This is what exposed that an explicitly requested enable_block_reuse=True was arriving as False. - trtllm_http_server.py: map RequestPerfMetrics onto the six nemo_vllm_* fields plus cached_tokens. Two details in the adapter are not obvious and were each arrived at from a failing run rather than from the API surface: **The epoch origin is anchored per request, not once globally.** TRT-LLM reports timing_metrics as steady_clock timedeltas and exposes no steady_clock_now() to Python, so an origin has to be estimated. A single global offset absorbs the first request's post-generation latency and then reports every later request's last_token that much too late; any later request that post-processes faster then reports last_token AFTER response_ready. That inverted the pair on 3621 of 11520 model calls on a 397B run. Anchoring on each request's own last_token_time makes the whole chain ordered by construction, because the frontend interval strictly contains the engine interval. Intra-request deltas stay exact; absolute cross-request alignment gains sub-millisecond jitter. **Prefix reuse comes from kv_cache_metrics and is clamped to the prompt length.** There is no `cached_tokens` attribute anywhere in tensorrt_llm 1.3.0rc23 -- reading one, as the vLLM path does, silently yields None on every request, which NeMo Gym then reports as zero reuse. On an agentic workload where every turn re-sends the whole conversation, a hardcoded zero reads as "the prefix cache is doing nothing", which is exactly the conclusion someone would act on. Reuse actually lives on RequestPerfMetrics.kv_cache_metrics and is counted in BLOCKS, so it is scaled by tokens_per_block and then clamped: the consumer validates 0 <= cached <= prompt and drops the whole cache-metric group when a partially-reused tail block rounds the value past the prompt -- deleting the metric on precisely the highest-reuse requests. Field-name notes for consumers: `queued_ts_us` is not a native TRT-LLM event and is mapped to arrival_time, so first_scheduled - queued is the TRT-LLM request queue; `generation_replica_idx` is a valid constant 0 in gen-only, not a missing field. Validated on the 397B gen-only workload (8 prompts x 16 generations x 3 steps): validate_timeline.py with --require-model-token-counts, --require-model-cache-metrics, --require-model-scheduling-timestamps and --require-generation-replica reports valid=true with zero errors, cache metrics present on 11519/11519 successful model calls, and zero timestamp inversions. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Three CacheTransceiverConfig knobs the recipe could not reach, each forwarded only when set so TRT-LLM's own defaults hold otherwise: - cache_transceiver_runtime: "auto" silently resolves to the C++ transceiver whenever it cannot confirm the model's preference, and a hybrid Mamba model under disaggregation needs the Python (v2) transceiver for its recurrent-state handoff -- so the recipe must be able to force it. - kv_cache_bounce_size_mb: coalesces a request's scattered per-block KV into one contiguous fabric-VMM buffer and a single multi-rail NIXL write, sidestepping per-block registration failures. - kv_transfer_timeout_ms: TRT-LLM's 60 s default is tuned for short prompts at low concurrency; at high rollout concurrency bulk ctx-side timeouts feed a cancel/retry churn that stresses the transceiver, so large multi-turn workloads want a much larger value. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
At high rollout concurrency the single DisaggServerActor (one uvicorn
process) saturates around 17 turns/s and becomes the replica's ceiling:
conversations pile up ahead of the engines while GPUs idle. This makes
the frontend horizontally scalable and, since the same investigation
needed to see where pre-generation time actually goes, adds permanent
end-to-end timing stamps.
Frontend sharding (num_frontend_workers, default 1 = old behavior):
- trtllm_generation: replica x frontend actor fan-out; NeMo-Gym receives
every frontend URL and its per-session client selection pins each
conversation to one frontend (sticky routing is per-process state, so
stickiness only holds within a frontend -- pinning sessions is what
makes N processes correct).
- Deterministic ports (frontend_base_port + idx on the pinned node) and
serve-from-constructor: a Ray-restarted actor re-binds the same port,
so the URL Gym holds stays valid across crashes.
- Snowflake node_id = replica_idx * num_frontends + frontend_idx keeps
ctx_request_id unique across frontends (process_id is hardwired 0
in-process and time.monotonic shares an origin per node).
Per-request work moved off the hot single processes:
- gen_tokids_ctxbytes / gen_strip_message_history: the gen leg carries
b64 int32 token ids instead of a 30k-int JSON array plus the full
message history it never reads.
- frontend_tokenize: frontends render the chat template and tokenize
via build_spliced_prompt_ids -- the exact pipeline the adapters use,
now hoisted to module level as the single source of truth -- and
attach prompt_token_ids_b64 to the ctx leg. Inbound payloads are
normalized through ChatCompletionRequest.model_dump(exclude_unset)
first: the raw Gym payload orders tool-JSON keys differently than the
pydantic-normalized form the adapters see, and the chat template is
sensitive to that order (turn-1 only; the splice covers later turns).
Guarded by ctx-side shadow validation
(NRL_TRTLLM_TOKENIZE_SHADOW_RATE) which re-tokenizes a sample on the
adapter and logs any divergence with a decoded token window.
- The supplied-ids short-circuit in the adapter route is hoisted before
template rendering, so the gen leg no longer renders a template whose
output it discards.
Full-path timeline stamps (all gated NRL_TRTLLM_EMIT_TIMELINE_FIELDS):
- Frontend ASGI middleware stamps receive/forward times as headers;
the response wrapper surfaces them as nemo_fe_recv/fwd_ts_us.
- The ctx adapter emits nemo_ctx_{arrival,queued,first_scheduled,
done}_ts_us on the context leg; the disagg service relays them onto
the final response (tekit carries the relay change), decomposing the
previously-opaque pre-generation leg into frontend / ctx submit /
ctx queue / prefill / KV-handoff segments per model call.
Validated end to end at conc-512 (2048 rollouts, 30-turn agentic
workload, 61k model calls): stamps present on 100% of calls, sub-leg
sum identical to the parent leg, tokenize shadow divergence zero.
Sharding the frontend (N=8) plus frontend tokenize took the equal-GPU
disagg configuration from well behind the aggregated baseline to ahead
of it.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The ctx-leg stamps previously required patching TRT-LLM's disagg service to copy four NeMo-named fields onto the generation response -- NeMo naming baked into the service's core method. TRT-LLM's ResponseHooks is the designed observability channel and on_ctx_resp already receives the full context response, so with the hooks class now substitutable upstream (response_hooks_cls), the whole relay lives here: _CtxStampRelayHooks parks the stamps on raw_req.state, and _attach_rollout_fields -- which already rewrites the final payload and holds the same raw_req -- moves them onto the response. Self-gating: when the ctx adapter does not emit the stamps (NRL_TRTLLM_EMIT_TIMELINE_FIELDS=0) nothing is stashed or attached. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The final response's cached_tokens comes from the generation leg, so under ctx-first disaggregation the per-call timeline had no view of how much of the prompt the CONTEXT engine actually reused -- and the engine-level prefill volume said reuse was far below the delta-only floor. The executor sets request.cached_tokens to the reused prefix length at the first chunk and bridges it onto the result; surface it as nemo_ctx_cached_tokens on the context response and relay it with the other ctx stamps. Zero is a meaningful value (a full miss), so it is not filtered like a timestamp. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…compute diagnostics The conversation-affine ADP router (kv_cache_routing_conversation_affinity) pins a conversation's turns to the rank that holds its prefix, but it only sees the id when the request carries ConversationParams. The HTTP adapter never forwarded one, so under attention-DP every turn landed on an effectively random rank: measured at 2P-DEP8 / conc 512 only 17-26% of turns found the previous turn's KV on the serving rank (about 1/DEP) and the ctx engine re-prefilled most of the history every turn. Read the id from the body's conversation_params (the Gym model proxy sends its session id there; one session per rollout) or, failing that, from the id the disagg service stamps onto disaggregated_params for its ctx/gen legs, and pass ConversationParams(conversation_id) to llm.generate_async. No id means no affinity, i.e. the previous behaviour. With the id: 96-97% of turns on the rank holding their prefix, ctx prefill tokens -66% to -88%. Also surface the context leg's compute accounting on the ctx response and relay it with the other ctx stamps (nemo_ctx_computed_tokens, nemo_ctx_first_begin, nemo_ctx_num_chunks): the engine's cached_tokens overshoots the real reuse point by a few blocks, so computed tokens and the first chunk's begin position are what the reuse analysis needs. These live on the RequestOutput (the per-choice CompletionOutput only forwards cached_tokens), hence the extra request_output argument. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…e ENV Comma-separated variable names get --container-env on every srun so the host value beats the image's baked ENV (pyxis lets the image win by default). First use: UCX_NET_DEVICES=all for NIXL KV transfer over the RDMA HCAs (the image bakes the management NIC os_p1s0) when running PD-disaggregated generation under this launcher.
…ontend middleware The MLPerf warmup sends required_prefix_token_ids (the vLLM server's prefix override extension) with every synthetic request. TRT-LLM's OpenAIDisaggServer validates bodies against extra="forbid" models, so the field turned each warmup request into a 400 and the warmup exception then took the driver down before run_start. The TRT-LLM engine adapter derives the on-policy prefix from the assistant messages and ignores the top-level field, so dropping it in the disagg frontend (like the Dynamo wrapper does) keeps both paths identical.
Moves [tool.trtllm].ref from 780b94e655 to 1c6e8f3dfb ("Leave generation-only
requests out of the post-refit KV recompute on disaggregated engines").
The other two pins that the file's own comment says move together stay put,
having been checked rather than assumed: tensorrt_llm/version.py reads
1.3.0rc23 at both refs, so `version` here and the `trtllm` extra in the root
pyproject.toml are already correct, and requirements.txt is unchanged between
the two refs, so [project].dependencies needs no re-sync.
This is not a fast-forward -- the refs have diverged, with 20+ commits on the
new side, mostly disaggregation work: in-flight weight updates on attention-DP
disagg engines, refit with piecewise CUDA graphs, the folded save-last prefill
that TLLM_MAMBA_FOLD_SAVE_LAST pins, and the FMHA trailing-argument fix. What
the old ref carried that the new one does not has not been audited.
The wheel cache is ref-scoped, so this forces a rebuild.
Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com>
Picks up "[None][fix] Ship the end-of-prompt recurrent state to the generation engine when the save-last fold is on" and everything else on user/zongfeij/rl since 1c6e8f3dfb.
…NeMo#3942) Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com>
Moves [tool.trtllm].ref in 3rdparty/TensorRT-LLM-workspace/pyproject.toml from 854bcea144 to b1cecb7c9d. build_trtllm compiles the engine from this ref, so the image picks the new TRT-LLM up on the next NEMO_RL_REVISION bump. The sha appears only here; nothing else in the tree references it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com>
Backport NVIDIA-NeMo#3984 (release_grads_before_refit) from optimized's patches 0002 and 0006 (originally by seonjinn <sna@nvidia.com>), then extend it to the trtllm generation backend: factory.py's _SUPPORTED_BACKENDS and validate_release_grads_before_refit only allowed vLLM. CollectiveWeightSynchronizer is backend-agnostic on the trainer side and TrtllmGeneration already implements the generation-side interface it needs, so trtllm just needed the gate relaxed. Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
max_batches was max_val_samples // val_batch_size, so a validation set that is not a multiple of the batch size silently dropped its tail (251 prompts in waves of 126 ran one wave of 126 and skipped the other 125). Round up so every validation prompt is scored; the per-wave metrics are aggregated across waves as before. Only the async GRPO validate() is changed (the sync/ppo paths keep the floor).
configure_fp8_moe_backend forced CUTLASS whenever precision="fp8" with is_mx=true, because MXFP8CutlassFusedMoEMethod used to be the only MXFP8 MoE implementation. TRT-LLM now also serves MXFP8 experts on Rubin through the fused FC1+FC2 CuTe DSL kernel (MXFP8CuteDslFusedMoEMethod, backend CUTEDSL), which inherits the CUTLASS weight storage and the bucket-by-bucket refit path, so the refit caster needs no change. Gen-only on Qwen3.5-397B TP8/EP8 the CuTe DSL kernel is 12-15 % faster per decode iteration and 5 % per step than CUTLASS. Allow CUTLASS or CUTEDSL for is_mx (case-insensitive, preserving the other MoeConfig fields), keep CUTLASS as the default when no backend is configured, and keep rejecting anything else. The error message now lists both accepted backends. Signed-off-by: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com>
Contributor
|
Directly cherry-picked to main...opt/dev-rubin-pin |
svcnvidia-nemo-ci
added
the
waiting-on-customer
hchings
force-pushed
the
opt/dev-rubin-pin
branch
from
September 17, 2026 18:46
99954a1 to
8c6f63f
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Add mxfp8 cutedsl moe backend for trtllm
Issues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information