Đây là những gì chúng tôi đã làm trong 9 tháng qua: làm cho việc đào tạo MoEs nhanh hơn ~2 lần và tiêu tốn bộ nhớ ít hơn ~2 lần! Những điểm nổi bật: - MoE thường mất nhiều thời gian và bộ nhớ nhất trong các mô hình hiện đại. Hóa ra, người ta có thể viết lại toán học quá trình ngược của MoE để giảm bộ nhớ kích hoạt mà bạn cần lưu trữ trong quá trình tiến lên khoảng ~2 lần, dẫn đến cùng một gradient mà không cần tính toán lại matmul. Tôi thực sự thích kết quả này, vì nó kết hợp cả cái nhìn thuật toán và hệ thống. - Phân tích các nút thắt trong lớp MoE dẫn đến một chiến lược tối ưu hóa tự nhiên: giảm thiểu việc đọc/ghi bộ nhớ càng nhiều càng tốt! Việc thu thập đầu vào cho quá trình tiến lên và gradient đầu ra cho quá trình ngược có thể đôi khi mất nhiều thời gian như các GEMM nhóm. Chúng tôi kết hợp thu thập với GEMM nhóm + chồng chéo truy cập bộ nhớ và tính toán để làm cho toàn bộ lớp nhanh hơn ~2 lần. - Tính toán top-k cho định tuyến chuyên gia có thể mất thời gian đáng ngạc nhiên, khoảng ~15-20% toàn bộ lớp MoE! Cài đặt top-k tiêu chuẩn sử dụng thuật toán top-k radix, rất tốt cho k lớn nhưng không tối ưu cho k nhỏ. Chúng tôi đã viết lại top-k bằng cách sử dụng thuật toán top-k bitonic, và đôi khi nó nhanh hơn 20-30 lần so với top-k của pytorch! Tất cả các kernel chính đều được viết bằng Cute-DSL nên chúng nên dễ dàng mở rộng (và cài đặt :D). Các kernel Hopper đã ra mắt, các kernel Blackwell gần như đã sẵn sàng. Các mô hình MoE trước đây thường kém hiệu quả phần cứng gấp 2 lần để đào tạo, hy vọng Sonic-MOE sẽ thay đổi điều đó.

Nvidia tiếp tục phát hành một số mô hình mở mạnh mẽ và nhanh nhất. Dữ liệu huấn luyện trước và sau khi huấn luyện cũng được phát hành, điều mà rất ít tổ chức đã làm.