器學(xué)習(xí):MLX 新手實戰(zhàn)指南)
如何在 Apple Silicon 上做機(jī)器學(xué)習(xí)MLX 新手實戰(zhàn)指南【免費下載鏈接】mlxMLX: An array framework for Apple silicon項目地址: https://gitcode.com/GitHub_Trending/ml/mlxMLX 是 Apple 機(jī)器學(xué)習(xí)研究團(tuán)隊推出的、面向 Apple silicon 的數(shù)組框架array framework。它的 Python API 與 NumPy 高度相似內(nèi)置延遲計算、統(tǒng)一內(nèi)存模型和可組合的函數(shù)變換自動微分、向量化可以在 Mac 上直接訓(xùn)練和推理模型。如果你熟悉 PyTorch 的寫法、想把代碼搬到 Apple Silicon 設(shè)備或者想在 MacBook 上跑本地模型這個項目值得花一小時上手。安裝命令怎么選三種場景根據(jù)你的運行環(huán)境三選一即可pip install mlx # Apple Silicon Mac需 macOS 14.0、原生 Python 3.10 pip install mlx[cuda12] # Linux NVIDIA GPUSM 7.5 及以上架構(gòu) pip install mlx[cpu] # Linux 純 CPU 版本在 Mac 上如果pip找不到匹配版本多半是 Python 不是原生的 arm 版本運行python -c import platform; print(platform.processor())輸出應(yīng)為arm如果是i386請換成原生 Python比如通過 Conda 安裝。5 分鐘跑通第一個數(shù)組先理解延遲計算下面這段代碼演示了 MLX 最核心的行為——操作先記錄、后執(zhí)行import mlx.core as mx a mx.array([1, 2, 3, 4]) b mx.array([1.0, 2.0, 3.0, 4.0]) c a b # 此時 c 只是計算圖上的一個配方尚未真正計算 mx.eval(c) # 顯式觸發(fā)計算 print(c) # array([2, 4, 6, 8], dtypefloat32)你會發(fā)現(xiàn)c a b之后并沒有立即產(chǎn)生數(shù)值MLX 只是在記錄一張動態(tài)計算圖只有在mx.eval、打印數(shù)組、調(diào)用.item()或轉(zhuǎn)成 numpy 時才會真正落盤計算。這個設(shè)計的直接好處是圖可以在執(zhí)行前被變換求導(dǎo)、向量化而且最終沒被使用的輸出不會白算。更多細(xì)節(jié)可以看倉庫里的說明docs/src/usage/lazy_evaluation.rst統(tǒng)一內(nèi)存在操作時指定設(shè)備而不是搬移數(shù)組Apple Silicon 的 CPU 和 GPU 直接共享同一個內(nèi)存池MLX 完全利用了這一點創(chuàng)建數(shù)組時不指定位置而是在調(diào)用操作時通過stream參數(shù)指定設(shè)備a mx.random.normal((100,)) b mx.random.normal((100,)) mx.add(a, b, streammx.cpu) # 讓 CPU 算 mx.add(a, b, streammx.gpu) # 讓 GPU 算與上一條并行執(zhí)行這兩個操作互不依賴MLX 會并行調(diào)度如果第二條依賴第一條的結(jié)果調(diào)度器會自動在兩個 stream 之間插入依賴不需要你手動同步。官方文檔里有個實測例子M1 Max矩陣乘法放 GPU、幾百次小的exp放 CPU總耗時約 1.4ms比全部堆在 GPU 上2.8ms快一倍??山M合的函數(shù)變換求導(dǎo)和向量化各一行MLX 提供了mx.grad、mx.vmap、mx.jvp、mx.vjp、mx.value_and_grad等變換并且可以任意嵌套組合。比如求二階導(dǎo)就是mx.grad(mx.grad(mx.sin))(x)mx.grad(mx.vmap(mx.grad(fn)))這樣的組合也是合法的。兩個最常用的模式訓(xùn)練時用mx.value_and_grad(fn)一次拿到函數(shù)值和梯度避免重復(fù)計算前向圖批量推理時給模型外層套一個mx.vmap(model)自動沿新維度向量化。保存與加載模型4 種序列化格式mx.load會根據(jù)文件擴(kuò)展名自動識別格式支持.npy單數(shù)組、.npz多數(shù)組、.safetensors和.gguf后兩者也常用于交換模型權(quán)重a, b mx.array([1.0]), mx.array([2.0]) mx.savez(arrays, a, bb) # 存為 arrays.npz mx.save_safetensors(arrays, {a: a, b: b}) # 存為 safetensors state mx.load(arrays.npz) # 返回 {名稱: 數(shù)組} 的字典動手示例跑倉庫自帶的訓(xùn)練腳本倉庫里有一個自包含的線性回歸訓(xùn)練腳本examples/python/linear_regression.py直接用python examples/python/linear_regression.py運行即可。它的流程值得逐行讀一遍生成隨機(jī)設(shè)計矩陣X和帶噪聲的標(biāo)簽y→ 定義損失0.5 * mx.mean(mx.square(X w - y))→ 用mx.grad(loss_fn)拿到梯度函數(shù) → 循環(huán) 10000 次做w w - lr * grad并mx.eval(w)固化結(jié)果 → 最后打印損失、與真實參數(shù)的距離和吞吐it/s。你可以把它當(dāng)模板替換成自己的數(shù)據(jù)就得到一個最小訓(xùn)練閉環(huán)。排坑與性能eval、內(nèi)存、GPU 分析忘記 eval 會導(dǎo)致計時失真訓(xùn)練循環(huán)里如果不 eval計算圖會一直累積time.perf_counter()測到的只是建圖時間。每步或每幾步eval 一次參數(shù)既是正確性也是性能問題。峰值內(nèi)存可以省一半得益于延遲計算先以 float32 實例化大模型、再換成 float16 權(quán)重中間不會真的算出 float32 那份數(shù)據(jù)峰值內(nèi)存接近減半。用 Metal 調(diào)試器分析 GPU以-DMLX_METAL_DEBUGON構(gòu)建后調(diào)用mx.metal.start_capture()/mx.metal.stop_capture()需配合MTL_CAPTURE_ENABLED1環(huán)境變量就能捕獲 GPU 追蹤文件在 Xcode 里回放。指南在 docs/src/dev/metal_debugger.rst。多設(shè)備張量并行要把大模型分片到多個設(shè)備時mlx.nn提供了AllToShardedLinear和ShardedToAllLinear等現(xiàn)成分層兩者前后銜接可以省掉中間通信。完整示例見 examples/python/distributed_tensor_parallel.py。去哪繼續(xù)學(xué)快速上手docs/src/usage/quick_start.rst更多示例examples/python/含邏輯回歸、分布式數(shù)據(jù)/張量并行、量化推理建議你先pip install mlx然后原樣跑一遍examples/python/linear_regression.py把延遲計算和mx.eval的節(jié)奏跑熟再動手替換成自己的數(shù)據(jù)。等你熟悉了這條主線函數(shù)變換和張量并行這兩塊會順理成章地接上?!久赓M下載鏈接】mlxMLX: An array framework for Apple silicon項目地址: https://gitcode.com/GitHub_Trending/ml/mlx創(chuàng)作聲明:本文部分內(nèi)容由AI輔助生成(AIGC),僅供參考