如何在 Apple Silicon 上做機(jī)器學(xué)習(xí)MLX 新手實(shí)戰(zhàn)指南【免費(fèi)下載鏈接】mlxMLX: An array framework for Apple silicon項(xiàng)目地址: https://gitcode.com/GitHub_Trending/ml/mlxMLX 是 Apple 機(jī)器學(xué)習(xí)研究團(tuán)隊(duì)推出的、面向 Apple silicon 的數(shù)組框架array framework。它的 Python API 與 NumPy 高度相似內(nèi)置延遲計(jì)算、統(tǒng)一內(nèi)存模型和可組合的函數(shù)變換自動(dòng)微分、向量化可以在 Mac 上直接訓(xùn)練和推理模型。如果你熟悉 PyTorch 的寫法、想把代碼搬到 Apple Silicon 設(shè)備或者想在 MacBook 上跑本地模型這個(gè)項(xiàng)目值得花一小時(shí)上手。安裝命令怎么選三種場(chǎng)景根據(jù)你的運(yùn)行環(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 版本運(yùn)行python -c import platform; print(platform.processor())輸出應(yīng)為arm如果是i386請(qǐng)換成原生 Python比如通過 Conda 安裝。5 分鐘跑通第一個(gè)數(shù)組先理解延遲計(jì)算下面這段代碼演示了 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 # 此時(shí) c 只是計(jì)算圖上的一個(gè)配方尚未真正計(jì)算 mx.eval(c) # 顯式觸發(fā)計(jì)算 print(c) # array([2, 4, 6, 8], dtypefloat32)你會(huì)發(fā)現(xiàn)c a b之后并沒有立即產(chǎn)生數(shù)值MLX 只是在記錄一張動(dòng)態(tài)計(jì)算圖只有在mx.eval、打印數(shù)組、調(diào)用.item()或轉(zhuǎn)成 numpy 時(shí)才會(huì)真正落盤計(jì)算。這個(gè)設(shè)計(jì)的直接好處是圖可以在執(zhí)行前被變換求導(dǎo)、向量化而且最終沒被使用的輸出不會(huì)白算。更多細(xì)節(jié)可以看倉(cāng)庫里的說明docs/src/usage/lazy_evaluation.rst統(tǒng)一內(nèi)存在操作時(shí)指定設(shè)備而不是搬移數(shù)組Apple Silicon 的 CPU 和 GPU 直接共享同一個(gè)內(nèi)存池MLX 完全利用了這一點(diǎn)創(chuàng)建數(shù)組時(shí)不指定位置而是在調(diào)用操作時(shí)通過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í)行這兩個(gè)操作互不依賴MLX 會(huì)并行調(diào)度如果第二條依賴第一條的結(jié)果調(diào)度器會(huì)自動(dòng)在兩個(gè) stream 之間插入依賴不需要你手動(dòng)同步。官方文檔里有個(gè)實(shí)測(cè)例子M1 Max矩陣乘法放 GPU、幾百次小的exp放 CPU總耗時(shí)約 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)))這樣的組合也是合法的。兩個(gè)最常用的模式訓(xùn)練時(shí)用mx.value_and_grad(fn)一次拿到函數(shù)值和梯度避免重復(fù)計(jì)算前向圖批量推理時(shí)給模型外層套一個(gè)mx.vmap(model)自動(dòng)沿新維度向量化。保存與加載模型4 種序列化格式mx.load會(huì)根據(jù)文件擴(kuò)展名自動(dòng)識(shí)別格式支持.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ù)組} 的字典動(dòng)手示例跑倉(cāng)庫自帶的訓(xùn)練腳本倉(cāng)庫里有一個(gè)自包含的線性回歸訓(xùn)練腳本examples/python/linear_regression.py直接用python examples/python/linear_regression.py運(yùn)行即可。它的流程值得逐行讀一遍生成隨機(jī)設(shè)計(jì)矩陣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í)參數(shù)的距離和吞吐it/s。你可以把它當(dāng)模板替換成自己的數(shù)據(jù)就得到一個(gè)最小訓(xùn)練閉環(huán)。排坑與性能eval、內(nèi)存、GPU 分析忘記 eval 會(huì)導(dǎo)致計(jì)時(shí)失真訓(xùn)練循環(huán)里如果不 eval計(jì)算圖會(huì)一直累積time.perf_counter()測(cè)到的只是建圖時(shí)間。每步或每幾步eval 一次參數(shù)既是正確性也是性能問題。峰值內(nèi)存可以省一半得益于延遲計(jì)算先以 float32 實(shí)例化大模型、再換成 float16 權(quán)重中間不會(huì)真的算出 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è)備張量并行要把大模型分片到多個(gè)設(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把延遲計(jì)算和mx.eval的節(jié)奏跑熟再動(dòng)手替換成自己的數(shù)據(jù)。等你熟悉了這條主線函數(shù)變換和張量并行這兩塊會(huì)順理成章地接上?！久赓M(fèi)下載鏈接】mlxMLX: An array framework for Apple silicon項(xiàng)目地址: https://gitcode.com/GitHub_Trending/ml/mlx創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考