nvidia--tensorrt/samples/sampleDistCollective/README.md

TensorRT Distributed Collective Sample

Table Of Contents

• Description
• How does this sample work?
  • TensorRT API layers and ops
• Prerequisites
• Running the sample
  • Sample --help options
• Additional resources
• License
• Changelog
• Known issues

Description

This sample, sampleDistCollective, demonstrates how to use TensorRT for multi-GPU inference by creating and running TensorRT networks with IDistCollectiveLayer. It tests a specific collective operation specified via the required --op argument by building a network for that operation and verifying the results.

How does this sample work?

The sample builds a TensorRT network containing a distributed collective layer, then runs inference across multiple GPUs using NCCL for GPU-to-GPU communication.

Specifically:

• INetworkDefinition::addDistCollective is called to add the collective layer (kALL_REDUCE, kALL_GATHER, kBROADCAST, kREDUCE, kREDUCE_SCATTER, kALL_TO_ALL, kGATHER, or kSCATTER).
• IDistCollectiveLayer::setNbRanks is called to set the number of ranks for the collective operation.
• The NCCL unique ID is coordinated via a shared file specified by TRT_NCCL_ID_FILE. Rank 0 generates the ID and writes it to the file; other ranks wait and read it.
• ncclCommInitRank is called to initialize the NCCL communicator on each rank.
• IExecutionContext::setCommunicator is called to set the NCCL communicator on the execution context.
• After inference, each rank verifies its output matches the expected result for the collective operation.

TensorRT API layers and ops

In this sample, the IDistCollectiveLayer is used for distributed collective operations across multiple GPUs. For more information, see the TensorRT Developer Guide: Layers documentation.

Prerequisites

1. Multiple GPUs: This sample requires at least 2 GPUs.

2. NCCL: Install NCCL library (version should be >= 2.19.0 and < 3.0):

   sudo apt-get install -y libnccl2 libnccl-dev
   

3. Process Launcher (one of the following):

   • SLURM: srun command (available on HPC clusters)
   • Open MPI: mpirun command

     sudo apt-get install -y openmpi-bin libopenmpi-dev
     

Running the sample

1. Compile this sample by following build instructions in TensorRT README. The binary named sample_dist_collective will be created in the <TensorRT root directory>/bin directory.

2. Run the sample with 2 processes. The sample requires the following environment variables:

   • TRT_MY_RANK: The rank of this process (0 to WORLD_SIZE-1).
   • TRT_WORLD_SIZE: The total number of processes.
   • TRT_NCCL_ID_FILE: Path to a shared file for NCCL ID coordination. Rank 0 writes the NCCL unique ID to this file, and other ranks read from it. The file should be empty or non-existent before starting.

   Using SLURM (srun):

   srun --ntasks=2 bash -lc 'export TRT_MY_RANK=$SLURM_PROCID; \
       export TRT_WORLD_SIZE=$SLURM_NTASKS; \
       export TRT_NCCL_ID_FILE=/tmp/nccl_id.txt; \
       ./sample_dist_collective --op all_reduce'
   

   Using Open MPI (mpirun):

   mpirun -np 2 bash -lc 'export TRT_MY_RANK=$OMPI_COMM_WORLD_RANK; \
       export TRT_WORLD_SIZE=$OMPI_COMM_WORLD_SIZE; \
       export TRT_NCCL_ID_FILE=/tmp/nccl_id.txt; \
       ./sample_dist_collective --op all_reduce'
   

   Note: Make sure to delete or clear the TRT_NCCL_ID_FILE before each run to ensure a fresh NCCL ID is generated.

   Available operations:

   • all_reduce - Reduces data across all ranks and distributes the result to all ranks
   • all_gather - Gathers data from all ranks and distributes the concatenated result to all ranks
   • broadcast - Broadcasts data from rank 0 to all other ranks
   • reduce - Reduces data across all ranks and sends the result to rank 0
   • reduce_scatter - Reduces data across all ranks and scatters portions of the result to each rank
   • all_to_all - Exchanges equally sized data chunks between all ranks
   • gather - Gathers data from all ranks to rank 0
   • scatter - Scatters data from rank 0 to all ranks

3. Verify that the sample ran successfully. If the sample runs successfully you should see output similar to the following:

   [I] Rank 0 - Generated NCCL ID and wrote to file: /tmp/nccl_id.txt
   [I] Rank 1 - Read NCCL ID from file: /tmp/nccl_id.txt
   [I] Rank 0 - ALL_REDUCE PASSED
   [I] Rank 1 - ALL_REDUCE PASSED
   [I] Rank 0 - ALL_REDUCE test completed successfully!
   [I] Rank 1 - ALL_REDUCE test completed successfully!
   

   This output shows that the sample ran successfully; PASSED.

Sample --help options

To see the full list of available options and their descriptions, use the -h or --help command line option.

Additional resources

The following resources provide a deeper understanding about distributed computing with TensorRT:

Documentation

• Introduction To NVIDIA's TensorRT Samples
• Working With TensorRT Using The C++ API
• NVIDIA's TensorRT Documentation Library
• NVIDIA NCCL Documentation

License

For terms and conditions for use, reproduction, and distribution, see the TensorRT Software License Agreement documentation.

Changelog

January 2026

• Initial release of sampleDistCollective.

Known issues

There are no known issues with this sample.