如何在 Llama-Factory 中自定义损失函数

这段代码的关键在于，它没有改动任何训练流程，只是替换了损失计算部分。你可以把它看作一个'插槽'——只要返回的是标量 loss，PyTorch 就能自动完成反向传播。这意味着你的自定义逻辑可以非常复杂，比如引入对比学习项、KL 散度约束，甚至是基于外部奖励的强化学习目标。

更重要的是，Llama-Factory 提供了配置驱动的加载机制。你不需要修改主程序，只需将上述类保存为 trainers/custom_trainer.py，然后在 YAML 配置中声明：

trainer_type: custom
custom_trainer_path: ./trainers/custom_trainer.py
label_smoothing: 0.1

框架会在初始化时动态导入并实例化你的 CustomTrainer，自动注入所有配置参数。这种插件式架构让实验迭代变得极其高效：换损失就像换电池一样简单。

但别忘了，灵活性也意味着责任。当你跳出默认路径时，有几个坑必须警惕。

首先是梯度稳定性。如果你在损失中加入了复杂的数学运算，比如除法、对数或指数操作，稍不注意就会导致 NaN 或梯度爆炸。建议始终用 torch.clamp 对输入做裁剪，并在调试阶段开启 torch.autograd.set_detect_anomaly(True) 来捕捉异常源头。

其次是分布式训练兼容性。在多 GPU 场景下，每个设备只看到一部分 batch。如果你在损失中做了全局归一化或统计量计算（如均值、方差），必须确保这些值是在所有设备上同步聚合过的。否则，梯度更新会不一致。好在 Llama-Factory 默认使用 DistributedDataParallel，你可以借助 torch.distributed.all_reduce 手动同步张量，或者干脆避免跨设备依赖。

再来看内存效率。长序列任务中，一次性展开所有 token 的 logits 和 labels 可能占用巨大显存。例如，一个 batch size 为 8、序列长度为 8192 的输入，展平后的形状是 (8*8192, vocab_size)，对于 32K 词表来说就是近 2GB 的中间张量。解决办法是分块计算或使用 reduction='none' 后按需降维，而不是盲目 .mean()。

还有一个常被忽视的点是日志可解释性。当你加了权重、平滑或多个损失项时，最终的 loss 值已经不能直接和原始交叉熵比较了。建议在训练日志中同时输出原始 loss 和加权后的 total loss，方便分析收敛行为。Llama-Factory 支持 TensorBoard，你可以轻松记录这些辅助指标：

if self.args.local_rank == 0:  # 主进程记录
    self.log({"base_loss": base_loss.item(), "weighted_loss": weighted_loss.item()})

说到实际应用，我们再回到那个金融客服的例子。假设你有一组标注好的问题类别，其中'退款政策'、'账户安全'等属于高优先级。与其靠数据过采样来提升曝光，不如直接在损失层面赋予它们更高权重：

def compute_loss(self, model, inputs):
    labels = inputs["labels"]
    category_ids = inputs.get("category_id", None)
    outputs = model(**inputs)
    logits = outputs["logits"]
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()
    loss_per_token = self.ce_loss(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
    
    # Reshape to [batch_size, seq_len] and average over sequence
    loss_per_sample = loss_per_token.view(labels.size(0), -1).mean(dim=1)
    
    if category_ids is not None:
        class_weights = {
            0: 1.0,   # login_issue
            1: 5.0,   # refund_policy
            2: 8.0,   # account_frozen
            3: 1.5    # feature_request
        }
        weights = torch.tensor([class_weights[cid.item()] for cid in category_ids], device=loss_per_sample.device)
        loss_per_sample = loss_per_sample * weights
    
    return loss_per_sample.mean()

这种方法的优势在于，它不改变数据分布，避免了因重复采样带来的噪声放大；同时又能精准地将业务意图转化为可优化目标。而且，权重参数完全可以从配置文件读取，做到代码与策略解耦。

类似的思路还能拓展到更多高级场景：

• Focal Loss：抑制易分类样本的贡献，聚焦难例；
• Contrastive Loss：在检索增强问答中拉近 query 与 positive passage 的表示距离；
• KL Div Loss：在蒸馏任务中对齐教师模型与学生模型的输出分布；
• Multi-task Learning：联合优化生成任务和分类任务，共享主干网络。

这些都不是理论设想，而是已经在推荐系统、医疗诊断、法律文书生成等领域落地的技术实践。关键在于，你是否拥有一个足够开放的框架来承载这些创新。

Llama-Factory 的真正优势，不只是支持 LoRA、QLoRA 这些热门技术，而是它把整个微调链条打开给你看，并告诉你：'这里也可以改。' 它的设计哲学很明确：通用性解决共性问题，可扩展性应对个性需求。

这也解释了为什么它能在众多微调工具中脱颖而出。相比 Alpaca-LoRA 这类脚本型项目，它提供了 WebUI 和模块化 API；相比纯 CLI 工具，它又保留了深度定制的空间。无论是想快速验证想法的研究者，还是需要稳定交付的企业开发者，都能找到自己的位置。

未来的大模型训练，不会停留在'喂数据、调 learning rate'的层面。随着应用场景越来越复杂，我们需要的是语义感知的优化目标、任务感知的损失结构，甚至是用户反馈驱动的动态调整机制。而这一切的起点，往往就是一个被重新定义的 compute_loss 方法。

当你能把'这个问题很重要'翻译成'这个样本的损失要翻倍'，你就不再只是在训练模型，而是在塑造它的价值观。这才是高级微调的真正意义。