在 comm_fusion=False 模式下，当前的实现将 all_reduce 的异步逻辑放在 pre_all_reduce_groups 内部处理，导致 wait_and_apply_grads 阻塞后续 reduce_scatter 的下发，无法实现跨模块的通信/计算重叠。
comm_fusion=False
pre_all_reduce_groups
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] =========================================================================================
class TorchHSDPStateV2(HSDPState): """Torch HSDP cell state""" pre_reduce_scatter_params: List = [] # 存储 AllReduceParamGroup 的 allreduce handles，在 root_backward_hook 中统一处理 pending_all_reduce_groups: List[AllReduceParamGroup] = []
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()
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()
@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 # 由调用者处理同步
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()
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
pending_all_reduce_groups
异步融合 AllReduce 流水线优化
Context
在
comm_fusion=False模式下，当前的实现将 all_reduce 的异步逻辑放在pre_all_reduce_groups内部处理，导致wait_and_apply_grads阻塞后续 reduce_scatter 的下发，无法实现跨模块的通信/计算重叠。分析
设计
修改方案
流水线设计
时序图
重叠效果
修改文件
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关键点
pre_all_reduce_groups: 存储刚发出 reduce_scatter 的组，等待下一次 post_backward 处理pending_all_reduce_groups: 存储已发出 allreduce 的组，等待 root_backward_hook 统一处理验证方案