插件化扩展教程：如何在ms-swift中自定义loss函数和optimizer
在大模型训练日益复杂的今天，一个“万能但僵硬”的框架已经难以满足多样化任务的需求。无论是做指令微调、人类偏好对齐（如DPO、KTO），还是尝试最新的低秩优化技术（如GaLore），研究人员和工程师都希望不改源码、快速验证新想法。
这正是插件化架构的价值所在——它让框架像乐高一样可拼装。而ms-swift作为魔搭社区推出的大规模模型训练与部署一体化平台，早已将这一理念贯彻到底。其对loss和optimizer的灵活替换能力，正是支撑算法创新与工程落地的关键底座。
我们不妨从一个问题出发：假设你正在训练一个对话模型，目标不是简单地预测下一个token，而是让模型学会区分“好回答”和“坏回答”。传统的交叉熵损失显然不够用了——你需要一种能感知偏好的损失函数。
怎么办？重写整个Trainer？当然不用。ms-swift允许你只写几行代码，定义一个新的compute_loss逻辑，然后把它“插”进训练流程里。这就是所谓的自定义Loss函数。
同理，当你面对百亿参数模型显存爆满的窘境时，也不必死磕AdamW。你可以引入像GaLore这样的轻量级优化器，通过对梯度做低秩投影来大幅降低内存占用——而且无需修改核心训练循环，只需换上另一个“插头”。
这种高度解耦的设计思路，使得ms-swift既能保持主干稳定，又能支持最前沿的科研探索。
自定义Loss：不只是计算差异，更是建模学习信号
在深度学习中，loss函数远不止是“算个误差”那么简单。它是引导模型学习方向的指挥棒。标准的交叉熵适用于分类任务，但在更复杂的场景下，我们需要更精细的控制。
比如，在知识蒸馏或偏好学习中，标签不再是非黑即白的类别，而是带有强度信息的连续值（例如用户打分0.8 vs 0.3）。这时如果还用普通BCELoss，就会忽略样本之间的相对质量差异。
于是我们可以设计一个类似KTO风格的加权损失：
class CustomKtoLoss: def __init__(self, beta: float = 0.1): self.beta = beta self.bce_loss = nn.BCEWithLogitsLoss(reduction='none') def compute_loss(self, model, inputs: Dict[str, torch.Tensor]) -> torch.Tensor: input_ids = inputs["input_ids"] labels = inputs["labels"] # 归一化后的偏好分数 [0,1] attention_mask = inputs["attention_mask"] outputs = model(input_ids=input_ids, attention_mask=attention_mask) logits = outputs.logits[:, -1] # 取最后一个token的预测值 probs = torch.sigmoid(logits) # 动态加权：高质量样本赋予更高权重 pos_weight = 1.0 + torch.exp(-self.beta * labels) neg_weight = 1.0 + torch.exp(self.beta * (1 - labels)) per_sample_loss = self.bce_loss(logits, labels) weighted_loss = pos_weight * per_sample_loss return weighted_loss.mean()这个小小的改动带来了显著变化：模型不再平均对待所有样本，而是更关注那些“明显更好”的回答。实践中你会发现，收敛速度更快，生成结果也更具一致性。
关键在于，这个类只需要实现一个方法——compute_loss(model, inputs)，返回一个标量tensor即可。ms-swift的Trainer会自动接管后续的反向传播和参数更新。你不需要操心分布式训练、混合精度或者梯度裁剪，这些都被封装好了。
小贴士：自定义loss中最容易出错的是设备不一致问题。务必确保所有张量都在同一设备（如GPU）上。另外，不要在loss中手动调用
.zero_grad()或.step()，这些由Trainer统一管理。
如果说loss决定了“学什么”，那么optimizer就决定了“怎么学”。
传统优化器如AdamW为每个参数维护动量和方差状态，导致显存消耗通常是模型本身的2~3倍。这对于几十亿甚至上百亿参数的模型来说，几乎不可承受。
有没有办法减少这部分开销？有——比如最近火出圈的GaLore（Gradient Low-Rank Projection）。
它的核心思想很简单：大多数全连接层的梯度具有低内在秩（intrinsic low rank），也就是说，可以用一个小得多的子空间来近似表示。于是我们可以在更新前先对梯度做一次投影，在低维空间中进行优化，再映射回原空间。
下面是一个简化版的实现：
import torch from torch.optim import Optimizer class SimpleGaloreOptimizer(Optimizer): def __init__(self, params: Iterable[torch.nn.Parameter], lr: float = 1e-3, rank: int = 128, alpha: float = 0.75): defaults = dict(lr=lr, rank=rank, alpha=alpha) super().__init__(params, defaults) self.W_resid = {} for group in self.param_groups: for p in group['params']: if p.requires_grad and p.dim() > 1: self.init_galore_projection(p, group['rank']) def init_galore_projection(self, param: torch.Tensor, rank: int): rows, cols = param.shape device = param.device dtype = param.dtype if rows >= cols: U = torch.empty(cols, cols, device=device, dtype=dtype) torch.linalg.qr(U, out=(U, _)) self.state[param]['projector'] = U[:, :rank].contiguous() else: U = torch.empty(rows, rows, device=device, dtype=dtype) torch.linalg.qr(U, out=(U, _)) self.state[param]['projector'] = U[:rank, :].contiguous() @torch.no_grad() def step(self, closure=None): loss = None if closure is not None: with torch.enable_grad(): loss = closure() for group in self.param_groups: lr = group['lr'] for p in group['params']: if p.grad is None or not p.requires_grad: continue grad = p.grad.data state = self.state[p] if 'projector' in state and grad.dim() > 1: proj = state['projector'] if grad.size(0) >= grad.size(1): update_flat = (grad @ proj) * lr update = update_flat @ proj.T else: update_flat = (proj @ grad) * lr update = proj.T @ update_flat p.data -= update else: p.data -= lr * grad return loss这段代码虽然短，却包含了GaLore的核心机制：构造正交投影矩阵、判断矩阵形状以决定左右乘顺序、仅对高维参数启用投影等。
更重要的是，它完全兼容PyTorch的Optimizer协议，因此可以直接传给ms-swift的Trainer：
trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, optimizers=(SimpleGaloreOptimizer(model.parameters(), lr=5e-5), None), compute_loss=CustomKtoLoss(beta=0.2), )就这么简单，你的训练就已经运行在一个显存更友好、收敛更稳定的优化路径上了。
当然，如果你不想自己实现，ms-swift也内置了对GaLore、Q-Galore等先进优化器的支持，只需通过配置文件一键开启：
# config.yaml optimizer_type: galore galore_rank: 128 galore_update_interval: 50 galore_scale: 0.1然后在初始化Trainer时不传optimizers参数，框架会自动根据配置加载对应优化器。
实际应用中的几个典型场景
场景一：偏好对齐任务中传统loss收敛慢
很多团队在做DPO或KTO时发现，模型很难稳定地区分优劣回答。原因就在于标准loss没有建模“差距程度”——两个回答哪怕差距很大，loss也只当作一对正负样本处理。
解决方案就是使用带隐式奖励建模的loss，比如上面提到的KTO-style加权loss，或者SimPO这类基于margin的设计。它们能让模型更敏感地捕捉到质量差异，从而加快收敛。
场景二：超大模型训练显存不足
当你训练一个百亿参数以上的模型时，AdamW带来的额外显存开销可能直接让你无法增大batch size。这时候切换到GaLore类优化器，往往能带来40%~60%的显存下降，相当于多出一张卡的容量。
尤其在多节点训练中，这种节省非常可观。而且由于投影本身是无损近似的，性能通常不会下降，有时反而因为更稳定的更新而略有提升。
场景三：稀疏更新或特定层定制策略
有些任务只需要微调部分层（如LoRA中的适配器），其他层冻结。这时你可以结合参数分组，在自定义optimizer中为不同层设置不同的学习率或更新方式。
例如：
# 只对包含'lora_'的参数启用GaLore filtered_params = (p for n, p in model.named_parameters() if 'lora_' in n and p.requires_grad) optimizer = SimpleGaloreOptimizer(filtered_params, lr=1e-4)这样既保证了关键模块的高效更新，又避免了不必要的计算开销。
设计背后的思考：为什么插件化如此重要？
一个好的训练框架，应该像操作系统一样：内核稳定可靠，外设自由扩展。ms-swift正是朝着这个方向演进。
通过开放compute_loss和optimizers这两个接口，它实现了真正的“策略与执行分离”：
- 科研人员可以专注于新算法的设计，而不必陷入工程细节；
- 工程师可以通过配置文件快速部署最优方案，提高复现性和可维护性；
- 企业用户可以在同一套流程下管理多种任务类型，降低运维复杂度。
而且这种扩展是安全的——自定义逻辑被隔离在独立组件中，即使出错也不会破坏主干流程。建议的做法是在关键位置添加日志和异常捕获：
def compute_loss(self, model, inputs): try: # your custom logic return loss except Exception as e: print(f"[Loss Error] {str(e)}") raise此外，强烈建议为自定义组件编写单元测试，尤其是检查梯度是否正常流动：
# 测试示例 def test_custom_loss(): model = YourModel() inputs = { "input_ids": torch.randint(0, 1000, (2, 10)), "labels": torch.rand(2, 1), "attention_mask": torch.ones(2, 10) } loss_fn = CustomKtoLoss() loss = loss_fn.compute_loss(model, inputs) assert loss.requires_grad loss.backward() # 确保能反向传播真正强大的框架，不是因为它功能最多，而是因为它允许别人让它变得更强。
ms-swift通过对loss和optimizer的插件化支持，把“创新能力”交还给了开发者。你可以用它跑通标准微调，也可以用来验证最新的论文方法；可以用于小规模实验，也能支撑超大规模训练。
这种灵活性的背后，是一套清晰的抽象：只要遵循compute_loss接口，任何损失都能接入；只要符合torch.optim.Optimizer规范，任何更新策略都能运行。
掌握这一点，你就不再只是一个使用者，而成了框架的共建者。而这，或许才是推动AI技术持续前进的真正动力。
