Đào tạo một LLM bằng Swift, Phần 2: Các framework tích hợp của macOS

Trong bài viết này, tôi sẽ xem xét một số framework được tích hợp sẵn trong macOS dành cho các thuật toán số. Theo chủ đề của loạt bài này, tôi chủ yếu sẽ xem xét các framework cũng có thể huấn luyện một mô hình ML. Tuy nhiên, có rất nhiều cách tiếp cận khác nhau – Accelerate (BLAS), BNNS, CoreML, MPSGraph – và thách thức thực sự là biết nên sử dụng cái nào – nếu chúng có thể sử dụng được để huấn luyện. Mục lục fchGelu là gì? Accelerate (BLAS) Accelerate (BNNS) MPSGraph CoreML? Tải xuống tất cả mã Kết luận fchGelu là gì? Giống như lần trước, tôi sẽ nói về mã trong các ví dụ nhưng tôi sẽ không thực sự giải thích cơ chế của LLM hoặc thuật ngữ. Tôi thực sự ở đây để nói về các framework của macOS, không phải các mô hình. Tôi biết, điều này khá khó hiểu nhưng đây không phải là khóa học dành cho người mới bắt đầu. Nếu bạn muốn tìm hiểu thêm về thuật ngữ, hãy thử xem video "Let's reproduce GPT-2 (124M)" của Andrej Karpathy, nơi anh ấy giải thích mọi thứ. Accelerate (BLAS) Thư viện macOS đầu tiên dành cho học máy mà tôi muốn nói đến là framework Accelerate. Accelerate thực sự là một framework tổng hợp cho một số thư viện nhỏ hơn. Accelerate là một thư viện quan trọng trong macOS nhưng bạn có thể dành cả sự nghiệp mà không cần sử dụng trực tiếp. Nó đã tồn tại từ Mac OS X 10.2 Jaguar dưới dạng các thư viện riêng biệt vecLib, BLAS và LAPACK, sau đó vào Mac OS X 10.3 Panther, nó được hợp nhất thành Accelerate. Bạn còn nhớ những cái tên mèo lớn không? Thời kỳ vui vẻ. Nói chung, Accelerate chứa các thuật toán có thể tái sử dụng được tối ưu hóa cho các lệnh vector SIMD. Trong Swift, chúng ta không nhất thiết cần framework Accelerate để vector hóa SIMD (trong bài viết trước, tôi đã sử dụng Relaxed.multiplyAdd và tính năng tự động vector hóa của Swift để đạt được khả năng vector hóa SIMD tuyệt vời) nhưng vẫn rất hữu ích khi sử dụng Accelerate khi bạn không muốn nhìn chằm chằm vào mã assembly của riêng mình. Ví dụ, tôi gần đây đã thêm tính năng hiển thị vào một thư viện phân tích cú pháp PDF đơn giản và sử dụng mã sau đây dựa trên vImage của Accelerate để áp dụng mặt nạ hình ảnh: if let matte { guard var matteBuffer = try? vImage_Buffer(width: width, height: height, bitsPerPixel: 8) else { return nil } defer { matteBuffer.free() } vImageBufferFill_ARGB8888(&matteBuffer, [1, matte.r, matte.g, matte.b], vImage_Flags(kvImageNoFlags)) vImageAlphaBlend_ARGB8888(&baseBuffer, &matteBuffer, &baseBuffer, vImage_Flags(kvImageNoFlags)) } kết quả là nhanh hơn khoảng 5 lần so với việc lặp lại pixel thô mà tôi đã sử dụng trước đây (và nhanh hơn khoảng 20 lần trong các bản dựng Debug). Bạn có thể nghĩ rằng bạn có thể làm điều này bằng cách vẽ vào một CGContext, và bạn sẽ đúng nhưng đoán xem điều đó sử dụng gì bên trong? Các hàm tương tự. Tất cả những gì tôi đang làm ở đây là loại bỏ người trung gian và cho phép bản thân kiểm soát trực tiếp hơn một chút. Quay trở lại chủ đề nhân ma trận từ bài viết trước, Accelerate cung cấp cho chúng ta triển khai sgemm của BLAS. BLAS là viết tắt của "Basic Linear Algebra Subprograms" (Chương trình con đại số tuyến tính cơ bản) và sgemm là viết tắt của "Single precision GEneral Matrix Multiplication" (Phép nhân ma trận tổng quát độ chính xác đơn). Việc có người khác tối ưu hóa phép nhân ma trận là tốt nhưng Accelerate BLAS còn mang lại một lợi thế quan trọng khác: nó cho phép chúng ta truy cập vào đơn vị AMX của Apple Silicon mà không cần đến những thủ thuật phức tạp mà tôi đã cần trong bài viết trước. Để minh họa cách thức hoạt động, hãy xem xét nhân tích ma trận cơ bản (ngây thơ) trong Swift từ lần trước: static func matmul_forward(out: inout [Float], inp: [Float], weight: [Float], bias: [Float]?, B: Int, T: Int, C: Int, OC: Int) { for b in 0..<B { for t in 0..<T { let bt = b * T + t for o in 0..<OC { var val = bias?[o] ?? 0 for i in 0..<C { val += inp[bt * C + i] * weight[o * C + i] } out[bt * OC + o] = val } } } } Thực hiện tương tự với BLAS sẽ như sau: static func matmul_forward(out: inout [Float], inp: [Float], weight: [Float], bias: [Float]?, B: Int, T: Int, C: Int, OC: Int) { cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, Int32(OC), Int32(B * T), Int32(C), 1, weight, Int32(C), inp, Int32(C), 0, &out, Int32(OC)) guard var bias else { return } out.withUnsafeMutableBufferPointer { outBuffer in guard let outBase = outBuffer.baseAddress else { return } for bt in 0..<(B * T) { cblas_saxpy(Int32(OC), 1, &bias, 1, outBase.advanced(by: bt * OC), 1) } } } Hàm matmul_forward gần như giống hệt một hàm sgemm điển hình, chỉ được tích hợp thêm bước bias. Bỏ qua tất cả các tối ưu hóa khác đã đề cập trong bài viết trước, chỉ sử dụng cblas_sgemm tại 9 vị trí trong triển khai "Basic Swift" nơi nó được áp dụng, chúng ta có kết quả: Model | Tokens/s | Training iterations/s ---|---|--- Basic Swift | 0.054 | 0.014 AMX | 5.884 | 1.678 Accelerate BLAS | 8.086 | 2.015 Thay đổi này (tại 9 vị trí) đã làm cho triển khai "Basic Swift" nhanh hơn 144 lần. Nó thậm chí còn nhanh hơn 1,2 lần so với nỗ lực sử dụng trực tiếp đơn vị AMX. Triển khai Accelerate BLAS thực sự ấn tượng. Và không chỉ là tốc độ thô. Dưới đây là biểu đồ sử dụng CPU khi chạy: Biểu đồ sử dụng CPU của Instruments cho các mô hình AMX (nửa trái) và BLAS (nửa phải) Framework Accelerate quản lý phép nhân ma trận vừa nhanh vừa có mức sử dụng CPU rất thấp. Điều duy nhất vượt qua nó là triển khai "Tiled Metal" của tôi chạy trên GPU. Model | Tokens/s | Training iterations/s ---|---|--- Accelerate BLAS | 8.086 | 2.015 Tiled Metal | 15.049 | 3.371 Điều đó không thực sự đáng ngạc nhiên. GPU thực tế được tạo ra cho các phép toán ma trận hiệu suất cao và tôi đang chạy trên M3 Max, có khá nhiều lõi GPU. Bạn sẽ phải điên rồ mới nghĩ rằng CPU có thể đánh bại điểm số GPU đó. Accelerate (BNNS) Hãy thử một chút điên rồ. Bạn có nhớ tôi đã nói rằng BLAS là viết tắt của "Basic Linear Algebra Subprograms" (Chương trình con Đại số tuyến tính cơ bản) không? Dựa trên cùng một tiêu chuẩn từ những năm 1970.