插件化扩展教程：如何在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技术持续前进的真正动力。