Whale根据分布式模式和用户配置的资源,自动对训练数据进行分片。如果训练数据无法均分,您可以通过环境变量配置数据分片策略。本文介绍Whale支持的数据分片策略、数据分片策略的参数列表及数据分片策略的配置方法。

数据分片策略

分布式训练时,每个Worker需要获取训练数据的不同分片。如果每个Worker处理的数据完全一样,则训练效果等同于单机,无法发挥分布式加速效果。Whale根据分布式模式和用户配置的资源,自动对训练数据进行分片。每个Worker依据分布式模式分配不同的角色,包括Master和Slave。Master负责构图和驱动执行,Slave被动执行分配到的任务。

Whale的训练数据分片策略如下:
  1. 根据分布式模式(例如数据并行、模型并行、流水并行或算子拆分)和资源,确定每个Worker的角色。
  2. 根据Master Worker切分训练数据。

数据分片场景实例

  • 数据并行
    在数据并行场景中,每个Worker都是Master角色,Whale按照Master Worker数量均匀切分训练数据。如下图所示,共3个Master Worker,Whale会将训练数据的文件列表均分为3份,每个Worker获取一个分片。Worker内的多个GPU按照Batch依次获取数据。数据并行的训练数据切分
  • 混合数据并行和模型并行

    在混合数据并行和模型并行的场景,模型会切分为多个Stage分配至不同Driver。如下图所示,模型分成2个Stages,Master Worker执行Stage 0,Slave Worker执行Stage 1。(4个Worker中有2个Master角色,2个Slave角色。)

    Whale根据Master Worker数量均匀切分训练数据。如下图所示,有2个Master Worker,即将训练数据的文件列表均分为2份,每个MasterWorker获取一个分片,Worker内的多个GPU按照Batch依次获取数据。

    混合数据并行和模型并行的数据切分
  • 训练数据无法均分
    当训练数据的文件数目不能整除Master Worker数量时,Whale提供三种处理方式:
    • 不切分,进行随机Shuffle。
      每个Master Worker获取完整的文件列表,不切分训练数据。为保证收敛,对文件列表进行随机Shuffle。
      说明
      • 训练数据无法均分时,Whale默认采用该处理方式。
      • 当文件数量小于Master Worker数量时,Whale只支持该处理方式。
    • 不均匀切分,为部分Master Worker多分配一个文件。

      尽可能均匀切分,容许部分Master Worker比其他的多分配一个文件。

    • 均匀切分,多余的文件直接丢弃。

      建议仅在文件数很多,且丢弃的文件不影响收敛时,使用该方式。

数据分片的参数列表

Whale支持的数据分片策略参数如下表所示,如果训练数据的文件数目不能整除Master Worker数量,您可以使用如下参数配置数据分片策略。
参数 描述 类型 默认值 备注
WHALE_UNBALANCED_IO_SLICING 是否允许切分数据时,每个Worker分配的文件数目不同,即为部分Worker多分配一个文件。取值范围如下:
  • False:不切分训练数据。
  • True:允许非均等切分训练数据,即为部分Worker多分配一个文件。
BOOL False 由于部分Worker会分到更多的训练数据文件,因此需要注意分布式训练的结束条件。如果以Epoch为结束条件,则训练可能挂起。建议使用StopAtStepHook控制模型训练是否结束。
WHALE_DROP_LAST_FILES 当训练数据的文件数目不能整除Master Worker数量时,是否依然均匀切分训练数据,将多余文件舍弃。该参数的优先级低于WHALE_UNBALANCED_IO_SLICING参数,如果这两个参数同时为True,则执行非均等切分策略。取值范围如下:
  • False:不切分训练数据。
  • True:均匀切分训练数据,将多余文件舍弃。
BOOL False 建议仅在训练数据文件数量远大于训练规模,且丢弃小部分数据不影响收敛时,才将该参数配置为True

数据分片策略配置

当训练数据的文件数目不能整除Master Worker数量时,Whale支持通过环境变量配置数据分片策略。您可以根据需求,在训练代码或执行脚本中配置参数:
  • 在训练代码中配置
    可以通过os.environ配置环境变量。例如,您可以使用如下命令,配置非均等切分(该参数为WHALE_UNBALANCED_IO_SLICING,详情请参见上面的数据分片参数列表部分)。
    import os
    os.environ["WHALE_UNBALANCED_IO_SLICING"]="True"
  • 在执行脚本中配置
    如果在PAI-DLC环境或手动开启训练,有Launch脚本,则可以在脚本中,配置环境变量:
    • 在执行脚本最开始,配置环境变量
      export WHALE_UNBALANCED_IO_SLICING=True
      ...
      python train.py
      其中train.py表示训练文件。
    • 在加载执行命令的语句前,配置环境变量。
      WHALE_UNBALANCED_IO_SLICING=True python train.py
      其中train.py表示训练文件。