dynamo: target native NCCL collectives at the op's process (sub)group

[pkisfaludi-nv]

Summary

The native TensorRT DistCollective converters (added for TRT 11 multi-device) always target the world group: they compute groups = np.arange(world_size) and layer.num_ranks = world_size, ignoring the process group the collective was actually issued on.

That makes a collective run over all ranks, which breaks 2-D device meshes where different collectives run over different subgroups — e.g. Ulysses / context-parallel all_to_all over a CP subgroup while tensor-parallel all_reduce uses a separate TP subgroup. TensorRT then rejects the layer, e.g.:

IDistCollectiveLayer ... All to All requires first input dimension to be divisible by nbRanks

because nbRanks was the world size (8) instead of the subgroup size (4).

Fix

The tensorrt::fused_nccl_* custom ops already carry group_name. This PR:

• adds _collective_group_ranks(group_name, world_size) which resolves the collective's process group to its global ranks via torch.distributed.distributed_c10d._resolve_process_group + get_process_group_ranks, falling back to the world group when it can't be resolved (single-program runs / group not created in this process);
• threads group_name from each fused-op converter (all_gather, reduce_scatter, all_reduce, all_to_all, scatter, gather) into the corresponding *_native impl;
• passes the resolved ranks to add_dist_collective and sets layer.num_ranks = len(groups).

World-group behavior is unchanged when there's a single (world) group.

Validation

On an 8×A100 (NVLink) CP4×TP2 run of the LTX-2.3 DiT, supplying the CP subgroup's ranks to add_dist_collective clears the All to All ... divisible by nbRanks build error for the Ulysses all_to_all (which previously used nbRanks=8 for a 4-rank subgroup). Files touched are Python-only (converters); no C++/runtime change.

Note for maintainers (separate issue)

With the subgroup fix in place, compiling a large graph (a 48-layer transformer, ~39k nodes) surfaces a separate TensorRT limitation: when a collective gets fused into a Myelin foreign node, the build fails with CollectiveOperation not supported for collective In toMyelinCommKind / Could not find any implementation for node {ForeignNode[...]}. Standalone collectives compile and run fine. This looks like a Myelin-side gap for collectives inside fused foreign nodes and is orthogonal to this converter change — flagging in case it's useful.

🤖 Generated with Claude Code

[meta-cla]

Hi @pkisfaludi-nv!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[narendasan]

@pkisfaludi-nv do you need me to send to the instructions to get on the CLA?

[apbose]

Thanks for the PR @pkisfaludi-nv!
So this PR bakes each collective subgroup into the engine at build.
I am thinking about the runtime scenario for the 2D Mesh case. Going by the ncclCommsplit changes for subgroup routing
Myelin per-subgroup NCCL communicators !19723
TensorRT standalone DistCollective subgroup routing !49040
TensorRT fused-Myelin subgroup routing !49062, if we take these two cases:

• C++ runtime (TRTEngine.cpp::bind_nccl_comm): today the auto-resolver bails when several NCCL groups exist. For CP4×TP2 do we pin the world group via distributed_context, or does bind need to prefer the world/parent group?
• Python runtime (_PythonTorchTensorRTModule, get_active_group()): same question — is setting the world group active enough?

[apbose]

had a doubt in this for order sensitive collectives on a permuted device mesh.
sorted() in the above would give a deterministic order which all the rank members in the group would agree on. I assume it would be a no-op for all_reduce/reduce_scatter where order wont affect the result.

But in cases like all_to_all and all_gather won't the ops be layout sensitive? eg: if DeviceMesh yields a group whose get_process_group_ranks is not ascending eg: [say rank 5, rank 2]. The subgroup comm {2,5} has two members and sorted would give the wrong order.