Atten res cache#79517
Conversation
… optimization Complete the backward path for the block-attention-residual (BlockAttnRes) pipeline communication optimization under virtual pipeline parallel (VPP): - _can_free: skip releasing tensors marked pp_block_cached so cached received blocks keep their data pointer for same-stage reuse. - _update_block_cache: split cached blocks into received (raw object, pp_block_cached) and produced (detach cc, pp_cc_ref) so producer-chunk autograd graph is not reused across chunks. - _backward_step: inject cc.grad accumulated by same-stage later chunks back into the producer block's grad_tensors before autograd.backward. - _forward_step_helper: tag output tuple with pp_mb_idx. - _backward_step_helper: free per-microbatch block cache / meta after the last logical chunk backward. Mathematically equivalent to the non-optimized path; differences are bounded bf16 accumulation-order noise.
|
你的PR提交成功，感谢你对开源项目的贡献! |
|
root seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account. You have signed the CLA already but the status is still pending? Let us recheck it. |
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
已复查，当前还有几处阻塞性问题，细节见 inline review comments。
| if not self._dynamic_shape: | ||
| if not self._send_recv_meta.has_send_meta: | ||
| self._send_recv_meta.set_send_message(output_tensor) | ||
| self._send_recv_meta.send_meta( | ||
| output_tensor, | ||
| _hcg.get_pipe_parallel_group(), | ||
| reverse=reverse, | ||
| block_cache_meta=block_cache_meta, | ||
| ) |
There was a problem hiding this comment.
这里把逐 microbatch 变化的 block_cache_meta 绑进了 has_send_meta/has_recv_meta 的缓存协议里，但 p2p_cache_shape=True 时这两个标志会把后续 step 直接短路掉。结果是第二个 microbatch 起，接收端拿到的还是第一次的 recv_block_cache_meta，_merge_block_cache() 会继续按旧 meta 计算块数。
处理要求：请针对该评论修复并提交新的 commit。
建议把 block meta 从 shape meta 里拆出去，至少要保证它每次都会独立发送/接收，而不是复用缓存分支：
if block_cache_meta is not None:
# 每次都单独传递 block_cache_meta
...
else:
# 继续沿用当前的 shape cache 逻辑
...| recv_prev=recv_prev, | ||
| batch_p2p_comm=self._use_batch_p2p_comm, | ||
| skip_check_meta=not self.training, | ||
| block_cache_meta=meta_to_send, |
There was a problem hiding this comment.
这里把 block_cache_meta 传下去了，但这个 scheduler 里没有像主 interleave 路径那样把刚收到的 meta 持久化到 _recv_block_cache_meta_for_vpp[next_virtual_pp_rank]。这样 _forward_step_helper() 里的 _merge_block_cache() 下一轮拿到的 received_meta 仍然是 None，收到的块计数不会合进来。
处理要求：请针对该评论修复并提交新的 commit。
函数开头那次初始 recv_forward 之后也要补同样的保存，和主 interleave 路径保持一致，例如：
self.input_tensors[0].append(self._p2p_helper.recv_forward(...))
if self._block_atten_res_opt:
self._recv_block_cache_meta_for_vpp[0] = self._p2p_helper.get_recv_block_cache_meta()
input_tensor = self._p2p_helper.send_forward_recv_forward(...)
if self._block_atten_res_opt:
self._recv_block_cache_meta_for_vpp[next_virtual_pp_rank] = (
self._p2p_helper.get_recv_block_cache_meta() if recv_prev else None
)| recv_prev=recv_prev, | ||
| batch_p2p_comm=self._use_batch_p2p_comm, | ||
| skip_check_meta=not self.training, | ||
| block_cache_meta=meta_to_send, |
There was a problem hiding this comment.
这里同样缺少把本次接收到的 block meta 持久化到 _recv_block_cache_meta_for_vpp[next_forward_virtual_pp_rank]。当前只有 input_tensor 入队，_forward_step_helper() 下一轮拿到的 received_meta 会是 None，块计数对齐逻辑就失效了。
处理要求：请针对该评论修复并提交新的 commit。
函数开头那次 recv_forward 之后也要补同样的保存，避免第一个 virtual rank 的 meta 直接丢失，例如：
self.input_tensors[0].append(self._p2p_helper.recv_forward(...))
if self._block_atten_res_opt:
self._recv_block_cache_meta_for_vpp[0] = self._p2p_helper.get_recv_block_cache_meta()
input_tensor = self._p2p_helper.send_forward_recv_forward(...)
if self._block_atten_res_opt:
self._recv_block_cache_meta_for_vpp[next_forward_virtual_pp_rank] = (
self._p2p_helper.get_recv_block_cache_meta() if recv_prev else None
)3 participants
PR Category
Performance Optimization
PR Types
new features
Description
残差注意力在pp下的通信优化和显存优化
是否引起精度变化
是 但是不影响收敛