DeepSpeed 用于 HPU

DeepSpeed 使您能够在 HPU 上拟合和训练更大的模型，这得益于 ZeRO 论文中描述的各种优化。特别是，您可以使用以下两种 ZeRO 配置，这些配置已验证与 Gaudi 完全兼容

ZeRO-1：跨进程划分优化器状态。
ZeRO-2：跨进程划分优化器状态 + 梯度。

这些配置与 Habana 混合精度完全兼容，因此可用于以 bf16 精度训练您的模型。

您可以在此处找到有关 DeepSpeed Gaudi 集成的更多信息。

设置

要在 Gaudi 上使用 DeepSpeed，您需要使用以下命令安装 Optimum Habana 和 Habana 的 DeepSpeed 分支：

pip install optimum[habana]
pip install git+https://github.com/HabanaAI/[email protected]

使用 Optimum Habana 中的 DeepSpeed

GaudiTrainer 允许像使用 Transformers Trainer 一样轻松地使用 DeepSpeed。这可以通过三个步骤完成

必须定义 DeepSpeed 配置。
deepspeed 训练参数可以用来指定 DeepSpeed 配置的路径。
必须使用 deepspeed 启动器运行您的脚本。

这些步骤在下面详细说明。关于如何在 Transformers Trainer 中使用 DeepSpeed 的全面指南也可在此处获得。

DeepSpeed 配置

要使用的 DeepSpeed 配置通过 JSON 文件传递，并允许您选择要应用的优化。以下是如何应用 ZeRO-2 优化和 bf16 精度的示例

{
    "steps_per_print": 64,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "bf16": {
        "enabled": true
    },
    "gradient_clipping": 1.0,
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": false,
        "reduce_scatter": false,
        "contiguous_gradients": false
    }
}

特殊值 "auto" 允许自动获取正确或最有效的值。您也可以自己指定值，但如果您这样做，则应注意不要与您的训练参数存在冲突的值。强烈建议阅读 Transformers 文档中的此部分，以完全了解其工作原理。

Habana 在此处提供了其他 HPU 配置示例。

Transformers 文档非常详细地解释了如何从头开始编写配置。所有配置可能性的更完整描述可在此处获得。

deepspeed 训练参数

要使用 DeepSpeed，您必须在 GaudiTrainingArguments 实例中指定 deespeed=path_to_my_deepspeed_configuration

training_args = GaudiTrainingArguments(
    # my usual training arguments...
    use_habana=True,
    use_lazy_mode=True,
    gaudi_config_name=path_to_my_gaudi_config,
    deepspeed=path_to_my_deepspeed_config,
)

此参数既表示应使用 DeepSpeed，又指向您的 DeepSpeed 配置。

启动您的脚本

最后，有两种可能的方式来启动您的脚本

使用 gaudi_spawn.py 脚本

python gaudi_spawn.py \
    --world_size number_of_hpu_you_have --use_deepspeed \
    path_to_script.py --args1 --args2 ... --argsN \
    --deepspeed path_to_deepspeed_config

其中 --argX 是要使用 DeepSpeed 运行的脚本的参数。

直接在代码中使用 DistributedRunner

from optimum.habana.distributed import DistributedRunner
from optimum.utils import logging

world_size=8 # Number of HPUs to use (1 or 8)

# define distributed runner
distributed_runner = DistributedRunner(
    command_list=["scripts/train.py --args1 --args2 ... --argsN --deepspeed path_to_deepspeed_config"],
    world_size=world_size,
    use_deepspeed=True,
)

# start job
ret_code = distributed_runner.run()

您应该在您的 Gaudi 配置中设置 "use_fused_adam": false，因为它与 DeepSpeed 还不兼容。

< > 更新在 GitHub 上