LlamaFactory
算力平台:

使用多个模型与 DeepSpeed ​

在使用 DeepSpeed 时，你可能会遇到需要同时加载和运行多个模型的情况。DeepSpeed 提供了多种方法来管理和优化多个模型的使用。以下是一些常见的场景和相应的解决方案。

1. 加载多个模型 ​

1.1 单个进程中的多个模型 ​

如果你需要在单个进程中加载多个模型，可以使用以下方法：

python
import deepspeed

# 定义模型
model1 = Model1()
model2 = Model2()

# 初始化 DeepSpeed 引擎
engine1, optimizer1, _, _ = deepspeed.initialize(model=model1, model_parameters=model1.parameters())
engine2, optimizer2, _, _ = deepspeed.initialize(model=model2, model_parameters=model2.parameters())

1.2 多个进程中的多个模型 ​

如果你需要在多个进程中加载多个模型，可以使用以下方法：

python
import deepspeed

# 定义模型
model1 = Model1()
model2 = Model2()

# 初始化 DeepSpeed 引擎
engine1, optimizer1, _, _ = deepspeed.initialize(model=model1, model_parameters=model1.parameters(), dist_init_required=True)
engine2, optimizer2, _, _ = deepspeed.initialize(model=model2, model_parameters=model2.parameters(), dist_init_required=True)

2. 优化多个模型的性能 ​

2.1 使用混合精度 ​

DeepSpeed 支持混合精度训练，可以显著提高训练速度和减少内存使用。你可以在配置文件中启用混合精度：

yaml
# deepspeed_config.json
{
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "initial_scale_power": 16,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1
  }
}

2.2 使用 ZeRO 优化 ​

DeepSpeed 的 ZeRO 优化可以显著减少内存使用，提高训练效率。你可以在配置文件中启用 ZeRO 优化：

yaml
# deepspeed_config.json
{
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "overlap_comm": true,
    "load_from_fp32_weights": true,
    "elastic_checkpoint": true
  }
}

3. 管理多个模型的资源 ​

3.1 资源分配 ​

在多模型场景中，合理分配资源非常重要。你可以使用 DeepSpeed 的资源管理功能来优化资源分配：

python
import deepspeed

# 定义模型
model1 = Model1()
model2 = Model2()

# 定义资源分配
resources = {
    "model1": {"gpu": 0},
    "model2": {"gpu": 1}
}

# 初始化 DeepSpeed 引擎
engine1, optimizer1, _, _ = deepspeed.initialize(model=model1, model_parameters=model1.parameters(), config_params=resources["model1"])
engine2, optimizer2, _, _ = deepspeed.initialize(model=model2, model_parameters=model2.parameters(), config_params=resources["model2"])

3.2 动态资源管理 ​

DeepSpeed 还支持动态资源管理，可以根据运行时的负载情况动态调整资源分配：

python
import deepspeed

# 定义模型
model1 = Model1()
model2 = Model2()

# 初始化 DeepSpeed 引擎
engine1, optimizer1, _, _ = deepspeed.initialize(model=model1, model_parameters=model1.parameters(), dynamic_resources=True)
engine2, optimizer2, _, _ = deepspeed.initialize(model=model2, model_parameters=model2.parameters(), dynamic_resources=True)

4. 调试和监控 ​

4.1 调试 ​

在多模型场景中，调试可能会更加复杂。你可以使用 DeepSpeed 的调试工具来帮助你定位问题：

python
import deepspeed

# 定义模型
model1 = Model1()
model2 = Model2()

# 初始化 DeepSpeed 引擎
engine1, optimizer1, _, _ = deepspeed.initialize(model=model1, model_parameters=model1.parameters(), debug=True)
engine2, optimizer2, _, _ = deepspeed.initialize(model=model2, model_parameters=model2.parameters(), debug=True)

4.2 监控 ​

DeepSpeed 提供了多种监控工具，可以帮助你监控多模型的运行状态：

python
import deepspeed

# 定义模型
model1 = Model1()
model2 = Model2()

# 初始化 DeepSpeed 引擎
engine1, optimizer1, _, _ = deepspeed.initialize(model=model1, model_parameters=model1.parameters(), monitor=True)
engine2, optimizer2, _, _ = deepspeed.initialize(model=model2, model_parameters=model2.parameters(), monitor=True)

通过以上方法，你可以有效地管理和优化多个模型的使用，提高训练效率和性能。

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

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

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

本教程将重点介绍两个常见的用例：

  1. 知识蒸馏，其中较小的学生模型被训练以模仿较大、性能更好的教师模型。如果学生模型可以放在单个 GPU 上，我们可以使用 ZeRO-2 进行训练，使用 ZeRO-3 进行教师模型的推理。这比对两个模型都使用 ZeRO-3 要快得多。
  2. 同时训练多个 不相交 的模型。

知识蒸馏 ​

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

通常情况下，你会为两个模型使用单个 [utils.DeepSpeedPlugin]。然而，在这种情况下，有两个独立的配置。Accelerate 允许你创建和使用多个插件 当且仅当 它们在一个 dict 中，这样你可以在需要时引用和启用正确的插件。

python
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 应该仅配置为推理模型，如下例所示。

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 配置示例如下所示。

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",
}

从这里开始，创建一个 [Accelerator] 并传入两个配置。

python
from accelerate import Accelerator

accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)

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

学生模型 ​

默认情况下，Accelerate 将 dict 中的第一个项目设置为默认或启用的插件（即 "student" 插件）。你可以使用 [utils.deepspeed.get_active_deepspeed_plugin] 函数来验证当前启用的插件。

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

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

python
assert active_plugin is accelerator.deepspeed_plugin

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

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

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

教师模型 ​

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

python
accelerator.state.select_deepspeed_plugin("teacher")

这会禁用 "student" 插件并启用 "teacher" 插件。Transformers 中的 DeepSpeed 状态配置将被更新，从而在使用 deepspeed.initialize() 时改变调用的插件配置。这允许你使用 Transformers 提供的自动 deepspeed.zero.Init 上下文管理器集成。

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

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

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

训练 ​

从这里开始，你的训练循环可以是任何你想要的形式，只要确保 teacher_model 从未被训练过。

python
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] 而无需额外的参数。

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

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

python
# 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
)

现在你可以训练：

python
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] 中的 相关测试。

Copyright © 2024-2025 llamafactory.cn