Unsloth và Nvidia hợp tác giúp tăng tốc độ huấn luyện LLM lên 25% trên GPU phổ thông

Fine-tuning (tinh chỉnh) hiện nay là một trong những tác vụ đòi hỏi tài nguyên tính toán khổng lồ, liên tục đẩy phần cứng đến giới hạn. Các GPU của NVIDIA được thiết kế riêng cho các khối lượng công việc này: chúng chia nhỏ các vấn đề phức tạp và xử lý song song. Unsloth hoạt động trên dải sản phẩm GPU của NVIDIA, từ các máy tính xách tay RTX cục bộ đến các siêu máy tính AI cá nhân DGX Spark.

Để giúp các nhà phát triển khai thác tối đa GPU của mình, Unsloth đã hợp tác với NVIDIA nhằm loại bỏ các nút thắt ẩn làm chậm quá trình huấn luyện. Các tối ưu hóa mới được triển khai này đã tăng tốc độ huấn luyện GPU lên khoảng 25% khi kết hợp lại. Dưới đây là chi tiết về cách họ đã thực hiện điều đó.

Unsloth và Nvidia hợp tác tối ưu hóa LLM

Tối ưu hóa quy trình huấn luyện: Vượt qua các nút thắt ẩn

Khi tối ưu hóa huấn luyện mô hình, các nhà phát triển thường bắt đầu với các nhân (kernels) có tác động cao như: phép nhân ma trận (matmuls), cơ chế chú ý (attention), các thao tác kết hợp (fused ops), nhóm GEMM, v.v. Mặc dù các nhân này thực hiện phần lớn các phép tính số học, nhưng khi các thành phần chính được tối ưu hóa, một lớp nút thắt khác sẽ xuất hiện. GPU bị đình trệ do các công việc phụ thuộc vào metadata. Thời gian chạy xây dựng lại các cấu trúc dữ liệu giống hệt nhau ở mỗi lần lặp, và các luồng sao chép/tính toán thực hiện tuần tự thay vì có thể chồng chéo lên nhau.

Bằng cách nhắm mục tiêu vào các nút thắt này, Unsloth và NVIDIA đã hợp tác thực hiện ba cải tiến chính. Mô hình chung của các tối ưu hóa này rất đơn giản: giảm thiểu các công việc quản lý lặp lại và để công việc sao chép diễn ra song song với các tính toán hữu ích.

1. Caching Packed-Sequence (Đóng gói chuỗi)

Thay vì đệm (padding) tất cả các chuỗi về cùng một độ dài và lãng phí tài nguyên tính toán trên các token đệm, chúng ta nối chúng thành một chuỗi đóng gói (packed sequence) dài hơn:

Padded vs Packed Sequence

Mô hình vẫn cần biết nơi mỗi chuỗi gốc bắt đầu và kết thúc. Vì vậy, cùng với các token đã đóng gói, chúng ta mang theo metadata chuỗi như độ dài, chỉ số chuỗi tích lũy (cu_seqlens), độ dài chuỗi tối đa và cấu trúc mặt nạ.

Điểm mấu chốt ở đây là: đối với một lô (batch) đóng gói cố định, metadata đó giống hệt nhau cho mọi lớp (layer). Nếu mô hình có L lớp, việc xây dựng hoặc đồng bộ hóa lại B một lần cho mỗi lớp không phải là công việc mới. Đó là cùng một thông tin được xây dựng lại nhiều lần.

Việc xây dựng B + xây dựng B + ... + xây dựng B (L lần) tạo ra chi phí overhead. Một số đường dẫn này có thể buộc đồng bộ hóa từ thiết bị sang máy chủ (device-to-host), hiệu quả là tạo ra một điểm đồng bộ hóa GPU-CPU. Một khi điều đó xảy ra trong đường dẫn mỗi lớp, chi phí overhead sẽ lặp lại ở mỗi lớp.

Thay đổi caching packed-sequence giúp giảm thiểu điều này. Thay vì xây dựng lại thông tin chuỗi đóng gói liên tục, nó lưu trữ (cache) metadata có thể tái sử dụng và các cấu trúc phía chú ý (attention-side) bắt nguồn từ nó, trên mỗi thiết bị, cho lô đóng gói hiện tại. Các cấu trúc được lưu trữ này sau đó được tái sử dụng trên các lớp.

2. Double Buffering cho Activation Checkpointing

Activation checkpointing là một kỹ thuật tiêu chuẩn để huấn luyện các mô hình lớn. Ý tưởng là tiết kiệm bộ nhớ bằng cách không giữ mọi activation trung bình sống qua quá trình truyền ngược (backward pass). Đổi lại, chúng ta phải trả thêm một chút công việc trong quá trình backward.

Tuy nhiên, điều này đặt ra một câu hỏi hệ thống: nếu một activation đã được chuyển ra ngoài (offload), làm thế nào để nó quay lại GPU cho backward?

Trong đường dẫn smart checkpointing của Unsloth, các activation có thể được dàn dựng trong bộ nhớ CPU được ghim (pinned CPU memory) và sao chép lại khi cần. Điều này tiết kiệm VRAM, nhưng có thể giới thiệu một nút thắt: đó là một mẫu tuần tự hóa (serialization pattern). Nếu một bộ đệm được sử dụng lại cho cả sao chép và tính toán, luồng sao chép và luồng tính toán sẽ lần lượt thay nhau hoạt động.

Một cách xử lý sạch hơn là sử dụng hai bộ đệm (double buffering). Trong khi quá trình truyền ngược đang chạy trên bộ đệm A, luồng sao chép có thể tải trước activation tiếp theo vào bộ đệm B. Sau đó, vai trò của chúng hoán đổi cho nhau. Điều này tạo ra sự chồng chéo đường ống (pipeline overlap), mặc dù không phải là sự chồng chéo hoàn hảo.

Double buffering không giảm lượng toán học. Nó che giấu độ trễ sao chép đằng sau các tính toán hữu ích. Loại tối ưu hóa này có xu hướng trở nên mạnh mẽ hơn khi mô hình đủ lớn để tính toán backward là đáng kể, nhưng không quá chiếm ưu thế以至于 chi phí sao chép biến mất trong tiếng ồn.

3. Tối ưu hóa MoE Routing

Thay đổi thứ ba chuyên biệt hơn, nhưng nó cho thấy cùng một mẫu trong MoE (Mixture of Experts) routing.

Trong đường dẫn MoE dựa trên PyTorch GPT-OSS mà chúng tôi kiểm tra, một phần tốn kém của việc định tuyến là figuring out các token đi đến expert nào. Một triển khai ngây thơ có thể thực hiện việc này bằng cách sử dụng torch.where cho từng expert.

Thoạt nhìn, điều này có vẻ vô hại. Nhưng torch.where phụ thuộc vào dữ liệu ở đây: số lượng token được định tuyến đến mỗi expert thay đổi từ lô này sang lô khác. Điều này có thể giới thiệu đồng bộ hóa CPU-GPU hoặc chi phí overhead thời gian chạy liên quan vì kích thước đầu ra phụ thuộc vào mẫu định tuyến.

Cách tiếp cận tốt hơn là nhóm mọi thứ một lần: nhóm tất cả các token cần định tuyến, sau đó tính toán các phần bù (offsets) một lần cho tất cả các expert. Về mặt toán học, lợi ích không phải là chúng ta thay đổi logic định tuyến. Chúng ta thay đổi tần suất chúng ta yêu cầu thời gian chạy trả lời một câu hỏi lập chỉ mục động.

Kết luận

Mặc dù ba tối ưu hóa này nằm ở các phần khác nhau của ngăn xếp (stack), chúng đang giải quyết cùng một vấn đề. Các cơ hội tối ưu hóa chính nằm trong mã "keo dán" (glue code) xung quanh các nhân chính.

Có một bài học kỹ thuật hữu ích ở đây. Khi các nhân toán học được tối ưu hóa, "nhanh hơn" thường có nghĩa là một trong hai điều: làm ít toán học hơn hoặc làm ít công việc quản lý hơn. Các tối ưu hóa này tập trung vào cái sau, giúp các nhà phát triển tận dụng tối đa phần cứng NVIDIA của họ để xây dựng các mô hình AI tốt hơn và nhanh hơn.