Transformers：Deepspeed场景下如何重写Trainer子类自定义优化器

问题背景 #

启用Deepspeed训练时，直接通过Trainer构造函数的optimizers参数传入自定义Adafactor优化器和AdafactorSchedule调度器，会触发以下报错：

RuntimeError: Passing optimizers is not allowed if Fairscale, Deepspeed or PyTorch FSDP is enabled.

You should subclass Trainer and override the create_optimizer_and_scheduler method.

需求是使用自定义配置的Adafactor优化器和对应的学习率调度器完成Deepspeed训练。

解决方案 #

核心思路是自定义Trainer子类，重写优化器和调度器的创建方法，让Deepspeed能够正确处理分布式环境下的参数封装。具体实现可以选择两种方式：

方式1：分别重写优化器与调度器创建方法 #

自定义CustomTrainer类，继承自Trainer，分别实现create_optimizer和create_scheduler方法：

from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments, DataCollatorWithPadding
from transformers.optimization import Adafactor, AdafactorSchedule
import torch

# 自定义Trainer子类
class CustomTrainer(Trainer):
    def create_optimizer(self):
        # 实例化自定义配置的Adafactor优化器
        optimizer = Adafactor(
            self.model.parameters(),
            scale_parameter=True,
            relative_step=True,
            warmup_init=True,
            lr=None  # Adafactor开启relative_step时无需手动指定学习率
        )
        return optimizer
    
    def create_scheduler(self, num_training_steps: int, optimizer=None):
        # 基于优化器生成Adafactor专属调度器
        lr_scheduler = AdafactorSchedule(optimizer if optimizer is not None else self.optimizer)
        return lr_scheduler

# 模型初始化
model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", 
    num_labels=2,
    torch_dtype=torch.float16,
)

# 训练参数配置
training_args = TrainingArguments(
    output_dir='./xd',
    num_train_epochs=4,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32*2,
    warmup_steps=0,  # Adafactor的warmup_init=True已包含预热逻辑，无需重复设置
    tf32=True,
    bf16=False,
    fp16=True,
    dataloader_num_workers=16,
    gradient_accumulation_steps=1,
    evaluation_strategy='steps',
    eval_steps=2500,
    save_steps=1000,
    deepspeed=r'ds_config_zero3.json',
    disable_tqdm=False,
    weight_decay=0.01,
    logging_dir='./xd_logs',
    logging_steps=1000,
)

# 使用自定义Trainer初始化，无需传入optimizers参数
trainer = CustomTrainer(
    model=model,
    args=training_args,
    # 可在此添加train_dataset、eval_dataset等参数
)

# 启动训练
trainer.train()

方式2：直接重写组合创建方法 #

如果需要更紧凑的逻辑，可以直接重写create_optimizer_and_scheduler方法，一次性完成优化器和调度器的创建：

class CustomTrainer(Trainer):
    def create_optimizer_and_scheduler(self, num_training_steps: int):
        self.optimizer = Adafactor(
            self.model.parameters(),
            scale_parameter=True,
            relative_step=True,
            warmup_init=True,
            lr=None
        )
        self.lr_scheduler = AdafactorSchedule(self.optimizer)
        return self.optimizer, self.lr_scheduler

关键注意事项 #

• 禁止在Trainer外部提前实例化优化器：Deepspeed需要对模型参数进行分布式封装（如Zero3的参数分片），提前实例化会导致优化器无法正确获取分布式后的参数。
• Adafactor参数适配：开启relative_step=True时，无需手动指定lr，调度器会自动根据训练步数调整学习率；warmup_init=True已包含预热逻辑，需将TrainingArguments中的warmup_steps设为0避免冲突。

内容的提问来源于stack exchange，提问作者John Stud