[MLX Framework] Apple đã bắt đầu cuộc chơi !

Giới thiệu MLX framework ( dành riêng cho việc huấn luyện mô hình học máy cho Apple Silicon)

Các nhóm học máy (ML) của Apple đã phát hành ra một framework chỉ dành riêng cho các dòng chip M (Apple Silicon) . Bạn có thể xem chi tiết hơn ở đây (link) Ở trong bài viết này, tôi cũng như các bạn thật sự cảm thấy rất tò mò về framework này, Apple nổi tiếng với việc đem tới những trải nghiệm tốt nhất về tối ưu phần cứng nhất trên các thiết bị như Iphone, Ipad, Macbook,.... Thì giờ đây với MLX framework, tôi cũng đang có cảm giác như vậy. MLX được lấy cảm hứng từ các framework đi trước như Pytorch, Jax, ... Tuy nhiên, MLX thêm hỗ trợ cho một mô hình bộ nhớ thống nhất ( a unified memory model), có nghĩa là các mảng sống trong bộ nhớ dùng chung và các thao tác có thể được thực hiện trên bất kỳ loại thiết bị được hỗ trợ nào mà không thực hiện các bản sao dữ liệu.

Ở link này ml-examples ( các ví dụ để developer các thể trải nghiệm về hiệu năng của MLX), như mô tả của repo thì có thể chạy mượt mà LLama2-7b hay một vài mô hình lớn khác như Stable_diffusion, mistral,.... Và đặc biệt trong repo có nói đó chính là việc chúng ta hoàn toàn có thể fine tuning lora trên dòng chip M, nghe khá thú vị tuy nhiên thực tế thì chắc cần phải kiểm chứng thêm. Trong repo còn đưa ra những dẫn chứng rằng MLX có thể chạy nhanh hơn ~15% so với PyTorch trong việc inference Stable_diffusion.

Trước đây mình có sử dụng thử device:mps của PyTorch ( tất nhiên là trải nghiệm rất tốt, mình có thể train cả model GAN🤣) và cũng đã có suy nghĩ về việc liệu Apple có cho ra mắt một framework chuyên dụng cho huấn luyện học máy hay không, vì các tiến trình công nghệ chip của Apple càng ngày càng tăng GPU và neural engine. Và cuối cùng ngày này cũng đã đến, ngày Apple chính thức cùng các ông lớn tham gia vào cuộc chạy đua công nghệ 🤩

Trải nghiệm thử MLX framework với bài toán kinh điển ( huấn luyện tập MNIST)

Đây là dòng lệnh rất quan trọng, nếu như hiện lên là i386 thì bạn phải thực hiện config đến khi output à arm thì lúc đó mới có thể trải nghiệm MLX một cách trọn vẹn.

Để thực hiện config chuyển đổi từ i386 thành arm, bạn có thể tham khảo ở link này

Huấn luyện tập MNIST trên MLX Framework

Việc quan trọng đầu tiên là cài đặt thư viện MLX :

pip install mlx

*Có một vài lưu ý nhỏ, đó là hệ điều hành MacOS >= 13.3

Sau đó tiến hành tải repo này :

git clone https://github.com/ml-explore/mlx-examples.git

Di chuyển tới thư mục mnist:

cd mnist

Và trải nghiệm sức mạnh của Apple Silicon thôi ( ở đây mình đang sử dụng chip M1 phiên bản base chỉ có 8 GPU):

python main.py --gpu

*Mình chỉ cần chạy 20 epochs và cho một kết quả khá là tích cực

Một epoch chạy chưa đến 1s 🤪. Đúng là Apple không bao giờ làm chúng ta thất vọng. Không những thế độ chính xác trong 20 epochs đầu tiên ~97%.

Khá hứng thú đúng không các bạn, vì mình cũng là một iFan🤣. Mình sẽ thử fine tuning với Llama2-7b để xem MLX có thật sự mạnh hay không !

Link tải file pretrained mình để ở đây nha (link)