[论文解读] Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training

Can Karakus, Rahul Huilgol+8 more|arXiv (Cornell University)|Nov 10, 2021

Topic Modeling参考文献 28被引用 8

一句话总结

本文介绍了 Amazon SageMaker 模型并行，这是一种灵活、通用的框架，通过 PyTorch 中的自动流水线并行和张量并行，实现高效的大模型训练。该框架无需代码修改即可支持多种架构——包括非 Transformer、异构及 MoE 模型——同时保留原生 PyTorch 特性，并在性能上与 DeepSpeed 竞争。

ABSTRACT

With deep learning models rapidly growing in size, systems-level solutions for large-model training are required. We present Amazon SageMaker model parallelism, a software library that integrates with PyTorch, and enables easy training of large models using model parallelism and other memory-saving features. In contrast to existing solutions, the implementation of the SageMaker library is much more generic and flexible, in that it can automatically partition and run pipeline parallelism over arbitrary model architectures with minimal code change, and also offers a general and extensible framework for tensor parallelism, which supports a wider range of use cases, and is modular enough to be easily applied to new training scripts. The library also preserves the native PyTorch user experience to a much larger degree, supporting module re-use and dynamic graphs, while giving the user full control over the details of the training step. We evaluate performance over GPT-3, RoBERTa, BERT, and neural collaborative filtering, and demonstrate competitive performance over existing solutions.

研究动机与目标

• 解决缺乏通用、灵活的大模型训练框架的问题，以支持多样化的模型架构和使用场景。
• 通过自动划分和极少代码修改，降低新训练脚本的集成负担。
• 在分布式训练系统中完整保留原生 PyTorch 特性，如动态计算图和模块复用。
• 为多 GPU 和多节点环境提供可扩展、可扩展的流水线并行与张量并行系统。
• 与现有方案（如 DeepSpeed）相比，实现具有竞争力的训练性能。

提出的方法

• 该框架采用模块-服务器架构，实现对任意模型架构在设备间的自动、负载均衡的自动划分。
• 提出一种通用的张量并行框架，支持统一的大模型、具有孤立大组件的模型，以及混合专家（MoE）架构。
• 设计专用的设备到设备（D2D）通信后端，处理设备和节点间动态生成的通信请求。
• 系统提供灵活的 API，通过钩子（init_hook、forward_hook、return_hook）将原生 PyTorch 模块映射到其分布式对应物，无需修改原始代码。
• 包含实用工具如 parameter_creation_scope 和分布式张量操作（例如 allgather、allreduce、reduce-scatter），用于细粒度控制张量并行。
• 提供上下文管理器 smp.delay_param_initialization，将参数初始化推迟到模型部署阶段，从而减轻 CPU 内存压力。

实验结果

研究问题

• RQ1 能否构建一个通用的模型并行框架，支持任意模型架构且仅需极少代码修改？
• RQ2 如何将张量并行泛化以支持非均匀模型、大嵌入层及混合专家架构？
• RQ3 在分布式训练系统中，原生 PyTorch 特性（如动态图和模块复用）能在多大程度上被保留？
• RQ4 所提框架在真实训练工作负载中的性能与现有方案（如 DeepSpeed）相比如何？
• RQ5 自动负载均衡划分能否有效应用于复杂、异构的模型架构？

主要发现

• 该框架成功训练了 GPT-3、RoBERTa、BERT 和神经协同过滤模型，其训练吞吐量与 DeepSpeed 相比具有竞争力。
• 自动划分在无需手动分析架构或重新配置的情况下，实现了设备间的有效负载均衡。
• 系统完整保留了 PyTorch 的动态计算图特性，支持条件执行流和参数复用。
• D2D 通信后端能高效处理动态生成的通信模式，降低复杂训练工作流中的开销。
• 使用 smp.delay_param_initialization 可在 CPU 内存受限时训练包含数十亿参数的模型。
• 张量并行框架支持的使用场景范围广于以往方案，包括 MoE 和非 Transformer 架构。