ms-swift实战：为Qwen2.5模型定制损失函数的进阶指南

在模型微调领域，真正掌握训练过程的核心在于理解如何根据特定需求调整优化目标。本文将带您深入探索ms-swift框架中自定义损失函数的实现细节，从基础概念到实战应用，帮助您在Qwen2.5模型微调中获得更精细的控制能力。

1. 理解自定义损失函数的必要性

当您已经熟悉基础的微调流程后，标准交叉熵损失可能无法满足所有场景需求。自定义损失函数在以下情况尤为关键：

  • 类别不平衡问题：当训练数据中某些类别的样本远多于其他类别时，标准损失函数可能导致模型偏向多数类
  • 特定任务需求：如需要强调某些token的重要性（如实体识别中的关键词）
  • 正则化需求：需要在损失函数中加入特定约束（如L1/L2之外的定制正则项）

ms-swift框架通过compute_loss_func参数提供了灵活的接口，让我们能够在不修改框架核心代码的情况下注入自定义逻辑。这种设计既保持了框架的稳定性，又为高级用户提供了足够的扩展空间。

提示：在开始自定义前，建议先使用默认损失函数完成一次基准训练，这将帮助您评估自定义改进的实际效果

2. 环境准备与基础配置

2.1 安装与依赖

确保您的环境已安装最新版ms-swift和相关依赖：

pip install ms-swift torch>=2.0.0 transformers>=4.40.0

2.2 基础模型加载

以下是加载Qwen2.5-3B-Instruct模型的标准流程：

from swift.llm import get_model_tokenizer
import torch

model_id = 'Qwen/Qwen2.5-3B-Instruct'
model, tokenizer = get_model_tokenizer(
    model_id_or_path=model_id,
    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
    model_kwargs={"device_map": "auto"}
)

2.3 LoRA配置要点

微调大型语言模型时，LoRA（Low-Rank Adaptation）是资源效率最高的方法之一。以下是经过优化的配置：

from swift.tuners import Swift, LoRAConfig

lora_config = LoRAConfig(
    r=16,
    lora_alpha=32,
    target_modules=[
        "q_proj", "k_proj", "v_proj", 
        "o_proj", "gate_proj", 
        "up_proj", "down_proj"
    ],
    lora_dropout=0.05,
    bias="none"
)
model = Swift.prepare_model(model, lora_config)

3. 自定义损失函数的设计与实现

3.1 基础损失函数结构

所有自定义损失函数都应继承自torch.nn.Module，并实现forward方法。以下是基本框架：

import torch.nn as nn
import torch.nn.functional as F

class CustomLoss(nn.Module):
    def __init__(self, alpha=0.5):
        super().__init__()
        self.alpha = alpha  # 自定义超参数
        self.ce_loss = nn.CrossEntropyLoss()
        
    def forward(self, outputs, labels, **kwargs):
        logits = outputs['logits'] if isinstance(outputs, dict) else outputs
        mask = kwargs.get('attention_mask', None)
        
        # 基础交叉熵计算
        loss = self.ce_loss(
            logits.view(-1, logits.size(-1)),
            labels.view(-1)
        )
        
        return loss

3.2 处理Attention Mask

当处理变长序列时，正确处理padding部分至关重要：

def forward(self, outputs, labels, **kwargs):
    logits = outputs['logits']
    mask = kwargs.get('attention_mask', None)
    
    if mask is not None:
        # 计算每个token的loss
        per_token_loss = F.cross_entropy(
            logits.view(-1, logits.size(-1)),
            labels.view(-1),
            reduction='none'
        ).view_as(labels)
        
        # 应用mask
        valid_tokens = mask.sum()
        loss = (per_token_loss * mask).sum() / (valid_tokens + 1e-8)
    else:
        loss = self.ce_loss(
            logits.view(-1, logits.size(-1)),
            labels.view(-1)
        )
    
    return loss

3.3 实现Focal Loss解决类别不平衡

对于类别不平衡问题，Focal Loss是有效的解决方案：

class FocalLoss(nn.Module):
    def __init__(self, gamma=2.0, alpha=None):
        super().__init__()
        self.gamma = gamma
        self.alpha = alpha
        
    def forward(self, outputs, labels, **kwargs):
        logits = outputs['logits']
        mask = kwargs.get('attention_mask', None)
        
        # 计算softmax概率
        probs = F.softmax(logits, dim=-1)
        # 获取目标类别的概率
        class_probs = probs.gather(-1, labels.unsqueeze(-1)).squeeze(-1)
        
        # 计算focal loss
        focal_loss = -((1 - class_probs) ** self.gamma) * torch.log(class_probs + 1e-8)
        
        if mask is not None:
            valid_tokens = mask.sum()
            loss = (focal_loss * mask).sum() / (valid_tokens + 1e-8)
        else:
            loss = focal_loss.mean()
            
        return loss

4. 高级技巧与实战应用

4.1 组合多个损失函数

实际项目中，我们经常需要组合多个损失目标：

class CompositeLoss(nn.Module):
    def __init__(self, loss_weights=None):
        super().__init__()
        self.loss_weights = loss_weights or {
            'ce': 1.0,
            'kl': 0.1,
            'l2': 0.01
        }
        
    def forward(self, outputs, labels, **kwargs):
        # 基础交叉熵
        ce_loss = F.cross_entropy(
            outputs['logits'].view(-1, outputs['logits'].size(-1)),
            labels.view(-1)
        )
        
        # KL散度正则项
        kl_loss = F.kl_div(
            F.log_softmax(outputs['logits'], dim=-1),
            F.softmax(outputs['teacher_logits'], dim=-1),
            reduction='batchmean'
        )
        
        # L2正则
        l2_loss = sum(p.pow(2.0).sum() for p in model.parameters())
        
        total_loss = (
            self.loss_weights['ce'] * ce_loss +
            self.loss_weights['kl'] * kl_loss +
            self.loss_weights['l2'] * l2_loss
        )
        
        return total_loss

4.2 调试与验证技巧

集成自定义损失函数时，以下调试方法非常有用：

  1. 梯度检查：

    # 在训练循环中添加
for name, param in model.named_parameters():
    if param.grad is not None:
        print(f"{name} grad mean: {param.grad.mean().item()}")

  2. 损失值监控：

    # 自定义Trainer回调
from transformers import TrainerCallback

class LossLoggingCallback(TrainerCallback):
    def on_step_end(self, args, state, control, **kwargs):
        if state.global_step % 100 == 0:
            print(f"Step {state.global_step}: Loss {state.log_history[-1]['loss']}")

  3. 数值稳定性检查：

    def forward(self, outputs, labels, **kwargs):
    logits = outputs['logits']
    if torch.isnan(logits).any():
        raise ValueError("NaN detected in logits")
    # ...其余计算

4.3 性能优化建议

大规模训练时，这些优化可以显著提升效率：

  • 使用混合精度训练：

    training_args = TrainingArguments(
    fp16=True,  # 或bf16=True
    # 其他参数...
)

  • 内存优化技巧：

    # 在自定义损失中避免不必要的中间变量
loss = F.cross_entropy(
    outputs['logits'].view(-1, outputs['logits'].size(-1)),
    labels.view(-1),
    reduction='none'
)
# 立即应用mask，避免存储完整矩阵
if mask is not None:
    loss = (loss.view_as(labels) * mask).sum() / mask.sum()

5. 完整训练流程示例

将以上组件整合为完整工作流：

from swift import Trainer, TrainingArguments
from swift.llm import get_template

# 准备模板
template = get_template(
    model.model_meta.template,
    tokenizer,
    max_length=512
)
template.set_mode('train')

# 训练参数
training_args = TrainingArguments(
    output_dir="./output",
    per_device_train_batch_size=4,
    learning_rate=3e-5,
    num_train_epochs=3,
    logging_steps=100,
    save_steps=1000,
    fp16=torch.cuda.is_available(),
)

# 启用梯度
model.enable_input_require_grads()

# 初始化Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    compute_loss_func=FocalLoss(gamma=2.0),
    template=template,
)

# 开始训练
trainer.train()

# 保存模型
model.save_pretrained("./final_model")

在实际项目中，我发现Focal Loss的参数γ=2.0和α=0.25的组合在大多数文本分类任务中表现良好。对于特别长的序列（超过512 tokens），建议在损失计算前检查attention mask的有效性，避免因大量padding导致损失计算不稳定。

# LoRA
Logo
亚马逊云科技技术品牌专区

更多推荐

cover

别再只聊大模型！AI的下一场革命，藏在“身体”里

avatar 亚马逊云科技技术品牌专区
cover

别再只聊大模型！人工智能的下一场革命，藏在“身体”里

avatar 亚马逊云科技技术品牌专区
cover

别再只聊大模型！AI的下一场革命，藏在“身体”里

avatar 亚马逊云科技技术品牌专区