平时用unsloth微调模型大多数人只碰训练脚本改改参数、换换数据集就完事了。但真到了生产环境或者做科研实验的时候总会有一些默认配置解决不了的需求。我这次就是被逼着把unsloth的源码翻出来改了几处整个过程踩了不少坑也弄明白了很多之前只知道用、不知道原理的东西。这篇就分享一下我改unsloth的实际经历、改动的思路和完整的验证过程给正准备对unsloth“下手”的朋友一个参考。1. 为什么要改unsloth三个绕不开的真实场景先说结论unsloth本身已经做得非常好了绝大多数情况下你不需要动它的源码。我这次改它是因为手头的项目撞上了三个默认设置解决不了的问题属于“不改就没法往下走”的情况。1.1 训练过程中需要即时评估而不是只盯loss曲线我之前的微调流程很简单配好数据集启动UnslothTrainer盯着loss降到某个值保存adapter完事。但这次的项目不太一样模型要输出特定格式的结构化内容比如JSON字段必须齐全、某些关键词必须出现。loss降得再漂亮也不代表格式就对了。我需要每训练几十步就让模型实际跑一次推理看看它在真实问题上的回答是否满足格式要求。默认的TrainingArguments只支持loss日志和常规的evaluate逻辑没法在训练中途方便地做生成式评估。我可以改训练脚本但每次都要写一堆回调、还要考虑会不会和unsloth的优化逻辑冲突。与其在外围打转不如直接进源码里看它到底怎么组织的。1.2 训练数据结构不是标准SFT格式大多数微调教程用的都是sharegpt格式或者alpaca格式。我这次的数据集带了复杂的结构化字段每条样本不只是“指令回答”还包含多个子任务、多段上下文以及一部分纯文本样本——它们不需要计算loss只需要作为上下文出现。默认的UnslothTrainerDataCollator在处理这类混合样本时很吃力。它的padding逻辑比较固定对每一条样本都按同等规则处理一旦样本类型不同出来的batch就很别扭。要么padding太多浪费显存要么不该计算loss的部分混进了loss计算。改训练脚本里的tokenize_function只能解决一部分问题真正的痛点在于数据整理阶段。1.3 checkpoint需要携带业务上下文训练完的adapter不是给自己用的后续要交给部署的同学去接推理服务。他们需要知道训练数据长什么样、用了哪个模板版本、特殊token的约定是什么。这些信息默认不会写进adapter_config.json里每次都得人工整理一份说明文档效率很低、还容易漏。如果可以改动保存逻辑把这些自定义字段自动写进adapter的配置文件后续加载的时候就能直接读取省掉一堆沟通成本。这就涉及到UnslothTrainer的保存路径了。三个需求叠加在一起结论就非常明确必须动手改。1.4 改之前先自问真的需要动源码吗在这里我得先给所有想改unsloth的朋友提个醒这个框架的设计很紧凑很多优化是环环相扣的乱改容易翻车。动手之前务必要先排除掉几个不需要动源码的路径大多数功能调整都可以通过UnslothTrainingArguments的参数解决比如embedding_learning_rate、use_gradient_checkpointing这些。数据集格式问题优先在tokenize_function里做适配实在不行再动collator。生成式评估可以用transformers的TrainerCallback来做完全不碰源码也能实现。我这次为什么最终还是改了源码因为三处需求同时存在而且每一处都要动到unsloth内部比较深的地方。如果只是单个需求我建议还是先用参数、用回调解决保持源码纯净后续升级不痛苦。2. 修改前的源码地图知道往哪里下刀改源码最怕的就是一头扎进去找不到方向。unsloth的代码量不算特别大但内部结构比较紧凑新版本迭代也快。我改之前先把整个包的结构粗读了一遍确认了三块最常动刀的位置。2.1 unsloth源码的骨架与关键文件我用的是当前比较主流的unsloth版本安装之后源码在site-packages/unsloth目录下。里面几个最关键的部分是文件/目录核心职责我这次是否改动unsloth/training.pyUnslothTrainer、UnslothTrainerDataCollator等训练相关类都在这改了动了collator和保存逻辑unsloth/models/llama.py针对LLaMA系列模型的monkey patch把原attention等模块替换成手写kernel版本没动这是显存优化的核心碰它风险太高unsloth/kernels/底层Triton kernel和C拓展完全没动出问题很难排查unsloth/chat_template_utils.py对话模板相关的处理没改这个版本已经直接走tokenizer的模板了对我这次的需求来说training.py是主战场。models/llama.py这类文件里是模型层的补丁除非你要改模型结构本身否则不要碰。2.2 补丁机制决定了很多东西不能乱动unsloth最大的特色是monkey patch机制。它在加载模型之后会用自己实现的高性能模块替换掉模型中原始的LlamaAttention、LlamaMLP等组件。这些组件在models/llama.py里被重新定义替换完之后你看到的模型结构已经和原始模型不一样了。这个机制对你改代码有非常实际的影响如果你想改模型正向传播过程中的某个逻辑不能只改model.py里的forward函数因为实际跑的时候调用的是patch之后的版本。你必须改patch里面的那个类。很多新手在这里就懵了改了一处没生效还以为是自己改错了。同时patch之后的模型配合了梯度检查点等优化里面的state处理是精心设计过的。如果你在forward里加了额外的tensor操作显存占用可能一下子就爆了所以尽量少往模型层里加东西。2.3 动手前的备份与隔离实践我强烈不建议直接改site-packages里的原文件。原因很简单unsloth升级很频繁你要是改了原文件一旦重新安装或者升级版本改动就全没了而且很难追溯。我的做法是把整个unsloth包复制一份到项目目录下命名为unsloth_custom然后在训练脚本里优先import unsloth_custom而不是原始的unsloth。这样改动都在自己的项目目录里源头可控也方便记录自己改了哪些地方。具体操作就是在项目根目录建一个unsloth_custom包把site-packages里对应的文件复制过来。导入的时候由于项目根目录在sys.path最前面会优先导入你自己那份。修改之后为了确保用的确实是自定义版本可以在包里加一行print([custom unsloth loaded])验证。这样的隔离方式我实测下来比较稳既能改源码又不污染环境多个项目还能各用各的版本。3. 第一处改动自定义数据collator适配混合任务先解决最麻烦的数据问题。我需要让collator支持混合类型样本并且让部分样本不计算loss。3.1 默认collator做了什么限制在哪里UnslothTrainerDataCollator继承自HuggingFace的DataCollatorForSeq2Seq核心工作是把样本padding到batch内最大长度生成attention_mask和labels处理label中的padding部分通常用-100屏蔽这个collator对标准SFT数据没有问题但有个隐含假设所有样本结构一致。它会用同一个padding逻辑处理每一条样本。当我的数据集里既有对话样本、又有纯文本上下文样本时问题就来了。纯文本样本的标签应该全部是-100而对话样本只有回复部分算loss。默认collator没法区分这两种情况。另外每轮对话中问题部分和回答部分的loss要不要屏蔽、模板里的系统提示算不算loss在不同任务里有不同讲究。默认collator不提供这么细的控制。3.2 继承重写一个混合任务collator的实例我的方案是继承UnslothTrainerDataCollator在__call__里覆盖核心逻辑。这里给出一段思路参考实际代码需要根据你自己的数据结构来调整from unsloth import UnslothTrainerDataCollator import torch class MixedTaskCollator(UnslothTrainerDataCollator): def __init__(self, model, tokenizer, max_seq_length, context_mask_prob0.0, **kwargs): super().__init__(modelmodel, tokenizertokenizer, max_seq_lengthmax_seq_length, **kwargs) self.context_mask_prob context_mask_prob def __call__(self, features): # 1. 先把混合样本拆成普通样本和上下文样本两组 normal_features [f for f in features if f.get(loss_type) supervised] context_features [f for f in features if f.get(loss_type) context_only] # 2. 上下文样本的所有labels设为-100不参与loss for f in context_features: f[labels] [-100] * len(f[input_ids]) # 3. 正常样本按原始逻辑处理这里沿用父类处理 # 注意父类内部会调用tokenizer.pad不会重复调用 batch super().__call__(normal_features context_features) # 4. 对上下文样本做随机屏蔽增强我项目里需要的功能 if self.context_mask_prob 0.0 and context_features: for i, f in enumerate(context_features): input_len len(f[input_ids]) if torch.rand(1).item() self.context_mask_prob: mask_start max(0, input_len - 4) batch[labels][i, mask_start:] -100 return batch这里有几个非常关键的细节必须调用父类的__call__来执行实际的padding操作不要自己重写padding逻辑。因为unsloth在padding时做了一些很底层的长度对齐优化绕过了等于放弃优化。loss_type这个字段是你自己在tokenize_function里加进去的。你需要在构造数据集时给每条样本打上这个标记。labels里的-100是HuggingFace loss屏蔽的约定值交叉熵计算会自动忽略这些位置。如果你要用不同的损失函数这个约定就要对应调整。改完之后我专门跑了一个小实验验证效果打了一大批样本检查输出的labels里上下文部分是否全为-100。结果符合预期而且显存占用没有明显增加。这说明collator层面的改动能平滑地融入unsloth的优化流程。4. 第二处改动在训练循环里插入生成式评估loss正常收敛之后我开始做生成式评估。这里要说一下其实transformers.Trainer提供了evaluate方法但它默认是为普通指标设计的要拿模型生成文本再逐条检查格式用起来比较别扭。我的选择是直接用callback。4.1 为什么不直接改evaluate方法UnslothTrainer继承自transformers.Trainer它确实有evaluate()。但这个方法的触发时机和评估逻辑比较固定而且unsloth在某些模式下会有一些内部的状态处理你在外面直接调用evaluate可能会干扰训练流程。更灵活的做法是注册一个TrainerCallback在on_step_end或者on_log的时候触发。这样评估逻辑完全独立不侵入trainer内部也不影响unsloth自身的优化。理论上说这也属于“不修改源码”的范畴。但我的情况是评估逻辑比较复杂要加载额外的数据处理模块而且要在多个训练脚本中复用。如果每次都写一大段回调代码脚本冗余太多了。所以我直接把一个自定义的callback类写进了unsloth_custom/training.py里作为内置回调这样所有基于自定义unsloth的训练脚本都能直接用。4.2 callback实现步骤与代码核心代码如下from transformers import TrainerCallback, TrainerControl, TrainerState from transformers import TrainingArguments import torch class GenerationEvalCallback(TrainerCallback): def __init__(self, eval_prompts, tokenizer, interval_steps50, max_new_tokens128): self.eval_prompts eval_prompts self.tokenizer tokenizer self.interval_steps interval_steps self.max_new_tokens max_new_tokens def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): if state.global_step % self.interval_steps ! 0: return model kwargs[model] model.eval() with torch.no_grad(): for i, prompt in enumerate(self.eval_prompts): messages [{role: user, content: prompt}] inputs self.tokenizer.apply_chat_template( messages, tokenizeTrue, add_generation_promptTrue, return_tensorspt, ).to(model.device) outputs model.generate( inputs, max_new_tokensself.max_new_tokens, do_sampleTrue, temperature0.7, pad_token_idself.tokenizer.eos_token_id, ) response self.tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokensTrue) print(f\n[Step {state.global_step}] Prompt {i}: {prompt[:50]}...) print(fResponse: {response[:200]}) # 这里可以做格式检查、关键词匹配等把指标log进去 model.train()在使用的地方简单注册from unsloth_custom.training import GenerationEvalCallback training_args UnslothTrainingArguments(...) trainer UnslothTrainer( modelmodel, tokenizertokenizer, argstraining_args, train_datasetdataset, callbacks[GenerationEvalCallback(eval_promptseval_prompts, tokenizertokenizer)], )要点解释一下这些都是我踩过的坑model.eval()和model.train()必须成对出现。漏掉train()会导致后续训练的dropout和batchnorm行为异常整个训练曲线都会飘。pad_token_id要显式指定。unsloth对部分模型会调整tokenizer配置有些模型的pad_token是None生成的时候不指定很容易报错。生成的输入用apply_chat_template处理保证和训练时的模板一致。如果训练时用的是自定义模板这里也必须保持一致否则评估结果完全没参考意义。callback里生成的输出不要直接抛给trainer.log()除非你把它转成标准的logs字典格式。我在第一版里就是直接print后续做指标聚合时再统一记录。这样改完之后训练日志里每隔50步就会打印出模型在几个固定问题上的实际表现。我不需要再手动中断训练去试模型效果整个训练过程的监控质量提升了一大截。5. 第三处改动让checkpoint自动携带业务元数据这一处改动最“温和”没有动任何训练逻辑只是扩展了保存时的输出信息。5.1 默认adapter配置里缺了什么unsloth保存LoRA adapter时会生成一个adapter_config.json。里面包含base_model_name_or_path、r、lora_alpha、target_modules这些参数。这些对模型加载来说是必要的但对业务方来说远远不够。部署的同学拿到adapter之后最常见的问题就是“这个模型是用什么数据训练的模板是什么需要用什么特殊token格式”这些信息在纯技术配置里完全找不到。每次都得找训练的人问问完还要写文档文档还经常和实际不一致。如果能在保存的时候把这些信息自动写进配置后续加载的时候直接读取能省掉一大半沟通成本。5.2 修改保存逻辑的完整思路UnslothTrainer在保存adapter时会调用model.save_pretrained()通过save_adapter之类的路径把LoRA权重和配置写到目标目录。我在unsloth_custom/training.py里找到保存相关的调用在它保存完之后追加一个元数据写入步骤。思路是在UnslothTrainer里增加一个metadata属性训练脚本设置好之后保存的时候自动合并进配置文件。核心逻辑类似这样class CustomUnslothTrainer(UnslothTrainer): def __init__(self, *args, metadataNone, **kwargs): super().__init__(*args, **kwargs) self.metadata metadata or {} def _save(self, output_dirNone, state_dictNone): super()._save(output_diroutput_dir, state_dictstate_dict) # 在原有保存之后追加业务元数据 config_path os.path.join(output_dir, adapter_config.json) if os.path.exists(config_path): with open(config_path, r, encodingutf-8) as f: adapter_config json.load(f) adapter_config[custom_metadata] self.metadata with open(config_path, w, encodingutf-8) as f: json.dump(adapter_config, f, indent2, ensure_asciiFalse)使用的时候trainer CustomUnslothTrainer( modelmodel, tokenizertokenizer, argstraining_args, train_datasetdataset, metadata{ task: structured_output_generation, data_version: v2.1, chat_template_version: custom_v3, excluded_fields_from_loss: [context, conversation_history], note: 上下文部分不参与loss计算只作为输入上下文 }, )这里有几个需要注意的地方_save方法只在调用trainer.save_model()时触发。TrainingArguments里的save_strategy触发的是_save_checkpoint走的是checkpoint路径逻辑略有不同。如果你想让每几千步保存的中间checkpoint也带上元数据需要把逻辑同时挂钩到_save_checkpoint。我为了省事只在最终save_model()的时候写入元数据中间checkpoint不管。不要把训练状态写进元数据。比如global_step、train_loss这些值在每次保存时都不一样写进adapter_config.json会造成配置漂移。meta信息最好只放和数据集、模板、任务定义相关的静态内容。修改保存逻辑要保证幂等性。也就是说重复保存多次不会重复叠加字段。因为_save可能在训练中多次触发我上面这段代码里的逻辑是覆盖写入所以是幂等的这也算一个设计上的小细节。改完之后部署同事加载adapter时顺手看一眼配置文件训练数据版本一目了然再也不用追着问人了。6. 修改后的回归验证与踩坑记录源码改完了不代表事情就结束了。最关键的其实是验证——怎么证明改动没有破坏unsloth原有的优化效果。6.1 回归验证的四步流程我花了两天时间做验证核心思路很简单改之前跑一遍基线改之后再跑一遍对比训练曲线和显存占用。具体分为四步第一步最小数据集烟雾测试。拿几百条样本开一个极小的batch跑两三个step检查能否正常启动、日志是否正常。这一步主要排除语法错误和import问题。我在这里就发现了一个问题——自定义collator在deepcopy特征时出现了类型不匹配的错误排查后发现是loss_type字段在copy过程中丢失了需要特殊处理。第二步完整训练曲线对比。用同一个数据集、同样的超参数分别跑原始unsloth和修改后的版本对比loss下降曲线。正常情况应该有微小的随机差异但整体趋势必须一致。如果修改后loss波动异常大多半是labels的mask逻辑有问题。第三步显存占用对比。unsloth最核心的价值就是省显存。我这次改动没有碰kernel层也没有给模型forward加东西理论上显存不应该有明显增加。实测下来效果符合预期训练速度和显存占用基本与改前持平。如果你改了模型层之后显存暴涨那你得回头想想是不是不小心破坏了什么优化。第四步保存—加载闭环测试。把修改后训练的adapter保存下来用一个全新的原版unsloth环境加载跑一次推理看模型行为是否正常。这一步能验证我写的元数据有没有把adapter_config.json弄坏。实测加载正常自定义字段被正确读取没有影响LoRA权重。6.2 几个容易翻车的地方这次改源码踩了不少坑总结几个我认为最有价值的版本兼容性问题是最大的坑。unsloth迭代很快不同版本的UnslothTrainer类内部结构有差异。如果你照网上的教程改源码很可能因为版本不同而找不到对应的代码位置。我的做法是在改之前先把自己的unsloth版本固定住不要让它在训练中途升级。-100标签屏蔽的位置要非常小心。有些模型的自定义损失函数可能不会读取labels里的-100而是直接忽略labels里所有的负值。如果你的代码里自定义了损失函数记得检查它是否遵循-100约定。callback里做生成评估时注意模型状态切换。如果模型用了torch.compile或者unsloth的优化状态在某些版本下model.eval()和model.train()的切换可能会和优化状态有交互。实测中遇到过一次奇怪的问题评估完回到训练模式后某个layer的缓存没有清理干净导致后续loss出现一次尖峰。最终的解决办法是在model.train()之后手动做一次model.zero_grad()。不要为了“改”而改。我这次改的三个点每一个都有明确的使用场景支撑。如果你的项目用默认unsloth就能跑得很顺真没必要动源码。维护一份自定义版本的unsloth是有成本的升级、依赖管理、多人协作都会变得稍微复杂。7. 一点经验最小改动原则与后续维护思路改完这一轮我的体会是对unsloth这种性能敏感、结构紧凑的库最稳妥的改造策略是“最小改动”能加一层就加一层绝不去动底层kernel。这次改动一共就集中在两个文件里training.py里加了collator的覆盖逻辑和一个callbackUnslothTrainer的保存逻辑里加了元数据注入。模型层的文件一个都没碰kernel层更不用说。这样一来unsloth升级的时候我可以清楚地看到自己改了什么手动从一个版本挪到另一个版本的工作量也可控。如果你也准备改unsloth我给一个比较实用的建议把自己所有改动都集中在一个自定义文件里比如unsloth_custom/patches.py统一用mixin或回调方式挂载而不是零散地改原文件。这样做的最大好处是升级unsloth时你只需要重新import这个文件不需要再逐个找改动点。最后再分享一个实用的小技巧改完源码后在训练脚本最前面写一个断言检查自定义版本是否真的被加载了。比如在一个自定义模块里放一个全局常量CUSTOM_UNSLOTH_FLAG v20240615_custom assert CUSTOM_UNSLOTH_FLAG in sys.modules.get(unsloth, ), custom unsloth not loaded!这个断言能在训练启动时快速暴露import路径错误省得你稀里糊涂用错了版本训练跑了一大半才发现改动没生效。这类“调试性断言”在维护自定义版本时非常实用尤其当你同时有好几个项目加载了不同版本的unsloth时。 
长河君 · 2026/10/23 3:33:24 · 阅读 1.2k unsloth源码改造实战：自定义collator、评估回调与checkpoint元数据
平时用unsloth微调模型大多数人只碰训练脚本改改参数、换换数据集就完事了。但真到了生产环境或者做科研实验的时候总会有一些默认配置解决不了的需求。我这次就是被逼着把unsloth的源码翻出来改了几处整个过程踩了不少坑也弄明白了很多之前只知道用、不知道原理的东西。这篇就分享一下我改unsloth的实际经历、改动的思路和完整的验证过程给正准备对unsloth“下手”的朋友一个参考。1. 为什么要改unsloth三个绕不开的真实场景先说结论unsloth本身已经做得非常好了绝大多数情况下你不需要动它的源码。我这次改它是因为手头的项目撞上了三个默认设置解决不了的问题属于“不改就没法往下走”的情况。1.1 训练过程中需要即时评估而不是只盯loss曲线我之前的微调流程很简单配好数据集启动UnslothTrainer盯着loss降到某个值保存adapter完事。但这次的项目不太一样模型要输出特定格式的结构化内容比如JSON字段必须齐全、某些关键词必须出现。loss降得再漂亮也不代表格式就对了。我需要每训练几十步就让模型实际跑一次推理看看它在真实问题上的回答是否满足格式要求。默认的TrainingArguments只支持loss日志和常规的evaluate逻辑没法在训练中途方便地做生成式评估。我可以改训练脚本但每次都要写一堆回调、还要考虑会不会和unsloth的优化逻辑冲突。与其在外围打转不如直接进源码里看它到底怎么组织的。1.2 训练数据结构不是标准SFT格式大多数微调教程用的都是sharegpt格式或者alpaca格式。我这次的数据集带了复杂的结构化字段每条样本不只是“指令回答”还包含多个子任务、多段上下文以及一部分纯文本样本——它们不需要计算loss只需要作为上下文出现。默认的UnslothTrainerDataCollator在处理这类混合样本时很吃力。它的padding逻辑比较固定对每一条样本都按同等规则处理一旦样本类型不同出来的batch就很别扭。要么padding太多浪费显存要么不该计算loss的部分混进了loss计算。改训练脚本里的tokenize_function只能解决一部分问题真正的痛点在于数据整理阶段。1.3 checkpoint需要携带业务上下文训练完的adapter不是给自己用的后续要交给部署的同学去接推理服务。他们需要知道训练数据长什么样、用了哪个模板版本、特殊token的约定是什么。这些信息默认不会写进adapter_config.json里每次都得人工整理一份说明文档效率很低、还容易漏。如果可以改动保存逻辑把这些自定义字段自动写进adapter的配置文件后续加载的时候就能直接读取省掉一堆沟通成本。这就涉及到UnslothTrainer的保存路径了。三个需求叠加在一起结论就非常明确必须动手改。1.4 改之前先自问真的需要动源码吗在这里我得先给所有想改unsloth的朋友提个醒这个框架的设计很紧凑很多优化是环环相扣的乱改容易翻车。动手之前务必要先排除掉几个不需要动源码的路径大多数功能调整都可以通过UnslothTrainingArguments的参数解决比如embedding_learning_rate、use_gradient_checkpointing这些。数据集格式问题优先在tokenize_function里做适配实在不行再动collator。生成式评估可以用transformers的TrainerCallback来做完全不碰源码也能实现。我这次为什么最终还是改了源码因为三处需求同时存在而且每一处都要动到unsloth内部比较深的地方。如果只是单个需求我建议还是先用参数、用回调解决保持源码纯净后续升级不痛苦。2. 修改前的源码地图知道往哪里下刀改源码最怕的就是一头扎进去找不到方向。unsloth的代码量不算特别大但内部结构比较紧凑新版本迭代也快。我改之前先把整个包的结构粗读了一遍确认了三块最常动刀的位置。2.1 unsloth源码的骨架与关键文件我用的是当前比较主流的unsloth版本安装之后源码在site-packages/unsloth目录下。里面几个最关键的部分是文件/目录核心职责我这次是否改动unsloth/training.pyUnslothTrainer、UnslothTrainerDataCollator等训练相关类都在这改了动了collator和保存逻辑unsloth/models/llama.py针对LLaMA系列模型的monkey patch把原attention等模块替换成手写kernel版本没动这是显存优化的核心碰它风险太高unsloth/kernels/底层Triton kernel和C拓展完全没动出问题很难排查unsloth/chat_template_utils.py对话模板相关的处理没改这个版本已经直接走tokenizer的模板了对我这次的需求来说training.py是主战场。models/llama.py这类文件里是模型层的补丁除非你要改模型结构本身否则不要碰。2.2 补丁机制决定了很多东西不能乱动unsloth最大的特色是monkey patch机制。它在加载模型之后会用自己实现的高性能模块替换掉模型中原始的LlamaAttention、LlamaMLP等组件。这些组件在models/llama.py里被重新定义替换完之后你看到的模型结构已经和原始模型不一样了。这个机制对你改代码有非常实际的影响如果你想改模型正向传播过程中的某个逻辑不能只改model.py里的forward函数因为实际跑的时候调用的是patch之后的版本。你必须改patch里面的那个类。很多新手在这里就懵了改了一处没生效还以为是自己改错了。同时patch之后的模型配合了梯度检查点等优化里面的state处理是精心设计过的。如果你在forward里加了额外的tensor操作显存占用可能一下子就爆了所以尽量少往模型层里加东西。2.3 动手前的备份与隔离实践我强烈不建议直接改site-packages里的原文件。原因很简单unsloth升级很频繁你要是改了原文件一旦重新安装或者升级版本改动就全没了而且很难追溯。我的做法是把整个unsloth包复制一份到项目目录下命名为unsloth_custom然后在训练脚本里优先import unsloth_custom而不是原始的unsloth。这样改动都在自己的项目目录里源头可控也方便记录自己改了哪些地方。具体操作就是在项目根目录建一个unsloth_custom包把site-packages里对应的文件复制过来。导入的时候由于项目根目录在sys.path最前面会优先导入你自己那份。修改之后为了确保用的确实是自定义版本可以在包里加一行print([custom unsloth loaded])验证。这样的隔离方式我实测下来比较稳既能改源码又不污染环境多个项目还能各用各的版本。3. 第一处改动自定义数据collator适配混合任务先解决最麻烦的数据问题。我需要让collator支持混合类型样本并且让部分样本不计算loss。3.1 默认collator做了什么限制在哪里UnslothTrainerDataCollator继承自HuggingFace的DataCollatorForSeq2Seq核心工作是把样本padding到batch内最大长度生成attention_mask和labels处理label中的padding部分通常用-100屏蔽这个collator对标准SFT数据没有问题但有个隐含假设所有样本结构一致。它会用同一个padding逻辑处理每一条样本。当我的数据集里既有对话样本、又有纯文本上下文样本时问题就来了。纯文本样本的标签应该全部是-100而对话样本只有回复部分算loss。默认collator没法区分这两种情况。另外每轮对话中问题部分和回答部分的loss要不要屏蔽、模板里的系统提示算不算loss在不同任务里有不同讲究。默认collator不提供这么细的控制。3.2 继承重写一个混合任务collator的实例我的方案是继承UnslothTrainerDataCollator在__call__里覆盖核心逻辑。这里给出一段思路参考实际代码需要根据你自己的数据结构来调整from unsloth import UnslothTrainerDataCollator import torch class MixedTaskCollator(UnslothTrainerDataCollator): def __init__(self, model, tokenizer, max_seq_length, context_mask_prob0.0, **kwargs): super().__init__(modelmodel, tokenizertokenizer, max_seq_lengthmax_seq_length, **kwargs) self.context_mask_prob context_mask_prob def __call__(self, features): # 1. 先把混合样本拆成普通样本和上下文样本两组 normal_features [f for f in features if f.get(loss_type) supervised] context_features [f for f in features if f.get(loss_type) context_only] # 2. 上下文样本的所有labels设为-100不参与loss for f in context_features: f[labels] [-100] * len(f[input_ids]) # 3. 正常样本按原始逻辑处理这里沿用父类处理 # 注意父类内部会调用tokenizer.pad不会重复调用 batch super().__call__(normal_features context_features) # 4. 对上下文样本做随机屏蔽增强我项目里需要的功能 if self.context_mask_prob 0.0 and context_features: for i, f in enumerate(context_features): input_len len(f[input_ids]) if torch.rand(1).item() self.context_mask_prob: mask_start max(0, input_len - 4) batch[labels][i, mask_start:] -100 return batch这里有几个非常关键的细节必须调用父类的__call__来执行实际的padding操作不要自己重写padding逻辑。因为unsloth在padding时做了一些很底层的长度对齐优化绕过了等于放弃优化。loss_type这个字段是你自己在tokenize_function里加进去的。你需要在构造数据集时给每条样本打上这个标记。labels里的-100是HuggingFace loss屏蔽的约定值交叉熵计算会自动忽略这些位置。如果你要用不同的损失函数这个约定就要对应调整。改完之后我专门跑了一个小实验验证效果打了一大批样本检查输出的labels里上下文部分是否全为-100。结果符合预期而且显存占用没有明显增加。这说明collator层面的改动能平滑地融入unsloth的优化流程。4. 第二处改动在训练循环里插入生成式评估loss正常收敛之后我开始做生成式评估。这里要说一下其实transformers.Trainer提供了evaluate方法但它默认是为普通指标设计的要拿模型生成文本再逐条检查格式用起来比较别扭。我的选择是直接用callback。4.1 为什么不直接改evaluate方法UnslothTrainer继承自transformers.Trainer它确实有evaluate()。但这个方法的触发时机和评估逻辑比较固定而且unsloth在某些模式下会有一些内部的状态处理你在外面直接调用evaluate可能会干扰训练流程。更灵活的做法是注册一个TrainerCallback在on_step_end或者on_log的时候触发。这样评估逻辑完全独立不侵入trainer内部也不影响unsloth自身的优化。理论上说这也属于“不修改源码”的范畴。但我的情况是评估逻辑比较复杂要加载额外的数据处理模块而且要在多个训练脚本中复用。如果每次都写一大段回调代码脚本冗余太多了。所以我直接把一个自定义的callback类写进了unsloth_custom/training.py里作为内置回调这样所有基于自定义unsloth的训练脚本都能直接用。4.2 callback实现步骤与代码核心代码如下from transformers import TrainerCallback, TrainerControl, TrainerState from transformers import TrainingArguments import torch class GenerationEvalCallback(TrainerCallback): def __init__(self, eval_prompts, tokenizer, interval_steps50, max_new_tokens128): self.eval_prompts eval_prompts self.tokenizer tokenizer self.interval_steps interval_steps self.max_new_tokens max_new_tokens def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): if state.global_step % self.interval_steps ! 0: return model kwargs[model] model.eval() with torch.no_grad(): for i, prompt in enumerate(self.eval_prompts): messages [{role: user, content: prompt}] inputs self.tokenizer.apply_chat_template( messages, tokenizeTrue, add_generation_promptTrue, return_tensorspt, ).to(model.device) outputs model.generate( inputs, max_new_tokensself.max_new_tokens, do_sampleTrue, temperature0.7, pad_token_idself.tokenizer.eos_token_id, ) response self.tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokensTrue) print(f\n[Step {state.global_step}] Prompt {i}: {prompt[:50]}...) print(fResponse: {response[:200]}) # 这里可以做格式检查、关键词匹配等把指标log进去 model.train()在使用的地方简单注册from unsloth_custom.training import GenerationEvalCallback training_args UnslothTrainingArguments(...) trainer UnslothTrainer( modelmodel, tokenizertokenizer, argstraining_args, train_datasetdataset, callbacks[GenerationEvalCallback(eval_promptseval_prompts, tokenizertokenizer)], )要点解释一下这些都是我踩过的坑model.eval()和model.train()必须成对出现。漏掉train()会导致后续训练的dropout和batchnorm行为异常整个训练曲线都会飘。pad_token_id要显式指定。unsloth对部分模型会调整tokenizer配置有些模型的pad_token是None生成的时候不指定很容易报错。生成的输入用apply_chat_template处理保证和训练时的模板一致。如果训练时用的是自定义模板这里也必须保持一致否则评估结果完全没参考意义。callback里生成的输出不要直接抛给trainer.log()除非你把它转成标准的logs字典格式。我在第一版里就是直接print后续做指标聚合时再统一记录。这样改完之后训练日志里每隔50步就会打印出模型在几个固定问题上的实际表现。我不需要再手动中断训练去试模型效果整个训练过程的监控质量提升了一大截。5. 第三处改动让checkpoint自动携带业务元数据这一处改动最“温和”没有动任何训练逻辑只是扩展了保存时的输出信息。5.1 默认adapter配置里缺了什么unsloth保存LoRA adapter时会生成一个adapter_config.json。里面包含base_model_name_or_path、r、lora_alpha、target_modules这些参数。这些对模型加载来说是必要的但对业务方来说远远不够。部署的同学拿到adapter之后最常见的问题就是“这个模型是用什么数据训练的模板是什么需要用什么特殊token格式”这些信息在纯技术配置里完全找不到。每次都得找训练的人问问完还要写文档文档还经常和实际不一致。如果能在保存的时候把这些信息自动写进配置后续加载的时候直接读取能省掉一大半沟通成本。5.2 修改保存逻辑的完整思路UnslothTrainer在保存adapter时会调用model.save_pretrained()通过save_adapter之类的路径把LoRA权重和配置写到目标目录。我在unsloth_custom/training.py里找到保存相关的调用在它保存完之后追加一个元数据写入步骤。思路是在UnslothTrainer里增加一个metadata属性训练脚本设置好之后保存的时候自动合并进配置文件。核心逻辑类似这样class CustomUnslothTrainer(UnslothTrainer): def __init__(self, *args, metadataNone, **kwargs): super().__init__(*args, **kwargs) self.metadata metadata or {} def _save(self, output_dirNone, state_dictNone): super()._save(output_diroutput_dir, state_dictstate_dict) # 在原有保存之后追加业务元数据 config_path os.path.join(output_dir, adapter_config.json) if os.path.exists(config_path): with open(config_path, r, encodingutf-8) as f: adapter_config json.load(f) adapter_config[custom_metadata] self.metadata with open(config_path, w, encodingutf-8) as f: json.dump(adapter_config, f, indent2, ensure_asciiFalse)使用的时候trainer CustomUnslothTrainer( modelmodel, tokenizertokenizer, argstraining_args, train_datasetdataset, metadata{ task: structured_output_generation, data_version: v2.1, chat_template_version: custom_v3, excluded_fields_from_loss: [context, conversation_history], note: 上下文部分不参与loss计算只作为输入上下文 }, )这里有几个需要注意的地方_save方法只在调用trainer.save_model()时触发。TrainingArguments里的save_strategy触发的是_save_checkpoint走的是checkpoint路径逻辑略有不同。如果你想让每几千步保存的中间checkpoint也带上元数据需要把逻辑同时挂钩到_save_checkpoint。我为了省事只在最终save_model()的时候写入元数据中间checkpoint不管。不要把训练状态写进元数据。比如global_step、train_loss这些值在每次保存时都不一样写进adapter_config.json会造成配置漂移。meta信息最好只放和数据集、模板、任务定义相关的静态内容。修改保存逻辑要保证幂等性。也就是说重复保存多次不会重复叠加字段。因为_save可能在训练中多次触发我上面这段代码里的逻辑是覆盖写入所以是幂等的这也算一个设计上的小细节。改完之后部署同事加载adapter时顺手看一眼配置文件训练数据版本一目了然再也不用追着问人了。6. 修改后的回归验证与踩坑记录源码改完了不代表事情就结束了。最关键的其实是验证——怎么证明改动没有破坏unsloth原有的优化效果。6.1 回归验证的四步流程我花了两天时间做验证核心思路很简单改之前跑一遍基线改之后再跑一遍对比训练曲线和显存占用。具体分为四步第一步最小数据集烟雾测试。拿几百条样本开一个极小的batch跑两三个step检查能否正常启动、日志是否正常。这一步主要排除语法错误和import问题。我在这里就发现了一个问题——自定义collator在deepcopy特征时出现了类型不匹配的错误排查后发现是loss_type字段在copy过程中丢失了需要特殊处理。第二步完整训练曲线对比。用同一个数据集、同样的超参数分别跑原始unsloth和修改后的版本对比loss下降曲线。正常情况应该有微小的随机差异但整体趋势必须一致。如果修改后loss波动异常大多半是labels的mask逻辑有问题。第三步显存占用对比。unsloth最核心的价值就是省显存。我这次改动没有碰kernel层也没有给模型forward加东西理论上显存不应该有明显增加。实测下来效果符合预期训练速度和显存占用基本与改前持平。如果你改了模型层之后显存暴涨那你得回头想想是不是不小心破坏了什么优化。第四步保存—加载闭环测试。把修改后训练的adapter保存下来用一个全新的原版unsloth环境加载跑一次推理看模型行为是否正常。这一步能验证我写的元数据有没有把adapter_config.json弄坏。实测加载正常自定义字段被正确读取没有影响LoRA权重。6.2 几个容易翻车的地方这次改源码踩了不少坑总结几个我认为最有价值的版本兼容性问题是最大的坑。unsloth迭代很快不同版本的UnslothTrainer类内部结构有差异。如果你照网上的教程改源码很可能因为版本不同而找不到对应的代码位置。我的做法是在改之前先把自己的unsloth版本固定住不要让它在训练中途升级。-100标签屏蔽的位置要非常小心。有些模型的自定义损失函数可能不会读取labels里的-100而是直接忽略labels里所有的负值。如果你的代码里自定义了损失函数记得检查它是否遵循-100约定。callback里做生成评估时注意模型状态切换。如果模型用了torch.compile或者unsloth的优化状态在某些版本下model.eval()和model.train()的切换可能会和优化状态有交互。实测中遇到过一次奇怪的问题评估完回到训练模式后某个layer的缓存没有清理干净导致后续loss出现一次尖峰。最终的解决办法是在model.train()之后手动做一次model.zero_grad()。不要为了“改”而改。我这次改的三个点每一个都有明确的使用场景支撑。如果你的项目用默认unsloth就能跑得很顺真没必要动源码。维护一份自定义版本的unsloth是有成本的升级、依赖管理、多人协作都会变得稍微复杂。7. 一点经验最小改动原则与后续维护思路改完这一轮我的体会是对unsloth这种性能敏感、结构紧凑的库最稳妥的改造策略是“最小改动”能加一层就加一层绝不去动底层kernel。这次改动一共就集中在两个文件里training.py里加了collator的覆盖逻辑和一个callbackUnslothTrainer的保存逻辑里加了元数据注入。模型层的文件一个都没碰kernel层更不用说。这样一来unsloth升级的时候我可以清楚地看到自己改了什么手动从一个版本挪到另一个版本的工作量也可控。如果你也准备改unsloth我给一个比较实用的建议把自己所有改动都集中在一个自定义文件里比如unsloth_custom/patches.py统一用mixin或回调方式挂载而不是零散地改原文件。这样做的最大好处是升级unsloth时你只需要重新import这个文件不需要再逐个找改动点。最后再分享一个实用的小技巧改完源码后在训练脚本最前面写一个断言检查自定义版本是否真的被加载了。比如在一个自定义模块里放一个全局常量CUSTOM_UNSLOTH_FLAG v20240615_custom assert CUSTOM_UNSLOTH_FLAG in sys.modules.get(unsloth, ), custom unsloth not loaded!这个断言能在训练启动时快速暴露import路径错误省得你稀里糊涂用错了版本训练跑了一大半才发现改动没生效。这类“调试性断言”在维护自定义版本时非常实用尤其当你同时有好几个项目加载了不同版本的unsloth时。