使用 Optimum Habana 进行分布式训练

随着模型变得越来越大，并行化已成为在有限硬件上训练更大模型并通过几个数量级加速训练速度的一种策略。

所有 PyTorch 示例和 GaudiTrainer 脚本都可以直接用于分布式训练。有两种方法可以启动它们

使用 gaudi_spawn.py 脚本

python gaudi_spawn.py \
    --world_size number_of_hpu_you_have --use_mpi \
    path_to_script.py --args1 --args2 ... --argsN

其中 --argX 是以分布式方式运行脚本的参数。这里提供了关于问答此处和文本分类此处的示例。

直接在代码中使用 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"],
    world_size=world_size,
    use_mpi=True,
)

# start job
ret_code = distributed_runner.run()

您可以为训练参数设置 --distribution_strategy fast_ddp 以获得更简单且通常更快的分布式训练管理。更多信息此处。

为了更深入地了解，我们邀请您阅读我们关于以下内容的指南：

加速训练
预训练
DeepSpeed 用于训练更大的模型
多节点训练以进一步加快您的分布式运行速度

< > 在 GitHub 上更新