使用苹果的 MLX 框架本地部署 LLMs
原文：
towardsdatascience.com/deploying-llms-locally-with-apples-mlx-framework-2b3862049a93
这是什么关于的？
2023 年 12 月，苹果发布了他们新的MLX 深度学习框架，这是一个适用于苹果硅上机器学习的数组框架，由他们的机器学习研究团队开发。本教程将探讨该框架，并演示在 MacBook Pro（MBP）上本地部署 Mistral-7B 模型。我们将设置一个本地聊天界面来与部署的模型交互，并测试其每秒生成的令牌数量以评估其推理性能。此外，我们还将深入研究 MLX API，以了解可调整模型行为和影响生成文本的可用工具。
与往常一样，代码可在公共 GitHub 仓库中找到：github.com/marshmellow77/mlx-deep-dive
为什么这很重要？
苹果的新机器学习框架 MLX，在苹果硅上机器学习的统一内存架构方面，相较于其他深度学习框架具有显著优势。与需要昂贵的 CPU 和 GPU 之间数据复制的传统框架（如 PyTorch 和 Jax）不同，MLX 维护的数据是可由两者访问的共享内存。这种设计消除了数据传输的开销，促进了更快的执行，尤其是在机器学习中常见的庞大数据集。对于在苹果设备上进行的复杂机器学习任务，MLX 的共享内存架构可能导致显著的加速。这一特性使得 MLX 对于希望在 iPhone 等设备上运行模型的开发者来说非常相关。
苹果在硅片设计方面的专业知识，MLX 预示着未来设备上 AI 应用可能集成到其芯片中的令人兴奋的功能。MLX 在加速和简化苹果平台上机器学习任务方面的潜力，使得它成为开发者应关注的框架。
初始设置
在部署模型之前，需要进行一些设置。首先，安装某些库是至关重要的。在继续安装之前，请记得创建一个虚拟环境：
pip install mlx-lm
这个库允许我们在本地部署大型语言模型（LLM），只需五行代码即可运行：
from mlx_lm import load, generate
model, tokenizer = load("mistralai/Mistral-7B-Instruct-v0.2")
prompt = """<s>[INST] Hello world! [/INST]"""
response = generate(model, tokenizer, prompt=prompt)
print(response)
第一次运行此脚本时，它将下载模型，这可能需要一些时间。在后续运行中，模型将从本地缓存中加载，显著加快处理速度。一旦模型下载完成，我们将收到以下响应：
图像由作者提供
这相当令人印象深刻，但幕后有许多事情在进行中。让我们揭开面纱，更好地了解实际上发生了什么：
图片由作者提供
一旦调用load()方法，它会检查模型是否在本地机器上可用。如果不可用，该方法将从 Hugging Face 模型库中下载模型，在我们的例子中是从huggingface.co/mistralai/Mistral-7B-Instruct-v0.2/tree/main下载。加载权重后，它们将被转换为 MLX 格式。此外，如果配置中指定了量化，模型将被量化。
量化及快速转换的模型
注意：Mistral-7B-Instruct-v0.2 模型的配置没有指定量化，这意味着我们将以完整权重加载模型。然而，可以将模型的原始权重转换为 MLX 格式（weights.npz）并同时进行量化，使用以下方法：
python convert.py --torch-path <path_to_torch_model> -q
在 Hugging Face 上还有一个活跃的 MLX 社区，他们已经将几个模型转换为 MLX 格式：
注意：截至目前，我们无法使用上述代码加载这些转换后的模型，因为
load()方法特别要求以.safetensors格式加载模型权重。要加载weights.npz格式的模型，我们可以从这个示例中获取灵感。
Mistral 指令模型
代码的另一个有趣方面是提示。大型语言模型（LLMs）越来越多地用于聊天机器人应用，以实现更自然的对话。与简单地生成连续文本不同，聊天机器人需要理解对话上下文，这包括“用户”和“助手”等不同角色之间的交流。模型输入由一系列消息组成，而不是单一段落。每个模型通常都有自己的聊天模板，我们可以从他们的模型页面看到 Mistral 的模型模板：
huggingface.co/mistralai/Mistral-7B-Instruct-v0.2
在上面的示例代码中，我们通过将特殊标记插入字符串来手动应用此模板。稍后，我们将探讨一种更有效的方法来实现这一点。
聊天机器人界面
由于我们已经下载了 Mistral-7B 指令版本模型，我们应该通过将其用作聊天机器人模型来最大化其潜力。幸运的是，这可以通过Streamlit轻松实现，这是一个开源框架，只需几行代码就可以构建聊天机器人界面。
首先，让我们安装 Streamlit：
pip install streamlit
我们不会涵盖整个应用（完整代码可在 GitHub 上找到），但一些值得注意的亮点包括：
加载模型
# Cache the model loading to avoid reloading every time
@st.cache_resource
def load_model():
model, tokenizer = load("mistralai/Mistral-7B-Instruct-v0.2")
return model, tokenizer
model, tokenizer = load_model()
就像我们的第一个例子一样，我们使用mlx_lm库来加载我们的模型。为了避免每次向聊天机器人发送新消息时都加载模型，我们利用 Streamlit 的缓存管理功能，通过使用st.cache_resource装饰器。
跟踪对话
与缓存模型类似，我们还需要跟踪整个对话。聊天机器人通过简单的方法记住对话的先前部分：每次接收到新的用户输入时，将整个先前的对话作为一个长提示发送到模型中。为了跟踪对话，我们初始化一个如下所示的状态变量：
if 'messages' not in st.session_state:
st.session_state['messages'] = []
我们将按如下方式添加用户和助手消息：
st.session_state['messages'].append({"role": "user", "content": prompt})
...
st.session_state['messages'].append({"role": "assistant", "content": response})
将对话历史转换为所需格式
如前所述，Mistral-7B-Instruct 模型需要非常特定的格式来反映先前的对话。我们可以使用分词器的apply_chat_template()方法将对话历史转换为符合此规范的单一长提示：
formatted_conversation = tokenizer.apply_chat_template(st.session_state['messages'], tokenize=False)
让我们看看这是如何工作的，并确保它正常工作。我们可以通过运行以下内容来测试它：
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
chat = [
{'role': 'user', 'content': 'Hello, how are you?'},
{'role': 'assistant', 'content': 'I'm good, thank you for asking. How can I assist you today?'},
{'role': 'user', 'content': 'Can you tell me the weather forecast for London?'},
{'role': 'assistant', 'content': 'Sure, the forecast for London is sunny with a high of 23 degrees.'},
{'role': 'user', 'content': 'Thank you, that's very helpful!'}
]
print(tokenizer.apply_chat_template(chat, tokenize=False))
我们得到的结果如下：
<s>[INST] Hello, how are you? [/INST]I'm good, thank you for asking. How can I assist you today?</s>[INST] Can you tell me the weather forecast for London? [/INST]Sure, the forecast for London is sunny with a high of 23 degrees.</s>[INST] Thank you, that's very helpful! [/INST]
看起来聊天模板方法工作得很好，我们可以用它来格式化我们的对话历史。
时间测量
我们还将测量接收响应所需的时间，并计算模型每秒生成的标记数（TPS）。从技术上讲，这个度量并不完全准确，因为它没有考虑到模型处理提示所需的时间（随着对话的进行而增加）。然而，从用户体验的角度来看，这种区别通常不太明显。因此，我将使用这种方法进行测量。
如果你感兴趣，想专门测量生成阶段每秒生成的标记数，即模型处理提示之后，你可能想查看这个示例。
测试聊天机器人
一旦代码开发令我们满意，我们可以通过输入以下内容来启动聊天机器人：
streamlit run mlx-chatbot.py
这将自动打开一个浏览器标签页localhost:8501/（如果没有，只需在浏览器中输入此地址）。
我们现在可以开始与模型交互：
图片由作者提供
图片由作者提供
在开始时，每秒生成的标记数大约为 15，随着对话的变长，这个数字会下降到 12-13。这是因为模型需要更多的时间来解析整个先前的对话，这在我们在本例中测量 TPS 的方式中得到了反映。
内存消耗
其中一个关键问题是当然这个聊天机器人和本地部署的 LLM 消耗了多少内存？当在初始拥有 64GB 内存的 MBP M1 Max 上运行时，脚本使用了 14GB 的内存（通过top -o mem）
图片由作者提供
然而，一旦对话变长，内存消耗迅速上升到 46GB，并保持在那里。我怀疑这是程序允许使用的最大内存量：
图片由作者提供
探索 MLX API
有时我们希望我们的聊天机器人有不同的行为，例如，创建更长或更短的回答。在上面的例子中，我用“你好”开始对话时，回复对我来说太冗长了。注意，这不同于设置 max_tokens 参数。如果模型以较长的消息回复，并且 max_tokens 参数设置得较低，例如 50 个新标记，模型的回复将被截断。
到目前为止，它只支持max_tokens、temp和seed。然而，我预计随着时间的推移，MLX 研究团队或社区将添加其他推理参数。
parser.add_argument(
"--prompt", default=DEFAULT_PROMPT, help="Message to be processed by the model"
)
parser.add_argument(
"--max-tokens",
"-m",
type=int,
default=DEFAULT_MAX_TOKENS,
help="Maximum number of tokens to generate",
)
parser.add_argument(
"--temp", type=float, default=DEFAULT_TEMP, help="Sampling temperature"
)
parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="PRNG seed")
所以到目前为止，它只支持max_tokens、temp和seed。然而，我预计随着时间的推移，MLX 研究团队或社区将添加其他推理参数。
结论
在本教程中，我们已经在拥有 64GB 内存的 Macbook Pro 上本地部署了非常强大的 Mistral-7B 模型。我们还创建了一个聊天机器人界面来与模型交互。这仅仅是开始，无疑指向了一个未来，其中本地 LLM 部署将变得更加普遍和重要。
有许多方法可以继续和改进这个教程，我将留给读者去探索。一些想法包括：
-
探索其他 LLMs 和/或使用量化版本
-
高级：克隆 MLX 框架并实现其他推理参数
-
更密切地测量和监控内存消耗，并测试极限
-
高级：克隆 MLX 框架并实现其他推理参数
Heiko Hotz
👋 关注我的Medium和LinkedIn，了解更多关于生成式 AI、机器学习和自然语言处理的内容。
👥 如果你在伦敦，加入我们的NLP London Meetups。
图片由作者提供
中国智能体开发者社区，聚焦智能体与大模型开发，提供前沿资讯、实用工具链、开源项目及行业案例。通过技术沙龙、开发者大赛等活动，促进经验交流与协作，助力开发者快速构建创新智能应用。
更多推荐
4
0- 0
已为社区贡献111条内容
所有评论(0)