使用 DeepSpeed 运行多个模型

本指南假设您已阅读并理解了DeepSpeed 使用指南。

使用 Accelerate 和 DeepSpeed 运行多个模型对以下场景很有用：

知识蒸馏
训练后技术，例如 RLHF（请参阅TRL库以获取更多示例）
同时训练多个模型

目前，Accelerate 提供了一个**非常实验性的 API** 来帮助您使用多个模型。

本教程将重点介绍两种常见用例

知识蒸馏，其中一个较小的学生模型被训练来模仿一个更大、性能更好的教师模型。如果学生模型适合单个 GPU，我们可以使用 ZeRO-2 进行训练，并使用 ZeRO-3 对教师模型进行分片推理。这比对两个模型都使用 ZeRO-3 快得多。
同时训练多个分离的模型。

知识蒸馏

知识蒸馏是使用多个模型的一个很好的例子，但只训练其中一个。

通常，您会对两个模型都使用单个utils.DeepSpeedPlugin。但是，在这种情况下，存在两个单独的配置。Accelerate 允许您创建和使用多个插件，**当且仅当**它们在 dict 中时，以便您可以在需要时引用和启用正确的插件。

from accelerate.utils import DeepSpeedPlugin

zero2_plugin = DeepSpeedPlugin(hf_ds_config="zero2_config.json")
zero3_plugin = DeepSpeedPlugin(hf_ds_config="zero3_config.json")

deepspeed_plugins = {"student": zero2_plugin, "teacher": zero3_plugin}

zero2_config.json 应配置为完整训练（因此，如果您没有使用自己的训练，请指定 scheduler 和 optimizer），而 zero3_config.json 应仅配置推理模型，如下面的示例所示。

{
    "bf16": {
        "enabled": "auto"
    },
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": true,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": "auto",
        "stage3_max_reuse_distance": "auto",
    },
    "train_micro_batch_size_per_gpu": 1
}

下面显示了一个 zero2_config.json 配置示例。

{
    "bf16": {
        "enabled": "auto"
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto",
            "torch_adam": true,
            "adam_w_mode": true
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
}

即使此特定模型没有被训练，如果未指定 train_micro_batch_size_per_gpu，DeepSpeed 也会引发错误。

在此基础上，创建一个单个Accelerator 并传入这两个配置。

from accelerate import Accelerator

accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)

现在让我们看看如何使用它们。

学生模型

默认情况下，Accelerate 将 dict 中的第一个项目设置为默认或启用的插件（"student" 插件）。使用utils.deepspeed.get_active_deepspeed_plugin() 函数查看哪个插件已启用，从而验证这一点。

active_plugin = get_active_deepspeed_plugin(accelerator.state)
assert active_plugin is deepspeed_plugins["student"]

AcceleratorState 还将活动 DeepSpeed 插件保存在 state.deepspeed_plugin 中。

assert active_plugin is accelerator.deepspeed_plugin

由于 student 是当前活动的插件，因此让我们继续准备模型、优化器和调度器。

student_model, optimizer, scheduler = ...
student_model, optimizer, scheduler, train_dataloader = accelerator.prepare(student_model, optimizer, scheduler, train_dataloader)

现在是处理教师模型的时候了。

教师模型

首先，您需要在Accelerator 中指定应使用 zero3_config.json 配置。

accelerator.state.select_deepspeed_plugin("teacher")

这将禁用 "student" 插件并改为启用 "teacher" 插件。Transformers 内部的 DeepSpeed 有状态配置会更新，并且它会更改使用 deepspeed.initialize() 时调用哪个插件配置。这允许您使用 Transformers 提供的自动 deepspeed.zero.Init 上下文管理器集成。

teacher_model = AutoModel.from_pretrained(...)
teacher_model = accelerator.prepare(teacher_model)

否则，您应该使用 deepspeed.zero.Init 手动初始化模型。

with deepspeed.zero.Init(accelerator.deepspeed_plugin.config):
    model = MyModel(...)

训练

在此基础上，您的训练循环可以是您想要的任何内容，只要 teacher_model 从未在上面进行训练。

teacher_model.eval()
student_model.train()
for batch in train_dataloader:
    with torch.no_grad():
        output_teacher = teacher_model(**batch)
    output_student = student_model(**batch)
    # Combine the losses or modify it in some way
    loss = output_teacher.loss + output_student.loss
    accelerator.backward(loss)
    optimizer.step()
    scheduler.step()
    optimizer.zero_grad()

训练多个分离的模型

训练多个模型是一个更复杂的场景。在其当前状态下，我们假设每个模型在训练期间都与其他模型**完全分离**。

这种情况仍然需要创建两个utils.DeepSpeedPlugin。但是，您还需要第二个Accelerator，因为不同的 deepspeed 引擎在不同的时间被调用。单个Accelerator 每次只能携带一个实例。

由于state.AcceleratorState 是一个有状态的对象，因此它已经知道两个可用的utils.DeepSpeedPlugin。您只需使用没有额外参数的第二个Accelerator 进行实例化。

first_accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)
second_accelerator = Accelerator()

您可以调用 first_accelerator.state.select_deepspeed_plugin() 来启用或禁用特定插件，然后调用 prepare。

# can be `accelerator_0`, `accelerator_1`, or by calling `AcceleratorState().select_deepspeed_plugin(...)`
first_accelerator.state.select_deepspeed_plugin("first_model")
first_model = AutoModel.from_pretrained(...)
# For this example, `get_training_items` is a nonexistent function that gets the setup we need for training
first_optimizer, first_scheduler, train_dl, eval_dl = get_training_items(model1)
first_model, first_optimizer, first_scheduler, train_dl, eval_dl = accelerator.prepare(
    first_model, first_optimizer, first_scheduler, train_dl, eval_dl
)

second_accelerator.state.select_deepspeed_plugin("second_model")
second_model = AutoModel.from_pretrained(...)
# For this example, `get_training_items` is a nonexistent function that gets the setup we need for training
second_optimizer, second_scheduler, _, _ = get_training_items(model2)
second_model, second_optimizer, second_scheduler = accelerator.prepare(
    second_model, second_optimizer, second_scheduler
)

现在您可以训练了

for batch in dl:
    outputs1 = first_model(**batch)
    first_accelerator.backward(outputs1.loss)
    first_optimizer.step()
    first_scheduler.step()
    first_optimizer.zero_grad()
    
    outputs2 = model2(**batch)
    second_accelerator.backward(outputs2.loss)
    second_optimizer.step()
    second_scheduler.step()
    second_optimizer.zero_grad()

资源

要查看更多示例，请查看目前在 [Accelerate] 中的相关测试。

< > 在 GitHub 上更新

加速

使用 DeepSpeed 运行多个模型

知识蒸馏

学生模型

教师模型

训练

训练多个分离的模型

资源