ms-swift 深度定制实战：调试、数据增强、新增 token 与自定义 loss

1. 为什么我最终选了 ms-swift 作为主力训练框架1.1 从“能跑通”到“跑得舒服”的选型逻辑我最早接触大模型微调的时候用的是最原始的那套方案自己写 Trainer、自己拼 collator、自己处理 padding一个模型一个脚本改一次结构就要重写一遍训练循环。后来试过几个封装框架要么抽象太重、想改底层结构得翻半天源码要么功能太薄、连 LoRA 和全参的切换都要自己写。直到用上 ms-swift才算找到一个平衡点——它把训练、微调、推理、评测、量化、部署这条链路都串起来了同时底层又足够开放想改模型结构、换 loss、加 token 都能直接下手。ms-swift 是魔搭社区推出的一套大模型与多模态模型的微调部署框架支持 450 多个纯文本大模型和 150 多个多模态大模型覆盖全参训练、LoRA、QLoRA、DoRA、LongLoRA 等主流微调方式也支持 DPO、ORPO、KTO、CPO 这类偏好对齐训练。它最吸引我的地方在于命令行一条swift sft就能起训但当你需要深度定制时它又不会拦着你。这一点对做研究或者做产品落地的人来说太重要了——前期快速验证用 CLI后期要改结构、改 loss、加自定义 token 时直接改 Python 代码就行。这篇文章我打算把最近这段时间用 ms-swift 踩过的坑和总结的经验一次性讲清楚包括 VSCode 调试环境怎么搭、数据集怎么注册、动态数据增强怎么接、新增 token 怎么处理、回归训练怎么配、模型结构怎么改、自定义 loss 怎么写。内容偏实操适合已经跑通过一次 SFT、想往深水区走的人。如果你还没装过 ms-swift建议先跟着官方 README 跑一个最小的 LoRA 例子再回来看这篇。1.2 ms-swift 的架构分层与可扩展点要改得动一个框架先得知道它的分层。ms-swift 大致可以分成这么几层最上层是 CLI 和Swift类负责参数解析和流程编排中间是Trainer层基于 HuggingFace Trainer 做了大量定制比如Seq2SeqTrainer的变体再往下是模型加载层通过get_model_tokenizer这类工厂函数把模型、tokenizer、config 组装起来最底层是数据层Dataset和Template负责把原始样本转成模型能吃的 input_ids 和 labels。这个分层决定了你改东西的位置想改训练流程就动 Trainer想改模型结构就动 model 加载逻辑想改数据格式就动 Template想加 loss 就动 Trainer 的compute_loss。我后面讲的每一个定制点都会对应到具体是哪一层。理解了这个你再看官方文档里那些参数就不会觉得是一盘散沙了。另外提一句ms-swift 的版本迭代挺快不同版本之间 API 会有差异。我下面写的代码基于较新的版本如果你用的是老版本某些函数签名可能对不上遇到报错先去看对应版本的源码别硬套。2. VSCode 调试环境搭建让断点真正停下来2.1 解释器与依赖的隔离调试大模型训练代码第一件事是把环境隔离干净。我习惯用 conda 建一个独立环境Python 版本跟 ms-swift 要求对齐一般是 3.10 或 3.11。装依赖的时候有个坑不要直接pip install ms-swift就完事因为训练不同模型需要的依赖不一样比如多模态模型要额外的视觉库量化训练要 bitsandbytes 或 auto-gptq。我的做法是先装基础包再按需装扩展conda create -n swift python3.10 -y conda activate swift pip install ms-swift -U # 需要多模态能力时 pip install ms-swift[all] -U装完之后在 VSCode 里按CtrlShiftP输入Python: Select Interpreter选中刚才那个 conda 环境的 python 路径。这一步看着简单但我见过太多人调试时断点不停最后发现是 VSCode 用的系统 python根本没装 ms-swift。2.2 launch.json 的关键配置VSCode 调试 Python 靠的是.vscode/launch.json。调试 ms-swift 训练脚本核心是把命令行参数通过args传进去同时把工作目录设对。我常用的配置长这样{ version: 0.2.0, configurations: [ { name: Swift SFT Debug, type: debugpy, request: launch, program: ${workspaceFolder}/train_sft.py, console: integratedTerminal, cwd: ${workspaceFolder}, env: { CUDA_VISIBLE_DEVICES: 0, NCCL_P2P_DISABLE: 1 }, args: [ --model_type, qwen2_7b, --dataset, my_dataset, --num_train_epochs, 1 ], justMyCode: false } ] }这里有几个点值得说。justMyCode一定要设成false否则你只能在自己写的代码里断点进不了 ms-swift 框架内部的函数而调试框架问题恰恰需要进到它内部。console用integratedTerminal而不是internalConsole因为训练过程有大量进度条输出内部控制台会卡。NCCL_P2P_DISABLE这个环境变量在单卡调试时能避免一些 NCCL 初始化的问题多卡训练时再按需去掉。2.3 断点该打在哪里环境搭好只是第一步关键是知道在哪下断点。我调试 ms-swift 一般会在这几个位置打数据加载阶段Template.encode或Dataset.__getitem__看一条样本经过模板处理后 input_ids 和 labels 到底长什么样这是排查“loss 不降”的第一现场。模型前向之前Trainer.compute_loss入口确认 batch 里的字段和 shape。loss 计算处如果你自定义了 loss这里必打。保存/加载 checkpoint 处排查 tokenizer 和模型权重是否同步保存。提示调试训练脚本时把per_device_train_batch_size调到 1、max_steps设成 5 左右能让你在几十秒内走完整个训练循环断点调试效率高很多。别一上来就用真实配置跑等半天才停到第一个断点。2.4 远程调试的场景如果你是在服务器上训练、本地用 VSCode 写代码那 Remote-SSH 插件基本是标配。连上远程之后解释器选远程环境里的 pythonlaunch.json 里的路径也要改成远程路径。远程调试有个细节训练脚本里如果有os.fork或者多进程 dataloader断点可能会在子进程里失效。解决办法是把dataloader_num_workers设成 0先保证单进程能断住调通了再开多进程。3. 数据集注册从原始文件到可训练样本3.1 数据格式的两种主流选择ms-swift 支持的数据格式主要有两类一类是标准格式比如query/response或者messages这种对话格式另一类是自定义格式你自己写加载函数。我一般优先用标准格式因为框架已经帮你处理好了模板拼接、label 掩码这些事。一个典型的对话格式样本长这样{ messages: [ {role: user, content: 帮我写一个快速排序}, {role: assistant, content: def quicksort(arr): ...} ] }如果你的数据是这种结构直接存成 jsonl然后在命令行里--dataset /path/to/data.jsonl就能用。但实际项目里数据往往没这么规整比如你要做回归任务、要做多模态、要控制哪些 token 参与 loss这时候就得自定义。3.2 自定义数据集注册的完整流程ms-swift 注册自定义数据集的标准做法是继承Dataset或者用register_dataset装饰器。我以注册一个本地 jsonl 数据集为例走一遍完整流程。第一步写一个加载函数返回Dataset对象from swift.llm import Dataset, register_dataset register_dataset def load_my_dataset(dataset_id_or_path, **kwargs): # 读取原始文件 rows [] with open(dataset_id_or_path, r, encodingutf-8) as f: for line in f: rows.append(json.loads(line)) # 转成 swift 的 Dataset dataset Dataset.from_list(rows) return dataset第二步在训练脚本里 import 这个模块让装饰器生效。很多人卡在这里函数写好了但框架找不到就是因为没 import装饰器没执行注册表里没这个数据集。第三步命令行里用你注册的名字swift sft --dataset my_dataset --dataset_path /path/to/data.jsonl这里有个容易混淆的点--dataset传的是数据集标识--dataset_path传的是实际路径。如果你注册时把路径写死了那--dataset_path可以不传但更灵活的做法是让加载函数接收路径参数。3.3 数据预处理与 label 构造数据注册进来只是第一步真正决定训练效果的是 label 怎么构造。ms-swift 的 Template 机制会自动把 query 部分 mask 掉label 设为 -100只对 response 部分计算 loss。但有些场景你需要更精细的控制比如只对答案的最后一段算 loss可以在 Template 里重写_encode方法手动控制 labels。回归任务label 不是 token 序列而是一个连续值这时候要改数据格式和 loss。多轮对话要决定是每轮都算 loss还是只算最后一轮。我踩过的一个坑是自定义 Template 时忘了处理 padding 的 label。padding 位置的 label 必须是 -100否则 loss 会被 padding 拉偏。框架默认会处理但你一旦重写编码逻辑就得自己保证这一点。3.4 数据集缓存与版本管理ms-swift 会对处理过的数据集做缓存缓存在~/.cache/modelscope之类的目录下。调试阶段我建议把缓存关掉或者每次清掉否则你改了数据预处理逻辑跑起来发现还是老结果会怀疑人生。命令行加--dataset_num_proc 1并手动删缓存目录能避免大部分“改了没生效”的问题。另外数据集版本管理别偷懒。我习惯在数据文件名里带上日期和版本号比如sft_data_v3_20240601.jsonl同时在训练配置里记录用了哪个版本。不然过两周你回头复现实验根本想不起来当时用的是哪份数据。4. 动态数据增强让每个 epoch 看到的样本都不一样4.1 为什么静态增强不够用传统的数据增强是在训练前把数据扩增好存成更大的文件。但大模型微调场景下这种方式有两个问题一是数据量本来就大扩增后存储和加载成本高二是不同 epoch 看到同样的增强结果模型容易过拟合到增强模式上。动态数据增强的思路是在__getitem__里实时做增强每个 epoch 同一份原始样本可能被增强成不同形式。4.2 在 Dataset 层实现动态增强实现位置就在你自定义的 Dataset 里。我以文本任务为例写一个带随机增强的 Datasetimport random class AugDataset(Dataset): def __init__(self, rows, tokenizer, template): self.rows rows self.tokenizer tokenizer self.template template def __len__(self): return len(self.rows) def __getitem__(self, idx): row self.rows[idx] query row[query] response row[response] # 动态增强随机决定是否做同义替换/截断/拼接 if random.random() 0.3: query self.random_truncate(query) if random.random() 0.2: response self.random_paraphrase(response) return self.template.encode(query, response)关键点是增强逻辑必须幂等且可复现——如果你用了随机数记得在训练开始时设好 seed否则实验没法复现。我一般会在__getitem__里用random模块然后在训练脚本开头random.seed(42)。4.3 增强策略的选择与配比增强不是越多越好配比很关键。我常用的几类增强和它们的适用场景增强类型适用场景建议概率风险同义替换文本分类、意图识别0.2-0.3替换后语义漂移随机截断长文本摘要0.1-0.2丢失关键信息回译翻译、改写任务0.1引入噪声模板扰动结构化输出任务0.3破坏格式注意增强概率超过 0.5 之后我实测下来 loss 曲线会变得很抖模型收敛变慢。建议从 0.1 开始逐步往上加每次加完看验证集指标。4.4 动态增强与多进程 dataloader 的冲突这是个隐蔽的坑。如果你开了dataloader_num_workers 0每个 worker 进程有自己独立的随机状态导致增强结果不可复现。解决办法有两个一是在 worker 初始化函数里设 seed通过worker_init_fn二是干脆在调试阶段把 workers 设成 0。生产训练时我一般用前者保证可复现的同时还能并行加速。5. 新增 token词表扩展的正确姿势5.1 什么时候需要新增 token不是所有场景都需要加 token。以下几种情况我会考虑新增领域专有名词频繁出现且被切得很碎比如医学术语、代码里的特殊符号、需要引入特殊控制符比如|reg|这种标记回归任务的 token、多模态场景下要加图像占位符。如果只是普通文本微调加 token 反而可能破坏预训练学到的词表分布得不偿失。5.2 新增 token 的完整操作在 ms-swift 里加 token核心是改 tokenizer 和模型 embedding。流程如下from transformers import AutoTokenizer, AutoModelForCausalLM tokenizer AutoTokenizer.from_pretrained(model_dir) model AutoModelForCausalLM.from_pretrained(model_dir) new_tokens [|reg|, |img|] num_added tokenizer.add_tokens(new_tokens) model.resize_token_embeddings(len(tokenizer))这里有个关键细节resize_token_embeddings之后新增的 embedding 是随机初始化的。如果你直接开始训练这些随机向量会带来不稳定。我的做法是先把新增 token 的 embedding 初始化为已有 token 的平均值或者用语义相近的 token 初始化import torch with torch.no_grad(): emb model.get_input_embeddings().weight mean_emb emb[:-num_added].mean(dim0) emb[-num_added:] mean_emb5.3 保存与加载的一致性加完 token 训练完保存的时候一定要把 tokenizer 和模型一起存。我见过有人只存了模型权重加载时用原始 tokenizer结果新增 token 的 id 对不上推理直接崩。正确做法是tokenizer.save_pretrained(save_dir)和model.save_pretrained(save_dir)都调用保证词表和权重同步。另外如果你用的是 LoRA新增 token 的 embedding 属于被训练的参数保存 LoRA 权重时要确认这部分有没有被包含进去。有些配置下 LoRA 只保存 attention 层的权重embedding 的改动会丢。稳妥起见新增 token 后我一般会做一次全参微调或者把 embedding 层也纳入 LoRA 的目标模块。5.4 新增 token 对训练的影响评估加 token 之后我建议做一次对比实验同一份数据一份加 token 一份不加看验证集 loss 和下游指标。如果加了之后没提升甚至变差说明这些 token 没必要加。我自己的经验是只有当某个 token 在数据里出现频率超过千分之一且被原始 tokenizer 切成 3 个以上子词时加它才有明显收益。6. 回归训练把生成模型改成打分模型6.1 回归任务的本质变化标准语言模型做的是分类任务——预测下一个 token 在词表上的概率分布。回归任务要输出一个连续值这意味着模型的输出头要从vocab_size维改成 1 维loss 要从交叉熵改成 MSE 或 Huber。在 ms-swift 里做回归核心就是改这两处。6.2 改输出头与 loss最直接的做法是继承模型类替换lm_headimport torch.nn as nn class RegressionModel(nn.Module): def __init__(self, base_model): super().__init__() self.base base_model hidden base_model.config.hidden_size self.reg_head nn.Linear(hidden, 1) def forward(self, input_ids, attention_mask, labelsNone): outputs self.base(input_ids, attention_maskattention_mask, output_hidden_statesTrue) last_hidden outputs.hidden_states[-1] # 取最后一个非 padding token 的表示 lengths attention_mask.sum(dim1) - 1 pooled last_hidden[torch.arange(last_hidden.size(0)), lengths] logits self.reg_head(pooled).squeeze(-1) loss None if labels is not None: loss nn.functional.mse_loss(logits, labels.float()) return {loss: loss, logits: logits}这里取 pooled 表示的方式很关键。用最后一个 token 还是用 mean pooling对回归效果影响很大。我实测下来对于序列级回归任务mean pooling 通常比 last token 更稳因为 last token 受 padding 和序列长度影响大。你可以两种都试看验证集 MAE。6.3 数据格式与 label 对齐回归任务的数据格式跟生成任务不一样label 是浮点数而不是 token 序列。所以你的 Dataset 返回的应该是input_ids、attention_mask和labelsfloat 类型而不是生成任务那种 shift 过的 token label。这一点如果搞混loss 会算得莫名其妙。6.4 回归训练的评估指标回归任务别只看 loss要看 MAE、RMSE 和相关系数。我一般会在验证阶段把预测值和真实值都存下来画个散点图看模型是不是在某些区间系统性偏高或偏低。这种可视化能发现 loss 看不出来的问题比如模型对某个数值范围完全不敏感。7. 改模型结构与自定义 loss7.1 改结构的常见需求改模型结构的场景很多加一个 adapter 分支、改 attention 的计算方式、在中间层插入一个分类头、把某几层冻结。ms-swift 的模型加载走的是get_model_tokenizer工厂你要改结构最干净的方式是拿到模型后直接改而不是去改框架源码。from swift.llm import get_model_tokenizer model, tokenizer get_model_tokenizer(model_type, model_id_or_path) # 冻结前 N 层 for i, layer in enumerate(model.model.layers): if i 4: for p in layer.parameters(): p.requires_grad False7.2 自定义 loss 的接入点自定义 loss 的接入点在 Trainer 的compute_loss。ms-swift 的 Trainer 继承自 HuggingFace你可以重写这个方法from swift.llm import Seq2SeqTrainer class MyTrainer(Seq2SeqTrainer): def compute_loss(self, model, inputs, return_outputsFalse): outputs model(**inputs) logits outputs.logits labels inputs[labels] # 自定义 loss交叉熵 一个正则项 ce_loss nn.functional.cross_entropy( logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index-100 ) reg_loss sum(p.pow(2).sum() for p in model.parameters()) * 1e-5 loss ce_loss reg_loss return (loss, outputs) if return_outputs else loss写自定义 loss 有几个坑一是 ignore_index 要设对生成任务里 padding 的 label 是 -100不忽略的话 loss 会被拉偏二是 shape 要对齐logits 是[batch, seq, vocab]labels 是[batch, seq]view 的时候别搞反三是数值稳定性自己写 softmax 相关的 loss 时记得用 logsumexp 技巧别直接 exp 再 log。7.3 结构改动后的权重加载如果你改了结构比如加了新层加载预训练权重时会有一部分参数匹配不上。这时候from_pretrained会报 missing keys 的警告。我的做法是先把不匹配的 key 打印出来确认哪些是新加的、哪些是名字对不上的然后决定是随机初始化还是手动映射。别直接忽略这个警告有时候它意味着你的改动把原有权重加载错了。8. 常见问题与排查技巧实录8.1 训练不收敛的排查顺序loss 不降是最常见的问题我一般按这个顺序排查看数据断点进__getitem__确认 input_ids 和 labels 是不是符合预期label 有没有全被 mask 成 -100。看学习率太大直接发散太小几乎不动。LoRA 一般 1e-4 到 5e-4全参一般 1e-5 到 5e-5。看 loss 计算断点进compute_loss手动算一遍确认没有 shape 错位。看梯度加一行print(grad_norm)如果一直是 0 或者 nan说明梯度断了或者爆炸了。8.2 常见报错速查表报错信息可能原因解决方向CUDA out of memorybatch 太大或序列太长降 batch、开梯度累积、开 gradient checkpointingKeyError: xxx数据集字段名不对检查 Template 期望的字段名size mismatch新增 token 后没 resize调 resize_token_embeddingsloss is nan学习率过大或数据有脏样本降 lr、检查数据断点不停justMyCode 为 true 或解释器选错改 launch.json、重选解释器8.3 我踩过的几个真实坑第一个坑是数据集缓存。改了预处理逻辑但结果没变查了半天发现是缓存没清。现在我养成习惯每次改数据相关代码先删缓存目录。第二个坑是新增 token 后忘了 resize。tokenizer 加了 token但模型 embedding 没扩训练时直接 index out of range。这个错误信息不明显容易误判成数据问题。第三个坑是自定义 loss 里 ignore_index 没设。padding 参与 loss 计算导致 loss 数值虚高模型学到的全是 padding 模式。这个坑很隐蔽因为 loss 确实在降只是降的是 padding 的 loss。第四个坑是多卡训练时随机种子没同步。动态增强在多卡下每个进程增强结果不一样导致梯度方向不一致。解决办法是在worker_init_fn里用 rank 和 epoch 组合设 seed。8.4 性能调优的几个实用开关训练速度慢的时候我会依次检查这几个开关gradient_checkpointing开了没省显存但慢一点、flash_attention用了没快很多、dataloader_num_workers设了没数据加载不拖后腿、bf16开了没比 fp16 稳。这几个组合调好训练速度能差出两三倍。我个人在实际操作中的体会是ms-swift 这套框架最大的价值不在于它帮你省了多少代码而在于它把大模型训练里那些琐碎但关键的环节——模板、label 掩码、tokenizer 同步、checkpoint 管理——都标准化了。你在这个标准化的基础上做定制比从零搭一套要省心得多。但标准化也意味着你得先理解它的约定不然改起来处处碰壁。我建议新手先把官方 example 跑通然后挑一个最简单的定制点比如换个 loss动手改一遍改通了再往深了走。最后再分享一个小技巧调试训练脚本时把max_steps设成 3 到 5配合断点能在几分钟内走完整个训练循环比等完整 epoch 高效太多。