AMSP: Reducing Communication Overhead of ZeRO for Efficient LLM Training
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2023
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866913263807102976 |
|---|---|
| author | Chen, Qiaoling Hu, Qinghao Wang, Guoteng Xiong, Yingtong Huang, Ting Chen, Xun Gao, Yang Yan, Hang Wen, Yonggang Zhang, Tianwei Sun, Peng |
| author_facet | Chen, Qiaoling Hu, Qinghao Wang, Guoteng Xiong, Yingtong Huang, Ting Chen, Xun Gao, Yang Yan, Hang Wen, Yonggang Zhang, Tianwei Sun, Peng |
| contents | Training large language models (LLMs) encounters challenges in GPU memory consumption due to the high memory requirements of model states. The widely used Zero Redundancy Optimizer (ZeRO) addresses this issue through strategic sharding but introduces communication challenges at scale. To tackle this problem, we propose AMSP, a system designed to optimize ZeRO for scalable LLM training. AMSP incorporates three flexible sharding strategies: Full-Replica, Full-Sharding, and Partial-Sharding, and allows each component within the model states (Parameters, Gradients, Optimizer States) to independently choose a sharding strategy as well as the device mesh. We conduct a thorough analysis of communication costs, formulating an optimization problem to discover the optimal sharding strategy. Additionally, AMSP optimizes distributed LLM training by efficiently overlapping communication with computation. Evaluations demonstrate up to 52\% Model FLOPs Utilization (MFU) when training the LLaMA-based model on 1024 GPUs, resulting in a 1.56 times improvement in training throughput compared to newly proposed systems like MiCS and ZeRO++. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2311_00257 |
| institution | arXiv |
| publishDate | 2023 |
| record_format | arxiv |
| spellingShingle | AMSP: Reducing Communication Overhead of ZeRO for Efficient LLM Training Chen, Qiaoling Hu, Qinghao Wang, Guoteng Xiong, Yingtong Huang, Ting Chen, Xun Gao, Yang Yan, Hang Wen, Yonggang Zhang, Tianwei Sun, Peng Distributed, Parallel, and Cluster Computing Training large language models (LLMs) encounters challenges in GPU memory consumption due to the high memory requirements of model states. The widely used Zero Redundancy Optimizer (ZeRO) addresses this issue through strategic sharding but introduces communication challenges at scale. To tackle this problem, we propose AMSP, a system designed to optimize ZeRO for scalable LLM training. AMSP incorporates three flexible sharding strategies: Full-Replica, Full-Sharding, and Partial-Sharding, and allows each component within the model states (Parameters, Gradients, Optimizer States) to independently choose a sharding strategy as well as the device mesh. We conduct a thorough analysis of communication costs, formulating an optimization problem to discover the optimal sharding strategy. Additionally, AMSP optimizes distributed LLM training by efficiently overlapping communication with computation. Evaluations demonstrate up to 52\% Model FLOPs Utilization (MFU) when training the LLaMA-based model on 1024 GPUs, resulting in a 1.56 times improvement in training throughput compared to newly proposed systems like MiCS and ZeRO++. |
| title | AMSP: Reducing Communication Overhead of ZeRO for Efficient LLM Training |
| topic | Distributed, Parallel, and Cluster Computing |
| url | https://arxiv.org/abs/2311.00257 |