【RFC】hyper_parallel.hsdp 兼容 torch.distributed.fsdp.fully_shard - HSDPParam

背景

使用数据并行进行训练时，各节点存储的模型参数是相同的，模型参数更新过程使用的优化器状态也是相同的。从存储的角度看，集群中模型参数和优化器状态存在冗余存储。从计算角度看，节点间更新模型参数的优化器计算完全一致，存在冗余计算。如果将模型参数和优化器状态按节点个数切分，每个节点存储不同的模型参数切片和优化器状态切片，不仅消除了冗余存储，由于通过优化器参与模型参数更新的是切片数据，还能消除优化器的冗余计算。这种并行方式，我们一般称为优化器并行，又称ZeRO（Zero Redundancy Optimizer）。
目前，业界中主要使用的是Torch的fully_shard接口，其本质上就是ZeRO-3的实现。

方案

基本原理

以业界常见的优化器并行方案ZeRO为例，有ZeRO-1，ZeRO-2，ZeRO-3三种优化级别：

ZeRO1

• 切分对象: 权重 (Weights) 和 优化器状态 (Optimizer State)。
• 行为: 模型参数和优化器状态被切分到 Sharding 组内的不同设备上。正向执行前通过AllGather获取完整模型参数，反向结束释放AllGather的显存并对梯度做ReduceScatter，获取梯度分片用于更新参数分片。

ZeRO2

• 切分对象: 权重 (Weights)、优化器状态 (Optimizer State) 以及 累加梯度 (Gradients)。
• 行为: 正向执行前通过AllGather获取完整模型参数，所有Micro Step结束后才释放AllGather的显存，相较于ZeRO-1，每一个Micro Step的梯度都经过ReduceScatter之后再进行累加，可以降低训练过程中的显存压力，但是会带来更大的通信开销。

ZeRO3

• 切分对象: 权重 (Weights)、优化器状态 (Optimizer State)、累加梯度 (Gradients)。
• 行为：正向执行前通过AllGather获取完整模型参数，正向结束后立刻释放AllGather所占用的空间，反向开始前再重新做AllGather。反向结束再释放AllGather的显存，完整的模型参数并未常驻内存，实现真正意义上的参数切分。梯度处理方面与ZeRO-2一致，每一个Micro Step的梯度经过ReduceScatter之后再进行累加。

接口设计

在接口设计上，我们选择兼容Torch的fully_shard接口，设计如下

def hsdp(
           cell,
           *,
           mesh,
           reshard_after_forward,
           shard_placement_fn,
           mp_policy,
           offload_policy,
           ignored_params
   )

2.1 fully_shard 接口入参分析

入参名称	功能含义
cell	需要进行分布式并行封装的原始网络模块或子层。
mesh	定义参数分片与副本复制通信域的设备拓扑网格（DeviceMesh）。
reshard_after_forward	控制前向计算结束后，是否立即释放完整参数并回归至分片状态。
shard_placement_fn	自定义参数切分逻辑。
mp_policy	设置模型在计算、存储与通信各环节中的混合精度策略。
offload_policy	定义参数或梯度从显存卸载（Offload）至 CPU 内存的调度策略。
ignored_params	明确指定不参与 HSDP 参数分片与通信调度的参数集合。

2.2 HSDPCell 核心方法分析

HSDPCell 提供以下方法实现对模型运行时状态的精细化调控。

方法名称	功能含义
set_requires_gradient_sync	配置当前模块在反向传播过程中是否触发梯度同步通信。
set_forward_prefetch_cells	正向预取。
set_backward_prefetch_cells	反向预取。
reshard	手动触发逻辑，强制将完整参数切分为分片状态并释放多余显存。
unshard	手动触发逻辑，立即通过通信补全分片参数至完整状态。
set_is_last_backward	标记当前反向过程是否为本轮训练步的最后一次计算迭代。用于资源清理
set_reshard_after_forward	动态设置前向计算结束后是否重新切分参数。
set_reshard_after_backward	动态设置反向计算结束后是否重新切分参数。
set_requires_all_reduce	HSDP场景是否进行AllReduce通信。

可以通过上述接口搭配，达到ZeRO-1， ZeRO-2，ZeRO-3级别的优化器并行切分策略：
ZeRO-1 = rehsard_after_forward=True + set_require_gradient_sync=False
ZeRO-2 = rehsard_after_forward=True + set_require_gradient_sync=True
ZeRO-3 = rehsard_after_forward=True + set_require_gradient_sync=True

2.3 HSDPParam

HSDPParam 的主要作用是维护参数状态，进行参数分片计算，对张量生命周期进行管理。

方法名称	功能含义
_init_sharded_param	参数分片
_init_sharded_post_forward_param_metadata	根据post-forward 元数据参数分片
init_all_gather_outputs	预分配 all-gather 输出缓冲区
init_unsharded_param	构建完整参数
to_sharded	切换到shard
to_sharded_post_forward	切换到 post-forward
to_unsharded	切换到完unshard
to_sharded_dtensor	转换为 sharded DTensor
to_sharded_post_forward_dtensor	转换为 post-forward DTensor
alloc_all_gather_outputs	分配 all-gather 输出内存
free_unsharded_param	释放完整参数内存
all_gather_inputs	获取 all-gather 输入
to_accumulated_grad_if_needed	梯度累积处理
accumulate_unsharded_grad_if_needed	累积梯度

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 14
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/14

Start by locating the hyper_parallel.hsdp and HSDPParam implementations, then compare their existing interfaces and state transitions with torch.distributed.fsdp.fully_shard. Review the listed HSDPCell and HSDPParam methods and determine the compatibility scope. Done means the agreed interface and ZeRO behavior are implemented and validated by the relevant tests, though no test files are named here.