chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Waiting to run

This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
+132
View File
@@ -0,0 +1,132 @@
# TensorRT Distributed Collective Sample
**Table Of Contents**
- [Description](#description)
- [How does this sample work?](#how-does-this-sample-work)
* [TensorRT API layers and ops](#tensorrt-api-layers-and-ops)
- [Prerequisites](#prerequisites)
- [Running the sample](#running-the-sample)
* [Sample `--help` options](#sample-help-options)
- [Additional resources](#additional-resources)
- [License](#license)
- [Changelog](#changelog)
- [Known issues](#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](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#layers) is used for distributed collective operations across multiple GPUs. For more information, see the [TensorRT Developer Guide: Layers](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#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):
```bash
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
```bash
sudo apt-get install -y openmpi-bin libopenmpi-dev
```
## Running the sample
1. Compile this sample by following build instructions in [TensorRT README](https://github.com/NVIDIA/TensorRT/). 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):**
```bash
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):**
```bash
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](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
- [Working With TensorRT Using The C++ API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#c_topics)
- [NVIDIA's TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
- [NVIDIA NCCL Documentation](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/index.html)
## License
For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
## Changelog
January 2026
- Initial release of `sampleDistCollective`.
## Known issues
There are no known issues with this sample.