Distributed Training

Huấn luyện phân tán phân chia công việc huấn luyện một model trên nhiều bộ xử lý, thường là GPU hoặc TPU trong cùng một máy hoặc trên một cụm đám mây hay cụm tại chỗ, để rút ngắn thời gian huấn luyện hoặc cho phép huấn luyện một model quá lớn để vừa bộ nhớ của một thiết bị. Các thiết bị trao đổi gradient, activation hoặc phân mảnh tham số qua các kết nối tốc độ cao bằng những thư viện truyền thông tin như NVIDIA NCCL. Đây là cách huấn luyện các model nền tảng và model ngôn ngữ lớn hiện đại, đồng thời cũng rút ngắn quá trình huấn luyện phát hiện đối tượng thường ngày trên các máy trạm nhiều GPU.

Cách thức Hoạt động của Huấn luyện Phân tán#

Song song hóa dữ liệu là chiến lược phổ biến nhất. Mỗi GPU lưu một bản sao đầy đủ của model, xử lý phần riêng trong mỗi batch dữ liệu huấn luyện rồi tham gia thao tác all-reduce để lấy trung bình gradient, nhờ đó mọi bản sao đều áp dụng cùng một lần cập nhật gradient descent. DistributedDataParallel (DDP) của PyTorch chồng lấp quá trình truyền thông tin này với lượt lan truyền ngược. Kích thước batch toàn cục bằng kích thước batch trên mỗi GPU nhân với số GPU, vì vậy learning rate thường phải tăng tương ứng: Goyal và cộng sự đã huấn luyện ResNet-50 trên ImageNet trong một giờ bằng 256 GPU với batch gồm 8,192 ảnh, sử dụng quy tắc tăng tuyến tính và giai đoạn warmup, đạt hiệu suất mở rộng khoảng 90% từ 8 đến 256 GPU.

Song song hóa dữ liệu phân mảnh giúp giảm chi phí bộ nhớ của các bản sao đó. ZeRO và Fully Sharded Data Parallel (FSDP) của PyTorch phân chia tham số, gradient và trạng thái optimizer trên nhiều GPU, rồi chỉ tập hợp đầy đủ trọng số của mỗi layer khi layer đó đang tính toán, nhờ vậy lượng bộ nhớ cần thiết trên mỗi GPU giảm khi bổ sung thêm GPU.

Song song hóa model phân chia chính mạng nơ-ron. Song song hóa tensor chia các phép nhân ma trận riêng lẻ trên nhiều GPU, song song hóa pipeline phân công các layer liên tiếp cho những GPU khác nhau rồi truyền các micro-batch qua chúng, song song hóa ngữ cảnh chia nhỏ các chuỗi dài, còn song song hóa chuyên gia đặt các chuyên gia của model hỗn hợp chuyên gia trên các GPU khác nhau. Các framework như Megatron-LM kết hợp những phương pháp này với song song hóa dữ liệu để huấn luyện các model có hàng trăm tỷ tham số.

Thực hành huấn luyện đa GPU#

Giao tiếp là phần chi phí chính: lưu lượng all-reduce tăng theo kích thước model, vì vậy GPU kết nối bằng NVLink trong cùng node mở rộng quy mô tốt hơn các node liên kết qua mạng chậm hơn. Thống kê batch normalization được tính riêng trên từng GPU, trừ khi các layer được chuyển sang synchronized batch normalization; điều này quan trọng khi phần batch trên mỗi GPU nhỏ. Ultralytics YOLO26 chạy huấn luyện đa GPU qua DDP: truyền nhiều ID thiết bị sẽ khởi chạy một process cho mỗi GPU bằng torch.distributed.run, sử dụng backend NCCL trên CUDA và chia đều giá trị batch cho các GPU.

from ultralytics import YOLO

# Tải model YOLO26 đã pretrained
model = YOLO("yolo26n.pt")

# Huấn luyện trên GPU 0 và 1 bằng DDP; batch=32 nghĩa là mỗi GPU xử lý 16 ảnh ở mỗi bước
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, batch=32, device=[0, 1])

Ứng dụng thực tế#

• Model ngôn ngữ lớn: Meta đã huấn luyện Llama 3 405B trên tối đa 16K GPU H100, kết hợp song song hóa tensor, pipeline, ngữ cảnh và dữ liệu theo phương pháp mà các tác giả gọi là song song hóa 4D.
• Xe tự hành: Các nhóm tri giác huấn luyện model phát hiện và phân đoạn trên kho lưu trữ lớn gồm bản ghi camera, lidar và radar từ đội xe; việc này chỉ khả thi trên các cụm GPU lớn.
• Phân tích ảnh y tế: Các thể tích CT và MRI 3D độ phân giải cao gây quá tải bộ nhớ của một GPU, còn việc phân tán batch hoặc các phân mảnh model trên nhiều GPU cho phép huấn luyện các thể tích và model lớn hơn.

Huấn luyện phân tán và các khái niệm liên quan#

• Học liên kết: Học liên kết cũng huấn luyện trên nhiều thiết bị, nhưng dữ liệu được lưu trên điện thoại, tại bệnh viện hoặc trên các máy khách khác để bảo vệ quyền riêng tư; chỉ các bản cập nhật model được gửi đến máy chủ trung tâm. Huấn luyện phân tán tập hợp dữ liệu trong một cụm để tối đa hóa tốc độ.
• Học phân tách: Học phân tách chia model giữa máy khách và máy chủ để dữ liệu thô không bao giờ rời khỏi máy khách, đánh đổi quyền riêng tư bằng việc tăng lượng truyền thông tin.
• Tích lũy gradient: Tích lũy gradient cộng dồn gradient qua nhiều batch nhỏ trước một lần cập nhật, mô phỏng phép tính của batch lớn trên một GPU với chi phí thời gian, trong khi song song hóa dữ liệu xử lý đồng thời các batch đó trên nhiều GPU.