title: AMSP: Super-Scaling LLM Training via Advanced Model States Partitioning
summary: Large Language Models (LLMs) have demonstrated impressive performance across
  various downstream tasks. When training these models, there is a growing
  inclination to process more tokens on larger training scales but with
  relatively smaller model sizes. Zero Redundancy Optimizer (ZeRO), although
  effective in conventional training environments, grapples with scaling
  challenges when confronted with this emerging paradigm. To this end, we propose
  a novel LLM training framework AMSP, which undertakes a granular partitioning
  of model states, encompassing parameters (P), gradient (G), and optimizer
  states (OS). Specifically, AMSP(1) builds a unified partitioning space,
  enabling independent partitioning strategies for P, G, and OS; (2)
  incorporates a scale-aware partitioner to autonomously search for optimal
  partitioning strategies: (3) designs a dedicated communication optimizer to
  ensure proficient management of data placement discrepancies arising from
  diverse partitioning strategies. Our evaluations show that AMSP achieves up to
  90.3% scaling efficiency across 1024 GPUs.
paper_authors._id: 6543072e5670dd57b76d66ff
paper_authors.name: Qiaoling Chen
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-02T13:28:04.329Z
paper_authors._id: 654372ee91a22f5a118fdeb9
paper_authors.avatarUrl: /avatars/9a93f7a8ee000cf7b02ab061a4756216.svg
paper_authors.fullname: QiaolingChen
paper_authors.type: user
paper_authors.user: QiaolingChen
paper_authors._id: 6543072e5670dd57b76d6700
paper_authors.name: Qinghao Hu
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-02T13:28:11.156Z
paper_authors._id: 62fb8412610dae1bcd036db7
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/62fb8412610dae1bcd036db7/LIaEoDrvrq8hDl8509ZuF.jpeg
paper_authors.fullname: Qinghao Hu
paper_authors.type: user
paper_authors.user: Qinghao
paper_authors._id: 6543072e5670dd57b76d6701
paper_authors.name: Zhisheng Ye
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-02T13:28:18.603Z
paper_authors._id: 6384e0d5b2906edaf835d56c
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/6384e0d5b2906edaf835d56c/tEo0rap-hd02bVy_u1cJb.jpeg
paper_authors.fullname: Zhisheng Ye
paper_authors.type: user
paper_authors.user: yzs981130
paper_authors._id: 6543072e5670dd57b76d6702
paper_authors.name: Guoteng Wang
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-02T13:28:29.284Z
paper_authors._id: 6542391f3fcd1aee202383d2
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/noauth/5O7dZAVGmwhsx1THSJ-gn.jpeg
paper_authors.fullname: wang
paper_authors.type: user
paper_authors.user: guoteng
paper_authors._id: 6543072e5670dd57b76d6703
paper_authors.name: Peng Sun
paper_authors._id: 6543072e5670dd57b76d6704
paper_authors.name: Yonggang Wen
paper_authors._id: 6543072e5670dd57b76d6705
paper_authors.name: Tianwei Zhang
paper_title: AMSP: Super-Scaling LLM Training via Advanced Model States Partitioning
paper_summary: Large Language Models (LLMs) have demonstrated impressive performance across
  various downstream tasks. When training these models, there is a growing
  inclination to process more tokens on larger training scales but with
  relatively smaller model sizes. Zero Redundancy Optimizer (ZeRO), although
  effective in conventional training environments, grapples with scaling
  challenges when confronted with this emerging paradigm. To this end, we propose
  a novel LLM training framework AMSP, which undertakes a granular partitioning
  of model states, encompassing parameters (P), gradient (G), and optimizer
  states (OS). Specifically, AMSP(1) builds a unified partitioning space,
  enabling independent partitioning strategies for P, G, and OS; (2)
  incorporates a scale-aware partitioner to autonomously search for optimal
  partitioning strategies: (3) designs a dedicated communication optimizer to
  ensure proficient management of data placement discrepancies arising from
  diverse partitioning strategies. Our evaluations show that AMSP achieves up to
  90.3% scaling efficiency across 1024 GPUs.
paper_ai_keywords: Zero Redundancy Optimizer (ZeRO) Large Language Models (LLMs) tokens model states parameters ($P$) gradient ($G$) optimizer states ($OS$) scaling unified partitioning space scale-aware partitioner communication optimizer scaling efficiency