LM Studio 0.3.4 发布，支持 Apple MLX | LM Studio 博客

LM Studio 0.3.4 发布，包含 MLX 引擎，可在 Apple Silicon Mac 上极其高效地运行本地大型语言模型 (LLM)。

从此处下载适用于 Apple Silicon 的 LM Studio。继续阅读以了解有关 LM Studio 中 MLX 的更多信息。

在 M3 Max 上，Llama 3.2 1B 的运行速度约为每秒 250 个 token

👾 对设计和构建系统感兴趣吗？ 我们正在招聘。查看此处的职位空缺。

LM Studio 0.3.4 中的 MLX 支持包括：

从 Hugging Face 搜索并下载任何受支持的 MLX LLM（就像您一直使用 GGUF 模型一样）
通过聊天 UI 或从您的代码中使用在 localhost 上运行的类似 OpenAI 的本地服务器来使用 MLX 模型
强制 LLM 以特定的 JSON 格式进行响应（感谢 Outlines）
使用 LLaVA 等视觉模型以及更多模型，并通过聊天或 API 使用它们（感谢 mlx-vlm）
加载并运行多个同步 LLM。您甚至可以混合搭配 llama.cpp 和 MLX 模型！

这篇博文的其余部分深入探讨了 LM Studio 中 MLX 的技术细节。

特别感谢 Awni Hannun (MLX) Rémi Louf (.txt/Outlines) 审阅这篇博文草稿以及 mlx-engine 代码。

什么是 MLX？

... 以及我为什么要关心？

MLX 是 Apple 推出的全新开源 AI/ML 软件堆栈，专门为 Apple Silicon 优化。它利用了 Apple M 芯片中强大的加速硬件。

MLX 由 Apple 的工程师开发，并得到不断壮大的开发者社区的支持，有望成为在 Mac 上运行本地 AI 的极具竞争力的选择。

MLX 核心库使用 C++ 编写，并辅以社区支持的 Python 以及 Swift 前端。

我们很高兴推出 LM Studio 中的 MLX 支持。这篇博文将介绍有关 MLX 的一些技术细节，以及 LM Studio 的 MLX 引擎的特别之处。

LM Studio 的 mlx-engine 是开源的

LM Studio 的 MLX 引擎是一个 Python 模块，使用以下软件包的组合构建

mlx-lm - Apple 的 MLX（Python 前端），由 Awni Hannun 和 Apple 团队开发。
Outlines - 来自 LLM 的结构化生成，由 Rémi Louf 和 .txt 团队开发
mlx-vlm - 适用于 Apple MLX 的视觉 LLM，由 Prince Canuma 开发。

mlx-engine 在 MIT 许可下开源。仓库：https://github.com/lmstudio-ai/mlx-engine。

LM Studio 中的 MLX 使用... Python？！

我们将 MLX 集成到 LM Studio 的旅程始于 Swift。虽然这种方法运行良好，但最终以下设计目标使 Python 成为更好的选择。

设计目标 1： 我们希望与社区一起迭代 MLX 引擎

更多开发者和研究人员熟悉 Python

设计目标 2： 我们希望能够在最新模型和技术发布后立即支持它们

Python 中的 MLX 往往更快地获得对新模型的支持

将 mlx-lm 支持添加到 LM Studio 需要能够在便携式、跨平台的方式中部署和运行 Python 组件。理想情况下，我们还希望能够将这些组件与主 LM Studio 应用程序中已有的 C/C++ 组件完全集成（这最终排除了某些潜在的候选解决方案，例如 conda 环境）。

LM Studio 的初始 Python 运行时支持构建在 python-build-standalone 项目和 Python 虚拟环境之上，使用即将发布的实用程序，该实用程序支持创建一组集成的、可独立下载的 Python 应用程序环境，这些环境共享通用的运行时和框架层（毕竟，如果可以合理避免，没有人愿意下载和安装多个 PyTorch 或 CUDA 副本）。

这个“堆叠式虚拟环境”实用程序使用 CPython 解释器的“站点自定义”功能，以及对虚拟环境内容的一些预发布和安装后调整，以允许这些虚拟环境在机器之间可靠传输，并使用 CPython 的 -m 命令行开关调用包含的应用程序启动模块。

请期待 10 月份晚些时候发布的有关这方面的更详细的技术公告。

小型深入探讨：`mlx-engine` 的一些功能

使用 MLX 运行文本生成模型

python MLX 生态系统的一个关键部分是 mlx_lm。该项目提供了一种使用 CLI 工具或几行 Python 代码轻松运行大型语言模型的方法，例如

from mlx_lm.utils import load, generate_step
import mlx.core as mx

def mlx_stream(prompt: str):
    model, tokenizer = load("/path/to/mlx/model")
    prompt_tokens = mx.array(tokenizer.encode(prompt))
    while True:
        yield generate_step(
            model=model,
            prompt=prompt_tokens
        )

for token in mlx_stream(prompt="Hello world!"):
    print(token, end="", flush=True)

让我们打开 generate_step 的引擎盖，以便更好地了解正在发生的事情

def generate_step(*args, **kwargs):

    # --snip--

    def sample(logits):
        logprobs = logits - mx.logsumexp(logits)
        if temp == 0:
            token = mx.argmax(logits, axis=-1)
        else:
            if top_p > 0 and top_p < 1.0:
                token = top_p_sampling(logits, top_p, temp)
            elif min_p != 0.0:
                token = min_p_sampling(logits, min_p, min_tokens_to_keep, temp)
            else:
                token = categorical_sampling(logits, temp)

        return token, logprobs

    y = prompt
    tokens = None

    def _step(y):
        logits = model(y[None], cache=cache)
        logits = logits[:, -1, :]

        nonlocal tokens
        tokens = mx.concat([tokens, y]) if tokens is not None else y
        for processor in logits_processor:
            logits = processor(tokens, logits)

        y, logprobs = sample(logits)
        return y, logprobs.squeeze(0)

    y, logprobs = _step(y)
    while True:
        next_y, next_logprobs = _step(y)
        yield y.item(), logprobs
        y, logprobs = next_y, next_logprobs

我们可以看到此处发生的重要操作

使用模型的 __call__ 方法评估模型。这将返回一个 logits 数组，其中每个元素对应于模型词汇表中的一个项目。logits 定义了词汇表中项目的概率分布。
使用用户提供的采样参数从 logits 数组中选择（即采样）一个 token。
采样的 token 将返回给调用者。

让我们看看如何向此生成循环添加我们的用户会喜欢的功能。

使用 Outlines 启用结构化生成

让我们向生成器添加一个功能：用户可以请求生成器输出有效的 json。我们可以使用 Outlines，来自 .txt。

Outlines 实现了来自 LLM 的结构化生成（例如，创建 json 输出）。这个包支持 mlx_lm 运行时，我们将利用它。Outlines 通过将用户提供的 json 模式转换为正则表达式来完成其工作。查看这个标题模式。

{
  "type": "object",
  "properties": {
    "title": {
      "type": "string",
      "minLength": 1
    }
  },
  "required": [
    "title"
  ]
}

Outlines 将该模式转换为这个正则表达式字符串

\{[ ]?"title"[ ]?:[ ]?"([^"\\\x00-\x1F\x7F-\x9F]|\\["\\]){1,}"[ ]?\}

这是一个更易于人类阅读（但不太精确）的正则表达式字符串版本：\{"title": ".{1,}"\}

使用此正则表达式字符串，Outlines 的生成循环如下

评估模型。即，处理提示并输出每个 token 的 logits。
对于每个 token，评估采样它是否会违反正则表达式。如果是，我们将其概率设置为零。我们说我们屏蔽了 logits。
使用屏蔽的 logits 采样一个 token。

mlx_lm 的 generate_step 允许我们定义 logits 处理器，因此让我们定义一个处理器来屏蔽 logits，以便输出与正则表达式匹配

from outlines.processors.structured import JSONLogitsProcessor

class OutlinesJSONLogitsProcessor:
    def __init__(self, json_schema, tokenizer):
        self.logits_processor = JSONLogitsProcessor(json_schema, tokenizer)

    def __call__(self, tokens: mx.array, logits: mx.array):
        logits_1d = logits.flatten()  # convert to 1-dimensional array
        logits_1d = self.logits_processor(tokens, logits_1d)
        logits = logits_1d[None]  # convert back to original shape
        return logits

我们可以使用此对象的实例化来调用 mlx 生成步骤

def mlx_stream(prompt: str):
    model, tokenizer = load("/path/to/mlx/model")
    prompt_tokens = mx.array(tokenizer.encode(prompt))
    json_schema='''{"type":"object","properties":{"title":{"type":"string","minLength":1}},"required":["title"]}'''  # define schema
    while True:
        yield generate_step(
            model=model,
            prompt=prompt_tokens,
            logits_processor=[OutlinesJSONLogitsProcessor(json_schema, tokenizer)]  # output valid json
        )

我们成功了！现在，只要提供 json 模式，我们就可以生成 json。

使用 MLX 运行视觉模型

MLX python 生态系统的另一部分是 mlx_vlm，这是一个用于运行视觉 LLM 的软件包。这是 mlx_vlm 中的 generate_step 方法，为了简洁而编辑过

def generate_step(*args, **kwargs):
    def sample(logits: mx.array) → Tuple[mx.array, float]:
        if temp == 0:
            token = mx.argmax(logits, axis=-1)
        else:
            if top_p > 0 and top_p < 1.0:
                token = top_p_sampling(logits, top_p, temp)
            else:
                token = mx.random.categorical(logits * (1 / temp))
        return token, logprobs

    # --snip--

    def _step(y):
        logits = model.language_model(y[None], cache=cache, mask=mask)
        logits = logits[:, -1, :]
        y, logprobs = sample(logits)
        return y, logprobs.squeeze(0)

    y = prompt
    logits = model(y, pixel_values, cache=cache, mask=mask)
    logits = logits[:, -1, :]
    y, logprobs = sample(logits)
    while True:
        next_y, next_logprobs = _step(y)
        yield y.item(), logprobs
        y, logprobs = next_y, next_logprobs

让我们比较和对比 mlx_vlm 实现与 mlx_lm 实现

mlx_vlm 评估使用 model.__call__ 方法。首次评估处理像素数据，随后的评估使用底层语言模型。
与 mlx_lm 相比，mlx_vlm 中的 sample 函数可用的采样方法更少。
mlx_vlm 中没有 logits_processor。

使用 mlx_lm 的 logits 处理和采样是有益的，同时也可以使用 mlx_vlm 的视觉模型。让我们实现它！

我们将编写一个类，该类将在首次调用时评估像素数据，并在后续调用中使用语言模型

class VisionModelWrapper:
    def __init__(self, vision_model, image_processor, pixel_values, mask):
        self.vision_model = vision_model
        self.image_processor = image_processor
        self.pixel_values = pixel_values
        self.mask = mask
        self.first_call = False

    def __call__(self, *args, **kwargs):
        if self.pixel_values is not None and not self.first_call:
            self.first_call = True
            return self.vision_model(self.input_ids, self.pixel_values, self.mask, **kwargs)
        else:
            return self.vision_model.language_model(*args, mask=self.mask, **kwargs)

现在，我们可以将其传递到 mlx_lm.generate_step 中

def mlx_stream(prompt: str):
    # load and wrap the vision model
    vision_model_dict, tokenizer = load_vision_model("/path/to/mlx/vision_model", "/path/to/image")
    vision_model_wrapper = VisionModelWrapper(**vision_model_dict)

    prompt_tokens = mx.array(tokenizer.encode(prompt))
    json_schema='''{"type":"object","properties":{"title":{"type":"string","minLength":1}},"required":["title"]}'''
    while True:
        yield generate_step(
            model=vision_model_wrapper,
            prompt=prompt_tokens,
            logits_processor=[OutlinesJSONLogitsProcessor(json_schema, tokenizer)]
        )

现在我们可以使用图像提示 LLM，并让它为我们生成标题！

使用 VLM 和结构化输出为图像添加标题

跨提示的 KV 缓存

跨提示的 KV（键值）缓存是一种优化技术，使 LLM 引擎能够重用先前交互的计算结果。这可以大大缩短模型响应时间或“首个 token 的时间”。

KV 缓存在聊天场景中尤其有价值，在聊天场景中，大部分提示（聊天记录）在对模型的生成请求中通常是相同的。

示例

时间步 1 (T1) - 用户发送提示 "总结这篇长文章：<此处为长文章...>"

{
    "User" : "Summarize this long article: <long article here...>"
}

时间步 2 (T2) - LLM 引擎对输入执行推理，计算模型权重和输入 token 嵌入之间的大型矩阵乘法，以生成输出 token："这篇文章讨论了...的影响"

{
    "User" : "Summarize this long article: <long article here...>",
    "AI" : "This article discusses the impact of..."
}

时间步 3 (T3) - 用户发送提示 "文章中是否提到了任何人？"。整个聊天记录将发送给 LLM，以便为其提供适当的上下文来继续对话。

{
    "User" : "Summarize this long article: <long article here...>",
    "AI" : "This article discusses the impact of...",
    "User" : "Are there any people mentioned in the article?"
}

时间步 4 (T4) - LLM 引擎对输入（来自 T1、T2 和 T3 的所有 token）执行推理，计算模型权重和输入 token 嵌入之间的大型矩阵乘法，以生成输出 token："是的，文章提到了几位关键人物，包括..."

{
    "User" : "Summarize this long article: <long article here...>",
    "AI" : "This article discusses the impact of...",
    "User" : "Are there any people mentioned in the article?",
    "AI" : "Yes, the article mentions several key figures, including..."
}

KV 缓存

KV 缓存利用了以下事实：当我们到达 T3 时，询问 LLM 关于 "文章中提到的人" 时，我们已经在 T1 和 T2 中执行了矩阵计算，这些计算与 T3 中需要计算的计算相同

{
    # START OF PREVIOUSLY COMPUTED
    "User" : "Summarize this long article: <long article here...>",
    "AI" : "This article discusses the impact of..."
    # END OF PREVIOUSLY COMPUTED
    "User" : "Are there any people mentioned in the article?"
}

因此，如果我们将 T1 和 T2 中计算的结果保存到 KV 缓存中，并在 T3 时向引擎提供对 KV 缓存的访问权限，则引擎只需对提示的新部分 "文章中是否提到了任何人？" 执行计算

{
    KV CACHE,
    "User" : "Are there any people mentioned in the article?"
}

这可以大大缩短 T4 中的响应时间。在我们的测试中，使用约 3000 个 token 的文章和 Meta-Llama-3.1-8B-Instruct-4bit，T4 响应时间从没有 KV 缓存时的约 10 秒降至有 KV 缓存时的仅 0.11 秒。

我们当前的 MLX KV 缓存实现

在实现时，mlx-lm 向其 generate_step 函数公开了一个 cache_history 参数

def generate_step(
    *args,
    cache_history: Optional[List[Tuple[mx.array, mx.array]]] = None,
    **kwargs
) → Generator[Tuple[mx.array, mx.array], None, None]:

通过传递适当的 cache_history（类似于上面的 KV 缓存），我们能够在我们的 MLX 引擎中实现 KV 缓存的初始版本。

我们通过改编 mlx-lm 的 PR 添加从文件加载 KV 缓存的功能来实现这一点，我们在其中通过缓存包装器预处理模型内部的提示

    def process_prompt(self, prompt_tokens, cache_wrapper, generate_args) → mx.array:
        """
        This method processes the prompt and adds its tokens to the cache history
        """

        # --snip--

        # prefill cache with prompt_tokens, except those that need to have a repetition penalty applied
        # (repetition penalty not currently possible for cached tokens)
        if "repetition_context_size" not in generate_args:
            generate_args["repetition_context_size"] = (
                20  # default value for mlx_lm.utils.generate_step
            )
        repetition_context_size = generate_args["repetition_context_size"]

        cache_history, generate_step_input = cache_wrapper.update_cache(
            prompt_tokens,
            num_tokens_to_exclude=repetition_context_size
        )

        generate_args["cache_history"] = cache_history

        return generate_step_input

上面看到的 cache_wrapper.update_cache，借鉴了 cache_prompt.py 以按块填充缓存

    # adapted from https://github.com/ml-explore/mlx-examples/blob/324184d670ec11916a5e92314171d497b312eefe/llms/mlx_lm/cache_prompt.py#L121-L137
    step_size = 512
    processed: int = 0
    while processed < len(tokens_to_process):
        # Here we evaluate the input prompt chunk by chunk to fill the cache
        chunk: mx.array = tokens_to_process[processed:processed+step_size]
        self.model(chunk[None], cache=self.cache)
        mx.eval([c.state for c in self.cache])
        self.tokens: mx.array = mx.concatenate([self.tokens, chunk]) if self.tokens is not None else chunk
        processed += chunk.size

现在缓存已创建并保存到 generate_args["cache_history"]，我们可以简单地将 generate_args 和 generate_step_input 传递给 mlx_lm.utils.generate_step

    # `process_prompt` function from above
    generate_step_input = process_prompt(prompt_tokens, cache_wrapper, generate_args)

    max_tokens = generate_args.pop("max_tokens")

    for (token, _), n in zip(
        # generate_step_input is now just the uncached repetition penalty tokens
        # generate_args has "cache_history" member, set in `process_prompt`
        mlx_lm.utils.generate_step(generate_step_input, model, **generate_args),
        range(max_tokens),
    ):

这使得 generate_step 函数能够利用先前计算的结果（存储在 cache_history 中）来大大缩短响应时间，与执行整个提示的原始处理相比。

然后，我们可以跨提示处理调用存储此 cache_history 对象，在其基础上构建，以保持聊天场景的响应性，即使在非常长的对话期间也是如此。但是，至关重要的是要确保处理到 cache_history 中的 token 在这样做时仍然对应于提示中的起始 token。有关此的更多信息，请查看 update_cache 函数中的缓存重置行为。

LM Studio 0.3.4 中的其他新功能

新增

任务控制：Cmd+Shift+M 搜索模型，Cmd+Shift+R 管理 LM 运行时
从 UI 设置结构化输出 (JSON 模式)

Bug 修复

修复长时间使用后出现黑屏的问题（参考：lmstudio-bug-tracker#98）
修复本地服务器端口只能使用 1234 的问题（参考：lms#80）
修复从 Obsidian 调用 embedding API 无法工作的问题（参考：tracker#142）
修复 RAG 有时在文档处理期间失败的问题

下载最新版本的 LM Studio 应用，适用于 macOS、Windows 或 Linux。
LM Studio 新手？请访问文档：文档：LM Studio 入门。
如需讨论和社区支持，请加入我们的 Discord 服务器。
如果您想在工作单位使用 LM Studio，请联系我们：LM Studio 工作版