Chạy mô hình Kimi K3 2.78T tham số trên laptop 64GB RAM nhờ engine WASTE

WASTE — Weight-Aware Streaming Tensor Engine (Công cụ Tensor truyền phát nhận biết trọng số)

Kimi K3 — 2,78 nghìn tỷ tham số — chạy trên một chiếc laptop phổ thông.

WASTE là một công cụ suy luận có thể nhúng, được viết bằng ngôn ngữ C và không có các phụ thuộc runtime từ bên thứ ba. Nó giữ phần thân (trunk) của mô hình trong bộ nhớ, truyền phát trực tiếp các chuyên gia (experts) được chọn từ ổ đĩa và sử dụng RAM còn lại làm bộ nhớ đệm chuyên gia có giới hạn.

Điểm chứng minh hiện tại của nó là mô hình Kimi K3 với trọng số mở hoàn chỉnh: 2,78 nghìn tỷ tham số, được chuyển đổi thành một container dung lượng 982 GiB và chạy trên MacBook Pro 64 GB với tốc độ 0,49–0,54 token mỗi giây. Đây không phải là một biến thể đã được chưng cất (distilled), cắt tỉa (pruned) hay rút gọn.

WASTE được viết cho đúng một mô hình đó và đúng một ràng buộc đó: K3 không vừa với RAM của các hệ thống máy tính phổ thông hiện nay. Nó có dung lượng 1,42 TB khi được công bố và 982 GB sau khi chuyển đổi. Tuy nhiên, một mô hình hỗn hợp chuyên gia (mixture of experts) chỉ kích hoạt khoảng 4% chính nó trên mỗi token, vì vậy gần như toàn bộ trọng số đó đều ở trạng thái nhàn rỗi tại bất kỳ thời điểm nào — và trọng số nhàn rỗi không cần phải nằm trong bộ nhớ, nó chỉ cần có thể truy cập kịp thời. WASTE giữ nó trên ổ đĩa theo bố cục mà mỗi chuyên gia chỉ tốn đúng một lần đọc, truyền phát những gì mỗi token thực sự cần và dành mọi byte RAM còn lại cho phần dữ liệu lặp lại.

Vị thế hiện tại

Công cụ này hoạt động chính xác: mọi lớp đều được kiểm chứng đối chiếu với PyTorch, các logit cuối cùng khớp đến 3,6e-06 và tháp thị giác (vision tower) khớp với oracle của chính nó đến 2,3e-06. Nó cũng chậm — nửa token mỗi giây, mất ba mươi giây cho câu trên.

Cả hai yếu tố đó đều quan trọng, và yếu tố thứ hai không nên được hiểu là một lời khước từ trách nhiệm. Chúng tôi không biết đến một bản trình diễn nào khác về việc truyền phát một mô hình có kích thước này từ ổ đĩa trên một máy tính phổ thông: chúng tôi không tìm thấy ví dụ nào về truyền phát NVMe ở quy mô nghìn tỷ, và các công thức được ghi chép tốt nhất cho lớp 671B đều giả định một máy chủ với một terabyte DDR5. Đây là báo cáo về những gì chúng tôi tìm thấy chứ không phải một cuộc khảo sát — kho lưu trữ này không chứa thư mục tham khảo hay bảng so sánh, vì vậy hãy đọc nó như một lời mời gửi phản hồi thay vì một kết quả cuối cùng. Phần thú vị không nằm ở tốc độ, mà ở chỗ toàn bộ mô hình nằm trong phạm vi có thể truy cập được trên một máy tính phổ thông duy nhất — và từ đây, vấn đề trở thành bài toán kỹ thuật thay vì tính khả thi.

Các đòn bẩy đã thay đổi vị trí. Việc chồng lấp các lượt đọc chuyên gia với các phép tính toán mang lại hiệu quả ~1,6x và đã được triển khai; hai phương án có vẻ lớn hơn — đọc ít byte hơn trên mỗi token và giữ nhiều byte hơn trong RAM — đều đã được đo lường và bị từ chối, một vì bộ định tuyến (router) của dòng mô hình này không có phần đuôi để hạ cấp và một vì bộ nhớ đệm mà máy tính không thể duy trì thường trú thì không thể mua được bằng bất cứ giá nào. Ngay cả khi các lượt đọc được chồng lấp, chúng vẫn chiếm 55% một bước giải mã (decode step) so với 27% của các phép tính, vì vậy những gì còn lại là một ổ đĩa nhanh hơn hoặc một máy tính có nhiều RAM hơn, chứ không phải là chạy thêm một lượt qua các kernel. docs/EFFICIENCY.md là bản tường trình về cách mỗi phương án đó được định giá, bao gồm cả hai phương án đã được xây dựng trước khi đo lường.

Những gì nó mở ra, một cách cụ thể: một mô hình quy mô tiên phong phản hồi mà không cần mạng, không cần hóa đơn theo mỗi token và không có dữ liệu nào rời khỏi máy — đó chính là sự khác biệt giữa "bạn không được gửi dữ liệu đó đến API" và "hãy chạy nó tại đây". Định dạng và công cụ này không phụ thuộc sâu vào K3; K3 đơn giản là trường hợp khó nhất hiện nay, và một mô hình truyền phát ở mức 2,78T cũng sẽ truyền phát thoải mái ở mức 48B.

Mọi con số trong tài liệu này đều được đo lường trên commit mà nó được công bố, và những con số sai đều được ghi lại là sai trong docs/LEARNED.md thay vì âm thầm sửa chữa.

Tại sao lại có cái tên này

Mỗi token được trả lời bởi một dịch vụ đám mây đều phải trả phí hai lần: một lần trên hóa đơn và một lần bằng điện năng của trung tâm dữ liệu đang chạy một mô hình lẽ ra có thể vừa — dù chật vật, vụng về, nhưng thực sự — trên phần cứng đang nằm ngay trên bàn làm việc. WASTE hướng tới việc trở thành bước đi cụ thể đầu tiên để chấm dứt sự lãng phí token đó. Từ viết tắt này ra đời sau.

Những gì bạn cần

Các kích thước ở đây là lũy thừa của hai, theo cách mà cả df và công cụ đều báo cáo: container có dung lượng 982 GiB, mà nhà cung cấp ổ đĩa sẽ gọi là 1,05 TB.

Mức RAM tối thiểu là mức mà công cụ từ chối khởi chạy nếu thấp hơn, và nó gần như hoàn toàn là 27,28 GB của phần thân thường trú. Thông lượng hữu ích bắt đầu ở mức cao hơn: trên máy 64 GB, công cụ tự cấp ngân sách 46 GB, trong đó 17,56 GB là bộ nhớ đệm chuyên gia, và đó là đỉnh của đường cong đo lường. Một máy 32 GB về mặt kỹ thuật có thể mở mô hình nhưng sẽ bị phân trang (page) rất tệ; hãy coi 64 GB là yêu cầu thực tế.

Tốc độ lưu trữ không phải là chi tiết nhỏ. Một token đọc 17 GB dữ liệu chuyên gia. Trên SSD nội bộ, tốc độ là 12,78 GB/s và mô hình truyền phát mượt mà; qua cổng USB, tốc độ là 0,94 GB/s và cùng một token đó mất mười ba giây. Hãy chuyển đổi sang NVMe nội bộ và chỉ sử dụng ổ đĩa ngoài để tải xuống.

Nếu không có sẵn một terabyte, cùng công cụ và định dạng đó có thể chạy Kimi-Linear-48B-A3B-Instruct từ một container 19 GB với mức tối thiểu 1,87 GB, ở tốc độ 10,7 tok/s. Đó là con đường tốt để thử nghiệm WASTE trước khi dành riêng một ổ đĩa cho K3.

Nó là gì

Đường dẫn là thư mục container mà bộ chuyển đổi đã ghi — không có mở rộng ~ ở đây, đó là công việc của shell.

Cách thức hoạt động

Vị trí quyết định tốc độ

Một mô hình được chuyển đổi một lần thành container.waste: một tệp manifest JSON, một phần thân thường trú và một ngân hàng chuyên gia cho mỗi lớp. Mỗi bản ghi chuyên gia được căn chỉnh 4 KiB với cổng (gate) của nó, các ma trận lên và xuống nằm liền kề, vì vậy việc định tuyến đến một chuyên gia tốn đúng một lệnh pread — không phải ba, không phải một lần tìm kiếm (seek) cho mỗi ma trận. Các phép tính toán chưa bao giờ là nút thắt cổ chai.

Các lượt đọc bỏ qua bộ nhớ đệm trang (F_NOCACHE trên macOS, O_DIRECT trên Linux, FILE_FLAG_NO_BUFFERING trên Windows). Điều đó là có chủ đích: với một container nhỏ hơn RAM, kernel sẽ lưu vào bộ nhớ đệm mọi thứ, và tỷ lệ hit được đo theo cách đó là một hư cấu không tồn tại được khi tiếp xúc với mô hình 982 GB.

Tiêu đề của mỗi bản ghi được kiểm tra ngay khi nhập vào — đúng magic, đúng chuyên gia mà chỉ mục yêu cầu, đúng các offset phù hợp — vì vậy một ngân hàng bị cắt bớt hoặc ghép nối sẽ dừng quá trình tạo và gọi tên bản ghi thay vì trả lời từ các byte sai. Điều đó không tốn kém gì có thể đo lường được. Bản ghi cũng mang một mã crc32 trên payload của nó, và việc kiểm tra đó là --verify, mặc định là tắt: nó là một lượt quét qua mọi bản ghi trên mỗi lần cache miss, khoảng 5% trên Kimi-Linear và 1% trên K3. Đáng giá cho một container bạn đã sao chép hoặc tải xuống và chưa đọc kể từ đó; không đáng giá trên mỗi token của container bạn tự chuyển đổi. Xem docs/FORMAT.md.

Ba bit cho mỗi trọng số chuyên gia

Các chuyên gia được lưu trữ dưới dạng lượng tử hóa vector dư (residual vector quantization) — ba giai đoạn của codebook 256 mục trên các vector 8 chiều, 3,00 bit cho mỗi trọng số — và ma trận không bao giờ được hiện thực hóa. Đối với mỗi token, công cụ xây dựng một bảng các tích vô hướng từng phần, một cho mỗi mục codebook trên mỗi vị trí vector, sau đó mỗi hàng chuyên gia là ba lượt đọc bảng và hai phép cộng.

Phần thân (trunk) giữ ở mức 4 và 8 bit. Mô hình được huấn luyện với phương pháp huấn luyện nhận biết lượng tử hóa (quantization-aware training) chỉ trên các chuyên gia, vì vậy nó không có khả năng chịu đựng được huấn luyện cho một phần thân bị nén: một phần thân 3-bit đã được xây dựng và đo lường, dự đoán bộ nhớ đệm vẫn giữ nguyên, nhưng thông lượng thì không, và đầu ra bị sụp đổ.

Mức tối thiểu của bộ nhớ đệm là tập hợp làm việc của một token

Con số có tính dự báo cao nhất trong dự án này. K3 chạm vào 16 chuyên gia trong mỗi 92 lớp trên mỗi token: 17,0 GB. Dưới mức đó, một chuyên gia được lưu vào bộ nhớ đệm cho một token sẽ bị đẩy ra (evicted) trước khi token tiếp theo yêu cầu nó, và tỷ lệ hit không phải là thấp — mà là bằng không. Trên mức đó, đường cong uốn cong mạnh.

Được đo theo thứ tự đó, trên một máy tính đang nhàn rỗi. Thứ tự rất quan trọng: chạy lại sau khi các hàng 52 và 58 GB đã đẩy máy tính vào tình trạng phân trang, 46 GB cho kết quả 0,22–0,25 thay vì 0,32 — trong khi báo cáo số lần hit và miss giống hệt nhau đến từng chữ số. Công cụ này có tính tất định; máy tính thì không, và nó không phục hồi hoàn toàn giữa các lần chạy. Hãy quét theo chiều tăng dần.

Cột giải mã (decode) có trước tính năng đọc trước (read-ahead) và chưa được quét lại: 46 GB hiện chạy ở tốc độ 0,51 thay vì 0,32. Hình dạng của đường cong là mục đích của bảng, và đọc trước không làm thay đổi nó — nó ẩn I/O đằng sau các phép tính, giúp mỗi hàng nhanh hơn nhưng không hàng nào có ngân sách khác biệt.

Mọi thứ trong thiết kế bộ nhớ đều tồn tại để vượt qua ranh giới đó, đó là lý do tại sao công cụ hoạt động để giải phóng RAM thay vì tiết kiệm nó.

Và có một giới hạn ở phía bên kia, gần hơn vẻ ngoài của nó. Hãy đọc bảng đó hai lần: tỷ lệ hit tăng dần xuống dưới. Ở mức 58 GB trên máy 64 GB, bộ nhớ đệm phục vụ 37% chuyên gia từ RAM và công cụ chậm hơn tám lần so với mức 46 GB, nơi nó phục vụ 13%. Công cụ nằm trong ngân sách của nó; máy tính thì không, vì vậy hệ điều hành đẩy bộ nhớ đệm chuyên gia ra ngoài, và một "hit" trở thành lỗi trang (page fault) thay vì lượt đọc ổ đĩa mà công cụ đang quản lý.

Vì vậy, cửa sổ sử dụng được rất hẹp. Nó mở ra ở mức ~46 GB, nơi bộ nhớ đệm cuối cùng cũng xóa được tập hợp làm việc của một token, và nó đã đóng lại ở mức 52 — trên một máy tính đang nhàn rỗi, với 49 GB trống trước khi chạy. Nó cũng đủ nhạy để thay đổi dưới một tác động có vẻ không liên quan: lấy 1,11 GB bảng nhúng (embedding table) ra khỏi tập hợp thường trú để đưa thẳng vào bộ nhớ đệm ở một ngân sách cố định, và điều đó đủ để đẩy mức 58 GB từ 0,32 tok/s xuống 0,04.

Vì vậy, mặc định không lấp đầy máy tính. Bộ nhớ đệm chuyên gia chỉ có giá trị ở các bội số nguyên của tập hợp làm việc đó, và phần dư trên một bội số chỉ mua được vài điểm tỷ lệ hit trong khi đẩy máy tính về phía phân trang. Khi tự chọn ngân sách, công cụ giảm xuống từng tập hợp làm việc một và lấy mức lớn nhất vừa với bảy phần tám RAM: K3 yêu cầu mức tối thiểu + 3× — 80,63 GB — và nhận được mức tối thiểu + 1× trên chiếc laptop này, ngân sách 46 GB và bộ nhớ đệm 17,56 GB. Đó là đỉnh của đường cong ở trên, đạt được mà không cần cờ (flag). Một máy 128 GB vẫn nhận được toàn bộ 3×.

Một phiên bản trước đó đã lấy mọi byte cho đến mức giới hạn, tạo ra bộ nhớ đệm 27 GB trên máy này — giữa hai ngân sách được đo ở mức 0,11 và 0,04 tok/s. Bài học thực sự là bộ nhớ đệm mà bạn không kiểm soát được thì không phải là bộ nhớ đệm, và hệ quả là một công cụ nên ngừng yêu cầu bộ nhớ trước khi hệ điều hành bắt đầu lấy lại nó.

Attention tuyến tính và bộ nhớ đệm KV được hấp thụ

Attention của K3 là một mô hình lai 3:1: Kimi Delta Attention, mang theo một trạng thái tái phát (recurrent state) kích thước cố định thay vì bộ nhớ đệm KV ngày càng tăng, và gated multi-head latent attention. Các lớp MLA lưu vào bộ nhớ đệm latent rộng 512 thay vì các key và value mở rộng theo từng head, với kv_b_proj được hấp thụ vào query và output: