ModelScope/SWIFT项目插件化开发指南

谭妲茹

于 2025-06-05 09:14:28 发布

阅读量345

点赞数 4

CC 4.0 BY-SA版权

本文链接：https://blog.youkuaiyun.com/gitblog_00485/article/details/148442994

ModelScope/SWIFT项目插件化开发指南

swift 魔搭大模型训练推理工具箱，支持LLaMA、千问、ChatGLM、BaiChuan等多种模型及LoRA等多种训练方式(The LLM training/inference framework of ModelScope community, Support various models like LLaMA, Qwen, Baichuan, ChatGLM and others, and training methods like LoRA, ResTuning, NEFTune, etc.) 项目地址: https://gitcode.com/gh_mirrors/swift1/swift

引言

ModelScope/SWIFT作为一个高效的深度学习训练框架，在3.0版本中引入了插件化(Pluginization)机制，这一创新设计极大地提升了框架的灵活性和可扩展性。本文将深入解析SWIFT框架中的插件化机制，帮助开发者掌握如何通过插件方式定制训练流程的各个环节。

插件化机制概述

插件化是SWIFT框架的核心设计理念之一，它允许开发者在不修改框架核心代码的情况下，通过实现特定接口来定制训练流程。这种设计带来了以下优势：

解耦核心代码与定制逻辑
便于功能扩展和维护
支持多种训练场景的灵活切换
降低二次开发的学习成本

核心插件类型详解

1. 回调机制(Callback)插件

回调机制是控制训练流程的重要方式，开发者可以通过实现TrainerCallback类来干预训练过程的关键节点。

class CustomCallback(TrainerCallback):
    def on_train_begin(self, args, state, control, **kwargs):
        """训练开始时执行的操作"""
        print("训练开始，准备初始化...")
        
    def on_epoch_end(self, args, state, control, **kwargs):
        """每个epoch结束时执行的操作"""
        if state.epoch % 2 == 0:
            print(f"第{state.epoch}个epoch完成")

典型应用场景包括：

实现早停机制(EarlyStopping)
自定义日志记录
训练过程监控
动态调整学习率

2. 损失函数(Loss)插件

SWIFT框架支持自定义损失函数，这对于特定任务优化至关重要。

@register_loss_func("focal_loss")
def focal_loss(outputs, labels, loss_scale=None, num_items_in_batch=None):
    """实现Focal Loss以处理类别不平衡问题"""
    ce_loss = F.cross_entropy(outputs, labels, reduction='none')
    pt = torch.exp(-ce_loss)
    loss = (1 - pt) ** 2 * ce_loss  # γ=2
    return loss.mean()

注意事项：

目前仅支持PT(预训练)和SFT(监督微调)任务
分类任务(seq_cls)和人类对齐任务(DPO/PPO)不支持自定义损失
需要确保损失函数与任务类型匹配

3. 损失权重(Loss Scale)插件

损失权重机制允许开发者对不同token赋予不同的训练权重，这在处理特定任务时非常有用。

class KeywordEmphasisLossScale(LossScale):
    """对特定关键词赋予更高权重"""
    
    KEYWORDS = ["重要", "关键", "必须"]
    
    def get_loss_scale(self, context, context_type, is_last_round, **kwargs):
        if context_type == ContextType.RESPONSE:
            tokens = jieba.lcut(context)
            scales = [2.0 if token in self.KEYWORDS else 1.0 for token in tokens]
            return tokens, scales
        return super().get_loss_scale(context, context_type, is_last_round)

应用场景：

强调对话系统中的关键回复
增强特定领域术语的学习
实现渐进式学习策略

4. 评估指标(Metric)插件

自定义评估指标可以帮助开发者更准确地衡量模型性能。

def compute_rouge_metrics(eval_preds):
    """计算ROUGE指标"""
    preds, labels = eval_preds
    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
    
    rouge = Rouge()
    scores = rouge.get_scores(decoded_preds, decoded_labels, avg=True)
    return {
        'rouge-1': scores['rouge-1']['f'],
        'rouge-2': scores['rouge-2']['f'],
        'rouge-l': scores['rouge-l']['f']
    }

METRIC_MAPPING = {
    'rouge': (compute_rouge_metrics, None),
}

5. 优化器(Optimizer)插件

SWIFT支持自定义优化器和学习率调度器，满足不同训练需求。

def create_adamw_with_warmup(args, model, dataset):
    """创建带warmup的AdamW优化器"""
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters() 
                      if not any(nd in n for nd in no_decay)],
            "weight_decay": args.weight_decay,
        },
        {
            "params": [p for n, p in model.named_parameters() 
                      if any(nd in n for nd in no_decay)],
            "weight_decay": 0.0,
        }
    ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate)
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=args.warmup_steps, 
        num_training_steps=args.max_steps
    )
    return optimizer, scheduler

optimizers_map = {
    'adamw_warmup': create_adamw_with_warmup,
}

6. 调谐器(Tuner)插件

调谐器是SWIFT的特色功能，支持开发者灵活定义模型微调策略。

class AdapterTuner(Tuner):
    """实现Adapter微调策略"""
    
    @staticmethod
    def prepare_model(args, model):
        """配置Adapter参数"""
        config = AdapterConfig(
            dim=model.config.hidden_size,
            reduction_factor=args.adapter_reduction_factor,
            non_linearity=args.adapter_non_linearity
        )
        return add_adapter(model, config)
    
    @staticmethod
    def save_pretrained(model, save_directory, **kwargs):
        """保存Adapter参数"""
        save_adapter(model, save_directory, **kwargs)
    
    @staticmethod
    def from_pretrained(model, model_id, **kwargs):
        """加载Adapter参数"""
        return load_adapter(model, model_id, **kwargs)

7. 奖励模型插件

SWIFT支持两种奖励模型，用于强化学习场景：

过程奖励模型(PRM)

class StepwisePRM(PRM):
    """分步推理过程奖励模型"""
    
    def __call__(self, infer_requests, **kwargs):
        rewards = []
        for request in infer_requests:
            response = request.messages[-1]['content']
            steps = response.split('\n')
            # 根据推理步骤质量计算奖励
            reward = min(len(steps) * 0.1, 1.0)  
            rewards.append(reward)
        return rewards

结果奖励模型(ORM)

class CodeCorrectnessORM(ORM):
    """代码正确性评估模型"""
    
    def __call__(self, infer_requests, ground_truths, **kwargs):
        rewards = []
        for req, gt in zip(infer_requests, ground_truths):
            # 执行代码并比较输出
            user_code = req.messages[-1]['content']
            gt_output = execute_code(gt)
            user_output = execute_code(user_code)
            rewards.append(1.0 if user_output == gt_output else 0.0)
        return rewards