[USENIX ATC] Whale, uma estrutura de treinamento distribuída eficiente para modelos de hiperescala que suportam clusters de GPU heterogêneos

Introdução: Whale (EPL), uma estrutura eficiente de treinamento de grandes modelos, foi selecionada pelo USENIX ATC

Autor: Zhang Jie, Jia Xianyan

Recentemente, o artigo do Alibaba Cloud Machine Learning PAI "Whale: Efficient Giant Model Training over Heterogeneous GPUs" em uma estrutura de treinamento distribuída eficiente para modelos de aprendizado profundo foi aceito pela USENIX ATC'22, a principal conferência acadêmica internacional na área de sistemas de computador.

Whale é uma estrutura de treinamento distribuída desenvolvida pela plataforma PAI de aprendizado de máquina da Alibaba Cloud. O nome de código aberto é EPL (Easy Parallel Library). Whale abstrai e encapsula diferentes estratégias de paralelização e oferece suporte a ela em uma estrutura de treinamento distribuída. Uma variedade de estratégias paralelas e otimização abrangente de memória de vídeo, computação, comunicação, etc., para fornecer uma estrutura de treinamento distribuída eficiente e fácil de usar. Whale fornece uma interface concisa e fácil de usar, e os usuários podem combinar várias estratégias paralelas híbridas adicionando algumas linhas de anotação. Ao mesmo tempo, Whale fornece uma estratégia paralela distribuída automatizada baseada na percepção de hardware, percebe recursos como poder de computação e armazenamento de hardware diferente e divide razoavelmente o modelo de acordo com esses recursos, equilibra a quantidade de computação em hardware diferente e maximiza eficiência computacional. Com o algoritmo de balanceamento de carga com reconhecimento de hardware da Whale, os modelos Bert-Large, Resnet50 e GNMT são 1,2 a 1,4 vezes mais rápidos no treinamento de GPU heterogêneo. Enquanto isso, usando a estrutura Whale, o modelo de trilhões de M6 é pré-treinado em 3 dias usando 480 V100s. Em comparação com modelos anteriores da mesma escala treinados no setor, os recursos de computação são economizados em mais de 80% e a eficiência do treinamento é aumentada em quase 11 vezes. O uso adicional de 512 GPUs treinou 10 trilhões de modelos com níveis utilizáveis em 10 dias.

Histórico e desafios

Com a popularidade do aprendizado profundo nos últimos anos, a escala de parâmetros do modelo também cresceu rapidamente. Os dados do OpenAI mostram que:

Antes de 2012, o tempo de cálculo do modelo dobrava a cada 2 anos, de acordo com a Lei de Moore;
Após 2012, o tempo de cálculo do modelo dobrou a cada 3,4 meses, superando em muito a velocidade de desenvolvimento de hardware;

Em particular, a escala dos parâmetros do modelo cresceu rapidamente nos últimos dois anos. Google, NVIDIA, Ali e Zhiyuan Research Institute lançaram modelos de trilhões de parâmetros, e alguns grandes fabricantes também lançaram dezenas de bilhões e centenas de bilhões de modelos de parâmetros. À medida que o tamanho dos parâmetros do modelo aumenta, o efeito do modelo é gradualmente melhorado, mas também traz maiores desafios para a estrutura de treinamento. Já existem alguns frameworks de treinamento distribuídos, como Horovod, Tensorflow Estimator, PyTorch DDP, etc. que suportam paralelismo de dados, Gpipe, PipeDream, PipeMare, etc. suportam paralelismo de pipeline, Mesh Tensorflow, FlexFlow, OneFlow, MindSpore, etc. divisão, mas quando Existem alguns desafios quando queremos treinar um modelo muito grande:

As estruturas de treinamento distribuídas atuais suportam apenas um pequeno número de estratégias paralelas e não possuem uma abstração unificada para suportar todas as estratégias paralelas e suas estratégias híbridas.
A implementação de estratégias paralelas complexas requer muita modificação do código do modelo e compreensão do sistema subjacente, o que aumenta muito a dificuldade de uso dos usuários.
Devido à diferença de poder de computação e memória de vídeo de GPUs heterogêneas no cluster, a estratégia de paralelização estática não pode fazer uso total de recursos heterogêneos para obter um treinamento eficiente.

Quebra

Para enfrentar os desafios do treinamento distribuído atual, desenvolvemos o Whale, uma estrutura de treinamento distribuído eficiente, de uso geral e com reconhecimento de hardware. Whale abstrai e define duas primitivas distribuídas (replicar e dividir). Os usuários podem expressar e transformar várias estratégias de paralelização e suas combinações adicionando algumas linhas de marcadores primitivos ao modelo, o que reduz bastante o limite para o uso de estruturas distribuídas. O tempo de execução Whale funde as informações de marcação do usuário no gráfico de computação e conclui automaticamente a otimização paralela do modelo. Ao mesmo tempo, o Whale fornece uma estratégia paralela distribuída automatizada com base no reconhecimento de hardware para otimizar o desempenho de treinamento distribuído em clusters de GPU heterogêneos. O design do Whale well equilibra a intervenção do usuário do modelo e as oportunidades de otimização do sistema, permitindo que cada engenheiro de algoritmo treine com facilidade e eficiência grandes tarefas de modelo distribuído.

Arquitetura de tecnologia

Conforme mostrado na figura abaixo, o framework Whale é dividido principalmente nos seguintes módulos:

Camada de interface: a interface de programação do modelo do usuário é baseada no TensorFlow, e o Whale fornece uma interface de expressão de estratégia de paralelização fácil de usar, permitindo que os usuários combinem várias estratégias paralelas híbridas;
Camada de expressão intermediária: Converta o modelo de usuário e estratégias paralelas em expressões internas e expresse várias estratégias paralelas por meio de TaskGraph, VirtualDevices e abstração de estratégia;
并行化引擎层：基于中间表达，Whale会对计算图做策略探索，结合硬件资源进行显存/计算/通信优化，并自动生成分布式计算图。
Runtime执行引擎：将分布式执行图转成TFGraph，再调用TF 的Runtime来执行；

并行化策略表达

Whale通过strategy annotation的方式来划分模型为多个TaskGraph，并在此基础上进行并行化。 Whale有两类strategy：replicate 和 split。通过这两种并行化接口，可以表达出各种不同的并行化策略，例如：

数据并行: 下面这个例子是一个数据并行的例子，每个模型副本用一张卡来计算。如果用户申请了8张卡，就是一个并行度为8的数据并行任务。

流水并行：在下面的例子里，模型被切分成2个 TaskGraph, "stage0"和"stage1"，用户可以通过配置pipeline.num_micro_batch参数来设定pipeline的micro batch数量。在这个例子里，"stage_0"和"stage_1"组成一个模型副本，共需要2张GPU卡。如果用户申请了8张卡，Whale会自动在pipeline外嵌套一层并行度为4的数据并行（4个pipeline副本并行执行）。

算子拆分并行：在以下例子中，Whale会对split scope下的模型定义做拆分，并放置在不同的GPU卡上做并行计算。

同时Whale支持对上述并行策略进行组合和嵌套，来组成各种混合并行策略，更多示例可以参考开源代码的文档和示例。

Parallel Planner

Paraller Planner是Whale runtime的重要一环，它的职责是生成一个高效的分布式执行plan。Parallel Planner的流程包含 (a) Paraller Planner的输入包含用户模型、用户标记(可选)、硬件计算资源和其他用户配置配置。 (b) 将物理计算资源映射成VirtualDevice, 用户无需担心如何将算子放置在分布式物理设备上。(c) 模型被分割成TaskGraph子图。因为Whale允许对不同的TaskGraph应用不同的分布式策略，所以在TaskGraph之间可能存在输入/输出shape不匹配。在这种情况下，Paraller Planner将自动在两个TaskGraphs之间插入相应的桥接层。

硬件感知的负载均衡算法

当模型训练资源包含异构硬件（比如混合了V100和T4），硬件感知的负载均衡算法可以提高在异构资源下的训练效率。Whale设计了两种平衡策略：Intra-TaskGraph和Inter-TaskGraph的平衡。

（1）对于Intra-TaskGraph的计算平衡，Whale会profile出模型的算力FLOP，并按照机器的计算能力按比例分配对应的计算负载，以此均衡每个step不同卡型上的模型计算时间。对于数据并行，我们可以通过调整不同副本上的batch大小来实现计算负载的均衡（保持全局batch不变）。对于算子拆分，我们可以通过不均匀的维度拆分来实现不同卡上子模块的计算负载均衡。

（2）当多个TaskGraph在多个异构GPU上执行流水并行，我们需要Inter-TaskGraph的计算平衡来提升整体的计算效率。由于流水并行的计算特点，前面的TaskGraph会比后面的TaskGraph多缓存一些前向计算的结果，因此对于Transformer类模型均衡layer切分的情况下，前面的TaskGraph会有更大的显存压力。因此在TaskGraph放置的时候，Whale会优先将前面的TaskGraph放置在显存容量更大的GPU卡上。与此同时，Whale会按照GPU卡的算力去切分模型，使得切分后的TaskGraph计算load和GPU算力成正比。

应用示例

借助Whale框架，我们4行代码实现M6模型数据并行+专家并行的混合并行训练。

如下图所示，MoElayer采用专家并行，其他layer采用数据并行：

并首次在480 V100 上，3天内完成万亿M6模型的预训练。相比此前业界训练同等规模的模型，此次仅使用480张V100 32G GPU就成功训练出万亿模型M6，节省算力资源超80%，且训练效率提升近11倍。进一步使用512 GPU在10天内即训练出具有可用水平的10万亿模型。

结语

Whale 通过统一的抽象支持各种并行化策略，并提出简洁的并行化原语来提高框架的易用性。同时，Whale提供了异构硬件感知的自动分布式计算图优化，实现高效的分布式模型训练。我们希望Whale能够成为一个大规模深度学习训练的基石，进一步推进模型算法创新和系统优化，使大模型训练技术能够广泛应用在实际生产环境中。Whale已经开源(github.com/alibaba/Eas…)，欢迎大家来试用和共建。

论文名称：Whale: Efficient Giant Model Training over Heterogeneous GPUs
论文作者：贾贤艳，江乐，王昂，肖文聪，石子骥，张杰，李昕元，陈浪石，李永，郑祯，刘小勇，林伟
论文链接：arxiv.org/pdf/2011.09…
开源链接：github.com/alibaba/eas…
参考文献：

[1] Junyang Lin, An Yang, Jinze Bai, Chang Zhou, Le Jiang, Xianyan Jia, Ang Wang, Jie Zhang, Yong Li, Wei Lin, et al. M6-10t: A sharing-delinking paradigm for effi- cient multi-trillion parameter pretraining. arXiv preprint arXiv:2110.03888, 2021.

[2] Junyang Lin, Rui Men, An Yang, Chang Zhou, Ming Ding, Yichang Zhang, Peng Wang, Ang Wang, Le Jiang, Xianyan Jia, Jie Zhang, Jianwei Zhang, Xu Zou, Zhikang Li, Xiaodong Deng, Jie Liu, Jinbao Xue, Huiling Zhou, Jianxin Ma, Jin Yu, Yong Li, Wei Lin, Jingren Zhou, Jie Tang, and Hongxia Yang. M6: A chinese multimodal pretrainer, 2021.

[3] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021.

[4] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R Devanur, Gregory R Ganger, Phillip B Gibbons, and Matei Zaharia. Pipedream: gen- eralized pipeline parallelism for dnn training. In Pro- ceedings of the 27th ACM Symposium on Operating Systems Principles, pages 1–15, 2019.

了解更多精彩内容，欢迎关注我们的阿里灵杰公众号，获取大数据+AI最新资讯

原文链接：click.aliyun.com/m/100034876…

本文为阿里云原创内容，未经允许不得转载。