Native DistCollective fails at build time when fused into a Myelin ForeignNode (large multi-device graph)

[pkisfaludi-nv]

Bug Description

The TRT-11 native multi-device collectives added in #4321 (add_dist_collective in py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py) work for small / standalone graphs, but fail during the TensorRT engine build on a large graph when a DistCollective layer is fused — together with surrounding reshape/unsqueeze/pointwise ops — into a Myelin ForeignNode. Myelin can't map the collective kind (toMyelinCommKind) once the IDistCollectiveLayer is absorbed into the fused node:

ERROR: Error Code 9: Skipping tactic 0x0 due to exception CollectiveOperation not supported for collective In toMyelinCommKind at /_src/optimizer/myelin/myelinDistCollectiveLayer.cpp:49
ERROR: Error Code 10: Internal Error (Could not find any implementation for node {ForeignNode[[UNSQUEEZE]...[SHUFFLE]-[.../proj_out/_reshape_copy]]} In computeCosts ... optimizer.cpp:4284)

Standalone/unfused collectives build and run fine (a 2-GPU all_reduce+all_to_all compiles + runs, no TRT-LLM); the failure only appears when the collective is fused into a ForeignNode, which happens naturally in a real 48-layer transformer (~39k nodes) with Ulysses all_to_all + tensor-parallel all_reduce.

To Reproduce (2 ranks)

A module with Linear → reshape/unsqueeze around a dist.all_to_all_single, so the collective is surrounded by shuffle/pointwise ops that Myelin fuses into one ForeignNode:

import torch, torch.nn as nn, torch.distributed as dist
import torch_tensorrt

class M(nn.Module):
    def __init__(self):
        super().__init__(); self.pin = nn.Linear(64,64); self.pout = nn.Linear(64,64)
    def forward(self, x):                       # [B,S,64]
        x = self.pin(x); b,s,h = x.shape
        x = x.reshape(b*s, h).unsqueeze(0).reshape(b*s, h).contiguous()
        y = torch.empty_like(x); dist.all_to_all_single(y, x)   # -> tensorrt::fused_nccl_all_to_all
        return self.pout(y.reshape(b, s, h))

def main():
    dist.init_process_group("nccl"); r = dist.get_rank(); torch.cuda.set_device(r)
    m = M().cuda().eval(); x = torch.randn(2, 8, 64, device="cuda")
    ep = torch.export.export(m, (x,))
    trt = torch_tensorrt.dynamo.compile(ep, inputs=[x], use_distributed_mode_trace=True)  # native collectives
    trt(x)

main()   # torchrun --nproc_per_node=2 repro.py

What works vs fails

• WORKS: standalone / unfused IDistCollectiveLayer (2-GPU all_reduce + all_to_all, no TRT-LLM), parity vs eager.
• FAILS: large graph where the collective is fused into a Myelin ForeignNode (errors above).

What we tried (did not help)

• optimization_level=0 — does not prevent the Myelin fusion.
• require_full_compilation=False — no effect; this is a build-time Myelin failure, not a converter-coverage partition fallback.
• There is no exposed torch-tensorrt setting to mark an IDistCollectiveLayer non-fusable, and no public TensorRT BuilderFlag / PreviewFeature / env var to disable Myelin fusion or keep collective layers standalone.

Expected behavior

Either (a) Myelin should support the collective kind inside a fused ForeignNode (extend toMyelinCommKind), or (b) the builder should not fuse an IDistCollectiveLayer into a ForeignNode (keep it standalone) so large multi-device graphs build, or (c) torch-tensorrt should expose a way to force a fusion boundary around add_dist_collective layers.

Environment

• torch-tensorrt 2.14.0.dev0 built from main@d2d459a (includes TRT 11 MD Ops #4321 "TRT 11 MD Ops")
• TensorRT 11.0.0.114 / 11.1.0.106 ; PyTorch 2.14 nightly (cu130) ; 8×A100 ; MULTIDEVICE_RUNTIME_10_16 preview enabled

Additional context

Flagged as a "Note for maintainers (separate issue)" in #4380 (which fixes the orthogonal subgroup/nbRanks divisibility problem — independent of this). Current workaround: add the fused collective ops (torch.ops.tensorrt.fused_nccl_all_to_all.default, ...fused_nccl_all_reduce.default) to torch_executed_ops so they run in PyTorch and create partition boundaries, keeping the DiT compute in TensorRT — at a graph-break/perf cost.

[apbose]

So I have been seeing this issue in Setup: Qwen3-30B-A3B (48 layers, OOTB EP → _c10d_functional.all_reduce), torch.compile backend, C++ runtime (use_python_runtime=False), TRT 11.0.0.114, 2×B200.] on torchTRT main

Symptom (runtime, not build):

IExecutionContext::setCommunicator: Error Code 3: mIsMdEngine == true failed — "The engine is not a multi-device engine"
IExecutionContext::enqueueV3: Error Code 1: Myelin [comm.cpp:48 get_nccl_comm] Invalid communicator ID: 0 In executeMyelinGraph
Engines build fine (48 fused_nccl_all_reduce layers created), but at execution setCommunicator is rejected on nearly every collective engine (2 of 96 bind attempts succeed; requires_native_multidevice=1 on all). The collectives then silently no-op → numerically wrong EP output.

What we found (from ProfilingVerbosity::DETAILED per-engine layer info): every per-layer engine does contain the native [DIST_COLLECTIVE]-[fused_nccl_all_reduce] (comm_allreduce) layer — but only engines that also get the signal/wait multi-stream sync layers are promoted to multi-device (mIsMdEngine=true). Comparing two:

• _run_on_acc_385 (collective from layers.47): 9 layers — comm_allreduce + signal×2 + wait×2 + a few kgen/shape_call, no GEMMs → mIsMdEngine=true → setCommunicator succeeds.
• _run_on_acc_9 (collective from layers.0): 21 layers — comm_allreduce + gemm×4 + kgen×13 + …, no signal/wait → mIsMdEngine=false → setCommunicator rejected. (This is the case for 47 of 48.)

So when fused_nccl_all_reduce is partition-fused into a large, compute-heavy TRT engine (collective + GEMMs), TRT does not emit the signal/wait scaffolding and does not set mIsMdEngine. Small, GEMM-free collective engines get promoted; large fused ones don't. This looks like the same "collective fused with surrounding ops" , manifesting as a runtime mIsMdEngine=false rather than the build-time toMyelinCommKind/ForeignNode failure. (No ForeignNode/Myelin-fusion in our layer info, FWIW — the collective is a distinct comm_allreduce layer in both engines.) @pkisfaludi-nv what is your opinion on this?