背景

使用数据并行进行训练时，各节点存储的模型参数是相同的，模型参数更新过程使用的优化器状态也是相同的。从存储的角度看，集群中模型参数和优化器状态存在冗余存储。从计算角度看，节点间更新模型参数的优化器计算完全一致，存在冗余计算。如果将模型参数和优化器状态按节点个数切分，每个节点存储不同的模型参数切片和优化器状态切片，不仅消除了冗余存储，由于通过优化器参与模型参数更新的是切片数据，还能消除优化器的冗余计算。这种并行方式，我们一般称为优化器并行，又称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 接口入参分析

2.2 HSDPCell 核心方法分析

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

可以通过上述接口搭配，达到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 的主要作用是维护参数状态，进行参数分片计算，对张量生命周期进行管理。