Accelerate 文档
在 DeepSpeed 中使用多个模型
并获得增强的文档体验
开始使用
在 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 上更新