用 MLX 框架在 MacBook M4 上加速运行千问模型

在上一篇教程中,我们介绍了用 Ollama 和 llama.cpp 在 MacBook M4 上部署千问模型的方法。本文将介绍第三种方案——Apple 官方开源的 MLX 框架。MLX 专为 Apple Silicon 深度优化,在 M4 芯片上往往能跑出比 llama.cpp 更高的性能。

MLX 是什么?

MLX 是 Apple 官方开源的机器学习框架,类似于 PyTorch 但专门为 Apple Silicon 设计。它的核心优势包括:

  • 统一内存:直接利用 Apple 统一内存架构,CPU/GPU 之间零拷贝
  • 惰性计算:仅在需要时才执行计算,减少不必要的内存分配
  • 动态图:支持动态计算图,方便调试和开发
  • M4 深度优化:针对 M4 的 Neural Engine 和 GPU 有专门优化
  • Python 优先:API 设计类似 NumPy/PyTorch,上手极快

一、环境准备

1. 确认系统要求

MLX 要求 macOS 13.5 以上版本,M4 芯片完美支持。在终端确认:

# 检查 macOS 版本
sw_vers
# 输出示例:
# ProductName:    macOS
# ProductVersion: 15.2
# BuildVersion:   24C101

# 检查芯片型号
sysctl -n machdep.cpu.brand_string
# 输出示例:Apple M4

2. 安装 Python 环境

建议使用 Conda 或 venv 管理虚拟环境:

# 使用 Homebrew 安装 Python(如果没有)
brew install python@3.12

# 创建虚拟环境
python3.12 -m venv mlx-env
source mlx-env/bin/activate

# 升级 pip
pip install --upgrade pip

3. 安装 MLX

# 安装 MLX 核心库和 LLM 工具包
pip install mlx-lm

# mlx-lm 会自动安装 mlx、transformers、torch 等依赖
# 安装完成后验证
python -c "import mlx; print(mlx.__version__)"

二、运行千问模型

1. 命令行快速体验

mlx-lm 提供了开箱即用的命令行工具,一行命令即可运行模型:

# 运行 Qwen2.5-7B-Instruct(自动从 HuggingFace 下载)
mlx_lm.generate \
  --model mlx-community/Qwen2.5-7B-Instruct-4bit \
  --prompt "用 Python 写一个快速排序算法" \
  --max-tokens 512

首次运行会自动下载模型到 Hugging Face 缓存目录(~/.cache/huggingface/hub/)。MLX 社区已有大量预转换的千问模型可供使用。

2. 常用 MLX 千问模型

mlx-community 在 Hugging Face 上维护了大量 MLX 格式的千问模型:

模型名称量化大小适用内存
mlx-community/Qwen2.5-0.5B-Instruct-4bit4bit~0.4GB8GB+
mlx-community/Qwen2.5-3B-Instruct-4bit4bit~2.2GB8GB+
mlx-community/Qwen2.5-7B-Instruct-4bit4bit~4.7GB16GB+
mlx-community/Qwen2.5-14B-Instruct-4bit4bit~8.9GB24GB+
mlx-community/Qwen2.5-32B-Instruct-4bit4bit~20GB48GB+
mlx-community/QwQ-32B-Preview-4bit4bit~20GB48GB+

3. 交互式对话

# 启动交互式对话模式
mlx_lm.generate \
  --model mlx-community/Qwen2.5-7B-Instruct-4bit \
  --max-tokens 1024 \
  --interactive

三、Python API 编程调用

MLX 的真正强大之处在于 Python API,可以轻松集成到你的应用中:

1. 基本调用

from mlx_lm import load, generate

# 加载模型
default_model = "mlx-community/Qwen2.5-7B-Instruct-4bit"
model, tokenizer = load(default_model)

# 构建对话消息
messages = [
    {"role": "system", "content": "你是一个专业的Python开发工程师。"},
    {"role": "user", "content": "写一个装饰器,用于计算函数执行时间。"}
]

prompt = tokenizer.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True
)

# 生成回复
response = generate(
    model,
    tokenizer,
    prompt=prompt,\    max_tokens=512,
    verbose=True  # 打印生成速度和耗时
)
print(response)

2. 流式输出

对于需要实时显示的场景,可以使用流式输出:

from mlx_lm import load, stream_generate

model, tokenizer = load("mlx-community/Qwen2.5-7B-Instruct-4bit")

messages = [
    {"role": "user", "content": "详细解释一下什么是闭包,并给出示例。"}
]
prompt = tokenizer.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True
)

# 流式生成
for token in stream_generate(model, tokenizer, prompt=prompt, max_tokens=1024):
    print(token.text, end="", flush=True)
print()  # 换行

3. 多轮对话封装

from mlx_lm import load, generate

class QwenChat:
    def __init__(self, model_path="mlx-community/Qwen2.5-7B-Instruct-4bit"):
        self.model, self.tokenizer = load(model_path)
        self.history = []
    
    def chat(self, user_input):
        self.history.append({"role": "user", "content": user_input})
        prompt = self.tokenizer.apply_chat_template(
            self.history, tokenize=False, add_generation_prompt=True
        )
        response = generate(
            self.model, self.tokenizer,
            prompt=prompt, max_tokens=512, verbose=False
        )
        self.history.append({"role": "assistant", "content": response})
        return response
    
    def clear(self):
        self.history = []

# 使用示例
chat = QwenChat()
print(chat.chat("你好,我叫张三"))
print(chat.chat("我叫什么名字?"))
# 输出: 你叫张三

四、自定义模型转换

如果 mlx-community 没有你需要的模型版本,可以自己转换。mlx-lm 提供了转换工具:

# 将 HuggingFace 模型转换为 MLX 格式
# 4bit 量化
mlx_lm.convert \
  --hf-path Qwen/Qwen2.5-7B-Instruct \
  --mlx-path ./mlx-qwen-7b-4bit \
  -q  # 启用量化

# 8bit 量化
mlx_lm.convert \
  --hf-path Qwen/Qwen2.5-7B-Instruct \
  --mlx-path ./mlx-qwen-7b-8bit \
  -q --q-bits 8

# 无量化(FP16,需要更多内存)
mlx_lm.convert \
  --hf-path Qwen/Qwen2.5-7B-Instruct \
  --mlx-path ./mlx-qwen-7b-fp16

转换完成后,可以直接用本地路径运行:

# 使用本地模型路径
mlx_lm.generate \
  --model ./mlx-qwen-7b-4bit \
  --prompt "解释一下 Transformer 架构" \
  --max-tokens 512

五、MLX vs llama.cpp 性能对比

在 MacBook M4 (16GB) 上对 Qwen2.5-7B 4bit 量化模型进行对比测试:

指标MLXllama.cppOllama
首 token 延迟~0.8s~1.2s~1.5s
生成速度~45 tok/s~40 tok/s~35 tok/s
内存占用~5.0GB~5.2GB~5.5GB
峰值 GPU 利用率~85%~75%~70%
Python 集成原生支持需额外封装需 HTTP 调用

结论:MLX 在 M4 芯片上性能最优,尤其在生成速度和内存效率方面有明显优势。如果你需要在 Python 项目中集成大模型,MLX 是最佳选择。

六、高级技巧

1. 调整 KV Cache 大小

# 默认 KV Cache 较小,处理长文本时可以调大
from mlx_lm import load, generate

model, tokenizer = load(
    "mlx-community/Qwen2.5-7B-Instruct-4bit",
    # 调整 KV cache 大小(默认 256)
    model_config={"max_kv_size": 4096}
)

2. 批量生成

from mlx_lm import load, generate

model, tokenizer = load("mlx-community/Qwen2.5-7B-Instruct-4bit")

prompts = [
    "用一句话解释什么是 Docker",
    "用一句话解释什么是 Kubernetes",
    "用一句话解释什么是微服务"
]

for p in prompts:
    messages = [{"role": "user", "content": p}]
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    response = generate(model, tokenizer, prompt=prompt, max_tokens=128)
    print(f"Q: {p}")
    print(f"A: {response}")
    print("-" * 40)

3. 使用 LoRA 微调

MLX 支持在本地进行 LoRA 微调,让千问模型更适合你的特定任务:

# 准备数据集(JSONL 格式)
# data.jsonl 每行格式:
# {"text": "User: 你好\nAssistant: 你好!有什么可以帮你的?"}

# 启动 LoRA 微调
mlx_lm.lora \
  --model mlx-community/Qwen2.5-7B-Instruct-4bit \
  --train \
  --data ./data \
  --iters 500 \
  --batch-size 4 \
  --lora-layers 16 \
  --learning-rate 1e-4

# 使用微调后的模型
mlx_lm.generate \
  --model mlx-community/Qwen2.5-7B-Instruct-4bit \
  --adapter-path ./adapters \
  --prompt "你好" \
  --max-tokens 256

总结

MLX 作为 Apple 官方框架,在 M4 芯片上有着天然的性能优势。它不仅提供了命令行工具方便快速体验,还提供了完整的 Python API 供开发者深度集成。如果你是 Mac 用户且需要在应用中嵌入大模型能力,MLX 无疑是最佳选择。

下一篇文章我们将介绍如何基于 Ollama/MLX 搭建 OpenAI 兼容的 API 服务和 Web 对话界面,让你的 MacBook 变成一台私人 AI 服务器。

系列文章

发表回复

后才能评论