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表示训练文件。