fix: update ray.sub - #3942
Open
hchings wants to merge 24 commits into
Open
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>
Carried downstream patch from mlperf optimized qwen35_397b_grpo (patches/nemo-rl/0001-build-add-Rubin-runtime-compatibility.patch), originally generated against NVIDIA-NeMo/RL main d5fb8d0. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A restored replay buffer can make the current step immediately trainable before the collector claims the next target. Starting training in that state lets refit advance the collector past the missing lookahead and can stall the following step indefinitely. Report reserved targets through collector status and hold the startup barrier only until the next target is complete or actively claimed. This prevents the resume deadlock while preserving training and rollout overlap. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Backport the required FP8 refit pieces from NVIDIA-NeMo/RL PR NVIDIA-NeMo#2744 to daf46ff. Resolve the wrapped Qwen3.5 parameter hierarchy before FP8 classification. Also expand packed Megatron-Bridge expert slabs into per-expert block-FP8 weights and scales; unlike the former 41352f7 source fork, this pin no longer splits them on the default collective-refit path. Initialize the FP8 configuration inside RayExecutorV2 workers before model construction. RayExecutorV2 has no pre-load collective_rpc hook, so pass the configuration through its inherited environment and consume it when the NeMo-RL worker extension is imported.
Squash of NVIDIA-NeMo/RL PR NVIDIA-NeMo#3545 (head 7180bff, 2026-08-11) onto main@d5fb8d04 with patches 0001-0003 applied. Replaces the carried PR NVIDIA-NeMo#3296 module (bf16_trtllm_moe.py): the vllm_backend refit path now natively reloads unquantized FlashInfer TRTLLM MoE weights, detects the realized refit backend, constrains native refit scope, defaults injected reload state safely, and hardens layerwise-refit failures; the ModelOpt quant backend opts out explicitly. Retire when the upstream PR merges.
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>
Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
shuyixiong
force-pushed
the
opt/dev-rubin-pin
branch
from
September 1, 2026 06:42
6eafdf0 to
b7ef775
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 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