wehub-resource-sync
902 lines
35 KiB
C++
902 lines
35 KiB
C++
/*
|
|
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#include <algorithm>
|
|
#include <cassert>
|
|
#include <cctype>
|
|
#include <chrono>
|
|
#include <cmath>
|
|
#include <cstdlib>
|
|
#include <fstream>
|
|
#include <iostream>
|
|
#include <memory>
|
|
#include <optional>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <string_view>
|
|
#include <thread>
|
|
#include <unistd.h>
|
|
#include <vector>
|
|
|
|
#include "NvInfer.h"
|
|
#include "NvInferRuntime.h"
|
|
|
|
#include "common.h"
|
|
#include "logger.h"
|
|
#include "sampleDevice.h"
|
|
#include "sampleEngines.h"
|
|
#include "sampleInference.h"
|
|
#include "sampleOptions.h"
|
|
#include "sampleReporting.h"
|
|
|
|
using namespace nvinfer1;
|
|
using namespace sample;
|
|
using namespace samplesCommon;
|
|
|
|
// MD code start
|
|
#include <nccl.h>
|
|
// MD code end
|
|
|
|
using namespace std;
|
|
|
|
//! Checks NCCL return codes and asserts on failure since NCCL errors are unrecoverable communication failures.
|
|
#define NCCLCHECK(cmd) \
|
|
do \
|
|
{ \
|
|
ncclResult_t const r = (cmd); \
|
|
ASSERT(r == ncclSuccess); \
|
|
} while (0)
|
|
|
|
#define CHECK_CUDA(status) \
|
|
if (status != cudaSuccess) \
|
|
{ \
|
|
fprintf(stderr, "CUDA Error: %s\n", cudaGetErrorString(status)); \
|
|
exit(EXIT_FAILURE); \
|
|
}
|
|
|
|
namespace
|
|
{
|
|
|
|
using LibraryPtr = std::unique_ptr<DynamicLibrary>;
|
|
|
|
#if !TRT_STATIC
|
|
#if defined(_WIN32)
|
|
std::string const kNVINFER_PLUGIN_LIBNAME
|
|
= std::string{"nvinfer_plugin_"} + std::to_string(NV_TENSORRT_MAJOR) + std::string{".dll"};
|
|
std::string const kNVINFER_LIBNAME = std::string{"nvinfer_"} + std::to_string(NV_TENSORRT_MAJOR) + std::string{".dll"};
|
|
std::string const kNVONNXPARSER_LIBNAME
|
|
= std::string{"nvonnxparser_"} + std::to_string(NV_TENSORRT_MAJOR) + std::string{".dll"};
|
|
std::string const kNVINFER_LEAN_LIBNAME
|
|
= std::string{"nvinfer_lean_"} + std::to_string(NV_TENSORRT_MAJOR) + std::string{".dll"};
|
|
std::string const kNVINFER_DISPATCH_LIBNAME
|
|
= std::string{"nvinfer_dispatch_"} + std::to_string(NV_TENSORRT_MAJOR) + std::string{".dll"};
|
|
#else
|
|
std::string const kNVINFER_PLUGIN_LIBNAME = std::string{"libnvinfer_plugin.so."} + std::to_string(NV_TENSORRT_MAJOR);
|
|
std::string const kNVINFER_LIBNAME = std::string{"libnvinfer.so."} + std::to_string(NV_TENSORRT_MAJOR);
|
|
std::string const kNVONNXPARSER_LIBNAME = std::string{"libnvonnxparser.so."} + std::to_string(NV_TENSORRT_MAJOR);
|
|
std::string const kNVINFER_LEAN_LIBNAME = std::string{"libnvinfer_lean.so."} + std::to_string(NV_TENSORRT_MAJOR);
|
|
std::string const kNVINFER_DISPATCH_LIBNAME
|
|
= std::string{"libnvinfer_dispatch.so."} + std::to_string(NV_TENSORRT_MAJOR);
|
|
#endif
|
|
#endif // !TRT_STATIC
|
|
std::function<void*(void*, int32_t)> pCreateInferRuntimeInternal{};
|
|
std::function<void*(void*, void*, int32_t)> pCreateInferRefitterInternal{};
|
|
std::function<void*(void*, int32_t)> pCreateInferBuilderInternal{};
|
|
std::function<void*(void*, void*, int)> pCreateNvOnnxParserInternal{};
|
|
std::function<void*(void*, void*, int)> pCreateNvOnnxRefitterInternal{};
|
|
|
|
//! Track runtime used for the execution of trtexec.
|
|
//! Must be tracked as a global variable due to how library init functions APIs are organized.
|
|
RuntimeMode gUseRuntime = RuntimeMode::kFULL;
|
|
|
|
#if !TRT_STATIC
|
|
template <typename FetchPtrs>
|
|
bool initLibrary(LibraryPtr& libPtr, std::string const& libName, FetchPtrs fetchFunc)
|
|
{
|
|
if (libPtr != nullptr)
|
|
{
|
|
return true;
|
|
}
|
|
try
|
|
{
|
|
libPtr.reset(new DynamicLibrary{libName});
|
|
fetchFunc(libPtr.get());
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
libPtr.reset();
|
|
sample::gLogError << "Could not load library " << libName << ": " << e.what() << std::endl;
|
|
return false;
|
|
}
|
|
catch (...)
|
|
{
|
|
libPtr.reset();
|
|
sample::gLogError << "Could not load library " << libName << std::endl;
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
#endif // !TRT_STATIC
|
|
|
|
bool initNvinfer()
|
|
{
|
|
#if !TRT_STATIC
|
|
static LibraryPtr libnvinferPtr{};
|
|
auto fetchPtrs = [](DynamicLibrary* l) {
|
|
pCreateInferRuntimeInternal = l->symbolAddress<void*(void*, int32_t)>("createInferRuntime_INTERNAL");
|
|
try
|
|
{
|
|
pCreateInferRefitterInternal
|
|
= l->symbolAddress<void*(void*, void*, int32_t)>("createInferRefitter_INTERNAL");
|
|
}
|
|
catch (const std::exception& e)
|
|
{
|
|
sample::gLogWarning << "Could not load function createInferRefitter_INTERNAL : " << e.what() << std::endl;
|
|
}
|
|
|
|
if (gUseRuntime == RuntimeMode::kFULL)
|
|
{
|
|
pCreateInferBuilderInternal = l->symbolAddress<void*(void*, int32_t)>("createInferBuilder_INTERNAL");
|
|
}
|
|
};
|
|
return initLibrary(libnvinferPtr, sample::getRuntimeLibraryName(gUseRuntime), fetchPtrs);
|
|
#else
|
|
pCreateInferRuntimeInternal = createInferRuntime_INTERNAL;
|
|
pCreateInferRefitterInternal = createInferRefitter_INTERNAL;
|
|
pCreateInferBuilderInternal = createInferBuilder_INTERNAL;
|
|
return true;
|
|
#endif // !TRT_STATIC
|
|
}
|
|
|
|
bool initNvonnxparser()
|
|
{
|
|
#if !TRT_STATIC
|
|
static LibraryPtr libnvonnxparserPtr{};
|
|
auto fetchPtrs = [](DynamicLibrary* l) {
|
|
pCreateNvOnnxParserInternal = l->symbolAddress<void*(void*, void*, int)>("createNvOnnxParser_INTERNAL");
|
|
pCreateNvOnnxRefitterInternal = l->symbolAddress<void*(void*, void*, int)>("createNvOnnxParserRefitter_INTERNAL");
|
|
};
|
|
return initLibrary(libnvonnxparserPtr, kNVONNXPARSER_LIBNAME, fetchPtrs);
|
|
#else
|
|
pCreateNvOnnxParserInternal = createNvOnnxParser_INTERNAL;
|
|
pCreateNvOnnxRefitterInternal = createNvOnnxParserRefitter_INTERNAL;
|
|
return true;
|
|
#endif // !TRT_STATIC
|
|
}
|
|
|
|
[[nodiscard]] std::string toString(CollectiveOperation op)
|
|
{
|
|
switch (op)
|
|
{
|
|
case CollectiveOperation::kALL_REDUCE: return "ALL_REDUCE";
|
|
case CollectiveOperation::kALL_GATHER: return "ALL_GATHER";
|
|
case CollectiveOperation::kBROADCAST: return "BROADCAST";
|
|
case CollectiveOperation::kREDUCE: return "REDUCE";
|
|
case CollectiveOperation::kREDUCE_SCATTER: return "REDUCE_SCATTER";
|
|
case CollectiveOperation::kALL_TO_ALL: return "ALL_TO_ALL";
|
|
case CollectiveOperation::kGATHER: return "GATHER";
|
|
case CollectiveOperation::kSCATTER: return "SCATTER";
|
|
}
|
|
throw std::runtime_error("Unknown CollectiveOperation");
|
|
}
|
|
|
|
[[nodiscard]] bool icharEquals(char a, char b)
|
|
{
|
|
return std::tolower(static_cast<unsigned char>(a)) == std::tolower(static_cast<unsigned char>(b));
|
|
}
|
|
|
|
//! Case-insensitive string equality:
|
|
[[nodiscard]] bool iequals(std::string_view lhs, std::string_view rhs)
|
|
{
|
|
return std::equal(lhs.begin(), lhs.end(), rhs.begin(), rhs.end(), icharEquals);
|
|
}
|
|
|
|
//! Parse operation string to CollectiveOperation enum
|
|
[[nodiscard]] std::optional<CollectiveOperation> parseCollectiveOp(std::string_view opStr)
|
|
{
|
|
if (iequals(opStr, "all_reduce"))
|
|
{
|
|
return CollectiveOperation::kALL_REDUCE;
|
|
}
|
|
if (iequals(opStr, "all_gather"))
|
|
{
|
|
return CollectiveOperation::kALL_GATHER;
|
|
}
|
|
if (iequals(opStr, "broadcast"))
|
|
{
|
|
return CollectiveOperation::kBROADCAST;
|
|
}
|
|
if (iequals(opStr, "reduce"))
|
|
{
|
|
return CollectiveOperation::kREDUCE;
|
|
}
|
|
if (iequals(opStr, "reduce_scatter"))
|
|
{
|
|
return CollectiveOperation::kREDUCE_SCATTER;
|
|
}
|
|
if (iequals(opStr, "all_to_all"))
|
|
{
|
|
return CollectiveOperation::kALL_TO_ALL;
|
|
}
|
|
if (iequals(opStr, "gather"))
|
|
{
|
|
return CollectiveOperation::kGATHER;
|
|
}
|
|
if (iequals(opStr, "scatter"))
|
|
{
|
|
return CollectiveOperation::kSCATTER;
|
|
}
|
|
return std::nullopt;
|
|
}
|
|
|
|
void printUsage(char const* programName)
|
|
{
|
|
std::cout << "Usage:" << std::endl;
|
|
std::cout << " Set TRT_MY_RANK, TRT_WORLD_SIZE, and TRT_NCCL_ID_FILE, then run " << programName
|
|
<< " --op <operation>" << std::endl;
|
|
std::cout << std::endl;
|
|
std::cout << "Options:" << std::endl;
|
|
std::cout << " --op <operation> Specify the collective operation to test (required)." << std::endl;
|
|
std::cout << " Valid operations: all_reduce, all_gather, broadcast, reduce, reduce_scatter, all_to_all, gather, scatter";
|
|
std::cout << std::endl;
|
|
std::cout << " --help, -h Show this help message." << std::endl;
|
|
std::cout << std::endl;
|
|
std::cout << "Environment Variables (required):" << std::endl;
|
|
std::cout << " TRT_MY_RANK The rank of this process (0 to WORLD_SIZE-1)." << std::endl;
|
|
std::cout << " TRT_WORLD_SIZE The total number of processes." << std::endl;
|
|
std::cout << " TRT_NCCL_ID_FILE Path to a shared file for NCCL ID coordination." << std::endl;
|
|
std::cout << " Rank 0 writes the NCCL ID to this file, other ranks read from it." << std::endl;
|
|
std::cout << " The file should be empty or non-existent before starting." << std::endl;
|
|
std::cout << std::endl;
|
|
std::cout << "Example commands:" << std::endl;
|
|
std::cout << " SLURM:" << std::endl;
|
|
std::cout << " srun --ntasks=2 bash -lc 'export TRT_MY_RANK=$SLURM_PROCID; \\" << std::endl;
|
|
std::cout << " export TRT_WORLD_SIZE=$SLURM_NTASKS; \\" << std::endl;
|
|
std::cout << " export TRT_NCCL_ID_FILE=/tmp/nccl_id.txt; \\" << std::endl;
|
|
std::cout << " " << programName << " --op all_reduce'" << std::endl;
|
|
std::cout << std::endl;
|
|
std::cout << " Open MPI:" << std::endl;
|
|
std::cout << " mpirun -np 2 bash -lc 'export TRT_MY_RANK=$OMPI_COMM_WORLD_RANK; \\" << std::endl;
|
|
std::cout << " export TRT_WORLD_SIZE=$OMPI_COMM_WORLD_SIZE; \\" << std::endl;
|
|
std::cout << " export TRT_NCCL_ID_FILE=/tmp/nccl_id.txt; \\" << std::endl;
|
|
std::cout << " " << programName << " --op all_reduce'" << std::endl;
|
|
}
|
|
|
|
//! Get rank from TRT_MY_RANK environment variable.
|
|
//! Users should set this variable via a launcher wrapper script.
|
|
[[nodiscard]] int32_t getRankFromEnv()
|
|
{
|
|
char const* rankStr = std::getenv("TRT_MY_RANK");
|
|
if (!rankStr)
|
|
{
|
|
sample::gLogError << "FATAL: TRT_MY_RANK environment variable is not set!" << std::endl;
|
|
sample::gLogError << "Please set TRT_MY_RANK to the rank of this process (0 to WORLD_SIZE-1)." << std::endl;
|
|
sample::gLogError << "Run with --help for example commands." << std::endl;
|
|
ASSERT(false && "TRT_MY_RANK environment variable must be set");
|
|
}
|
|
return std::stoi(rankStr);
|
|
}
|
|
|
|
//! Get world size from TRT_WORLD_SIZE environment variable.
|
|
//! Users should set this variable via a launcher wrapper script.
|
|
[[nodiscard]] int32_t getWorldSizeFromEnv()
|
|
{
|
|
char const* worldSizeStr = std::getenv("TRT_WORLD_SIZE");
|
|
if (!worldSizeStr)
|
|
{
|
|
sample::gLogError << "FATAL: TRT_WORLD_SIZE environment variable is not set!" << std::endl;
|
|
sample::gLogError << "Please set TRT_WORLD_SIZE to the total number of processes." << std::endl;
|
|
sample::gLogError << "Run with --help for example commands." << std::endl;
|
|
ASSERT(false && "TRT_WORLD_SIZE environment variable must be set");
|
|
}
|
|
return std::stoi(worldSizeStr);
|
|
}
|
|
|
|
//! Convert a hex character to its integer value
|
|
[[nodiscard]] int32_t hexCharToInt(char c)
|
|
{
|
|
if (c >= '0' && c <= '9')
|
|
{
|
|
return c - '0';
|
|
}
|
|
if (c >= 'a' && c <= 'f')
|
|
{
|
|
return c - 'a' + 10;
|
|
}
|
|
if (c >= 'A' && c <= 'F')
|
|
{
|
|
return c - 'A' + 10;
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
//! Convert NCCL unique ID bytes to hex string
|
|
[[nodiscard]] std::string ncclIdToHex(ncclUniqueId const& id)
|
|
{
|
|
constexpr char kHEX_CHARS[] = "0123456789abcdef";
|
|
std::string hexStr;
|
|
hexStr.reserve(sizeof(ncclUniqueId) * 2);
|
|
for (size_t i = 0; i < sizeof(ncclUniqueId); ++i)
|
|
{
|
|
auto const byte = static_cast<uint8_t>(id.internal[i]);
|
|
hexStr += kHEX_CHARS[byte >> 4];
|
|
hexStr += kHEX_CHARS[byte & 0x0F];
|
|
}
|
|
return hexStr;
|
|
}
|
|
|
|
//! Parse hex string to NCCL unique ID
|
|
[[nodiscard]] ncclUniqueId hexToNcclId(std::string const& hexStr)
|
|
{
|
|
constexpr size_t kNCCL_UNIQUE_ID_BYTES = sizeof(ncclUniqueId);
|
|
constexpr size_t kEXPECTED_HEX_LEN = kNCCL_UNIQUE_ID_BYTES * 2;
|
|
|
|
if (hexStr.length() != kEXPECTED_HEX_LEN)
|
|
{
|
|
throw std::runtime_error("NCCL ID hex string has invalid length: " + std::to_string(hexStr.length())
|
|
+ " (expected " + std::to_string(kEXPECTED_HEX_LEN) + ")");
|
|
}
|
|
|
|
ncclUniqueId id;
|
|
for (size_t i = 0; i < kNCCL_UNIQUE_ID_BYTES; ++i)
|
|
{
|
|
int32_t const high = hexCharToInt(hexStr[2 * i]);
|
|
int32_t const low = hexCharToInt(hexStr[2 * i + 1]);
|
|
if (high < 0 || low < 0)
|
|
{
|
|
throw std::runtime_error(
|
|
"NCCL ID hex string contains invalid character at position " + std::to_string(2 * i));
|
|
}
|
|
id.internal[i] = static_cast<char>((high << 4) | low);
|
|
}
|
|
return id;
|
|
}
|
|
|
|
//! Get the NCCL ID file path from TRT_NCCL_ID_FILE environment variable.
|
|
[[nodiscard]] std::string getNcclIdFilePath()
|
|
{
|
|
char const* filePath = std::getenv("TRT_NCCL_ID_FILE");
|
|
if (!filePath)
|
|
{
|
|
sample::gLogError << "FATAL: TRT_NCCL_ID_FILE environment variable is not set!" << std::endl;
|
|
sample::gLogError << "Please set TRT_NCCL_ID_FILE to a shared file path accessible by all ranks." << std::endl;
|
|
sample::gLogError << "Run with --help for example commands." << std::endl;
|
|
ASSERT(false && "TRT_NCCL_ID_FILE environment variable must be set");
|
|
}
|
|
return std::string(filePath);
|
|
}
|
|
|
|
//! Get NCCL unique ID using file-based coordination.
|
|
//! Rank 0 generates the ID and writes it to the file.
|
|
//! Other ranks wait for the file to be written and read the ID from it.
|
|
[[nodiscard]] ncclUniqueId getNcclIdViaFile(int32_t rank)
|
|
{
|
|
std::string const filePath = getNcclIdFilePath();
|
|
constexpr size_t kEXPECTED_HEX_LEN = sizeof(ncclUniqueId) * 2;
|
|
constexpr int32_t kPOLL_INTERVAL_MS = 10;
|
|
constexpr int32_t kTIMEOUT_MS = 30000; // 30 seconds timeout
|
|
|
|
if (rank == 0)
|
|
{
|
|
// Rank 0: Check if stale file exists from a previous run
|
|
std::ifstream checkFile(filePath);
|
|
if (checkFile)
|
|
{
|
|
std::string content;
|
|
std::getline(checkFile, content);
|
|
if (!content.empty())
|
|
{
|
|
throw std::runtime_error(
|
|
"NCCL ID file already exists with content: " + filePath + "\n"
|
|
"This may be stale data from a previous run. Please delete it first:\n"
|
|
" rm -f " + filePath);
|
|
}
|
|
}
|
|
|
|
// Generate NCCL ID and write to file
|
|
ncclUniqueId id;
|
|
NCCLCHECK(ncclGetUniqueId(&id));
|
|
|
|
std::string const hexStr = ncclIdToHex(id);
|
|
|
|
std::ofstream outFile(filePath, std::ios::trunc);
|
|
if (!outFile)
|
|
{
|
|
throw std::runtime_error("Failed to open NCCL ID file for writing: " + filePath);
|
|
}
|
|
outFile << hexStr;
|
|
outFile.close();
|
|
|
|
sample::gLogInfo << "Rank 0 - Generated NCCL ID and wrote to file: " << filePath << std::endl;
|
|
return id;
|
|
}
|
|
else
|
|
{
|
|
// Other ranks: Wait for file to be written and read the ID
|
|
int32_t elapsedMs = 0;
|
|
std::string hexStr;
|
|
|
|
while (elapsedMs < kTIMEOUT_MS)
|
|
{
|
|
std::ifstream inFile(filePath);
|
|
if (inFile)
|
|
{
|
|
std::getline(inFile, hexStr);
|
|
if (hexStr.length() == kEXPECTED_HEX_LEN)
|
|
{
|
|
sample::gLogInfo << "Rank " << rank << " - Read NCCL ID from file: " << filePath << std::endl;
|
|
return hexToNcclId(hexStr);
|
|
}
|
|
}
|
|
// File not ready yet, wait and retry
|
|
std::this_thread::sleep_for(std::chrono::milliseconds(kPOLL_INTERVAL_MS));
|
|
elapsedMs += kPOLL_INTERVAL_MS;
|
|
}
|
|
|
|
throw std::runtime_error("Timeout waiting for NCCL ID file to be written by rank 0");
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
IRuntime* createRuntime()
|
|
{
|
|
if (!initNvinfer())
|
|
{
|
|
return {};
|
|
}
|
|
ASSERT(pCreateInferRuntimeInternal != nullptr);
|
|
return static_cast<IRuntime*>(pCreateInferRuntimeInternal(&gLogger.getTRTLogger(), NV_TENSORRT_VERSION));
|
|
}
|
|
|
|
IBuilder* createBuilder()
|
|
{
|
|
if (!initNvinfer())
|
|
{
|
|
return {};
|
|
}
|
|
ASSERT(pCreateInferBuilderInternal != nullptr);
|
|
return static_cast<IBuilder*>(pCreateInferBuilderInternal(&gLogger.getTRTLogger(), NV_TENSORRT_VERSION));
|
|
}
|
|
|
|
IRefitter* createRefitter(ICudaEngine& engine)
|
|
{
|
|
if (!initNvinfer())
|
|
{
|
|
return {};
|
|
}
|
|
ASSERT(pCreateInferRefitterInternal != nullptr);
|
|
return static_cast<IRefitter*>(pCreateInferRefitterInternal(&engine, &gLogger.getTRTLogger(), NV_TENSORRT_VERSION));
|
|
}
|
|
|
|
nvonnxparser::IParser* createONNXParser(INetworkDefinition& network)
|
|
{
|
|
if (!initNvonnxparser())
|
|
{
|
|
return {};
|
|
}
|
|
ASSERT(pCreateNvOnnxParserInternal != nullptr);
|
|
return static_cast<nvonnxparser::IParser*>(
|
|
pCreateNvOnnxParserInternal(&network, &gLogger.getTRTLogger(), NV_ONNX_PARSER_VERSION));
|
|
}
|
|
|
|
nvonnxparser::IParserRefitter* createONNXRefitter(IRefitter& refitter)
|
|
{
|
|
if (!initNvonnxparser())
|
|
{
|
|
return {};
|
|
}
|
|
ASSERT(pCreateNvOnnxRefitterInternal != nullptr);
|
|
return static_cast<nvonnxparser::IParserRefitter*>(
|
|
pCreateNvOnnxRefitterInternal(&refitter, &gLogger.getTRTLogger(), NV_ONNX_PARSER_VERSION));
|
|
}
|
|
|
|
//! Helper struct to hold test configuration for each collective operation
|
|
struct CollectiveTestConfig
|
|
{
|
|
CollectiveOperation op;
|
|
std::vector<float> rank0Input;
|
|
std::vector<float> rank1Input;
|
|
std::vector<float> rank0ExpectedOutput;
|
|
std::vector<float> rank1ExpectedOutput; // Different from rank0 for REDUCE_SCATTER
|
|
int32_t outputElementCount; // Number of output elements per rank
|
|
};
|
|
|
|
//! Get test configuration for a specific collective operation
|
|
CollectiveTestConfig getTestConfig(CollectiveOperation op, int32_t worldSize)
|
|
{
|
|
// Input: 12 elements per rank [3, 4]
|
|
// After transpose: [4, 3]
|
|
constexpr int32_t kINPUT_SIZE = 12;
|
|
|
|
switch (op)
|
|
{
|
|
case CollectiveOperation::kALL_GATHER:
|
|
{
|
|
// ALL_GATHER: Each rank contributes data, all ranks receive concatenated result
|
|
// Output: 12 * worldSize = 24 elements
|
|
std::vector<float> const expected = {0.0F, 1.0F, 2.0F, 3.0F, 100.0F, 101.0F, 102.0F, 103.0F, 4.0F, 5.0F, 6.0F,
|
|
7.0F, 104.0F, 105.0F, 106.0F, 107.0F, 8.0F, 9.0F, 10.0F, 11.0F, 108.0F, 109.0F, 110.0F, 111.0F};
|
|
return {CollectiveOperation::kALL_GATHER,
|
|
{0.0F, 1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F, 7.0F, 8.0F, 9.0F, 10.0F, 11.0F},
|
|
{100.0F, 101.0F, 102.0F, 103.0F, 104.0F, 105.0F, 106.0F, 107.0F, 108.0F, 109.0F, 110.0F, 111.0F}, expected,
|
|
expected, // Both ranks get same result
|
|
kINPUT_SIZE * worldSize};
|
|
}
|
|
|
|
case CollectiveOperation::kALL_REDUCE:
|
|
{
|
|
// ALL_REDUCE: Sum across all ranks, all ranks receive same result
|
|
std::vector<float> const expected
|
|
= {11.0F, 22.0F, 33.0F, 44.0F, 55.0F, 66.0F, 77.0F, 88.0F, 99.0F, 110.0F, 121.0F, 132.0F};
|
|
return {CollectiveOperation::kALL_REDUCE,
|
|
{1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F, 7.0F, 8.0F, 9.0F, 10.0F, 11.0F, 12.0F},
|
|
{10.0F, 20.0F, 30.0F, 40.0F, 50.0F, 60.0F, 70.0F, 80.0F, 90.0F, 100.0F, 110.0F, 120.0F}, expected,
|
|
expected, // Both ranks get same result
|
|
kINPUT_SIZE};
|
|
}
|
|
|
|
case CollectiveOperation::kBROADCAST:
|
|
{
|
|
// BROADCAST: Rank 0 sends data to all ranks
|
|
std::vector<float> const expected = {1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F, 7.0F, 8.0F, 9.0F, 10.0F, 11.0F, 12.0F};
|
|
return {CollectiveOperation::kBROADCAST,
|
|
{1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F, 7.0F, 8.0F, 9.0F, 10.0F, 11.0F, 12.0F},
|
|
{99.0F, 99.0F, 99.0F, 99.0F, 99.0F, 99.0F, 99.0F, 99.0F, 99.0F, 99.0F, 99.0F, 99.0F}, expected,
|
|
expected, // Both ranks get same result
|
|
kINPUT_SIZE};
|
|
}
|
|
|
|
case CollectiveOperation::kREDUCE:
|
|
// REDUCE: Sum across all ranks, only root (rank 0) receives result
|
|
// rank1's output is undefined, use empty vector
|
|
return {CollectiveOperation::kREDUCE,
|
|
{1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F, 7.0F, 8.0F, 9.0F, 10.0F, 11.0F, 12.0F},
|
|
{10.0F, 20.0F, 30.0F, 40.0F, 50.0F, 60.0F, 70.0F, 80.0F, 90.0F, 100.0F, 110.0F, 120.0F},
|
|
{11.0F, 22.0F, 33.0F, 44.0F, 55.0F, 66.0F, 77.0F, 88.0F, 99.0F, 110.0F, 121.0F, 132.0F},
|
|
{}, // rank1's output is undefined
|
|
kINPUT_SIZE};
|
|
|
|
case CollectiveOperation::kREDUCE_SCATTER:
|
|
// REDUCE_SCATTER: Reduce then scatter - each rank gets a different portion
|
|
// Input [3,4] -> transpose [4,3] -> reduce_scatter [2,3] -> transpose [3,2] = 6 elements
|
|
// After reduce: [[11,55,99],[22,66,110],[33,77,121],[44,88,132]]
|
|
// rank0 gets first half [[11,55,99],[22,66,110]], after transpose: [[11,22],[55,66],[99,110]]
|
|
// rank1 gets second half [[33,77,121],[44,88,132]], after transpose: [[33,44],[77,88],[121,132]]
|
|
return {CollectiveOperation::kREDUCE_SCATTER,
|
|
{1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F, 7.0F, 8.0F, 9.0F, 10.0F, 11.0F, 12.0F},
|
|
{10.0F, 20.0F, 30.0F, 40.0F, 50.0F, 60.0F, 70.0F, 80.0F, 90.0F, 100.0F, 110.0F, 120.0F},
|
|
{11.0F, 22.0F, 55.0F, 66.0F, 99.0F, 110.0F}, // rank0 expected
|
|
{33.0F, 44.0F, 77.0F, 88.0F, 121.0F, 132.0F}, // rank1 expected
|
|
kINPUT_SIZE / worldSize};
|
|
|
|
case CollectiveOperation::kALL_TO_ALL:
|
|
{
|
|
// ALL_TO_ALL: Input [3,4] -> transpose [4,3] -> all_to_all (count=6) -> [4,3] -> transpose [3,4]
|
|
// Rank 0 transposed: [0,4,8,1,5,9,2,6,10,3,7,11]; sends first 6 to self, last 6 to rank 1.
|
|
// Rank 1 transposed: [100,104,108,101,105,109,102,106,110,103,107,111]; sends first 6 to rank 0.
|
|
// Rank 0 output [4,3]: rows=[0,4,8],[1,5,9],[100,104,108],[101,105,109]; after transpose [3,4]:
|
|
// [0,1,100,101, 4,5,104,105, 8,9,108,109]
|
|
// Rank 1 output [4,3]: rows=[2,6,10],[3,7,11],[102,106,110],[103,107,111]; after transpose [3,4]:
|
|
// [2,3,102,103, 6,7,106,107, 10,11,110,111]
|
|
return {CollectiveOperation::kALL_TO_ALL,
|
|
{0.0F, 1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F, 7.0F, 8.0F, 9.0F, 10.0F, 11.0F},
|
|
{100.0F, 101.0F, 102.0F, 103.0F, 104.0F, 105.0F, 106.0F, 107.0F, 108.0F, 109.0F, 110.0F, 111.0F},
|
|
{0.0F, 1.0F, 100.0F, 101.0F, 4.0F, 5.0F, 104.0F, 105.0F, 8.0F, 9.0F, 108.0F, 109.0F}, // rank0
|
|
{2.0F, 3.0F, 102.0F, 103.0F, 6.0F, 7.0F, 106.0F, 107.0F, 10.0F, 11.0F, 110.0F, 111.0F}, // rank1
|
|
kINPUT_SIZE};
|
|
}
|
|
|
|
case CollectiveOperation::kGATHER:
|
|
{
|
|
// GATHER: All ranks send to root. Input [3,4] -> transpose [4,3] -> gather [8,3] -> transpose [3,8]
|
|
// Root recvbuf = [rank0_data || rank1_data] transposed to [3,8]:
|
|
// [0,1,2,3,100,101,102,103, 4,5,6,7,104,105,106,107, 8,9,10,11,108,109,110,111]
|
|
// Non-root output buffer is allocated but NCCL does not write to it.
|
|
std::vector<float> const rank0Expected = {0.0F, 1.0F, 2.0F, 3.0F, 100.0F, 101.0F, 102.0F, 103.0F, 4.0F,
|
|
5.0F, 6.0F, 7.0F, 104.0F, 105.0F, 106.0F, 107.0F, 8.0F, 9.0F, 10.0F, 11.0F, 108.0F, 109.0F, 110.0F,
|
|
111.0F};
|
|
return {CollectiveOperation::kGATHER,
|
|
{0.0F, 1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F, 7.0F, 8.0F, 9.0F, 10.0F, 11.0F},
|
|
{100.0F, 101.0F, 102.0F, 103.0F, 104.0F, 105.0F, 106.0F, 107.0F, 108.0F, 109.0F, 110.0F, 111.0F},
|
|
rank0Expected,
|
|
{}, // rank1's output is undefined (non-root does not receive)
|
|
kINPUT_SIZE * worldSize};
|
|
}
|
|
|
|
case CollectiveOperation::kSCATTER:
|
|
{
|
|
// SCATTER: Root scatters data to all ranks. Input [3,4] -> transpose [4,3].
|
|
// Root sendBuf=[0,4,8,1,5,9,2,6,10,3,7,11]; recvCount=6 per rank.
|
|
// Rank 0 gets first 6 -> [2,3] -> transpose [3,2]: [0,1,4,5,8,9]
|
|
// Rank 1 gets next 6 -> [2,3] -> transpose [3,2]: [2,3,6,7,10,11]
|
|
// Rank 1 input is ignored by NCCL (non-root has no sendBuf).
|
|
return {CollectiveOperation::kSCATTER,
|
|
{0.0F, 1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F, 7.0F, 8.0F, 9.0F, 10.0F, 11.0F},
|
|
{99.0F, 99.0F, 99.0F, 99.0F, 99.0F, 99.0F, 99.0F, 99.0F, 99.0F, 99.0F, 99.0F, 99.0F}, // ignored
|
|
{0.0F, 1.0F, 4.0F, 5.0F, 8.0F, 9.0F}, // rank0 expected
|
|
{2.0F, 3.0F, 6.0F, 7.0F, 10.0F, 11.0F}, // rank1 expected
|
|
kINPUT_SIZE / worldSize};
|
|
}
|
|
}
|
|
throw std::runtime_error("Unknown CollectiveOperation");
|
|
}
|
|
|
|
//! Build and execute a network with a specific collective operation
|
|
void testCollectiveOperation(
|
|
int32_t rank, int32_t worldSize, CollectiveTestConfig const& config, ncclComm_t comm, cudaStream_t stream)
|
|
{
|
|
sample::gLogInfo << "Rank " << rank << " - Testing " << toString(config.op) << std::endl;
|
|
|
|
// Create builder and network
|
|
auto builder = std::unique_ptr<IBuilder>(createInferBuilder(sample::gLogger.getTRTLogger()));
|
|
ASSERT(builder != nullptr);
|
|
|
|
auto network = std::unique_ptr<INetworkDefinition>(
|
|
builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kSTRONGLY_TYPED)));
|
|
ASSERT(network != nullptr);
|
|
|
|
// Create input tensor
|
|
constexpr int32_t kINPUT_ROWS = 3;
|
|
constexpr int32_t kINPUT_COLS = 4;
|
|
auto* input = network->addInput("input", DataType::kFLOAT, Dims2{kINPUT_ROWS, kINPUT_COLS});
|
|
ASSERT(input != nullptr);
|
|
|
|
auto* firstShuffle = network->addShuffle(*input);
|
|
ASSERT(firstShuffle != nullptr);
|
|
firstShuffle->setFirstTranspose({{1, 0}});
|
|
|
|
ReduceOperation reduceOp = ReduceOperation::kNONE;
|
|
if (config.op == CollectiveOperation::kALL_REDUCE || config.op == CollectiveOperation::kREDUCE
|
|
|| config.op == CollectiveOperation::kREDUCE_SCATTER)
|
|
{
|
|
reduceOp = ReduceOperation::kSUM;
|
|
}
|
|
int64_t root = -1;
|
|
if (config.op == CollectiveOperation::kBROADCAST || config.op == CollectiveOperation::kREDUCE
|
|
|| config.op == CollectiveOperation::kGATHER || config.op == CollectiveOperation::kSCATTER)
|
|
{
|
|
root = 0;
|
|
}
|
|
auto* collectiveLayer
|
|
= network->addDistCollective(*firstShuffle->getOutput(0), config.op, reduceOp, root, nullptr, 0);
|
|
ASSERT(collectiveLayer != nullptr);
|
|
|
|
// Set the number of ranks for the collective operation
|
|
if (!collectiveLayer->setNbRanks(worldSize))
|
|
{
|
|
throw std::runtime_error("Failed to set the number of ranks for the collective layer");
|
|
}
|
|
|
|
auto* secondShuffle = network->addShuffle(*collectiveLayer->getOutput(0));
|
|
ASSERT(secondShuffle != nullptr);
|
|
secondShuffle->setFirstTranspose({{1, 0}});
|
|
|
|
// Mark the reshape layer's output as the network output
|
|
network->markOutput(*secondShuffle->getOutput(0));
|
|
|
|
// Build engine
|
|
auto builderConfig = std::unique_ptr<IBuilderConfig>(builder->createBuilderConfig());
|
|
ASSERT(builderConfig != nullptr);
|
|
auto serializedEngine = std::unique_ptr<IHostMemory>(builder->buildSerializedNetwork(*network, *builderConfig));
|
|
ASSERT(serializedEngine != nullptr);
|
|
|
|
// Create runtime and deserialize engine
|
|
auto runtime = std::unique_ptr<IRuntime>(createInferRuntime(sample::gLogger.getTRTLogger()));
|
|
ASSERT(runtime != nullptr);
|
|
|
|
// Deserialize the CUDA engine
|
|
auto engine = std::unique_ptr<ICudaEngine>(
|
|
runtime->deserializeCudaEngine(serializedEngine->data(), serializedEngine->size()));
|
|
ASSERT(engine != nullptr);
|
|
|
|
// Create execution context for the engine
|
|
auto context = std::unique_ptr<IExecutionContext>(engine->createExecutionContext());
|
|
ASSERT(context != nullptr);
|
|
|
|
// Prepare input and output buffers
|
|
char const* inputName = engine->getIOTensorName(0);
|
|
char const* outputName = engine->getIOTensorName(1);
|
|
|
|
std::vector<float> const& inputChunk = (rank == 0) ? config.rank0Input : config.rank1Input;
|
|
std::vector<float> outputChunk(config.outputElementCount, 0.0F);
|
|
|
|
size_t const inputBytes = inputChunk.size() * sizeof(float);
|
|
size_t const outputBytes = outputChunk.size() * sizeof(float);
|
|
|
|
void* dInput = nullptr;
|
|
void* dOutput = nullptr;
|
|
CHECK_CUDA(cudaMalloc(&dInput, inputBytes));
|
|
CHECK_CUDA(cudaMalloc(&dOutput, outputBytes));
|
|
|
|
// Copy input data to GPU asynchronously
|
|
CHECK_CUDA(cudaMemcpyAsync(dInput, inputChunk.data(), inputBytes, cudaMemcpyHostToDevice, stream));
|
|
|
|
// Set input/output tensor addresses in the execution context
|
|
context->setInputTensorAddress(inputName, dInput);
|
|
context->setTensorAddress(outputName, dOutput);
|
|
context->setInputShape(inputName, Dims2{kINPUT_ROWS, kINPUT_COLS});
|
|
|
|
// Set NCCL communicator
|
|
if (!context->setCommunicator(comm))
|
|
{
|
|
cudaFree(dInput);
|
|
cudaFree(dOutput);
|
|
throw std::runtime_error("Failed to set communicator for " + toString(config.op));
|
|
}
|
|
|
|
// Run inference
|
|
if (!context->enqueueV3(stream))
|
|
{
|
|
cudaFree(dInput);
|
|
cudaFree(dOutput);
|
|
throw std::runtime_error("Inference failed for " + toString(config.op));
|
|
}
|
|
|
|
CHECK_CUDA(cudaStreamSynchronize(stream));
|
|
|
|
// Copy output data back to host asynchronously
|
|
CHECK_CUDA(cudaMemcpyAsync(outputChunk.data(), dOutput, outputBytes, cudaMemcpyDeviceToHost, stream));
|
|
CHECK_CUDA(cudaStreamSynchronize(stream));
|
|
|
|
// Get the expected output for this rank
|
|
std::vector<float> const& expectedOutput = (rank == 0) ? config.rank0ExpectedOutput : config.rank1ExpectedOutput;
|
|
|
|
// Determine if this rank should verify output
|
|
// REDUCE: only rank 0 receives valid result
|
|
// All other ops: both ranks can verify (same or different expected values)
|
|
bool const shouldVerify = !expectedOutput.empty();
|
|
|
|
if (shouldVerify)
|
|
{
|
|
constexpr float kEPS = 1e-5F;
|
|
for (size_t i = 0; i < outputChunk.size() && i < expectedOutput.size(); ++i)
|
|
{
|
|
if (std::abs(outputChunk[i] - expectedOutput[i]) > kEPS)
|
|
{
|
|
cudaFree(dInput);
|
|
cudaFree(dOutput);
|
|
throw std::runtime_error("Output mismatch for " + toString(config.op) + " at index " + std::to_string(i)
|
|
+ ": expected " + std::to_string(expectedOutput[i]) + ", got " + std::to_string(outputChunk[i]));
|
|
}
|
|
}
|
|
sample::gLogInfo << "Rank " << rank << " - " << toString(config.op) << " PASSED" << std::endl;
|
|
}
|
|
|
|
// Cleanup
|
|
cudaFree(dInput);
|
|
cudaFree(dOutput);
|
|
}
|
|
|
|
//! Main test function that runs a specific collective operation test
|
|
void runCollectiveTest(int32_t rank, int32_t worldSize, CollectiveOperation op)
|
|
{
|
|
// Check GPU availability
|
|
int32_t deviceCount = 0;
|
|
CHECK_CUDA(cudaGetDeviceCount(&deviceCount));
|
|
|
|
if (deviceCount < worldSize)
|
|
{
|
|
throw std::runtime_error("Not enough GPUs available. Need " + std::to_string(worldSize) + " but found "
|
|
+ std::to_string(deviceCount));
|
|
}
|
|
|
|
// Use rank to select GPU
|
|
CHECK_CUDA(cudaSetDevice(rank));
|
|
|
|
// Create CUDA stream
|
|
cudaStream_t stream;
|
|
CHECK_CUDA(cudaStreamCreate(&stream));
|
|
|
|
// Set up NCCL - rank 0 generates ID and writes to file, others read from file
|
|
ncclUniqueId const id = getNcclIdViaFile(rank);
|
|
ncclComm_t comm;
|
|
NCCLCHECK(ncclCommInitRank(&comm, worldSize, id, rank));
|
|
|
|
// Get test configuration for the specified operation
|
|
CollectiveTestConfig const config = getTestConfig(op, worldSize);
|
|
|
|
// Run the collective operation test
|
|
testCollectiveOperation(rank, worldSize, config, comm, stream);
|
|
|
|
sample::gLogInfo << "Rank " << rank << " - " << toString(op) << " test completed successfully!" << std::endl;
|
|
|
|
NCCLCHECK(ncclCommDestroy(comm));
|
|
CHECK_CUDA(cudaStreamDestroy(stream));
|
|
}
|
|
|
|
int main(int argc, char* argv[])
|
|
{
|
|
constexpr int32_t kREQUIRED_WORLD_SIZE = 2;
|
|
|
|
for (int32_t i = 1; i < argc; ++i)
|
|
{
|
|
std::string const arg = argv[i];
|
|
if (arg == "--help" || arg == "-h")
|
|
{
|
|
printUsage(argv[0]);
|
|
return 0;
|
|
}
|
|
}
|
|
|
|
// Get rank and world size from TRT_MY_RANK and TRT_WORLD_SIZE environment variables.
|
|
int32_t const rank = getRankFromEnv();
|
|
int32_t const worldSize = getWorldSizeFromEnv();
|
|
|
|
// Parse command line arguments
|
|
CollectiveOperation selectedOp{};
|
|
bool hasSelectedOp = false;
|
|
|
|
for (int32_t i = 1; i < argc; ++i)
|
|
{
|
|
std::string arg = argv[i];
|
|
if (arg == "--op" && i + 1 < argc)
|
|
{
|
|
++i;
|
|
auto parsedOp = parseCollectiveOp(argv[i]);
|
|
if (!parsedOp)
|
|
{
|
|
if (rank == 0)
|
|
{
|
|
sample::gLogError << "Invalid operation: " << argv[i] << std::endl;
|
|
printUsage(argv[0]);
|
|
}
|
|
return 1;
|
|
}
|
|
selectedOp = *parsedOp;
|
|
hasSelectedOp = true;
|
|
}
|
|
}
|
|
|
|
// --op is required
|
|
if (!hasSelectedOp)
|
|
{
|
|
if (rank == 0)
|
|
{
|
|
sample::gLogError << "Error: --op argument is required." << std::endl;
|
|
printUsage(argv[0]);
|
|
}
|
|
return 1;
|
|
}
|
|
|
|
// We need exactly 2 processes for this test
|
|
if (worldSize != kREQUIRED_WORLD_SIZE)
|
|
{
|
|
if (rank == 0)
|
|
{
|
|
sample::gLogError << "This sample requires exactly 2 processes, but " << worldSize << " were provided."
|
|
<< std::endl;
|
|
sample::gLogError << "Please set TRT_WORLD_SIZE=2 and launch 2 processes." << std::endl;
|
|
sample::gLogError << "Run with --help for example commands." << std::endl;
|
|
}
|
|
return 1;
|
|
}
|
|
|
|
try
|
|
{
|
|
runCollectiveTest(rank, worldSize, selectedOp);
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
sample::gLogError << "Rank " << rank << " - Exception: " << e.what() << std::endl;
|
|
return 1;
|
|
}
|
|
|
|
return 0;
|
|
}
|