Yandex giới thiệu YaFSDP: Công cụ AI mã nguồn mở hứa hẹn cách mạng hóa đào tạo LLM

Mục lục:

1. Khó khăn trong đào tạo LLM phân tán
2. Giới thiệu YaFSDP
3. Tiêu thụ bộ nhớ
4. Kiểm tra kích hoạt
5. Tối ưu hóa giao tiếp
6. Kết quả thực nghiệm và hiệu suất tăng cường
7. Kết luận

Khó khăn trong đào tạo LLM phân tán

Huấn luyện các mô hình ngôn ngữ lớn (LLM) trên nhiều GPU liên quan đến các hoạt động phức tạp dẫn đến hiệu quả thấp và tiêu thụ bộ nhớ cao. Một trong những vấn đề chính là cần gửi và nhận khối lượng dữ liệu khổng lồ giữa các GPU. Ví dụ: trong một hoạt động all_reduce điển hình, lượng dữ liệu gradient gấp đôi số lượng tham số mạng cần được truyền đi. Trong trường hợp mô hình Llama 70B, điều này có nghĩa là truyền 280 GB dữ liệu mỗi lần lặp.

Hơn nữa, trọng số, gradient và trạng thái bộ tối ưu hóa được nhân bản trên các GPU, dẫn đến tải bộ nhớ khổng lồ. Mô hình Llama 70B và bộ tối ưu hóa Adam yêu cầu hơn 1 TB bộ nhớ, vượt xa dung lượng bộ nhớ 80 GB điển hình của hầu hết các GPU. Sự dư thừa này làm chậm đáng kể quá trình đào tạo và thường khiến việc đưa ngay cả các mô hình có kích thước vừa phải vào bộ nhớ GPU trở nên không thực tế.

Giới thiệu YaFSDP

YaFSDP của Yandex cung cấp một giải pháp hiệu quả cao cho những thách thức này. Bằng cách tập trung vào việc tối ưu hóa tiêu thụ bộ nhớ và loại bỏ các điểm nghẽn giao tiếp, YaFSDP nâng cao hiệu quả của đào tạo LLM. Nó hoạt động bằng cách phân chia các lớp thay vì các tham số riêng lẻ, duy trì giao tiếp hiệu quả và tránh các hoạt động dư thừa. Ngoài ra, YaFSDP phân bổ trước bộ đệm cho tất cả dữ liệu cần thiết, đảm bảo rằng bộ phân bổ Torch không gây ra bất kỳ sự thiếu hiệu quả nào.

YaFSDP hoạt động bằng cách sử dụng hai bộ đệm cho trọng số và gradient trung gian, với các lớp lẻ sử dụng một bộ đệm và các lớp chẵn sử dụng bộ đệm còn lại. Trọng số từ các lớp khác nhau được lưu trữ trong cùng một bộ nhớ. Nếu các lớp có cùng cấu trúc, chúng sẽ luôn giống nhau. Điều quan trọng là phải đảm bảo rằng khi bạn cần lớp X, bộ đệm chứa trọng số cho lớp X. Tất cả các tham số sẽ được lưu trữ trong phần bộ nhớ tương ứng trong bộ đệm.

Tiêu thụ bộ nhớ

Trong quá trình đào tạo, những người tiêu thụ bộ nhớ chính là trọng số, gradient, trạng thái bộ tối ưu hóa, bộ đệm và kích hoạt. YaFSDP làm giảm đáng kể tiêu thụ bộ nhớ bằng cách tối ưu hóa cách các yếu tố này được lưu trữ và truy cập.

Trọng số, Gradient và Trạng thái Bộ tối ưu hóa: Những yếu tố này phụ thuộc vào số lượng quy trình, và tiêu thụ bộ nhớ của chúng có xu hướng tiến gần đến 0 khi số lượng quy trình tăng lên. Bằng cách phân chia các thành phần này trên các GPU, YaFSDP giảm thiểu sự nhân bản và do đó giảm mức sử dụng bộ nhớ.

Bộ đệm tiêu thụ một lượng bộ nhớ cố định và lưu trữ các giá trị trung gian trong quá trình tính toán.

Kích hoạt phụ thuộc vào kích thước mô hình và số lượng mã thông báo được xử lý trên mỗi GPU.

Kiểm tra kích hoạt

Kiểm tra kích hoạt là một kỹ thuật chỉ lưu trữ các kích hoạt cần thiết trong quá trình chuyển tiếp xuôi và tính toán lại chúng trong quá trình chuyển tiếp ngược. Điều này làm giảm đáng kể dấu chân bộ nhớ, vì chỉ dữ liệu cần thiết được lưu trữ. Ví dụ: khi đào tạo mô hình Llama 2 70B với kích thước lô là 8192 mã thông báo, lưu trữ kích hoạt có thể được giảm từ hơn 110 GB xuống chỉ 5 GB. Tuy nhiên, phương pháp này tạo ra chi phí tính toán bổ sung, mà YaFSDP cho phép tránh bằng cách không sử dụng kiểm tra kích hoạt cho một số lớp, điều này là có thể do tối ưu hóa bộ nhớ.

Tối ưu hóa giao tiếp

YaFSDP cải thiện hiệu quả giao tiếp GPU bằng cách đảm bảo rằng dữ liệu chỉ được truyền khi cần thiết và bằng cách chồng giao tiếp với tính toán. Nó sử dụng các luồng CUDA để quản lý các tính toán và giao tiếp đồng thời một cách hiệu quả.

Công cụ sử dụng hai luồng: một luồng tính toán và một luồng giao tiếp. Các sự kiện đồng bộ hóa các luồng này, đảm bảo rằng các hoạt động được thực thi theo thứ tự chính xác mà không gây ra điểm chết.

Quá trình chuyển tiếp xuôi trên lớp thứ ba không bắt đầu cho đến khi hoạt động all_gather hoàn tất (điều kiện 1). Tương tự, hoạt động all_gather trên lớp thứ ba sẽ không bắt đầu cho đến khi quá trình chuyển tiếp xuôi trên lớp đầu tiên sử dụng cùng một bộ đệm hoàn tất (điều kiện 2). Vì không có chu kỳ nào trong sơ đồ này, nên điểm chết là không thể xảy ra.

Kết quả thực nghiệm và hiệu suất tăng cường

Việc triển khai YaFSDP đã cho thấy những cải thiện đáng kể về hiệu quả đào tạo. Trong một kịch bản đào tạo trước với mô hình có 70 tỷ tham số, YaFSDP đã có thể tiết kiệm tài nguyên của khoảng 150 GPU. Điều này dịch sang tiết kiệm chi phí hàng tháng đáng kể, dao động từ $0,5 đến $1,5 triệu, tùy thuộc vào nhà cung cấp hoặc nền tảng GPU ảo.

YaFSDP làm giảm thời gian đào tạo lên tới 26% so với các phương pháp hiện có như FSDP và tối ưu hóa việc sử dụng bộ nhớ, cho phép đào tạo các mô hình lớn hiệu quả hơn.

Kết luận

YaFSDP đại diện cho một bước tiến đáng kể trong đào tạo LLM. Việc giải quyết các thách thức nghiêm trọng về tiêu thụ bộ nhớ và thiếu hiệu quả giao tiếp cho phép đào tạo các mô hình ngôn ngữ lớn nhanh hơn và hiệu quả hơn.