Huấn luyện LLM bằng Swift: Hành trình tối ưu hóa nhân ma trận từ Gflop/s lên Tflop/s

Huấn luyện LLM bằng Swift: Hành trình tối ưu hóa nhân ma trận từ Gflop/s lên Tflop/s

Trong bối cảnh bùng nổ của các Mô hình Ngôn ngữ Lớn (LLM), phần lớn các công việc huấn luyện và suy luận thường được thực hiện bằng Python kết hợp với các thư viện C++ tối ưu hóa ở mức thấp. Tuy nhiên, một câu hỏi thú vị được đặt ra là: Liệu ngôn ngữ Swift, vốn nổi tiếng với hiệu năng cao trên nền tảng Apple, có thể tự mình thực hiện nhiệm vụ này mà không cần dựa vào các framework có sẵn không?

Bài viết này sẽ đi sâu vào hành trình tối ưu hóa mã nguồn nhân ma trận (matrix multiplication) viết bằng Swift để huấn luyện LLM trên chip Apple Silicon. Mục tiêu là đưa hiệu suất từ mức Gflop/s lên Tflop/s, đồng thời so sánh khả năng của các đơn vị tính toán khác nhau như CPU, SIMD, AMX và GPU.

Bối cảnh và Thách thức

Tác giả đã bắt đầu bằng việc viết lại mã nguồn tham khảo llm.c của Andrej Karpathy (một phiên bản GPT-2 đơn giản viết bằng C thuần) sang ngôn ngữ Swift. Thách thức lớn nhất không phải là thuật toán mạng nơ-ron, mà là hiệu năng của phép nhân ma trận – hoạt động chiếm phần lớn thời gian tính toán trong quá trình huấn luyện.

Ban đầu, phiên bản Swift "cơ bản" chạy cực kỳ chậm, chỉ đạt 0,054 tokens/giây, kém xa rất nhiều so với phiên bản C gốc. Đây là điểm khởi đầu cho một quá trình tối ưu hóa dài hơi và thú vị.

So sánh hiệu suất các phương pháp tối ưu hóa

Tối ưu hóa trên CPU: Từ chậm đến nhanh hơn

Bước đầu tiên là loại bỏ các chi phí phụ (overhead) không cần thiết của Swift. Việc sử dụng Array mặc định trong Swift gây ra chi phí kiểm tra tham chiếu và copy-on-write. Bằng cách chuyển sang sử dụng UnsafeMutableBufferPointer và tính toán trực tiếp trên bộ nhớ, tốc độ đã được cải thiện đáng kể.

Tiếp theo, tác giả tận dụng các lệnh SIMD (Single Instruction, Multiple Data) thông qua thư viện Numerics của Apple. Thay vì thực hiện phép tính từng số một, SIMD cho phép xử lý nhiều số liệu cùng lúc, giúp tăng tốc độ đáng kể. Kết hợp với việc tái cấu trúc vòng lặp để tối ưu hóa pipeline của CPU, mã nguồn Swift đã bắt đầu tiệm cận với tốc độ của C.

Cuối cùng, việc sử dụng đa luồng (multithreading) thông qua DispatchQueue.concurrentPerform đã giúp tận dụng toàn bộ các nhân vật lý của CPU. Phiên bản Swift đa luồng đã vượt trội hơn so với llm.c đơn luồng, đạt tốc độ gấp 5,5 lần.

Vũ khí bí mật của Apple Silicon: AMX

Apple Silicon (như các chip M-series) sở hữu một đơn vị tính toán đặc biệt gọi là AMX (Apple Matrix Extensions), được thiết kế riêng cho các tác vụ nhân ma trận lớn. AMX hoạt động dựa trên kỹ thuật "tiling" – chia nhỏ ma trận thành các ô (tiles) nhỏ để phù hợp với thanh ghi của bộ xử lý.

Việc viết mã cho AMX phức tạp hơn nhiều, đòi hỏi phải thao tác trực tiếp trên bộ nhớ và sử dụng các lệnh assembly đặc thù. Tuy nhiên, kết quả nhận được rất xứng đáng: phiên bản sử dụng AMX nhanh hơn gấp 1,6 lần so với phiên bản đa luồng chỉ dùng CPU, đạt gần 1,7 Tflop/s trong các bài kiểm tra.

Kéo GPU vào cuộc với Metal

Khi CPU đã đạt đến giới hạn, bước tiếp theo là chuyển sang GPU thông qua API Metal của Apple. GPU có khả năng xử lý song song hàng nghìn luồng, lý tưởng cho các tác vụ đồ họa và AI.

Ban đầu, một kernel Metal cơ bản không mang lại cải thiện vượt trội so với AMX do chi phí truyền dữ liệu. Tuy nhiên, sau khi áp dụng kỹ thuật "threadgroup tiling" (tương tự như AMX nhưng trên GPU) để tối ưu hóa việc truy cập bộ nhớ, hiệu suất đã bùng nổ.

Phiên bản Metal tối ưu hóa cuối cùng đã đạt mức 1,1 Tflop/s thực tế, nhanh hơn gấp 382 lần so với mã Swift ban đầu và nhanh hơn gấp 22 lần so với mã C gốc (llm.c).

Kết quả và Bài học

Dưới đây là bảng so sánh tóm tắt hiệu suất của các phương pháp:

• llm.c (C gốc): 100% (mốc tham chiếu)
• Swift đa luồng: 558,5%
• Swift với AMX: 958,8%
• Metal với Tiling: 3057,7%

Quá trình tối ưu hóa này đã chỉ ra rằng Swift hoàn toàn có thể đạt được hiệu năng tương đương hoặc vượt qua C nếu người lập trình viên hiểu rõ về phần cứng và sẵn sàng viết mã ở mức thấp (low-level). Tuy nhiên, đây cũng là một lời nhắc nhở rằng việc viết mã tối ưu hóa thủ công là vô cùng phức tạp và dễ sai sót.

"Rõ ràng, chúng ta vẫn còn chặng đường dài phía trước. Apple có rất nhiều framework khác để huấn luyện mạng nơ-ron và bản triển khai cuối cùng trong loạt bài này sẽ hoàn thành nhiều hơn 20 vòng lặp huấn luyện trước khi hầu hết các triển khai trong bài này kịp hoàn thành vòng lặp đầu tiên."

Đối với các ứng dụng thực tế (production), người dùng nên sử dụng các thư viện đã được tối ưu hóa sẵn của Apple như Accelerate (BLAS), BNNS, CoreML hoặc MPSGraph thay vì tự viết lại nhân ma trận từ đầu. Tuy nhiên, hành trình này mang lại cái nhìn sâu sắc về cách thức hoạt động của phần cứng Apple Silicon và khả năng tiềm ẩn của ngôn ngữ Swift trong lĩnh vực AI.