训练 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.DeepSpeedEngine.forward(*args, **kwargs)

定义每次调用时执行的计算。

应由所有子类覆盖。

注意

虽然前向传递的配方需要在此函数中定义，但应随后调用Module 实例，而不是此函数，因为前者会处理运行注册的钩子，而后者会静默地忽略它们。

反向传播

deepspeed.DeepSpeedEngine.backward(*args, **kwargs)

优化器步骤

deepspeed.DeepSpeedEngine.step(self, lr_kwargs=None)

在对 effective_train_batch 执行前向和反向传播后，执行权重更新步骤。

梯度累积

deepspeed.DeepSpeedEngine.is_gradient_accumulation_boundary(self)

查询当前微批次是否处于梯度累积的边界，因此将触发梯度缩减和优化器步骤。

返回

如果当前步骤是梯度累积边界。

返回类型

bool

模型保存

deepspeed.DeepSpeedEngine.save_16bit_model(self, save_dir, save_filename='pytorch_model.bin', exclude_frozen_parameters=False)

保存 16 位模型权重

此方法将 16 位模型权重保存在所需的目标位置。

参数

• save_dir – 必需。保存模型的目录

• save_filename – 可选。要保存的文件名。默认为 pytorch_model.bin

• exclude_frozen_parameters – 可选。从检查点状态中排除冻结的参数。

返回

True 当模型已保存时，False 否则。如果 stage3_gather_16bit_weights_on_model_save 为 False，则不会保存。

重要：所有进程必须调用此方法，而不仅仅是排名为 0 的进程。这是因为进程需要同步工作以收集权重。如果只针对排名为 0 的进程调用此方法，则此方法将挂起等待与其他进程同步。

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