本文针对大规模分类场景存在的问题,介绍Whale的并行化设计和方案。通过结合算子拆分和数据并行,优化通信拓扑结构,以解决大规模分类任务无法单机训练或分布式训练性能较差的问题。在Whale中,您可以通过模型划分、资源划分及映射三个步骤,实现大规模分类任务的Fusion模式。

背景信息

分类问题是机器学习、数据挖掘的基础问题。针对大数据场景,分类数据通常达到数百万、千万或亿级别。随着分类规模增长,模型参数和模型规模本身呈超线性增长,超大规模分类是业务上普遍存在的问题,例如人脸识别分类数达到亿级别。

问题

以ResNet50分类任务为例,其模型结构如下图所示。ResNet50分类模型结构ResNet50分类的原始模型代码如下。
features = ResNet50(inputs)
logits = FC(features)
predictions = Softmax(logits)

本文选择可以使用数据并行的10 W分类数,以便与算子拆分进行性能对比。分类数为10 W时,ResNet50部分的模型权重大小89.6 MB,全连接层FC(Fully Connected Layer)部分的模型权重大小为781.6 MB,是前者的8.7倍。当采用数据并行进行分布式训练时,在后向阶段(Back Propagation过程),FC部分计算得到梯度后,立刻通过AllReduce进行梯度同步,同时ResNet50部分的梯度会继续进行后向计算。然而由于FC部分同步的梯度过大,当ResNet50部分的梯度计算完成时,FC部分的梯度通常还在通信过程中。因此,FC部分的梯度同步无法良好地与ResNet50的后向计算Overlap,导致整轮迭代过程中的通信占比十分高,性能低下。

当分类数更大时,分布式计算模式对显存有更高的要求。以Nvidia 16 GB V100为例,假设隐层数为2048,仅能支持存储190 W的FC权重。再加上中间输出、计算得到的梯度及Optimizer的中间变量,16 GB的V100卡能够支持的分类数约为20 W(未考虑卷积层的显存占用)。

如上所述,即使将分类层(FC部分)单独放到一张GPU卡上也无法支持百万规模的分类任务,现有数据并行方法无法支持超大规模分类。针对超大规模分类任务,必须对分类层进行OP内的算子拆分,将单个OP拆分为多分片,在不同设备中进行分片计算。

在此场景中,需要考虑分类规模和数据规模,超大的分类数必然伴随着大规模的训练数据。算子拆分可以解决模型本身很大的问题,例如上述的大规模分类层或BertLarge模型。对于卷积层(基于TF-Slim库实现的Resnet50模型Backbone,是通过卷积函数实现的),其模型本身较小,需要通过数据并行解决数据规模大的问题。因此,目前需要解决的问题是如何结合算子拆分和数据并行,进行分布式训练,同时提高通信性能。

实现方案

针对模型不同部分,采用不同的并行化策略加速分布式训练,以提高分布式性能:
  • FC部分

    权重大,使得梯度同步开销大。为避免数据并行同步过大引起性能问题,也为解决单卡无法存放超大规模分类的问题,先将FC部分的Weight权重按照列进行拆分,再使用多卡存储分片。因此,针对FC和Softmax计算,每张卡只负责不同部分的计算。

  • ResNet50部分

    权重小,计算耗时长,该部分梯度同步开销小。对ResNet50部分进行数据并行处理,每张卡读取不同的数据分片,以加速数据处理速度。

综上所述,针对ResNet50的大规模分类任务,将模型分为两个Stage。将Stage 0的ResNet50部分通过数据并行策略复制N份至不同的卡中,进行数据并行。将Stage 1的FC和Softmax部分通过算子拆分策略分片至不同卡中,进行模型并行。假设共六张卡(GPU0、GPU1、GPU2、GPU3、GPU4及GPU5),将GPU分为两组,分别进行数据并行和算子拆分,如下图所示。混合并行上图中,将Stage 0部分分别拷贝至[GPU0, GPU1, GPU2, GPU3]四张卡中进行数据并行,每张卡读取不同的数据分片进行训练。将Stage 1部分的计算拆分为两片放至[GPU4, GPU5]中进行算子拆分,并行化的计算图如下所示。并行化的计算图上图中,包括以下处理过程:
  1. [GPU0, GPU1, GPU2, GPU3]完成ResNet50的前向过程后,[GPU4, GPU5]收到[GPU0, GPU1, GPU2, GPU3]输出的Feature特性。
  2. [GPU4, GPU5]先将Feature特性按照行进行Concat,再进入FC部分。FC详细的计算图如下所示。FC计算图
  3. [GPU4, GPU5]分别得到输出的Logits后,分别计算Softmax部分(该过程需要进行跨卡通信,以获得全局最大Logits信息和每行的Sum值)。
通过上述方案虽然可以完成大规模分类任务,但是运行过程中任然存在以下问题:
  • 数据并行部分需要点对点发送Feature至算子拆分部分,存在热点问题。
  • 迭代过程中存在GPU空闲。当数据并行计算时,算子拆分GPU空闲等待。当算子拆分部分计算时,数据并行GPU空闲等待。
该问题严重影响性能,尤其是热点问题会导致分布式规模无法良好的扩展。因此,Whale对此进行了进一步优化,即采用Fusion模式计算。
Whale将Stage 0的ResNet50部分映射至所有卡中进行数据并行,将Stage 1的FC和Softmax部分映射至所有卡中进行分片计算。模型与硬件资源的映射如下图所示。模型与硬件映射Stage 0的输出从原来的点对点发送优化为去中心的AllGather通信,避免了热点问题。同时,每张卡中Stage 0计算结束,进行AllGather通信后,立刻计算Stage1部分,不会出现GPU空闲等待。具体的计算图如下所示。优化后的计算图

Whale实现多分类任务

在Whale中实现大规模分类任务Fusion模式需要三个步骤(标准的Whale分布式编程范式):模型划分、资源划分及映射。

  1. 模型划分。
    将模型各部分按照不同并行化需求划分为不同部分,即Scopes划分,详情请参见whale.scopes。上述的实现方案中将模型划分为两部分:
    • Stage 0:包括Resnet50部分。
    • Stage 1:包括FC和Softmax部分。
    针对ResNet50模型的原始代码,在Whale中通过两行Scopes代码(whale.replicawhale.split)即可实现模型划分,示例代码如下。
    with whale.replica():    # Stage 0
        features = ResNet50(inputs)
        
    with whale.split():      # Stage 1
        logits = FC(features)
        predictions = Softmax(logits)
  2. 资源划分。

    对硬件资源进行分组,将其划分为一个Virtual Devices(包含 [GPU0, GPU1, GPU2, GPU3, GPU4, GPU5]),即Cluster划分,详情请参见whale.cluster

    Whale提供的Cluster工具可以对申请的Worker进行划分,whale.cluster默认即可将所有Device虚拟为一个Virtual Device,示例代码如下。
    cluster = whale.cluster()
  3. 映射。
    将模型各部分分别映射至Virtual Device:
    • 对Resnet50部分采用数据并行,放至Virtual Device,即[GPU0, GPU1, GPU2, GPU3, GPU4, GPU5]每张卡都有一个Resnet50的模型副本。
    • FC和Softmax部分采用算子拆分,放至Virtual Device,即[GPU0, GPU1, GPU2, GPU3, GPU4, GPU5]每张卡先分别计算FC内的Matmul的一个分片,再分片进行Softmax计算。
    对于映射部分,在Whale中只需要对已生成的Cluster执行with语法,即可轻松完成模型到硬件资源的映射。通过Whale完成并行化的模型核心代码如下(包含算子拆分的大规模分类任务可执行代码请参见large_scale_classification.py)。
    cluster = whale.cluster() # 资源划分。
    
    with cluster:             # 映射。
        with whale.replica(): # Stage 0。
            features = ResNet50(inputs)
    
        with whale.split():   # Stage 1。
            logits = FC(features)
            predictions = Softmax(logits)

性能

以Whale的数据并行性能数据为Baseline进行对比,测试环境如下。
测试环境项 描述
GPU型号 ecs.gn6v-c10g1.20xlarge(V100 * 8)
网络 VPC-35 GB
NCCL_MAX_NRINGS NVIDIA官方参数,测试时取值为4。
NCCL_MIN_NRINGS NVIDIA官方参数,测试时取值为4。
在算子拆分的性能测试过程中,采用10 W分类的ResNet50模型。将FC部分通过算子拆分放至分布式服务器上进行分片计算。另外,卷积部分和分片的FC部分采用Fusion模式复用Worker。由于进行了算子拆分,FC部分的显存会分摊至其他卡,因此在算子拆分场景下数据并行的最大BatchSize可以变得更大。本文中的算子拆分对比测试场景分别为:
  • 数据并行:Whale Data Parallelism Batch Size=16
  • 固定Batch Size的算子拆分:Whale Model Parallelism Batch Size=16
  • 动态Batch Size的算子拆分,即Whale Model Parallelism + Dynamic Batch Size
对比结果如下图所示。性能对比
说明 横坐标表示GPU卡数,纵坐标表示加速比。
从上图中,可以得到以下结论:
  • 对于单机多卡或跨机任务,算子拆分的性能远优于数据并行任务。其原因主要是减少了FC部分的梯度同步通信量。
  • 根据GPU卡数动态调整Batch Size的算子拆分性能,相比数据并行的最大Batch Size性能更优。在64卡时,Batch Size可以调整到160,是数据并行的10倍。
  • 性能加速结果呈超线性加速。其原因主要因为动态调大Batch Size,使Apply阶段(参数更新)的计算量和通信量均减少了90%。