231206 MLX 框架

飞书用户4220

飞书用户8453

飞书用户9286

飞书用户1153

2024年1月10日修改

前言

在 231026 llama.cpp 极简中我介绍了使用 llama.cpp 来推理模型，今天一个新的框架横空出世。MLX 是专为苹果芯片（即你的 MacBook 或 Mac Studio！）而设计的高效机器学习框架。在短短的一周时间里已经斩获了 8000 Stars。

common.docs_name - LarkCCM_Docs_Menu_Image

MLX 与 PyTorch 的区别

MLX 和其他框架的一个显著区别对统一内存的原生支持。在 MLX 中，数组存在于共享内存中，比如对于下面这段代码。​

代码块

a = mx.random.normal((100,))​
b = mx.random.normal((100,))​
mx.add(a, b, stream=mx.cpu)​
mx.add(a, b, stream=mx.gpu)​

变量 a 和 b 是没有副本的，虽然是在 cpu 和 gpu 这两个不同的设备上运行。而这个改进在 PyTorch 里较难实现。​

MLX 的数组 API 与 Numpy 是一致的，有点像 Jax，这让调试变得十分容易。​

框架总是越变越复杂，MLX 提供了一个很小的起点，对于想要探索新的想法的人，是一个很好的起点，比如对于没有丰富框架经验的研究者来说，通过 MLX 查询 Stack 是很容易的。同样，你也可以轻松的添加 Metal（GPU）内核来尝试新的优化。​

案例

使用苹果芯片微调与推理 Llama-7B 模型

在 https://github.com/ml-explore/mlx-examples/tree/main/lora 这个案例中，作者使用 M2 Ultra LoRA 微调 Llama-7B。使用的数据集是 WikiSQL 在使用 MLX 框架时模型 1s 能看 475 个 Tokens 这是个较为可观的速度。考虑到苹果芯片的功率远远低于 N 卡。

敝人也在第一时间使用了 M2 Max LoRA 微调 Llama-2-7B，此时：

1.
可训练的参数量为 2,090,000，大约为模型的 1%。​

2.
使用 28m24s 分钟学习 426,000 个 Tokens。​

3.
数据来源于 wikisql。​

训练时的内存占用。

231206 MLX 框架​

231206 MLX 框架