XTuner项目中自定义优化器与学习率分组的实践指南

背景介绍

XTuner作为基于InternLM的训练框架，在模型微调场景中发挥着重要作用。在实际应用中，我们经常需要对模型不同部分设置差异化的训练参数，特别是当模型包含视觉编码器(ViT)和语言模型(LLM)两个部分时，设置不同的学习率是常见的需求。

优化器配置的进阶用法

在XTuner框架中，优化器配置通过optim_wrapper参数进行控制。标准的配置方式如下：

optim_wrapper = dict(
    type=AmpOptimWrapper,
    optimizer=dict(type=AdamW, lr=1e-4, weight_decay=0.01),
    clip_grad=dict(max_norm=1.0),
    accumulative_counts=4,
    loss_scale='dynamic',
    dtype='float16'
)

实现差异化学习率配置

对于包含多模态组件的模型(如同时包含ViT和LLM的LLaVA)，我们可以通过paramwise_cfg参数实现不同模块的学习率差异化设置：

optim_wrapper = dict(
    ...,
    paramwise_cfg=dict(
        custom_keys={
            'visual_encoder.': dict(lr_mult=0.1),  # 视觉编码器学习率降为1/10
            'language_model.': dict(lr_mult=1.0)   # 语言模型保持原学习率
        }
    )
)

这种配置方式不需要重写OptimWrapperConstructor，框架会自动根据模块名称前缀匹配并应用对应的学习率乘数。

梯度累积的实现机制

XTuner基于MMEngine实现梯度累积功能，其工作机制与HuggingFace Trainer有所不同：

1. 迭代次数(iter)不会因为梯度累积而减少
2. 仅在累积次数达到设定值时执行参数更新
3. 梯度会在每次反向传播时累积

这种设计保持了训练流程的连续性，便于日志记录和监控。

使用Deepspeed时的注意事项

当启用Deepspeed时，优化器配置会被Deepspeed接管，此时需要：

1. 在Deepspeed配置文件中指定优化器参数
2. 确保paramwise_cfg设置能够正确传递到Deepspeed优化器
3. 可能需要调整train.py中的相关逻辑以确保配置生效

最佳实践建议

1. 对于多模态模型，建议视觉部分使用较低学习率(通常为语言模型的1/5到1/10)
2. 梯度累积次数应根据显存情况合理设置，通常为2的幂次方
3. 使用混合精度训练时，注意loss scaling策略的选择
4. 定期检查梯度更新情况，确保各模块按预期进行训练

通过合理配置优化器参数，可以显著提升模型微调的效果和训练稳定性。XTuner提供的灵活配置选项使得针对不同模块的精细化调参成为可能。