HSDP下通信完全掩盖

异步融合 AllReduce 流水线优化

Context

在 comm_fusion=False 模式下，当前的实现将 all_reduce 的异步逻辑放在 pre_all_reduce_groups 内部处理，导致 wait_and_apply_grads 阻塞后续 reduce_scatter 的下发，无法实现跨模块的通信/计算重叠。

分析

设计

post_backward 里面：
  1. 先执行上一次 reduce_scatter 的 wait
  2. 然后下发当前模块的 reduce_scatter
  3. 然后执行上一个模块的融合后的 allreduce（发出异步）
  4. 收集所有 allreduce 的 handle

_root_backward_hook 的 apply_final_reduce 最后：
  - 执行 delay_apply_reduce_grads 函数
  - 该函数对所有 allreduce 执行 wait_and_apply_grads

修改方案

流水线设计

post_backward 里面：
  1. 先执行上一次 reduce_scatter 的 wait
  2. 然后下发当前模块的 reduce_scatter
  3. 然后执行上一个模块的融合后的 allreduce（发出异步）
  4. 收集所有 allreduce 的 handle

_root_backward_hook 的 apply_final_reduce 最后：
  - 执行 delay_apply_reduce_grads 函数
  - 该函数对所有 allreduce 执行 wait_and_apply_grads

时序图

时间 →
=========================================================================================
Layer N-2 post_backward:
  │
  ├─ 1. wait 上一次 reduce_scatter (无)
  ├─ 2. issue 当前模块 reduce_scatter ──────────────────────────────┐
  ├─ 3. issue 上一个模块 allreduce (无)                               │ RS 在通信
  └─ 4. 收集 allreduce handle (无)                                    │
                                                                       ▼
Layer N-1 backward compute ────────────────────────────────────────────────────────────
                                                                       │
Layer N-1 post_backward:                                               │
  │                                                                    │
  ├─ 1. wait 上一次 reduce_scatter (Layer N-2) ←───────────────────────┘
  ├─ 2. issue 当前模块 reduce_scatter ──────────────────────────────┐
  ├─ 3. issue 上一个模块 allreduce (Layer N-2) ─────────────────────┼──┐
  └─ 4. 收集 allreduce handle                                        │  │ AR 在通信
                                                                      │  │
Layer N backward compute ────────────────────────────────────────────────▼───
                                                                      │  │
Layer N post_backward:                                                │  │
  │                                                                   │  │
  ├─ 1. wait 上一次 reduce_scatter (Layer N-1) ←──────────────────────┘  │
  ├─ 2. issue 当前模块 reduce_scatter ──────────────────────────────┐    │
  ├─ 3. issue 上一个模块 allreduce (Layer N-1) ─────────────────────┼──┐ │
  └─ 4. 收集 allreduce handle                                        │  │ │
                                                                      │  │ │
...                                                                   │  │ │
                                                                      │  │ │
root_backward_hook:                                                   │  │ │
  │                                                                   │  │ │
  └─ delay_apply_reduce_grads():                                      │  │ │
       ├─ wait_and_apply_grads(Layer N-2 allreduce) ←────────────────┘  │ │
       ├─ wait_and_apply_grads(Layer N-1 allreduce) ←───────────────────┘ │
       ├─ wait_and_apply_grads(Layer N allreduce)                         │
       └─ ...                                                            │
=========================================================================================

重叠效果

Layer N-2: [RS_issue] ───────────────────────────────────────────────────────────────────────────┐
                                                                                                  │
Layer N-1:               [backward_compute] ──────────────────────────────── [RS_wait]            │
                                                                                     [RS_issue] ───┼──┐
                                                                                     [AR_issue] ───┼──┼──┐
Layer N:                                [backward_compute] ───────────────────────────────────────┘  │  │
                                                                                                     │  │
Layer N-1:                                                                              [RS_wait]    │  │
                                                                                        [RS_issue] ──┼──┼──┐
                                                                                        [AR_issue] ──┼──┼──┼──┐
Layer N:                                                                                               │  │  │  │
...                                                                                                    │  │  │  │
                                                                                                       ▼  ▼  ▼  ▼
root_backward_hook:                                                                          [wait_all_AR_and_apply]
=========================================================================================

修改文件

1. state.py - 新增类变量存储 allreduce handles

class TorchHSDPStateV2(HSDPState):
    """Torch HSDP cell state"""
    pre_reduce_scatter_params: List = []
    # 存储 AllReduceParamGroup 的 allreduce handles，在 root_backward_hook 中统一处理
    pending_all_reduce_groups: List[AllReduceParamGroup] = []

2. state.py - 修改 post_backward

def post_backward(self, *unused):
    for hsdp_param in self.hsdp_params:
        hsdp_param.accumulate_unsharded_grad_if_needed()

    if not self.reduce_grads:
        if self.reshard_after_backward:
            self.shard()
        for hsdp_param in self.hsdp_params:
            hsdp_param.to_accumulated_grad_if_needed()
        return

    if not self.comm_fusion:
        # Step 1: wait 上一次的 reduce_scatter，issue 当前模块的 reduce_scatter
        self._wait_prev_and_issue_current_reduce_scatter()

        # Step 2: issue 上一个模块的融合 allreduce（异步）
        self._issue_prev_fused_allreduce()

        # Step 3: 处理不需要 all_reduce 的参数 (FSDP / 单副本)
        self._apply_no_allreduce_params()
    else:
        # comm_fusion=True 的现有逻辑
        comm_ctx = get_comm_ctx()
        if comm_ctx.all_reduce_param_group is not None:
            comm_ctx.all_reduce_param_group.wait_all_reduce_and_apply_grad()
            comm_ctx.all_reduce_param_group = None
        if comm_ctx.pre_param_group is not None:
            comm_ctx.pre_param_group.wait_reduce_scatter_and_issue_all_reduce()
            comm_ctx.pre_param_group = None
        self.param_group.foreach_reduce(reduce_scatter_reduce_op=self.reduce_op_type)

    if self.reshard_after_backward:
        self.shard()

3. state.py - 新增辅助方法

def _wait_prev_and_issue_current_reduce_scatter(self):
    """Step 1: wait 上一次 reduce_scatter，issue 当前模块 reduce_scatter"""
    # Wait 上一次的 reduce_scatter
    if TorchHSDPStateV2.pre_all_reduce_groups:
        prev_group = TorchHSDPStateV2.pre_all_reduce_groups[-1]
        for hsdp_param in prev_group.hsdp_params:
            hsdp_param.reduce_scatter_output()
            hsdp_param.clear_reduce_scatter_output()

    # Issue 当前模块的 reduce_scatter
    self._issue_reduce_scatter_for_current_module()

def _issue_prev_fused_allreduce(self):
    """Step 2: issue 上一个模块的融合 allreduce"""
    if TorchHSDPStateV2.pre_all_reduce_groups:
        prev_group = TorchHSDPStateV2.pre_all_reduce_groups[-1]
        prev_group.issue_async_allreduce()
        # 收集到 pending 队列，在 root_backward_hook 中统一处理
        TorchHSDPStateV2.pending_all_reduce_groups.append(prev_group)
        # 从 pre 队列移除
        TorchHSDPStateV2.pre_all_reduce_groups.pop()

def _apply_no_allreduce_params(self):
    """Step 3: 处理不需要 all_reduce 的参数"""
    need_synchronize = False
    while TorchHSDPStateV2.pre_reduce_scatter_params:
        pre_hsdp_param, pre_orig_dtype, _ = TorchHSDPStateV2.pre_reduce_scatter_params.pop(0)
        reduced_grad = pre_hsdp_param.reduce_scatter_output()
        pre_hsdp_param.clear_reduce_scatter_output()
        need_synchronize = pre_hsdp_param.apply_reduced_grad(reduced_grad, pre_orig_dtype) or need_synchronize

    if need_synchronize:
        if self.device.type == "npu":
            torch.npu.current_stream().synchronize()
        elif self.device.type == "cuda":
            torch.cuda.current_stream().synchronize()

4. state.py - 新增 delay_apply_reduce_grads

@classmethod
def delay_apply_reduce_grads(cls):
    """在 root_backward_hook 最后调用，统一处理所有 allreduce 的 wait 和梯度应用"""
    need_synchronize = False

    for group in cls.pending_all_reduce_groups:
        need_synchronize = group.wait_and_apply_grads() or need_synchronize

    cls.pending_all_reduce_groups.clear()

    if need_synchronize:
        # 这里需要获取 device，可能需要从 hsdp_state 传入
        pass  # 由调用者处理同步

5. scheduler.py - 修改 _root_backward_hook

def _root_backward_hook(self):
    apply_final_reduce = self.scheduler_state != FSDPSchedulerState.BACKWARD
    self._backward_hook()

    if apply_final_reduce:
        TorchHSDPSchedulerV2.root_bp_state = False
        with torch.profiler.record_function(f"root_backward reduce:{self.hsdp_state.module_name}"):
            comm_ctx = get_comm_ctx()

            # 排空 comm_fusion=True 的流水线
            if comm_ctx.all_reduce_param_group is not None:
                comm_ctx.all_reduce_param_group.wait_all_reduce_and_apply_grad()
                comm_ctx.all_reduce_param_group = None
            if comm_ctx.pre_param_group is not None:
                comm_ctx.pre_param_group.apply_fusion_reduced_grad()
                comm_ctx.pre_param_group = None

            # 处理最后一个模块的 reduce_scatter 和 allreduce
            if TorchHSDPStateV2.pre_all_reduce_groups:
                last_group = TorchHSDPStateV2.pre_all_reduce_groups[-1]
                # Wait reduce_scatter
                for hsdp_param in last_group.hsdp_params:
                    hsdp_param.reduce_scatter_output()
                    hsdp_param.clear_reduce_scatter_output()
                # Issue allreduce
                last_group.issue_async_allreduce()
                TorchHSDPStateV2.pending_all_reduce_groups.append(last_group)
                TorchHSDPStateV2.pre_all_reduce_groups.clear()

            # 处理非融合模式的剩余参数
            self.hsdp_state.reduce_params()

            # 最后统一处理所有 allreduce 的 wait 和梯度应用
            TorchHSDPStateV2.delay_apply_reduce_grads()

6. param_group.py - AllReduceParamGroup 新增方法

def issue_async_allreduce(self) -> None:
    """发出异步 allreduce（不等待）"""
    if self.fused_buffer is None:
        raise RuntimeError("Fused buffer not allocated.")

    self.all_reduce_handle = dist.all_reduce(
        self.fused_buffer,
        op=dist.ReduceOp.SUM,
        group=self.replicate_group,
        async_op=True,
    )

def wait_and_apply_grads(self) -> bool:
    """等待 allreduce 完成并应用梯度"""
    if self.all_reduce_handle is not None:
        self.all_reduce_handle.wait()
        self.all_reduce_handle = None

    need_synchronize = False
    for idx, hsdp_param in enumerate(self.hsdp_params):
        reduced_grad = self.get_param_grad_view(idx, hsdp_param.sharded_size)

        if self.reduce_op == dist.ReduceOp.AVG and self.replicate_world_size > 1:
            reduced_grad = reduced_grad / self.replicate_world_size

        need_synchronize = hsdp_param.apply_reduced_grad(
            reduced_grad, self.orig_dtypes[idx]
        ) or need_synchronize

    self.fused_buffer = None
    return need_synchronize

关键点

1. pre_all_reduce_groups: 存储刚发出 reduce_scatter 的组，等待下一次 post_backward 处理
2. pending_all_reduce_groups: 存储已发出 allreduce 的组，等待 root_backward_hook 统一处理
3. post_backward 顺序: wait_prev_RS → issue_current_RS → issue_prev_AR
4. root_backward_hook: 统一处理所有 pending allreduce 的 wait 和梯度应用

验证方案

1. 单元测试：验证 AllReduceParamGroup 的方法正确性
2. 集成测试：在 HSDP 配置下运行训练，验证梯度正确性
3. 性能测试：对比修改前后的通信时间，确认流水线重叠效果

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

Start by reading the existing post_backward flow in state.py, _root_backward_hook in scheduler.py, and AllReduceParamGroup in param_group.py, then compare the comm_fusion paths. Run the stated unit and HSDP integration tests before changing behavior. Done means gradients remain correct and performance testing confirms improved communication/compute overlap.