训练API

deepspeed.initialize() 在其第一个参数中返回一个类型为 DeepSpeedEngine 的训练引擎。此引擎用于推进训练。

for step, batch in enumerate(data_loader):
    #forward() method
    loss = model_engine(batch)

    #runs backpropagation
    model_engine.backward(loss)

    #weight update
    model_engine.step()

前向传播

反向传播

优化器步进

梯度累积

模型保存

此外，当创建 DeepSpeed 检查点时，会添加一个脚本 zero_to_fp32.py，可用于将 fp32 主权重重构为单个 PyTorch state_dict 文件。

训练多个模型

DeepSpeed 支持训练多个模型，这在知识蒸馏和训练后RLHF等场景中非常有用。核心方法是为每个模型创建单独的 DeepSpeedEngine。

训练独立模型

以下代码片段演示了在同一数据集上独立训练多个模型。

model_engines = [engine for engine, _, _, _ in [deepspeed.initialize(m, ...,) for m in models]]
for batch in data_loader:
   losses = [engine(batch) for engine in model_engines]
   for engine, loss in zip(model_engines, losses):
      engine.backward(loss)

除了创建多个 DeepSpeedEngine（每个模型一个）之外，上述用法与典型的 DeepSpeed 用法类似。

共享损失的多模型联合训练

以下代码片段演示了在共享损失值上联合训练多个模型。

model_engines = [engine for engine, _, _, _ in [deepspeed.initialize(m, ...,) for m in models]]
for batch in data_loader:
    losses = [engine(batch[0], batch[1]) for engine in model_engines]
    loss = sum(l / (i + 1) for i, l in enumerate(losses))
    loss.backward()

    for engine in model_engines:
        engine._backward_epilogue()

    for engine in model_engines:
        engine.step()

    for engine in model_engines:
        engine.optimizer.zero_grad()

除了使用多个 DeepSpeedEngine 之外，上述用法与典型用法在两个关键方面有所不同：

反向传播调用是使用共同损失值而非各自模型引擎进行的。
在 loss.backward() 之后，会在模型引擎上调用 _backward_epilogue。