feat(trtllm): TRT-LLM prefill/decode disaggregation for GRPO rollouts - #4095
Draft
shuyixiong wants to merge 7 commits into
Draft
feat(trtllm): TRT-LLM prefill/decode disaggregation for GRPO rollouts#4095shuyixiong wants to merge 7 commits into
shuyixiong wants to merge 7 commits into
Conversation
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> (cherry picked from commit a8f1ad3)
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> (cherry picked from commit 407f383)
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>
(cherry picked from commit 769aab4)
…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> (cherry picked from commit 14690fb)
…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. (cherry picked from commit f68ebd8)
…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. (cherry picked from commit 46cd1a1)
The disagg base commit bundled a second, independent change: NRL_SKIP_FAILING_EMPTY_ROLLOUT, which turns "NeMo Gym returned a rollout with no assistant turn" from a hard failure into a prompt-only sample masked out of the loss. That is an error-handling policy change, not part of prefill/decode disaggregation, and it changes behaviour for every backend. Take it out and restore the original raise. Removed: - nemo_rl/environments/nemo_gym.py -- restored verbatim; its entire delta over the branch point was this feature. An empty rollout raises ValueError again, with the original message (no NRL_SKIP_FAILING_EMPTY_ROLLOUT hint). - nemo_rl/experience/rollouts.py -- restored verbatim; likewise, the only delta was carrying the `empty_rollout` flag onto the batch. - nemo_rl/algorithms/grpo.py -- _apply_empty_rollout_filter and its two call sites (grpo_train, async_grpo_train) plus the train/num_masked_seqs_by_empty_rollout metric. Kept in grpo.py: the disagg-aware gpus_per_instance sizing, which is the part of that commit's grpo.py delta that does belong to disaggregation -- under PD the NVLink-domain unit is the replica, not the engine, and sizing by the engine skipped domain pinning entirely. Verified: no empty_rollout / NRL_SKIP_FAILING_EMPTY_ROLLOUT / num_masked_seqs_by_empty references remain, and the three modules compile. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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
2 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 a one line overview of what this PR aims to accomplish.
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