Whale可以帮助您便捷地进行分布式并行训练,支持各种并行策略混合训练,同时提供多种通信优化方法。本文介绍在Whale中实现分布式并行策略的详细步骤,包括初始化、资源划分、模型划分及硬件资源到逻辑资源的映射,并提供了完整示例。同时,介绍如何开启Whale的通信优化功能。

背景信息

Whale是灵活、易用、高效、统一的分布式训练框架,提供简单且易用的接口以实现数据并行、模型并行、流水并行、算子拆分及多种并行方式组合的混合并行。Whale基于TensorFlow开发,完全兼容TensorFlow接口,您只需要在已有的TensorFlow模型上增加几行分布式并行策略描述代码,即可进行分布式混合并行训练。

版本配套关系

Whale支持的PAI-TensorFlow及Python版本如下:
  • PAI-TensorFlow版本:PAI-TensorFlow 1.12、PAI-TensorFlow 1.15
  • Python版本:Python 2.7、Python 3.6

使用Whale添加分布式策略

  1. 初始化。
    使用init()接口进行Whale初始化,详情请参见whale.init。示例如下。
    import whale as wh
    wh.init()
  2. 划分资源(Cluster)。
    Whale中通过Cluster抽象和定义资源,完成硬件资源到逻辑资源的映射。同时,Whale通过不同策略将资源进行分组,详情请参见whale.cluster
    cluster = wh.cluster()
  3. 划分模型(Scope)。
    Whale通过Scope指定分布式训练方式,且支持多种Scope组合成各种混合并行方式,详情请参见whale.scopes。Whale支持的Scopes类型包括:
    • replica:数据并行。
    • stage:模型并行。
    • pipeline:流水并行。
    • split:算子拆分
    • 以上Scope的组合。
    以下列举一些分布式训练在Whale中的实现方式,以供参考:
    • 数据并行
      with wh.replica():
          model_fn()
    • 模型并行
      假设模型有20个Layers,前10个Layers划分为一个Stage,后10个Layers划分为另一个Stage,示例如下。
      with wh.stage():
          # layers 1~10 of model
          model_fn1()
      with wh.stage():
          # layers 11~20 of model
          model_fn2()
    • 数据并行、模型并行及流水并行
      在模型并行基础上进行数据并行和流水并行,您只需要组合使用replicapipeline的Scope,示例如下。
      with wh.replica():
          with wh.pipeline(num_micro_batch=5):
              with wh.stage():
                  # layers 1~10 of model
                  model_fn1()
              with wh.stage():
                  # layers 11~20 of model
                  model_fn2()
    • 数据并行和算子拆分

      假设对一个模型前10个Layers进行数据并行,对后面的Layers进行算子拆分,示例如下。

      with wh.replica():
          # layers 1~10 of model
          model_fn1()
      with wh.split():
          # layers 11~12 of model
          model_fn2()

完整示例

  1. 对模型进行分布式改写。
    使用Whale将一个单机模型改造为数据并行,此处仅展示主体代码改写过程,示例如下。完整代码请参见simple_data_parallel.py,更多分布式并行策略的使用方法请参见Whale分布式范式剖析
    import whale as wh
    def main(_):
      wh.init()  # init whale
      with wh.cluster(): # cluster
        with wh.replica(): # data parallelism
           your_model_definition # construct your model here.
    主体代码的改写过程主要包括三个调用:
    1. Whale初始化:wh.init()
    2. Cluster定义及硬件资源到逻辑资源的映射:with wh.cluster()
    3. 将并行方式指定为数据并行:with wh.replica()
    将单机模型改造为数据并行方式进行训练后,Whale会自动感知任务资源,并根据数据并行方式划分资源。
  2. 切分训练数据。
    上一步中读取的Mock数据。如果使用真实数据进行分布式训练,则需要对训练数据进行切分,从而确保每个Worker获取不同的数据。

    Whale会自动对训练数据进行切分,训练数据切分策略及配置方法请参见训练数据分片

  3. 提交作业。
    您可以通过以下任意一种方式提交作业:
    • PAI-Studio

      支持算法组件TensorFlow 1120,提交作业的方法请参见PAI-TF调用方式

    • PAI-DLC
      • 如果登录服务器,手动执行训练脚本,则执行如下命令。
        python simple_data_parallel.py
        其中simple_data_parallel.py表示分布式训练任务的代码文件,需要根据实际情况替换。
      • 如果通过Arena提交任务,请参见使用Arena管理任务(推荐)
      • 如果是使用PAI-DLC Dashboard提交任务,请参见使用PAI-DLC Dashboard管理任务

Whale其他功能

Whale支持多种通信优化,您可以根据需求开启各种优化功能,配置方法请参见通信参数