Mở rộng quy mô Inference AI trên nhiều GPU bằng NVIDIA TensorRT với Hỗ trợ Multi-Device Inference

by vanito

Phá vỡ giới hạn bộ nhớ và tính toán của GPU đơn lẻ – Triển khai các pipeline Generative AI khổng lồ với hiệu suất tối đa nhờ TensorRT 11.0.

Các khối lượng công việc Generative AI đang nhanh chóng phát triển vượt quá ngân sách về bộ nhớ và khả năng tính toán của một GPU đơn lẻ. Đối với các nhà phát triển (inference developer) xây dựng các pipeline tạo phương tiện (media generation), thách thức đặt ra là việc mở rộng quy mô trên nhiều thiết bị mà không phải hy sinh các tối ưu hóa then chốt — chẳng hạn như kernel fusions, memory planning, và lượng tử hóa (quantization) — những yếu tố cốt lõi mà NVIDIA TensorRT mang lại cho các môi trường triển khai thực tế (production).

Hỗ trợ suy luận đa thiết bị (Multi-device inference), một tính năng mới được giới thiệu trong TensorRT 11.0, mang khả năng suy luận multi-GPU hiệu suất cao nguyên bản (native) vào runtime của TensorRT, cho phép triển khai production trên nhiều thiết bị, kể cả các thiết bị biên (edge devices).

Kết hợp hỗ trợ multi-device inference trong TensorRT với Torch-TensorRT, các developer có thể chuyển đổi và triển khai các mô hình PyTorch khổng lồ bên ngoài framework, phá vỡ hoàn toàn các giới hạn về bộ nhớ và tính toán của thiết bị đơn.

👉 Hãy tải xuống TensorRT 11.0 với tính năng multi-device inference từ NVIDIA Developer Portal để mở khóa khả năng tăng tốc đa thiết bị nguyên bản cho các mô hình của bạn.

NVIDIA NCCL: Lớp giao tiếp cho Distributed Inference

NVIDIA Collective Communications Library (NCCL) cung cấp các hoạt động giao tiếp tập thể (collective operations) hiệu suất cao trên nhiều GPU và nhiều node, hỗ trợ quá trình huấn luyện các mô hình quy mô lớn trên hàng ngàn GPU. NCCL tự động lựa chọn giao thức vận chuyển tối ưu cho một cấu trúc mạng cụ thể, trừu tượng hóa các kết nối như NVIDIA NVLink, NVIDIA NVSwitch, PCIe và InfiniBand đằng sau một giao diện đồng nhất.

Bằng cách tích hợp trực tiếp với NCCL, TensorRT kế thừa khả năng tối ưu hóa giao thức này cho các khối lượng công việc inference khi chạy ở chế độ đa thiết bị. Tính năng mới này hỗ trợ đầy đủ bộ distributed collectives của NVIDIA NCCL, bao gồm: AllReduce, Broadcast, Reduce, AllGather, ReduceScatter, AlltoAll, Gather, và Scatter.

Các chiến lược song song hóa cho Distributed Inference

Distributed inference có thể được thể hiện qua một số chiến lược song song hóa (parallelism), mỗi chiến lược mang lại các sự đánh đổi khác nhau giữa tiết kiệm bộ nhớ, mở rộng tính toán và chi phí giao tiếp. Các chiến lược phổ biến nhất là Tensor parallelism và Context parallelism.

Tensor parallelism (Song song hóa Tensor)

Trong tensor parallelism, trọng số (weights) của một layer đơn lẻ được phân chia trên nhiều GPU. Mỗi GPU tính toán một phần (shard) của phép nhân ma trận của layer đó, sau đó kết hợp các kết quả cục bộ thông qua một collective operation để tạo ra kết quả đầu ra hoàn chỉnh. Điều này giúp giảm dung lượng bộ nhớ cần thiết trên mỗi thiết bị, biến nó thành lựa chọn tự nhiên (và thường là duy nhất) khi trọng số của một layer vượt quá bộ nhớ của một GPU, bất kể độ dài chuỗi đầu vào hay kích thước batch.

Trong một khối Transformer, các phép chiếu song song theo cột (column-parallel projections) được ghép nối với các phép chiếu song song theo hàng (row-parallel projections), do đó mỗi khối chỉ yêu cầu một thao tác AllReduce duy nhất, giữ cho chi phí giao tiếp trong giới hạn kiểm soát.

An image showing the side-by-side comparison of column-wise and row-wise parallel projections.

Context parallelism (Song song hóa Ngữ cảnh)

Trong context parallelism, chuỗi đầu vào (input sequence) được phân chia trên các GPU dọc theo chiều chuỗi. Mỗi GPU chỉ xử lý một phần của chuỗi, trong khi các collective operations làm cho toàn bộ chuỗi toàn cục (global sequence) khả dụng ở những nơi cần thiết, chẳng hạn như trong quá trình attention. Context parallelism đặc biệt hiệu quả cho các khối lượng công việc có chuỗi dài (long-sequence workloads), nơi mà chi phí tính toán attention tăng theo cấp số nhân với độ dài chuỗi, biến nó thành yếu tố tiêu tốn nhiều compute và bộ nhớ nhất.

Điều này cũng đặc biệt phù hợp với các mô hình Diffusion và DiT, nơi attention hai chiều (bidirectional attention) giúp tránh các vấn đề mất cân bằng tải thường gặp với causal masks.

NVIDIA TensorRT 11.0 giới thiệu hỗ trợ cho các nguyên hàm IDistCollectiveLayer cần thiết cho nhiều chiến lược song song hóa khác nhau. Phần còn lại của bài viết này tập trung vào context parallelism, trực tiếp giải quyết vấn đề tốn kém nhất trong các pipeline generative media hiện đại: long-sequence attention.

Context Parallelism cho Generative Media

Các pipeline tạo ảnh và video dựa trên Diffusion tiêu tốn phần lớn ngân sách tính toán và bộ nhớ bên trong các khối attention hoạt động trên các chuỗi token dài. Một latent ảnh độ phân giải cao hoặc một video clip nhiều khung hình có thể tạo ra các chuỗi hàng chục ngàn token trên mỗi khối, và chi phí attention mở rộng theo cấp số nhân với độ dài chuỗi.

AllGather KV

Một cách tiếp cận đơn giản để triển khai context parallelism là AllGather KV. Tại đây, các rank trao đổi các mảnh Key (K) và Value (V) của chúng thông qua thao tác AllGather trước khi tính toán attention cục bộ, cho phép mỗi rank tham chiếu (attend) trên toàn bộ chuỗi đầu vào. Kết quả là tạo ra đầu ra attention bao phủ toàn bộ chuỗi cho mỗi rank, đánh đổi bằng một collective operation bổ sung cho mỗi khối attention, trong khi phép nhân ma trận cục bộ giảm tỷ lệ thuận với số lượng rank.

An image showing that context parallelism partitions a token sequence across GPUs along the sequence dimension. Each rank processes a slice of Q locally. Before attention, ranks run AllGather on K and V so every slice can attend over the full sequence.

Ring Attention

Một cải tiến tiềm năng so với AllGather KV là Ring Attention, nơi quá trình giao tiếp và tính toán được xếp chồng lên nhau (overlapped). GPU xử lý Query (Q) cục bộ đồng thời với việc K và V liên tục di chuyển qua một cấu trúc mạng hình vòng (ring topology). Ring Attention cũng giúp giảm thiểu việc sử dụng bộ nhớ: bằng cách sử dụng online softmax, các tensor K và V kích thước đầy đủ không cần phải được hiện thực hóa (materialized) trên bất kỳ GPU nào.

Hình ảnh này thể hiện một cách triển khai khác của song song ngữ cảnh, sử dụng Ring Attention. Hình ảnh cho thấy chỉ có các cặp K và V có kích thước một phần xuất hiện trên bất kỳ GPU nào, dẫn đến việc giảm thiểu tối đa lượng bộ nhớ cần thiết.

DeepSpeed Ulysses

Đối với ngữ cảnh siêu dài (long context hàng chục nghìn token), một cách tiếp cận khác là DeepSpeed Ulysses. Nó bắt đầu bằng cách phân vùng các mẫu riêng lẻ dọc theo chiều chuỗi trên các GPU. Trước khi tính toán attention, nó sử dụng all-to-all communication cho các Q, K, và V đã phân vùng. Điều này đảm bảo mỗi GPU nhận được toàn bộ độ dài chuỗi nhưng chỉ cho một tập hợp con không chồng chéo của các attention heads, cho phép tính toán song song.

Hình ảnh này thể hiện một cách triển khai khác của song song hóa ngữ cảnh, sử dụng DeepSpeed ​​Ulysses. Hình ảnh minh họa các bước giao tiếp kép "tất cả-đến-tất cả" hoán đổi chiều song song hóa từ độ dài chuỗi sang các đầu chú ý ngay trước khối chú ý, và sau đó quay lại độ dài chuỗi ngay lập tức sau đó.

Benchmarks: Tạo Media với Context Parallelism trong C++

Các pipeline tiêu biểu dưới đây (NVIDIA Cosmos 3 cho video và FLUX.1 cho hình ảnh) ban đầu được viết bằng PyTorch, sau đó được convert qua Torch-TensorRT để tạo ra các TensorRT engine phục vụ triển khai production bằng C++.

Video generation với NVIDIA Cosmos 3

Cosmos3-Nano có thể tạo ảnh, video, âm thanh dựa trên các input đa phương thức. Theo kết quả benchmark, chiến lược Ulysses là người chiến thắng rõ ràng khi mô hình diffusion có độ dài ngữ cảnh cực lớn.

Biểu đồ cột thể hiện độ trễ của các chiến lược song song ngữ cảnh khác nhau trên 1, 2, 4 và 8 GPU. Hình ảnh cho thấy độ trễ có khả năng mở rộng tốt khi sử dụng cả chiến lược Ulysses và AllGather KV.

Biểu đồ cột thể hiện khả năng mở rộng của các chiến lược song song hóa ngữ cảnh khác nhau trên 1, 2, 4 và 8 GPU.

Hai hình ảnh đặt cạnh nhau cho thấy một cánh tay robot đang lau sạch một chiếc đĩa bẩn bằng miếng bọt biển, với một số bát đĩa khác có thể nhìn thấy ở phía sau. Hai hình ảnh trông khá giống nhau.

Image generation với Flux.1

Mô hình FLUX.1-dev có khả năng tạo ảnh từ mô tả văn bản. Ulysses tiếp tục là chiến lược tối ưu nhất cho tạo hình ảnh, nhưng Ring Attention cũng cho thấy khả năng mở rộng (scaling) rất tốt lên tới 4 GPU.

Biểu đồ cột thể hiện độ trễ của các chiến lược song song ngữ cảnh khác nhau trên 1, 2, 4 và 8 GPU. Hình ảnh cho thấy độ trễ thấp nhất đạt được bằng cách áp dụng chiến lược Ulysses.

Biểu đồ cột thể hiện khả năng mở rộng của các chiến lược song song hóa ngữ cảnh khác nhau trên 1, 2, 4 và 8 GPU.

So sánh song song kết quả đầu ra của các chiến lược AllGather KV, Ring Attention và Ulysses CP. Các hình ảnh đầu ra trông rất giống nhau. Cả 3 hình ảnh đầu ra đều cho thấy núi Fuji với một cành cây anh đào ở phía trước.

Bắt đầu sử dụng TensorRT với Multi-device

Workflow cốt lõi vẫn tương tự như TensorRT trên một thiết bị đơn, điểm khác biệt là mạng lưới giờ đây có thể bao gồm các lớp giao tiếp phân tán (distributed communication layers).

Yêu cầu chuẩn bị (Prerequisites):

📥 Tải xuống TensorRT 11 từ NVIDIA Developer Portal.
⚙️ Cài đặt TensorRT 11 theo các tài liệu hướng dẫn chính thức.
🖥️ Chuẩn bị một máy chủ single-node có nhiều GPU.
🛠️ Cài đặt OpenMPI trong môi trường phát triển (bare metal hoặc container).

  1. Tạo mạng lưới để suy luận đa thiết bị

Ở cấp độ mạng, suy luận đa thiết bị được kích hoạt thông IDistCollectiveLayerqua giao tiếp giữa các GPU. Các phép toán tập thể có thể được thêm trực tiếp vào mạng TensorRT bằng cách sử dụng INetworkDefinition::addDistCollective:

using namespace nvinfer1;
// create empty network
auto network =
  std::unique_ptr<INetworkDefinition>(builder->createNetworkV2(
  1U << static_cast<uint32_t>(kSTRONGLY_TYPED)));
auto* input =
  network->addInput("input", DataType::kFLOAT, Dims2{3, 4});
ITensor& inputTensor = *network->getInput(0);
auto* collectiveLayer = network->addDistCollective(
    inputTensor,
    CollectiveOperation::kALL_REDUCE,
    ReduceOperation::kSUM,
    -1,       // root: -1 for collectives without a root rank
    nullptr,  // groups: nullptr means all ranks participate
    0         // groupSize
);
// set the world size aka total number of GPUs
collectiveLayer->setNbRanks(8);

Đối với các tập hợp rút gọn như ALL_REDUCE, REDUCE, và REDUCE_SCATTER, hãy chỉ định một hợp lệ ReduceOperation, chẳng hạn như kSUM. Đối với các tập hợp không rút gọn như ALL_GATHER, BROADCAST, ALL_TO_ALL, GATHER, và SCATTER, hãy sử dụng ReduceOperation::kNONE. Các phép toán dựa trên gốc, bao gồm BROADCAST, REDUCE, GATHER, và SCATTER, yêu cầu một hạng gốc hợp lệ.

  1. Build một engine
// create builder config
auto builderConfig = std::unique_ptr<IBuilderConfig>(builder->createBuilderConfig());
// build engine
auto serializedEngine =  std::unique_ptr<IHostMemory>(builder->buildSerializedNetwork
(*network, *builderConfig));
  1. Tạo ngữ cảnh thực thi
auto runtime = std::unique_ptr<IRuntime>(createInferRuntime(
sample::gLogger.getTRTLogger()));
  1. Liên kết các tensor IO
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});
  1. Thiết lập bộ communicator và hàng đợi chờ suy luận
context->setCommunicator(comm);
context->enqueueV3(stream);

Lưu ý: Bộ truyền thông NCCL cũng phải duy trì hiệu lực trong suốt thời gian tồn tại của ngữ cảnh thực thi sử dụng nó.

  1. Bắt đầu suy luận

Chạy ứng dụng với OpenMPI trên 8 GPU. Mỗi rank sẽ chọn thiết bị CUDA cục bộ, khởi tạo NCCL, tạo công cụ TensorRT riêng, tạo ngữ cảnh thực thi riêng và gắn bộ giao tiếp NCCL.

mpirun -np 8 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'

Tìm hiểu thêm

Nếu bạn muốn tìm hiểu thêm về các chủ đề được giới thiệu trong bài viết này, chúng tôi đã cung cấp một số liên kết hữu ích để bạn tham khảo thêm.

NCCL: Thư viện truyền thông tập thể của NVIDIA (NCCL)

Xử lý song song:

NVIDIA TensorRT:

NVIDIA Torch-TensorRT: Tài liệu hướng dẫn về Torch-TensorRT

Sẵn sàng phá vỡ giới hạn mở rộng quy mô AI của doanh nghiệp?

Việc triển khai và tối ưu hóa các pipeline Generative AI khổng lồ trên nhiều GPU không chỉ đòi hỏi các công cụ phần mềm tinh chỉnh như TensorRT, mà còn cần một nền tảng hạ tầng tính toán được thiết kế chuẩn xác, liền mạch.

Nếu bạn đang tìm kiếm lời giải cho bài toán hiệu suất và muốn xây dựng một hệ thống phần cứng đáp ứng tối đa khối lượng công việc AI phức tạp, hãy để NTC AI đồng hành cùng bạn. Đội ngũ kỹ sư AI của chúng tôi luôn sẵn sàng hỗ trợ, cung cấp các dịch vụ tư vấn chuyên sâu về các giải pháp điện toán AI, máy chủ GPU chuyên dụng hiệu suất cao và kiến trúc hạ tầng toàn diện, giúp doanh nghiệp của bạn tối ưu cả về chi phí lẫn năng lực vận hành.

👉 Liên hệ ngay với đội ngũ NTC AI để được tư vấn chi tiết và kiến tạo hạ tầng AI vững chắc cho tương lai.

____
Bài viết liên quan
TAG: , ,