Tích hợp và sử dụng Kernels Hub

Thư viện kernels cho phép nạp các kernel tính toán đã được tối ưu trực tiếp từ Hub.
Bạn có thể tìm thấy kernels trong các tổ chức chuyên biệt hoặc bằng cách tìm kiếm thẻ kernel trong Hub.

Kernels là những mảnh mã được tối ưu hóa hỗ trợ cho việc phát triển, huấn luyện và suy luận mô hình. Tại đây, chúng ta sẽ tập trung vào việc tích hợp với TRL, nhưng hãy xem các tài nguyên ở trên để tìm hiểu thêm về chúng.

Cài đặt

Để sử dụng kernels với TRL, bạn cần cài đặt thư viện này trong môi trường Python của mình:

pip install kernels

Sử dụng Kernels từ Hub trong TRL

Kernels có thể thay thế trực tiếp các triển khai attention, loại bỏ nhu cầu biên dịch thủ công các backend attention như Flash Attention và tăng tốc độ huấn luyện chỉ bằng cách lấy kernel attention tương ứng từ Hub.

Bạn có thể chỉ định một kernel khi tải một model:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "your-model-name",
    attn_implementation="kernels-community/flash-attn2"  # other options: kernels-community/vllm-flash-attn3, kernels-community/paged-attention
)

Hoặc khi chạy một script huấn luyện TRL:

python sft.py ... --attn_implementation kernels-community/flash-attn2

Hoặc sử dụng TRL CLI:

trl sft ... --attn_implementation kernels-community/flash-attn2

Bây giờ bạn có thể tận dụng các backend attention nhanh hơn với kernel đã được tối ưu hóa sẵn cho cấu hình phần cứng của bạn từ Hub, giúp tăng tốc cả quá trình phát triển và huấn luyện.

Chọn và cố định phiên bản Kernel

Các kho kernel trên Hub được phân bản theo nhánh (v1, v2, v3, …), và Transformers chọn một phiên bản mặc định cho từng kho. Phiên bản mặc định đó có thể thay đổi từ bản phát hành Transformers này sang bản khác, nên cùng một giá trị attn_implementation không phải lúc nào cũng trỏ đến cùng một bản build.

Để kiểm soát bản build nào được tải, hãy thêm một revision vào id của kho, có thể là một nhánh phiên bản hoặc một SHA của commit:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "your-model-name",
    attn_implementation="kernels-community/flash-attn2@v2",  # or a commit SHA to pin an exact build
)

Pinning (cố định phiên bản) hữu ích để đảm bảo một phiên chạy huấn luyện có thể tái lập được, và để duy trì một bản build đã được xác nhận hoạt động tốt khi một bản mới hơn gặp lỗi hồi quy. Lưu ý rằng mỗi phiên bản chỉ cung cấp các bản build cho một phạm vi các phiên bản Torch và CUDA nhất định, do đó một phiên bản được cố định có thể không có biến thể phù hợp với môi trường của bạn.

attn_implementation="flash_attention_2" sẽ tự động chuyển sang kernel kernels-community/flash-attn2 trên Hub nếu gói flash-attn chưa được cài đặt, do đó bạn có thể đang sử dụng một kernel từ Hub mà không hề yêu cầu điều đó một cách rõ ràng.

Các bản build v3 cho CUDA 12.8 hiện đang gặp lỗi ở bước truyền ngược (backward pass) đối với các mô hình sử dụng grouped-query attention. Nếu môi trường của bạn đang dùng phiên bản CUDA đó, hãy cố định (pin) @v2 cho đến khi huggingface/kernels-community#1085 được đóng.

So sánh các triển khai Attention

Chúng tôi đã đánh giá các triển khai attention khác nhau có sẵn trong transformers, cùng với các backend kernel khác nhau, bằng cách sử dụng TRL và SFT.
Các thí nghiệm được chạy trên một H100 GPU duy nhất với CUDA 12.9, tận dụng Qwen3-8B với batch size là 8, gradient accumulation là 1 và độ chính xác bfloat16.
Lưu ý rằng các kết quả được trình bày ở đây chỉ áp dụng cho cấu hình này và có thể thay đổi với các cấu hình huấn luyện khác.

Hình sau minh họa cả độ trễ (thời gian cho mỗi bước huấn luyện) và bộ nhớ cấp phát đỉnh cho các cách triển khai attention và các backend kernel khác nhau.
Các triển khai dựa trên kernel cho hiệu suất tương đương với attention được cài đặt tùy chỉnh, và việc tăng max_length của mô hình sẽ tiếp tục cải thiện hiệu suất. Mức tiêu thụ bộ nhớ tương tự nhau giữa tất cả các cách triển khai, không có sự khác biệt đáng kể. Chúng ta đạt được hiệu suất tương tự nhưng với ít trở ngại hơn, như được mô tả trong mục tiếp theo.

Flash Attention so với Hub Kernels

Việc xây dựng Flash Attention từ mã nguồn có thể mất nhiều thời gian, thường dao động từ vài phút đến vài giờ, tùy thuộc vào phần cứng, cấu hình CUDA/PyTorch và việc có sẵn các bánh xe (wheels) đã biên dịch trước hay không.

Ngược lại, Hugging Face Kernels cung cấp quy trình làm việc nhanh hơn và đáng tin cậy hơn nhiều. Nhà phát triển không cần lo lắng về các cấu hình phức tạp—mọi thứ đều được xử lý tự động. Trong các bài kiểm tra hiệu năng của chúng tôi, các kernel đã sẵn sàng để sử dụng chỉ trong khoảng 2,5 giây, không cần biên dịch. Điều này cho phép bạn bắt đầu huấn luyện gần như ngay lập tức, đẩy nhanh đáng kể quá trình phát triển. Chỉ cần chỉ định phiên bản mong muốn, và kernels sẽ lo phần còn lại.

Update on GitHub