如何在Huggingface Trainer中使用两个优化器进行迭代训练？

是的，完全可以在Trainer中使用两个（或多个）优化器，核心思路是重写Trainer类的关键方法，自定义优化器的创建和训练步骤的执行逻辑。下面是具体实现方案：

方法1：重写create_optimizer_and_scheduler和training_step #

步骤1：拆分模型参数 #

先将模型的两个部分参数分离，假设你的模型包含part_a和part_b两个子模块：

# 拆分模型参数组
params_a = list(model.part_a.parameters())
params_b = list(model.part_b.parameters())

步骤2：自定义Trainer子类 #

继承Trainer类，重写两个核心方法：

• create_optimizer_and_scheduler：跳过父类默认逻辑，直接使用自定义的多个优化器和调度器
• training_step：在每步训练中分别用两个优化器更新对应模型部分的参数

示例代码：

from transformers import Trainer

class MultiOptimizerTrainer(Trainer):
    def __init__(self, optimizer_a, optimizer_b, scheduler_a, scheduler_b, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.optimizer_a = optimizer_a
        self.optimizer_b = optimizer_b
        self.scheduler_a = scheduler_a
        self.scheduler_b = scheduler_b

    def create_optimizer_and_scheduler(self, num_training_steps: int):
        # 父类要求必须赋值optimizer和lr_scheduler属性，随便绑定一个不影响逻辑
        self.optimizer = self.optimizer_a
        self.lr_scheduler = self.scheduler_a
        return self.optimizer_a, self.scheduler_a

    def training_step(self, model, inputs):
        model.train()
        inputs = self._prepare_inputs(inputs)

        # 计算损失
        with self.compute_loss_context_manager():
            loss = self.compute_loss(model, inputs)

        if self.args.n_gpu > 1:
            loss = loss.mean()  # 多GPU环境下平均损失

        # 更新第一部分模型参数
        self.optimizer_a.zero_grad()
        loss.backward(retain_graph=True)  # 保留计算图供第二个优化器使用
        self.optimizer_a.step()
        self.scheduler_a.step()

        # 更新第二部分模型参数
        self.optimizer_b.zero_grad()
        loss.backward()
        self.optimizer_b.step()
        self.scheduler_b.step()

        return loss.detach() / self.args.gradient_accumulation_steps

步骤3：初始化并启动训练 #

from torch.optim import AdamW
from transformers import get_scheduler

# 创建两个独立优化器
optimizer_a = AdamW(params_a, lr=1e-4)
optimizer_b = AdamW(params_b, lr=5e-5)

# 计算训练总步数
num_training_steps = len(train_dataset) // training_args.per_device_train_batch_size * training_args.num_train_epochs

# 创建对应学习率调度器
scheduler_a = get_scheduler(
    name="linear",
    optimizer=optimizer_a,
    num_warmup_steps=0,
    num_training_steps=num_training_steps
)
scheduler_b = get_scheduler(
    name="linear",
    optimizer=optimizer_b,
    num_warmup_steps=0,
    num_training_steps=num_training_steps
)

# 初始化自定义Trainer
trainer = MultiOptimizerTrainer(
    optimizer_a=optimizer_a,
    optimizer_b=optimizer_b,
    scheduler_a=scheduler_a,
    scheduler_b=scheduler_b,
    model=model,
    args=training_args,
    train_dataset=train_dataset
)

# 启动训练
trainer.train()

方法2：简化版——利用optimizers参数批量传入 #

Trainer支持直接传入optimizers参数，结合重写training_step可以更灵活地处理多优化器场景：

from transformers import Trainer

class MultiOptimizerTrainer(Trainer):
    def __init__(self, optimizers, schedulers, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.optimizers = optimizers
        self.schedulers = schedulers

    def training_step(self, model, inputs):
        model.train()
        inputs = self._prepare_inputs(inputs)

        with self.compute_loss_context_manager():
            loss = self.compute_loss(model, inputs)

        if self.args.n_gpu > 1:
            loss = loss.mean()

        # 遍历所有优化器执行更新
        for idx, optimizer in enumerate(self.optimizers):
            optimizer.zero_grad()
            # 最后一个优化器不需要保留计算图
            if idx != len(self.optimizers) - 1:
                loss.backward(retain_graph=True)
            else:
                loss.backward()
            optimizer.step()
            self.schedulers[idx].step()

        return loss.detach() / self.args.gradient_accumulation_steps

# 批量创建优化器和调度器
optimizers = [AdamW(params_a, lr=1e-4), AdamW(params_b, lr=5e-5)]
schedulers = [
    get_scheduler("linear", optimizer=optimizers[0], num_training_steps=num_training_steps),
    get_scheduler("linear", optimizer=optimizers[1], num_training_steps=num_training_steps)
]

# 初始化Trainer
trainer = MultiOptimizerTrainer(
    optimizers=optimizers,
    schedulers=schedulers,
    model=model,
    args=training_args,
    train_dataset=train_dataset
)
trainer.train()

注意事项 #

• 如果两个模型部分需要不同的损失函数，可以在compute_loss方法中分别计算，再在training_step中对应执行反向传播。
• 启用梯度累积时，需要调整training_step逻辑，确保累积到指定步数后再执行优化器更新。
• 多GPU训练场景下，Trainer父类已处理大部分参数同步逻辑，但自定义步骤时需注意梯度的正确传播。

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