自动并行策略搜索 FSDP/HSDP/TP 显存与性能估算能力支持

本文档描述自动并行策略搜索中 FSDP/HSDP/TP 显存与性能估算能力的设计需求。不包含具体代码实现，侧重参数分片、优化器状态分片、张量并行切分、通信开销、接口契约、测试设计与验收标准。

1. 背景

Data Parallelism（DP）和 Tensor Parallelism（TP）是大模型训练中最基础的并行维度。DP 通过复制模型并切分数据 batch 来扩展训练规模，TP 通过切分模型参数、计算过程和中间激活来降低单卡计算与显存压力。近年来，DP 进一步演进出 Fully Sharded Data Parallelism (FSDP) / Hybrid Sharded Data Parallelism (HSDP) 等参数与优化器状态分片方案，在保持数据并行扩展能力的同时，显著降低了单卡显存开销。

自动并行策略搜索需要在不同 fsdp_degree、hsdp_degree 和 tp_degree 下估算显存占用、计算量和通信开销，用于过滤显存不可行策略并对候选策略进行性能排序。对于保留普通 DP 复制语义的场景，也应支持将 dp_degree 作为 HSDP 中的 replicate 维度或外层数据并行维度进行建模。

2. 目标与非目标

2.1 目标

• 支持 fsdp_degree、hsdp_degree、tp_degree 进入自动并行策略搜索空间；
• 支持 FSDP/HSDP/TP 组合策略的显存估算；
• 支持 FSDP/HSDP/TP 组合策略的性能估算；
• 支持 FSDP/HSDP 对参数、梯度、optimizer state 的分片建模；
• 支持 TP 对参数、计算过程和中间激活的切分建模；
• 支持 TP 相关整除约束校验；
• 支持 FSDP/HSDP 相关 shard / replicate 维度约束校验；
• 支持 global batch size 与数据并行相关维度的约束校验；
• 输出候选策略的 memory cost、communication cost 和 estimated step time；
• 补充单元测试和公式对齐测试。

2.2 非目标

• 本 issue 不实现 PP/EP/CP 的估算；
• 本 issue 不实现新的 TP runtime；
• 本 issue 不处理 MoE Expert Tensor Parallel 的完整建模；

3. 建模语义

3.1 显存估算

FSDP/HSDP/TP 显存估算应覆盖：

• 参数显存；
• 梯度显存；
• optimizer state 显存；
• activation 显存；
• FSDP/HSDP 对参数、梯度、optimizer state 的分片效果；
• HSDP 中 shard 维度与 replicate 维度对单卡显存的影响；
• TP 切分后的参数与中间 tensor 显存；
• TP collective 可能引入的临时 buffer；
• FSDP/HSDP 参数 all-gather、reduce-scatter 等通信相关 buffer；
• 数据并行梯度同步相关 buffer。

3.2 性能估算

FSDP/HSDP/TP 性能估算应覆盖：

• 计算量随 tp_degree 的切分变化；
• TP 组内 collective 通信开销，例如 all-gather、reduce-scatter、all-reduce；
• FSDP/HSDP 参数 all-gather、reduce-scatter、梯度同步等通信开销；
• HSDP 中 shard 组与 replicate 组对应的通信开销；
• 数据并行相关梯度同步通信开销；
• FSDP/HSDP/TP 组合下的端到端 step time 估算；
• 不同通信带宽参数下的性能变化。

3.3 约束

应检查：

• hidden_size、num_attention_heads、ffn_hidden_size 等是否可被 tp_degree 整除；
• global_batch_size、micro_batch_size、数据并行相关 degree 是否满足 batch size 关系；
• fsdp_degree、hsdp_shard_degree、hsdp_replicate_degree、tp_degree 的乘积不超过可用设备数；
• HSDP 的 shard 维度和 replicate 维度配置合法；
• 显存估算结果是否超过单卡 memory limit。

4. 实现要点

• FSDP/HSDP/TP 估算模块应能独立接收 normalized_model_config、cluster_config 和 strategy_config；
• FSDP/HSDP/TP 估算结果应包含 memory breakdown 和 performance breakdown；
• memory breakdown 中应区分参数、梯度、optimizer state、activation 和通信 buffer；
• performance breakdown 中应区分 TP collective、FSDP/HSDP 参数通信和数据并行梯度同步通信；
• 对不可行策略应返回 infeasible 状态和具体原因；
• 对通信开销应至少支持基于通信量和带宽的 symbolic 估算；
• 估算结果应可被自动并行策略搜索器统一排序和过滤。

5. 测试设计

6. 验收标准

• fsdp_degree、hsdp_degree、tp_degree 可进入自动并行策略搜索空间；
• FSDP/HSDP/TP 组合策略可输出 memory cost 和 performance cost；
• 显存估算包含参数、梯度、optimizer state、activation 和通信 buffer；
• 显存估算能够体现 FSDP/HSDP 对参数、梯度、optimizer state 的分片效果；
• 性能估算包含 TP collective、FSDP/HSDP 参数通信和数据并行梯度同步开销；
• 不合法 TP 整除关系、HSDP 维度配置和 batch size 关系可被校验并过滤；
• AP-FT-01～AP-FT-08 测试通过；
• 文档中包含 FSDP/HSDP/TP 建模范围、约束条件和典型示例。

7. 参考

• Megatron-LM：Tensor Parallelism 与 Data Parallelism；
• TorchTitan / PyTorch DTensor：DeviceMesh 与 TP parallel style；
• HyperParallel：tensor parallel、dtensor、自动并行策略搜索相关模块。