Amazon SageMaker模型并行:一种通用且灵活的的大模型训练框架
机器学习
2021-11-12 v1 人工智能
分布式、并行与集群计算
摘要
随着深度学习模型规模迅速增长,需要系统级的大模型训练解决方案。我们提出Amazon SageMaker模型并行,一个与PyTorch集成的软件库,利用模型并行及其他节省内存的特性轻松训练大模型。与现有方案相比,SageMaker库的实现更加通用和灵活:它能以极小代码改动自动切分任意模型架构并运行流水线并行,同时还提供一个通用且可扩展的张量并行框架,支持更广泛的用例,并且模块化程度足以轻松应用于新的训练脚本。该库还在更大程度上保留了原生PyTorch用户体验,支持模块复用和动态图,同时让用户对训练步骤的细节拥有完全控制。我们在GPT-3、RoBERTa、BERT和神经协同过滤上评估了性能,并展示出相对于现有方案的竞争力。
引用
@article{arxiv.2111.05972,
title = {Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training},
author = {Can Karakus and Rahul Huilgol and Fei Wu and Anirudh Subramanian and Cade Daniel and Derya Cavdar and Teng Xu and Haohan Chen and Arash Rahnama and Luis Quintela},
journal= {arXiv preprint arXiv:2111.05972},
year = {2021}
}
备注
24 pages. Submitted for review