学习LLM微调 @2 - 在Macbook Air上使用mlx进行微调

在上一篇文章中笔者初步尝试了通过线上的微调服务来对大模型进行微调。但很显然，线上服务难以用在生产环境中。因为一方面，涉及到企业信息安全问题，不能直接将数据送到线上；另一方面，在线上微调的模型只能作为API调用而不能下载微调的权重。

所以接下来笔者就要尝试在本地通过微调框架来执行微调任务。由于最近在出差，暂时手头上只有一台笔记本(2023年款Macbook Air M3，24G内存)，所以笔者计划通过mlx工具执行微调任务，借此机会也上手一下苹果的mlx框架。

安装`mlx`

mlx是苹果官方开发，在M系列芯片上运行的模型训练与推理框架。根据官方教程，直接安装mlx-lm 这个包。除了mlx基础框架外，mlx-lm会额外提供一系列直接在命令行中可以使用的模型调用，模型微调等功能，我们后续使用的命令主要都来自这个包。

$ python3 -m pip install mlx-lm

下载原始模型

mlx框架所支持的所有模型都可以在huggingface的mlx community页面中找到。

这里我们以Qwen2.5-3B-Instruct-bf16为例（实际上更大的模型我的Macbook也调不动了！），首先我们执行下面的命令来下载模型。完成下载后会看到模型给出的简单回复，证明模型已经下载完成。

$ python3 -m mlx_lm.generate --model mlx-community/Qwen2.5-3B-Instruct-bf16 --prompt "hello"
==========
Hello! How can I assist you today?
==========
Prompt: 30 tokens, 170.685 tokens-per-sec
Generation: 10 tokens, 15.950 tokens-per-sec
Peak memory: 6.216 GB

数据准备

上一篇文章中详细介绍了从modelscope获取带思维链的微调数据，并且进行初步数据处理的过程。在硅基流动提供的微调服务中，服务会自动将上传的数据分为训练集和验证集，但mlx需要我们手动完成这个拆分。

首先按照上一篇文章中的方法整理数据并得到对话格式的训练数据distill_r1_110k_sft_messages.jsonl, 并将如下脚本保存为separate.py，

from os import path

import sys
import json


fin = open(sys.argv[1], 'r', encoding='utf8')
dir = sys.argv[2]
ftrain = open(path.join(dir, "train.jsonl"), "w+")
ftest = open(path.join(dir, "valid.jsonl"), "w+")

for i, line in enumerate(fin.readlines()):
    if i % 10 == 0:
        ftrain.write(line)
    else:
        ftest.write(line)

fin.close()
ftrain.close()
ftest.close()

并执行

$ python3 separate.py distill_r1_110k_sft_messages.jsonl ./data

然后在./data中就会出现按照9:1拆分好的训练集train.jsonl和验证集valid.jsonl （mlx微调过程中会直接根据这两个文件名来获取数据）。如果你需要不一样的训练集和验证集配比，可以自己修改上述的脚本。

执行微调

下面就到了将大象装进冰箱的时间。我们执行下面的命令，对下载好的Qwen2.5-3B-Instruct-bf16模型进行微调。

$ python3 -m mlx_lm.lora \
    --model mlx-community/Qwen2.5-3B-Instruct-bf16 \
    --train --data ./data \ # 指定数据所在目录
    --learning-rate 0.0001 \  # 设置学习率
    --iters 180 \  # 设置迭代次数 (可以自由调整，取决于要训练多少时间)
    --adapter-path ./mlx-adapter \  # 将微调结果保存在给定目录
    --steps-per-report 10 --save-every 30 \  # 每10个迭代显示一次日志，每30个迭代保存一次模型快照
    --max-seq-length 2048  # 设置训练过程中使用的最大序列长度

这个执行过程取决于你的电脑配置，笔者使用的丐版Air大概用了8个小时左右完成了180个迭代的微调，过程中loss从最初的1.938逐渐降低到1.0左右（在之前用同样的方法微调0.5B的模型的时候，Loss值是完全降不下来的），可见Qwen2.5-3B的基础模型能力尚可，具备支持简单思维链能力的前提条件。

使用微调所得的adapter进行推理

如果你的微调过程正常结束，微调结果保存在./mlx-adapter中，那么在执行推理命令的时候添加一个--adapter-path参数，就可以推理微调模型了。例如，

$ python3 -m mlx_lm.generate \
    --model mlx-community/Qwen2.5-3B-Instruct-bf16 \
    --adapter-path ./mlx-adapter \
    -m 1000 \
    --prompt "hello!"

但如果你希望使用中间过程生成的快照(--save-every 30就会每30步生成一个快照)，那么你应当将./mlx-adapter中的快照复制出来并且重新建立一个adapter文件夹。

首先我们来看./mlx-adapter的目录结构:

$ ll mlx-adapter
total 130416
-rw-r--r--@ 1 liyi  staff   7.0M  2 26 16:25 0000030_adapters.safetensors
-rw-r--r--@ 1 liyi  staff   7.0M  2 26 17:26 0000060_adapters.safetensors
……
-rw-r--r--@ 1 liyi  staff   2.1M  2 26 15:00 0000600_adapters.safetensors
                                             # xx_adapters.safetensors就是第xx轮迭代的快照
-rw-r--r--@ 1 liyi  staff   704B  2 26 15:13 adapter_config.json
-rw-r--r--@ 1 liyi  staff   7.0M  2 26 22:26 adapters.safetensors
                                             # 这个文件是最终的微调结果

在使用python3 mlx_lm.generate命令进行推理时，载入adapter实际上看的是adapters.safetensors这个文件。因此我们只要建立一个新目录并把对应的快照文件复制出来，就可以用中间过程快照进行推理。例如：

$ mkdir ./mlx-adapter-180
$ cp ./mlx-adapter/0000180_adapters.safetensors ./mlx-adapter-180/adapters.safetensors
$ cp ./mlx-adapter/adapter_config.json ./mlx-adapter-180/adapter_config.json
$ python3 -m mlx_lm.generate \
    --model mlx-community/Qwen2.5-3B-Instruct-bf16 \
    --adapter-path ./mlx-adapter-180 \
    -m 1000 \
    --prompt "hello!"

微调结果展示

运行以下命令，让微调后的模型写一首奇怪的诗（狗头）：

python3 -m mlx_lm.generate --model mlx-community/Qwen2.5-3B-Instruct-bf16 \
    --prompt "以马桶堵了为题目写一首诗" \
    --adapter-path ./mlx-adapter -m 1000

首先模型产生了如下思考过程。

<think> 嗯，用户让我以“马桶堵了”为题目写一首诗。首先，我需要理解用户的需求。他们可能是一个写作者，想要创作一首诗歌，主题是马桶堵塞。或者，他们可能是在寻找灵感，或者需要为某个项目写诗，比如学校作业、写作练习或者创意写作课。接下来，我需要考虑用户可能的背景。他们可能是学生，需要完成写作任务；或者他们对写作感兴趣，想要尝试不同的主题，比如日常生活中的小事情。也可能他们对马桶堵塞有个人经历，想通过诗歌表达出来。然后，我需要确定诗歌的结构和风格。可能用户希望诗歌既有诗意，又能反映现实中的小烦恼。需要考虑如何将马桶堵塞这个常见但尴尬的问题转化为艺术表达，同时保持诗歌的节奏和韵律。接下来，我需要构思诗歌的内容。马桶堵塞通常与家庭、清洁、尴尬的场景有关。可能需要包含对堵塞的描述，比如水龙头滴答声、堵塞物的堆积，以及解决堵塞的方法，比如疏通剂、工具等。同时，可以加入情感元素，比如焦虑、无奈，或者幽默感。还要考虑诗歌的节奏和押韵。用户可能希望诗歌有韵律感，比如四行诗，每行四个字，或者五行诗，每行五个字。这样更容易形成韵律，读起来朗朗上口。另外，用户可能希望诗歌有隐喻或象征意义。马桶堵塞可能象征着生活中的堵塞，比如工作、人际关系的障碍。需要通过马桶堵塞来表达更深层次的情感，比如沟通、解决冲突的必要性。最后，我需要检查是否有其他潜在需求。用户可能希望诗歌具有普遍性，能够引起共鸣，无论是成年人还是青少年。同时，他们可能希望诗歌既有幽默成分，又不失深度。总结一下，用户需要一首以马桶堵塞为题的诗，可能希望它既有诗意，又能反映现实中的小烦恼，同时通过堵塞来表达更深层次的情感。需要结构合理，韵律优美，内容贴近生活，情感真实，隐喻恰当。 </think>

之后是回复的正文：

马桶堵了
滴答声敲打夜的寂静
堵塞物在管道里
堆积成沉默的山

疏通剂喷射
像愤怒的火焰
烧灼着堵塞的
每一寸空间

水龙头滴答
像是时间的低语
提醒着堵塞的
不可逆

我试图用工具
撬动沉默的山
但每一次尝试
都像是徒劳的挣扎

直到马桶终于
恢复了流水的节奏
滴答声变成了
生命流动的乐章

马桶堵了
是生活中的小插曲
但每一次疏通
都让我更加珍视
生命流动的奇迹

相比于微调之前，文采有了极大的提升。以下是微调之前的模型对同一个prompt的输出。

马桶堵了，水位渐高，
心中烦忧，如堵心头。
水声潺潺，似诉衷肠，
等待疏通，心亦彷徨。

忽闻门外，脚步匆匆，
是那熟悉，是那身影。
他手握工具，神情专注，
一通疏通，水归原流。

心中感激，如沐春风，
感谢那及时，如救火。
生活琐事，虽小亦烦，
但有良人，便无惧。

从此以后，不再轻视，
生活琐事，亦需重视。
疏通马桶，虽小亦需，
生活之理，亦需疏通。

愿此诗，能解人烦忧，
亦能提醒，生活琐事，
虽小亦需重视，
生活之理，亦需疏通

安装mlx