[论文解读] Whale: Efficient Giant Model Training over Heterogeneous GPUs

Xianyan Jia, Le Jiang+10 more|arXiv (Cornell University)|Nov 18, 2020

Advanced Neural Network Applications被引用 10

一句话总结

Whale 是一个高级深度学习框架，通过引入两种模型注解原语，统一表达混合并行策略，实现了在异构 GPU 上对万亿参数模型的高效、可扩展训练。它自动化了分布式图优化和硬件感知负载均衡，仅需极少的代码修改，便在 512 块 V100 GPU 上实现了 10 万亿参数的训练，效率极高。

ABSTRACT

The scaling up of deep neural networks has been demonstrated to be effective in improving model quality, but also encompasses several training challenges in terms of training efficiency, programmability, and resource adaptability. We present Whale, a general and efficient distributed training framework for giant models. To support various parallel strategies and their hybrids, Whale generalizes the programming interface by defining two new primitives in the form of model annotations, allowing for incorporating user hints. The Whale runtime utilizes those annotations and performs graph optimizations to transform a local deep learning DAG graph for distributed multi-GPU execution. Whale further introduces a novel hardware-aware parallel strategy, which improves the performance of model training on heterogeneous GPUs in a balanced manner. Deployed in a production cluster with 512 GPUs, Whale successfully trains an industry-scale multimodal model with over ten trillion model parameters, named M6, demonstrating great scalability and efficiency.

研究动机与目标

• 解决在异构 GPU 集群上训练具有万亿参数的巨大深度学习模型所面临的挑战。
• 减少用户在手动管理复杂并行策略（如数据并行、模型并行和流水线并行）方面的工作量。
• 通过运行时感知优化，弥合模型开发与异构硬件环境之间的差距。
• 在无需代码更改的情况下，实现从 1000 亿到 1 万亿参数的无缝扩展。
• 通过自动图重写，以最小的性能开销支持混合并行策略。

提出的方法

• 通过模型注解引入两种高级原语，以表达所有并行策略及其混合形式，将编程与底层分发逻辑解耦。
• 利用这些注解指导自动化图优化，将本地计算图转换为高效的分布式执行计划。
• 采用硬件感知负载均衡算法，根据计算和内存容量将模型分区映射到 GPU，提升资源利用率。
• 应用图替换和模式匹配技术，对计算图进行重写，以在具有异构能力的多块 GPU 上实现高效执行。
• 与 TensorFlow 集成，支持模型组件的自动分发，包括张量并行和流水线并行，无需手动设备放置。
• 利用张量卸载和内存优化技术，突破 GPU 内存限制，实现 10 万亿参数的训练。

实验结果

研究问题

• RQ1 如何在最小用户干预下，使深度学习框架高效支持混合并行策略（如数据并行、模型并行、流水线并行）？
• RQ2 高级注解与自动化图优化能否在异构 GPU 集群上实现万亿参数模型的高效训练？
• RQ3 如何有效管理硬件异构性，以平衡工作负载并避免大规模模型训练中的瓶颈？
• RQ4 统一框架在从 1000 亿到 1 万亿参数的模型规模范围内，其性能和可扩展性如何？
• RQ5 图优化技术能否在大规模分布式训练中有效应用，特别是在复杂的混合并行场景下？

主要发现

• Whale 在由 512 块异构 V100 GPU 组成的集群上成功训练了一个 10 万亿参数的多模态模型（M6），展示了强大的可扩展性。
• 通过在不同模型组件上结合数据并行与模型并行，梯度同步开销降低了 89.7%。
• 训练规模从 1000 亿参数扩展到 1 万亿参数，无需代码更改，凸显了其灵活性与易用性。
• 硬件感知负载均衡算法通过有效平衡计算与内存工作负载，显著提升了在异构 GPU 上的性能。
• Whale 通过高级抽象显著降低了用户工作量，同时保持了高效率与可扩展性，优于现有框架。
• 该框架在 480 块 V100 GPU 上实现了对 1 万亿参数模型的高效训练，性能可与使用 2048 个 TPU 核心的最先进方案相媲美。