使用 Apple 的 MLX 框架在本地部署 LLM

2023年12月，Apple发布了新的MLX深度学习框架，这是一个用于在Apple芯片上进行机器学习的阵列框架，由其机器学习研究团队开发。本教程将探讨该框架，并演示如何在MacBookPro（MBP）上本地部署Mistral-7B模型。我们将设置一个本地聊天界面来与已部署的模型进行交互，并根据每秒生成的令牌测试其推理性能。此外，我们将深入研究MLXAPI，以了解用于改变模型行为和影响生成文本的可用杠杆。

和往常一样，相关代码可以在公开的GitHub存储库中找到：。

Apple的新机器学习框架MLX具有在Apple芯片上进行机器学习的统一内存架构，与其他深度学习框架相比具有显着优势。与PyTorch和Jax等传统框架不同，这些框架需要在CPU和GPU之间进行成本高昂的数据复制，而MLX将数据维护在共享内存中，两者都可以访问。这种设计消除了数据传输的开销，促进了更快的执行，特别是对于机器学习中常见的大型数据集。对于Apple设备上的复杂ML任务，MLX的共享内存架构可以显著加快速度。此功能使MLX与希望在设备（例如iPhone）上运行模型的开发人员高度相关。

在部署模型之前，需要进行一些设置。首先，必须安装特定的库。在进行安装之前，请记得创建一个虚拟环境：

pipinstallmlx-lm

这个库允许我们在本地部署一个大型语言模型（LLM），并只需五行代码就能运行它：

frommlx_lmimportload,generatemodel,tokenizer=load("mistralai/")prompt="""s[INST]Helloworld![/INST]"""response=generate(model,tokenizer,prompt=prompt)print(response)

第一次执行此脚本时，它将下载模型，可能需要一些时间。在后续运行中，模型将从本地缓存加载，大大加快了这个过程。一旦模型被下载，我们就会收到响应：

根据下面的图片，我们深入了解下背后的原理：

一旦调用load()方法，它会检查模型是否在本地机器上可用。如果不存在，该方法会从HuggingFaceModelHub下载模型。加载权重后，它们将被转换为MLX格式。此外，如果在配置中指定，模型将进行量化。

请注意，模型的配置未指定量化，这意味着我们将加载模型的完整权重。但是，可以使用以下命令将模型的原始权重转换为MLX格式（）并同时进行量化：

_to_torch_model-q

HuggingFace上还有一个活跃的MLX社区，该社区已经将多个模型转换为MLX格式：

HuggingFaceMLX社区：

加载格式的模型示例：

代码的另一个重要功能是提示。大型语言模型（LLM）越来越多地用于聊天机器人应用程序，以实现更自然的对话。与简单地生成连续文本不同，聊天机器人需要理解对话上下文，其中包括“用户”和“助手”等不同角色之间的交流。模型输入由一系列消息组成，而不是单个段落。每个模型通常都有自己的聊天模板，我们可以从他们的模型页面看到Mistral的模型模板是如何显示的：

Mistral的模型聊天模板：

在上面的示例代码中，我们通过将特殊标记插入到字符串中来手动应用此模板。稍后，我们将探索一种更有效的方法来实现这一目标。

我们已经下载了Mistral-7B模型的指导版本，我们应该将其用作聊天机器人模型来最大限度地发挥其潜力。幸运的是，这可以通过Streamlit轻松完成，这是一个开源框架，允许我们只需几行代码即可构建聊天机器人界面。

pipinstallstreamlit

我们不会涵盖整个应用程序（完整的代码可以在GitHub上找到），但其中一些值得注意的亮点包括：

#缓存模型加载，以避免每次重新加载@_resourcedefload_model():model,tokenizer=load("mistralai/")returnmodel,tokenizermodel,tokenizer=load_model()

就像在我们的第一个示例中一样，我们利用mlx_lm库来加载我们的模型。为了防止每次向聊天机器人发送新消息时重新加载模型，我们利用Streamlit的缓存管理，使用_resource修饰符。

类似于缓存模型，我们还需要跟踪整个对话。聊天机器人使用一种简单的方法记住先前对话的部分：每次接收到新的用户输入时，整个先前的对话以一个长提示的形式发送到模型。为了跟踪对话，我们初始化一个有状态的变量，如下所示：

if'messages'_state:_state['messages']=[]

然后，我们将添加用户和助手的消息，如下所示：

_state['messages'].app({"role":"user","content":prompt})_state['messages'].app({"role":"assistant","content":response})

Mistral-7B-Instruct模型需要一种非常具体的格式来反映之前的对话。我们可以使用tokenizer的apply_chat_template()方法将对话历史转换为符合此规范的一个长提示：

formatted_conversation=_chat_template(_state['messages'],tokenize=False)

通过运行以下命令来测试，看看它是如何工作的，并确保它正常运行：

fromtransformersimportAutoTokenizertokenizer=_pretrained("mistralai/")chat=[{'role':'user','content':'Hello,howareyou?'},{'role':'assistant','content':'I\\'mgood,?'},{'role':'user','content':'CanyoutellmetheweatherforecastforLondon?'},{'role':'assistant','content':'Sure,theforecastforLondonissunnywithahighof23degrees.'},{'role':'user','content':'Thankyou,that\\'sveryhelpful!'}]print(_chat_template(chat,tokenize=False))

我们得到的结果是：

s[INST]Hello,howareyou?[/INST]I'mgood,?/s[INST]CanyoutellmetheweatherforecastforLondon?[/INST]Sure,theforecastforLondonissunnywithahighof23degrees./s[INST]Thankyou,that'sveryhelpful![/INST]

结果显示聊天模板方法正常运作，接下来可以用它来格式化我们的对话历史。

我们还将测量接收响应所需的时间，并计算模型每秒生成的令牌数（TPS）。从技术上讲，这个度量并不完全准确，因为它没有考虑模型处理提示所需的时间（随着对话的进行而增加）。然而，从用户体验的角度来看，这种区别通常不太明显。因此，我将在此使用此方法进行测量。

计算令牌数示例：

通过键入以下内容来启动聊天机器人：

启动后将自动在http://localhost:8501/上打开一个浏览器选项卡（如果没有，请在浏览器中键入此地址）。

现在我们可以开始与模型进行交互：

一开始，每秒生成令牌约为15个，后来随着对话时间的延长，它下降到12-13个。这是因为模型需要更多时间来解析整个上一个对话，这在本例中我们测量TPS的方式中有所体现。

当然，一个关键问题是这个聊天机器人和本地部署的LLM消耗了多少内存？在MBPM1Max上运行，具有64GB内存的初始脚本使用了14GB内存（通过top-omem查看）。

然而，一旦对话变得更长，内存消耗迅速上升到46GB，并保持在那里。我怀疑这是程序被允许使用的内存的最大值：

有时我们希望我们的聊天机器人表现不同，例如，生成更长或更短的答案。在上面的例子中，当我以“你好”开始对话时，响应对我来说太冗长了。请注意，这与设置max_tokens参数是分开的。如果模型以更长的消息响应，并且max_tokens参数设置得很低，例如在50个新令牌的情况下，模型的响应将被截断。

通常，还有其他模型推理参数可用于引导模型的行为，例如重复惩罚或长度惩罚。但是，MLXAPI还不支持这些参数。要了解支持哪些参数，我们可以深入研究MLX代码并检查generate()API：

_argument("--prompt",default=DEFAULT_PROMPT,help="Messagetobeprocessedbythemodel")_argument("--max-tokens","-m",type=int,default=DEFAULT_MAX_TOKENS,help="Maximumnumberoftokenstogenerate",)_argument("--temp",type=float,default=DEFAULT_TEMP,help="Samplingtemperature")_argument("--seed",type=int,default=DEFAULT_SEED,help="PRNGseed")

到目前为止，它只支持'max_tokens'、'temp'和'seed'。但是，我预计随着时间的推移，MLX研究团队或社区将添加其他推理参数。

我们将功能非常强大的Mistral-7B模型本地部署到具有64GB内存的MacbookPro上。我们还创建了一个聊天机器人界面来与模型进行交互。这仅仅是一个开始，它无疑预示着未来，本地LLM部署将变得更加普遍和重要。

有很多方法可以继续和改进本教程，我将留给读者。一些想法包括：

探索其他LLM和/或使用量化版本

改进了添加文件上传的界面，或用于更改推理参数（例如温度、最大令牌数等）的字段

更密切地测量和监控内存消耗，并测试限制

高级：克隆MLX框架并实现其他推理参数

MLX深度学习框架：

项目地址：

Mistral-7B：

HuggingFaceMLX社区：

加载格式的模型示例：

Mistral的模型聊天模板：

计算令牌数示例：