Tensor Parallelism

Tensor Parallelism là một kỹ thuật distributed training nâng cao được sử dụng trong machine learning để chia nhỏ các cấu trúc toán học cá nhân lớn, hay tensor, trên nhiều phần cứng tăng tốc như GPUs hoặc TPU. Khi huấn luyện các deep learning model khổng lồ, số lượng tham số có thể dễ dàng vượt quá dung lượng bộ nhớ của một thiết bị duy nhất. Thay vì đặt toàn bộ một neural network layer trên một GPU, tensor parallelism phân đoạn các ma trận trọng số và chia nhỏ các phép toán toán học (như phép nhân ma trận) trên nhiều thiết bị trong một cluster. Điều này cho phép model tận dụng sức mạnh tính toán và bộ nhớ kết hợp của toàn bộ thiết bị phần cứng, thực thi các phép tính song song theo mô hình Single-Program Multiple-Data (SPMD) đồng thời đồng bộ hóa kết quả thông qua các kết nối liên thông tốc độ cao như NVIDIA NVLink.

Cách thức hoạt động của Tensor Parallelism#

Cốt lõi của một neural network là các phép nhân ma trận. Tensor parallelism phân phối các hoạt động này bằng cách chia các ma trận theo hàng hoặc theo cột. Ví dụ, trong một fully connected layer hoặc cơ chế chú ý transformer, một GPU có thể tính toán nửa bên trái của ma trận trong khi GPU khác tính toán nửa bên phải. Sau khi các phép tính song song hoàn tất, các thiết bị sẽ giao tiếp—thường sử dụng All-Reduce collective operations nhanh chóng—để tổng hợp các kết quả một phần trước khi chuyển tensor hoàn chỉnh sang layer tiếp theo. Các tiến bộ học thuật gần đây vào năm 2025 đang tối ưu hóa hơn nữa quy trình này bằng cách giới thiệu các kích hoạt được đồng bộ hóa một phần nhằm giảm chi phí giao tiếp vốn thường gây nghẽn các compute cluster lớn.

Phân biệt các kỹ thuật song song liên quan#

Việc hiểu cách tensor parallelism phù hợp với bối cảnh rộng lớn hơn của tính toán phân tán đòi hỏi phải phân biệt nó với các chiến lược phổ biến khác:

• Tensor Parallelism so với Model Parallelism: Tensor parallelism là một danh mục phụ rất cụ thể của model parallelism. Trong khi model parallelism nói chung đề cập đến việc chia tách một mô hình trên các thiết bị theo bất kỳ cách nào, thì tensor parallelism chỉ đề cập nghiêm ngặt đến việc sharding các tensor riêng lẻ bên trong một lớp đơn lẻ.
• Tensor Parallelism so với Pipeline Parallelism: Pipeline parallelism là một dạng model parallelism khác giúp phân vùng mạng theo chiều sâu—đặt các lớp đầu tiên trên GPU 0, các lớp tiếp theo trên GPU 1, v.v. Điều này tạo ra các phụ thuộc tuần tự được gọi là các pipeline bubble. Tensor parallelism chia nhỏ chính các lớp đó, thực thi chúng đồng thời mà không có độ trễ tuần tự, nhưng yêu cầu băng thông mạng cao hơn nhiều.
• Tensor Parallelism và Data Parallelism: Trong data parallelism, toàn bộ model được sao chép hoàn toàn trên mọi GPU, và chỉ có tập dữ liệu huấn luyện được chia nhỏ trên các thiết bị. Đối với các kiến trúc được tối ưu hóa cao như Ultralytics YOLO26, dễ dàng khớp trên các GPU hiện đại, data parallelism thông qua DistributedDataParallel của PyTorch là phương pháp mặc định. Tensor parallelism thường chỉ cần thiết khi các tham số của một layer duy nhất vượt quá VRAM của phần cứng, gây ra lỗi hết bộ nhớ (OOM).

Các ứng dụng trong thực tế#

Tensor parallelism là không thể thiếu trong cơ sở hạ tầng AI hiện đại, đặc biệt là đối với các kiến trúc tiên tiến đòi hỏi quy mô tính toán khổng lồ:

• Training Large Language Models (LLMs): Massive foundation models like Meta's Llama 3 và DeepSeek V3 tận dụng các framework như NVIDIA Megatron-LM để triển khai tensor parallelism. Bởi vì các chiều ẩn và các attention head của các model này rất lớn, việc chia chúng trên một node 8-GPU là bắt buộc để huấn luyện hiệu quả và duy trì độ trễ thấp trong quá trình real-time inference.
• Large Vision Models (LVMs) and 3D Generation: As computer vision mở rộng quy mô hướng tới các hệ thống suy luận đa phương thức lớn, các nhà nghiên cứu sử dụng tensor parallelism kết hợp với pipeline parallelism trên các dịch vụ như AWS SageMaker để huấn luyện các vision transformer (ViT) khổng lồ. Kỹ thuật này cho phép xử lý hình ảnh độ phân giải cao và tạo video đòi hỏi các khối bộ nhớ liên tục khổng lồ.

Triển khai Tensor Parallelism trong PyTorch#

Trước đây, các kỹ sư phải viết logic phân tán tùy chỉnh phức tạp để phân đoạn tensor. Gần đây, PyTorch đã giới thiệu DTensor (Distributed Tensor), đơn giản hóa quy trình làm việc này một cách nguyên bản. Dưới đây là ví dụ về việc tạo một tensor được phân đoạn theo hàng bằng cách sử dụng official PyTorch Distributed Tensor API:

import torch
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import Shard, distribute_tensor

# Initialize a 1D device mesh across 2 GPUs
mesh = init_device_mesh("cuda", (2,))

# Create a standard PyTorch tensor (e.g., representing a layer's weights)
local_tensor = torch.randn(1024, 1024)

# Distribute the tensor across the GPUs by sharding along the first dimension (row-wise)
# Each GPU now holds a (512, 1024) chunk of the original tensor
distributed_tensor = distribute_tensor(local_tensor, mesh, [Shard(0)])

print(f"Global shape: {distributed_tensor.shape}, Local shape: {distributed_tensor.to_local().shape}")

Đối với các tác vụ thị giác được tối ưu hóa cho biên và model deployment nhanh chóng, các nhà phát triển thường dựa vào Ultralytics Platform để tự động xử lý việc tận dụng phần cứng tối ưu. Mặc dù các foundation model có hàng tỷ tham số yêu cầu cấu hình tensor parallelism thủ công, bạn vẫn có thể mở rộng quy mô huấn luyện cho các model như YOLO26 một cách hiệu quả bằng các CLI commands đơn giản ngay khi mở hộp. Điều này đảm bảo thông lượng tối đa bằng cách sử dụng liền mạch các kỹ thuật data parallelism nguyên bản cùng với các model training tips mạnh mẽ.