Đằng sau hậu trường của huấn luyện phân tán và lý do tại sao hệ thống dây GPU của bạn lại quan trọng như chiến lược của bạn

Trí Để tinh chỉnh Mistral-7B bằng Adam với độ chính xác hỗn hợp BF16, các tham số chiếm 14 GB (7 tỷ tham số, mỗi tham số 2 byte). Các gradient chiếm thêm 14 GB. Trạng thái bộ tối ưu hóa là phần tốn kém nhất, Adam duy trì hai giá trị trung bình động cho mỗi tham số (động lượng và phương sai), cả hai đều được lưu trữ ở định dạng FP32, chiếm khoảng 58 GB. Tổng cộng, chúng ta cần 87 GB trước khi mô hình xử lý bất kỳ lô dữ liệu hoặc tính toán bất kỳ kích hoạt nào. Đối với GPU A100 có 80 GB VRAM, mô hình này không thể chứa vừa. Do đó, chúng ta cần phân tán. Câu hỏi đặt ra là: mỗi GPU sẽ lưu trữ những gì? **Distributed Data Parallel (DDP): Mọi GPU đều lưu trữ mọi thứ** Câu trả lời đơn giản là mọi GPU đều lưu trữ toàn bộ mô hình. Đây là DDP, phương pháp mà hầu hết mọi người thường gặp đầu tiên. Mỗi GPU giữ một bản sao hoàn chỉnh của mô hình: tất cả các tham số, tất cả các gradient, tất cả các trạng thái bộ tối ưu hóa. Phần được chia sẻ là dữ liệu. Mỗi GPU xử lý một phần khác nhau của lô dữ liệu, thực hiện quá trình truyền xuôi và truyền ngược riêng, và tính toán gradient từ phần dữ liệu của mình. Vấn đề là các gradient này đều khác nhau vì mỗi GPU xử lý dữ liệu khác nhau. Nếu mỗi GPU tự cập nhật bộ tối ưu hóa của mình, các bản sao sẽ lệch nhau và chúng ta sẽ không còn một mô hình duy nhất mà có bốn mô hình. Do đó, trước khi cập nhật, các GPU sẽ tính trung bình các gradient của chúng. Mỗi GPU gửi gradient của mình đi, nhận gradient của các GPU khác, và tất cả đều có cùng một gradient trung bình. Sau đó, chúng thực hiện một bước tối ưu hóa giống hệt nhau và các bản sao vẫn đồng bộ. Việc tính trung bình này là phần giao tiếp duy nhất mà DDP cần, và nó xảy ra một lần mỗi bước. Nó được gọi là all-reduce, và chúng ta sẽ quay lại nó sau khi bắt đầu tính toán chi phí của từng chiến lược. Hình 1: Cách DDP hoạt động (Ảnh của tác giả) DDP nhanh và đơn giản chính xác vì nó giao tiếp rất ít. Một thao tác all-reduce mỗi bước, và PyTorch thậm chí còn chồng chéo nó với quá trình truyền ngược, vì vậy trên phần cứng tốt, bạn hầu như không nhận thấy nó. Nhưng hạn chế của nó nằm trong thiết kế. Mỗi GPU lưu trữ mọi thứ, điều đó có nghĩa là việc thêm GPU không giải quyết được vấn đề bộ nhớ. Nếu 87 GB không vừa trên một A100, nó cũng không vừa trên bốn A100, bởi vì mỗi trong số bốn GPU vẫn đang lưu trữ toàn bộ 87 GB. Thêm GPU vào DDP giúp tăng thông lượng, nhưng nó không làm giảm bộ nhớ mà mỗi GPU phải lưu trữ. **Fully Sharded Data Parallel (FSDP): Không GPU nào lưu trữ mọi thứ** Nếu vấn đề là mỗi GPU lưu trữ mọi thứ, giải pháp rõ ràng là đảm bảo không có GPU nào làm điều đó. Đây là FSDP. Thay vì mỗi GPU giữ một bản sao đầy đủ, mô hình được chia thành các phần, và mỗi GPU chỉ sở hữu một phần của các tham số.