MD-TRT Support, Compile/Export, C++ and Python by narendasan · Pull Request #4183 · pytorch/TensorRT

narendasan · 2026-04-12T19:08:47Z

Description

Opening this to test the CI

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
This change requires a documentation update

Checklist:

My code follows the style guidelines of this project (You can use the linters)
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas and hacks
I have made corresponding changes to the documentation
I have added tests to verify my fix or my feature
New and existing unit tests pass locally with my changes
I have added the relevant labels to my PR in so that relevant reviewers are notified

- C++ runtime: NCCL communicator init via c10d, rank/world_size serialization, DynamicOutputAllocator, ABI version bump to 8 - Python runtime: distributed support in PythonTorchTensorRTModule and TorchTensorRTModule, NCCL library auto-detection - Conversion: native TRT DistCollective API (AllGather, ReduceScatter, AllReduce) with TRT-LLM plugin fallback - Graph lowering: fuse c10d_functional collectives + wait_tensor into single ops - Feature detection: native_trt_collectives flag, platform validation, graceful fallback chain - Build: conditional NCCL compilation via torch_nccl toolchain - Examples: tensor_parallel_simple_example.py, tensor_parallel_llama_llm.py

…, ABI v9

…g and enable DTensor decomposition

…hapes Five interconnected fixes: 1. fold_get_attr_item_calls: fold scalar param .item() calls into Python scalars before AOT tracing. Inside FakeTensorMode, even real-tensor .item() calls raise DataDependentOutputException. 2. backends.py: three changes: - call fold_get_attr_item_calls before entering FakeTensorMode - detect vmap/higher-order ops and route them through aot_autograd instead of aot_export_joint_simple (which doesn't handle HOPs) - on TRT build failure, strip TRT-only kwargs (use_fp32_acc) from the fallback graph before returning it to PyTorch 3. _decompositions.py: prevent SDPA from leaking back into the decomp table via Core ATen Interchange ops even after being removed from TORCH_TRT_DECOMPOSITIONS. 4. partitioning/common.py: lower the default max dynamic shape from min*2^16 to min*2^12 — 65536 is too large for TRT to find kernel implementations for attention ops. 5. _TorchTensorRTModule.py: move CPU scalar inputs to CUDA before execution — aot_autograd lifts scalar attributes (e.g. head_dim^-0.5) as explicit graph inputs; TRT requires all inputs on CUDA. Also fixes remove_sym_nodes to match tensor sources by equality rather than local_name so that GetItemSource bases (from torch.compile dynamic=True) are matched correctly, and updates register_sdpa.py to handle aten.scaled_dot_product_attention.default (the form produced after aot_autograd) in addition to the flash/efficient variants.

…a 3.2 1B

narendasan · 2026-04-13T16:22:54Z

core/runtime/execute_engine.cpp


 std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intrusive_ptr<TRTEngine> compiled_engine) {
+  // All inputs are expected to be on CUDA. Warn and move any that are not.
+  for (auto& inp : inputs) {


I would like to remove this but didnt have time to check if the device operations in python suppress this correctly

narendasan · 2026-04-13T16:24:49Z

core/runtime/execute_engine.cpp

+  // the constructor-time bind was deferred (e.g. no collective had been issued
+  // at construction time, or for serialized programs loaded inline where there
+  // is no Python _TorchTensorRTModule.forward wrapper).
+  if (compiled_engine->is_md && !compiled_engine->nccl_initialized) {


Not entirely sure this is necessary

narendasan · 2026-04-13T16:25:16Z

core/runtime/register_jit_hooks.cpp

+        .def(
+            "set_group_name",
+            [](c10::intrusive_ptr<TRTEngine> self, std::string group_name) {
+              self->group_name = group_name;


We should have a test where we switch engines between process groups

narendasan · 2026-04-13T16:26:59Z

core/runtime/TRTEngine.cpp

+  // process group from the c10d registry.  PyTorch assigns sequential
+  // numeric names ("0", "1", ...) to process groups; probe until we
+  // find one with an NCCL backend.
+  if (this->group_name.empty() && this->is_md) {


We should only do this if there is one available group. If there are multiple NCCL groups available we should tell the user to manually select

narendasan · 2026-04-13T16:28:18Z

docsrc/tutorials/deployment/distributed_inference.rst

+
+        def forward(self, x):
+            out = self.linear(x)
+            out = torch.ops._c10d_functional.all_reduce(out, "sum", self.group_name)


Lets dig into this more after the PR lands

narendasan · 2026-04-13T18:45:30Z

py/torch_tensorrt/distributed/run/__init__.py

+logger = logging.getLogger("torchtrtrun")
+
+
+def _get_nccl_lib_dir() -> Optional[str]:


Move into its own file

apbose · 2026-04-14T07:13:11Z

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py


+        self._nccl_comm: Optional[Any] = None
+        self._has_nccl_ops: bool = False
+


this should be set before the self.setup_engine()

apbose · 2026-04-14T07:14:10Z

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

+        inspector = self.engine.create_engine_inspector()
+        engine_json = inspector.get_engine_information(trt.LayerInformationFormat.JSON)
+        self._has_nccl_ops = "NCCL" in engine_json or "AllReduce" in engine_json
+


something like this works
engine_json_lower = engine_json.lower()
self._has_nccl_ops = "dist_collective" in engine_json_lower or "nccl" in engine_json_lower or "allreduce" in engine_json_lower

apbose and others added 11 commits April 12, 2026 11:41

removing the try-except block in TRTengine.cpp and correcting the typis

aaa6557

Redesign distributed inference API: auto-detect rank, lazy NCCL setup…

4c1e68d

…, ABI v9

remove nccl.h dependancy

7cfa40b

clean up import and add comment

ac96255

moving setup_nccl_library call to example script

fe1c6f4

work on the save/load export part-add is_md flag, guard export tracin…

b658c7a

…g and enable DTensor decomposition

refactor: Adjusting how we use NCCL

a35dfe6

test: add torch.compile(backend='tensorrt') integration test for Llam…

2aa8f14

…a 3.2 1B

feat: llama3.2 working with MD-TRT

6f81a66

meta-cla bot added the cla signed label Apr 12, 2026

github-actions bot requested a review from zewenli98 April 12, 2026 19:09

This comment was marked as outdated.

Sign in to view

narendasan force-pushed the push-vqqzkszwrvyx branch from 67134da to b5b1f5f Compare April 12, 2026 19:15

This comment was marked as outdated.

Sign in to view

narendasan force-pushed the push-vqqzkszwrvyx branch from b5b1f5f to 1957cc4 Compare April 12, 2026 19:18

This comment was marked as outdated.

Sign in to view

feat: Support exported and serialization workflows for MD-TRT

0d2d61c

narendasan force-pushed the push-vqqzkszwrvyx branch 5 times, most recently from 473cff9 to 9022e03 Compare April 13, 2026 01:14

ci: fix nccl builds in CI

e08b0c5

narendasan force-pushed the push-vqqzkszwrvyx branch from 9022e03 to e08b0c5 Compare April 13, 2026 03:49

narendasan commented Apr 13, 2026

View reviewed changes

apbose reviewed Apr 14, 2026

View reviewed changes

		logger = logging.getLogger("torchtrtrun")


		def _get_nccl_lib_dir() -> Optional[str]:


		self._nccl_comm: Optional[Any] = None
		self._has_nccl_ops: bool = False

Conversation

narendasan commented Apr 12, 2026

Description

Type of change

Checklist:

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants