fix(distributed): subgroup NCCL regression + nccl_utils TRT 11 cleanup#4407
Open
apbose wants to merge 1 commit into
Open
fix(distributed): subgroup NCCL regression + nccl_utils TRT 11 cleanup#4407apbose wants to merge 1 commit into
apbose wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
There are some changes that do not conform to Python style guidelines:
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/distributed/test_native_nccl.py 2026-07-15 05:32:11.252442+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/distributed/test_native_nccl.py 2026-07-15 05:32:39.841891+00:00
@@ -1365,10 +1365,11 @@
This class is only reached when ENABLED_FEATURES.torch_tensorrt_runtime=False
(C++ extension not installed), so setup_engine() naturally instantiates
_TRTEngine without any patching needed.
"""
+
class LinearModel(nn.Module):
def forward(self, x: torch.Tensor) -> torch.Tensor:
return x + 1
model = LinearModel().cuda().eval()
@@ -2149,12 +2150,16 @@
tests = [
_multirank_all_reduce_correctness,
_multirank_all_gather_correctness,
_multirank_reduce_scatter_all_reduce_ops,
_multirank_all_to_all_correctness,
- lambda r, ws, dev: [_multirank_scatter_correctness(i, r, ws, dev) for i in range(ws)],
- lambda r, ws, dev: [_multirank_gather_correctness(i, r, ws, dev) for i in range(ws)],
+ lambda r, ws, dev: [
+ _multirank_scatter_correctness(i, r, ws, dev) for i in range(ws)
+ ],
+ lambda r, ws, dev: [
+ _multirank_gather_correctness(i, r, ws, dev) for i in range(ws)
+ ],
_multirank_distributed_mode_tp_model,
_multirank_distributed_mode_subgroup,
_multirank_cpp_runtime_bind_nccl,
_multirank_distributed_mode_context_switch,
_multirank_pg_migration,test_native_nccl.py: - Fix test_distributed_mode_subgroup and test_distributed_mode_context_switch regressions introduced by PR #4222 (Python runtime rework): add dist.barrier(group=subgroup) before first TRT forward so PyTorch's lazy ncclComm_t is initialized and bind_nccl_comm() sees a non-zero getCommPtr(). - Remove all use_python_runtime=True flags (deprecated no-op in 2.13). - Add _cpp_runtime_available() guard to TestPythonRuntimePickle so it skips in C++ builds where _TRTEngine's custom op is never registered. - Fix _multirank_scatter_correctness / _multirank_gather_correctness in run_multirank_tests(): wrap with lambdas that loop over root ranks to match their (root, rank, world_size, device) signature (broken since PR #4321). - Remove _multirank_distributed_mode_subgroup_python and _multirank_distributed_mode_context_switch_python (cannot work when C++ extension is loaded); collapse _multirank_pg_migration to cpp-only. _nccl_utils.py: - setup_nccl_for_torch_tensorrt(): drop ensure_nccl_symlink() call and pre-load libnccl.so.2 directly instead of via the unversioned symlink. TRT 11.0+ (ncclWrapper.cpp) tries dlopen("libnccl.so.2") first, matching PyTorch's soname exactly; glibc reuses the already-loaded handle so no symlink is needed. ensure_nccl_symlink() is kept for TRT < 11.0 users. - Update module docstring to reflect confirmed TRT 11.0+ behavior.
eaced4d to
5a37370
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
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.
test_native_nccl.py:
_nccl_utils.py: