Skip to content

【训练营】Pipeline 并行自定义布局 - #223

Open
CuiLingyunCrispy wants to merge 4 commits into
InfiniTensor:masterfrom
CuiLingyunCrispy:feat/pipeline-custom-layout
Open

CuiLingyunCrispy wants to merge 4 commits into
InfiniTensor:masterfrom
CuiLingyunCrispy:feat/pipeline-custom-layout

Conversation

@CuiLingyunCrispy

Copy link
Copy Markdown

背景

为 InfiniTrain 增加统一的 Pipeline 并行自定义布局(Pipeline Layout)能力:用户可通过命令行显式配置各 Pipeline Stage 的 Transformer 层范围及 Embedding / FinalNorm / LMHead 的归属,并支持自动负载均衡布局建议与 Pipeline 负载分析。

主要改动

统一 PipelineLayout 数据结构,按用户提供的每层计算代价或按每层参数量等自动建议布局;对布局合法性进行校验以及分析Pipeline 负载分析

对标标准

  • 全部通过标准
  • 支持根据各层参数量、Profiler 统计或用户提供的计算代价,自动生成近似负载均衡的布局建议。
  • 给出默认均匀布局和自定义布局的 Pipeline bubble、各 Stage 执行时间及吞吐对比,证明自定义布局能够改善负载不均衡场景。

测试

(1)单元测试:tests/distributed/test_pipeline_layout.cc(12 用例)+ test_pipeline_layout_suggest.cc(22 用例),全部通过
(2)端到端:scripts/verify_pipeline_layout_correctness.sh(PP=1 单卡 vs PP=2 7,5,逐 step loss 对比)

文档

  • 使用指导:docs/pipeline_layout_guide.md
  • 离线对比演示:docs/pipeline_layout_demo.cc

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant