[DRAFT] feat: enable native multi-device TensorRT on TensorRT-RTX builds - #4462
Open
apbose wants to merge 2 commits into
Open
[DRAFT] feat: enable native multi-device TensorRT on TensorRT-RTX builds#4462apbose wants to merge 2 commits into
apbose wants to merge 2 commits into
Conversation
TensorRT-RTX 1.5 ships the full multi-device API (IDistCollectiveLayer,
addDistCollective, setNbRanks, IExecutionContext::setCommunicator), but
Torch-TensorRT compiled all of it out of RTX builds.
core/runtime/TRTEngine.h gated TRT_HAS_NATIVE_NCCL on
NV_TENSORRT_MAJOR/MINOR >= 10.16. TensorRT-RTX defines TRT_MAJOR_RTX /
TRT_MINOR_RTX and then aliases NV_TENSORRT_MAJOR/MINOR to them, so
NV_TENSORRT_MAJOR is 1 on RTX and that comparison never matched.
ENABLE_TRT_NCCL_COLLECTIVES was therefore never defined, and the whole MD
runtime -- bind_nccl_comm, set_group_name, release_nccl_comm, the lazy
bind in execute_engine.cpp, and NATIVE_TRT_COLLECTIVES_AVAIL -- was absent
from RTX builds. The two release lines use incompatible version schemes,
so detect the RTX package first and version-check against its own
numbering, mirroring is_tensorrt_version_supported() on the Python side.
Nothing else needed changing: the Bazel NCCL detection probes PyTorch
(RTX-agnostic on Linux), USE_C10D_NCCL and nccl_headers already reach the
RTX configs, the tensorrt_rtx bindings expose add_dist_collective /
CollectiveOperation / ReduceOperation / set_communicator / num_ranks, and
MD is not behind a PreviewFeature on RTX so the existing hasattr guard in
_TRTInterpreter correctly skips it.
CI: the `distributed` suite was variants=("standard",) and so never ran
against RTX. Add the rtx variant, overriding away test_nccl_ops.py and
USE_TRTLLM_PLUGINS -- both are TensorRT-LLM-only and would no-op on RTX.
Also add a per-suite `runner` field to the manifest. The suite's
--multirank follow-ups need 2 GPUs, but #4397 dropped the explicit
multi-GPU runner the old build-test-linux-x86_64.yml set for this job, so
it had been falling back to the single-GPU validation_runner. Restore
g4dn.12xlarge for standard; rtx uses g5.12xlarge (A10G/SM 8.6) since the
TensorRT docs describe DistCollective as requiring Ampere or newer.
Not yet validated on hardware; needs a 2-GPU RTX Linux box.
apbose
changed the title
github-actions
Bot
added
component: tests
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
1 participant
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.
No description provided.