自定义Loss函数与Optimizer扩展教程，打造专属训练流程-平芜编程栈

自定义Loss函数与Optimizer扩展教程，打造专属训练流程

在大模型时代，通用训练范式正逐渐让位于“按需定制”的精细化调优策略。无论是金融领域的合规对话系统，还是医疗场景中的专业问答引擎，标准交叉熵损失和AdamW优化器往往难以满足复杂任务的需求。开发者越来越需要一种灵活机制：既能引入前沿对齐算法如DPO、KTO，又能集成低秩优化技术如GaLore，在不牺牲效率的前提下实现性能跃升。

魔搭社区推出的ms-swift框架正是为此而生。它不仅支持600+纯文本大模型与300+多模态模型的端到端训练部署，更通过其插件化架构，允许用户无缝扩展核心组件——从Loss、Metric到Trainer、Callback乃至Optimizer。这种设计打破了传统框架“黑箱化”的局限，让研究人员可以在统一平台上构建真正意义上的“专属训练流程”。

本文将聚焦于两个最具影响力的可扩展点：自定义Loss函数与自定义Optimizer，深入剖析其实现逻辑，并结合代码示例展示如何将其应用于真实训练场景。

灵活建模训练目标：自定义Loss函数的设计与实践

当面对偏好对齐、多任务学习或稀疏监督等复杂任务时，标准损失函数常常显得力不从心。例如，在人类反馈强化学习（RLHF）中，我们并不直接拥有标签，而是知道一个回答优于另一个。此时，传统的分类损失无法建模这种相对排序关系，必须依赖专门设计的目标函数。

ms-swift 的解决方案是将损失计算抽象为模块化组件，允许用户通过继承nn.Module并注册的方式注入自定义逻辑。整个流程嵌入在Trainer的训练步中：

DataLoader 提供一批结构化样本（如chosen_input_ids,rejected_input_ids）；
模型分别前向传播得到优选与劣选响应的logits；
自定义Loss接收这些输出，执行特定计算并返回标量损失；
反向传播基于该loss更新参数。

这种方式既保持了PyTorch原生的灵活性，又确保了与LoRA、QLoRA、FSDP等微调与分布式技术的兼容性。

以下是一个简化版 DPO（Direct Preference Optimization）损失的实现：

import torch import torch.nn as nn from typing import Dict, Any class CustomDpoLoss(nn.Module): """ 自定义DPO风格损失函数：直接优化偏好对数据中的胜者与败者响应 ref: https://arxiv.org/abs/2305.18290 """ def __init__(self, beta: float = 0.1): super().__init__() self.beta = beta # 温度系数，控制KL约束强度 def forward(self, policy_chosen_logps: torch.Tensor, policy_rejected_logps: torch.Tensor) -> Dict[str, torch.Tensor]: """ Args: policy_chosen_logps: 模型对优选回答的对数概率 (batch_size,) policy_rejected_logps: 模型对劣选回答的对数概率 (batch_size,) Returns: 包含loss、accuracy等指标的字典 """ losses = -torch.log( torch.sigmoid(self.beta * ( policy_chosen_logps - policy_rejected_logps )) ) reward_accuracies = (policy_chosen_logps > policy_rejected_logps).float().mean() return { "loss": losses.mean(), "reward_accuracy": reward_accuracies } # 注册到 ms-swift 框架 def register_custom_loss(): from swift.torchkit.trainer import SwiftTrainer SwiftTrainer.register_loss('dpo', CustomDpoLoss(beta=0.1))

这个实现的关键在于利用sigmoid函数将偏好判断转化为概率形式，同时通过beta调节KL散度项的影响程度，防止策略偏离参考模型过远。值得注意的是，返回值不仅包含主损失，还附带了reward_accuracy这一监控指标，便于实时观察模型是否正在学会区分优劣回答。

⚠️ 实践建议：
- 避免在计算过程中使用.detach()中断梯度流；
- 多卡训练下应确保 loss 为全局平均后的标量，必要时使用all_reduce同步；
- 建议添加 NaN 检查，如assert not torch.isnan(losses).any()，防止训练崩溃。

一旦完成注册，在配置文件中只需设置"loss_type": "dpo"即可启用该损失函数，无需修改任何主干代码。

降低显存开销：自定义Optimizer的工程实现路径

随着模型规模突破百亿甚至千亿参数，优化器状态（如Adam中的动量和方差）所占用的显存已成为训练瓶颈。以AdamW为例，每个参数需额外存储4字节（float32）的动量 + 4字节的方差，整体内存开销可达模型权重本身的两倍以上。

GaLore（Gradient Low-Rank Projection）提出了一种极具启发性的思路：并非所有参数都需要全精度更新。对于线性层权重这类高维张量，其梯度往往具有低内在秩特性，可通过SVD分解投影至低维子空间进行更新，从而大幅压缩优化器状态。

ms-swift 支持此类先进优化器的核心在于其参数分组机制与优化器工厂模式。用户只需实现标准优化器接口（step,zero_grad），即可被 Trainer 动态加载和调用。

下面是一个简化的 GaLore 实现：

import torch from torch.optim import Optimizer from typing import Iterable, Dict class GaLoreProjector: def __init__(self, rank: int = 128, update_proj_gap: int = 50, scale: float = 1.0): self.rank = rank self.update_proj_gap = update_proj_gap self.scale = scale self.proj_p = None self.step = 0 def project(self, grad: torch.Tensor) -> torch.Tensor: if self.step % self.update_proj_gap == 0: if grad.dim() >= 2: U, S, Vt = torch.linalg.svd(grad, full_matrices=False) self.proj_p = Vt[:self.rank].T # 保留右奇异向量 self.step += 1 return torch.mm(grad.flatten(start_dim=1), self.proj_p).detach() def project_back(self, grad_proj: torch.Tensor) -> torch.Tensor: return torch.mm(grad_proj, self.proj_p.T).view_as(grad_proj) class GaLoreOptimizer(Optimizer): """ 实现 GaLore (Gradient Low-Rank Projection) 优化器 ref: https://arxiv.org/abs/2310.04417 """ def __init__(self, params: Iterable[Dict], base_optimizer: str = 'adamw', **kwargs): defaults = {k: v for k, v in kwargs.items()} super().__init__(params, defaults) self.base_optimizer_name = base_optimizer self.projectors = {} def step(self, closure=None): loss = None if closure is not None: loss = closure() for group in self.param_groups: for p in group['params']: if p.grad is None: continue # 初始化投影器 if p not in self.projectors: self.projectors[p] = GaLoreProjector(rank=128) projector = self.projectors[p] grad_proj = projector.project(p.grad.data) # 使用基础优化器更新低维表示 state = self.state[p] if len(state) == 0: state['grad_proj'] = torch.zeros_like(grad_proj) state['grad_proj'].copy_(grad_proj) # 模拟AdamW更新（简化示例） state['grad_proj'] -= 1e-4 * state['grad_proj'] # dummy update # 映射回原始空间 p.data.add_(projector.project_back(state['grad_proj']), alpha=-group['lr']) return loss # 注册接口 def register_galore(): from swift.torchkit.optimizer import SwiftOptimizer SwiftOptimizer.register('galore', GaLoreOptimizer)

该实现中，GaLoreProjector负责将原始梯度投影到由右奇异向量张成的低秩空间，随后在低维空间执行参数更新，最后再映射回原形。实验表明，即使将秩设为128，也能在多数任务上保持接近全秩更新的性能，但显存消耗可降低数十倍。

⚠️ 注意事项：
- 仅适用于nn.Linear等矩阵型参数，不推荐用于偏置或 LayerNorm；
- 投影频率不宜过高，一般每50步更新一次投影矩阵即可维持稳定性；
- 若在多卡环境下共享投影矩阵，需注意跨设备同步问题；
- 当前不兼容GPTQ等静态量化方案，建议关闭相关功能。

注册后，只需在YAML配置中指定"optimizer": "galore"，即可自动应用该优化策略。

构建完整训练闭环：从理论到落地的应用整合

在实际项目中，自定义Loss与Optimizer往往协同工作，共同构成高效的训练流水线。以一个多模态人类对齐任务为例，系统架构如下所示：

graph TD A[Dataset] --> B[Data Collator] B --> C[Model Forward → Logits] C --> D{Custom Loss Module} D -->|policy_chosen_logps| E[DPO Loss] D -->|policy_rejected_logps| E E --> F[Loss Scalar] F --> G{Custom Optimizer} G --> H[GaLore Update] H --> I[Parameter Update] I --> C J[Trainer Control Loop] --> C J --> G

具体流程如下：

用户准备图文混合的偏好数据集（图像 + 问题 + 优选/劣选回答）；
加载 Qwen-VL-Chat 模型并启用 LoRA 微调；
在配置文件中声明：
yaml loss_type: "dpo" optimizer: "galore"
框架自动查找已注册的CustomDpoLoss与GaLoreOptimizer；
训练过程中：
- 模型分别编码优选与劣选回答，提取其对数概率；
- DPO损失引导模型增强优选响应的生成倾向；
- GaLore将注意力层的梯度投影至低秩空间，减少优化器内存占用；
最终输出具备偏好对齐能力且体积轻量的模型。

这种组合策略已在多个实际场景中验证有效：

实际痛点	解决方案
百亿级以上模型微调显存不足	GaLore + QLoRA联合压缩，优化器状态降至原大小1/10以下
标准CE损失无法捕捉人类偏好	引入DPO/SimPO类损失，直接建模排序行为
框架封闭难集成新算法	插件化机制支持热插拔，无需Fork仓库
多任务目标冲突	设计加权复合Loss（如语言建模 + 视觉对比），动态调整权重

此外，在工程实践中还需关注一些关键细节：