[TP] 基于 PyTorch AsyncCollectiveTensor 实现 Col/Row/Seq redistribute 异步

一、背景

PyTorch ColwiseParallel / RowwiseParallel / SequenceParallel 在模块 I/O 边界调用 DTensor.redistribute(async_op=True)。collective 提交后不立即 wait，而是返回 torch.distributed._functional_collectives.AsyncCollectiveTensor（ACT），在下游算子（如 Linear）第一次读取 local tensor 时才 wait_tensor()，从而实现通信与计算重叠。

Hyper parallelize_module 的 Col/Row/Seq 当前走 DTensor.redistribute() 全同步路径：TensorRedistribution → differentiable_all_gather_concat / differentiable_all_reduce 等均在返回前完成通信。功能正确，但无 TP 边界重叠。

本 Issue 目标：在 Torch 后端，基于 PyTorch AsyncCollectiveTensor + wait_tensor()，为 Col/Row/Seq 的 redistribute 增加 async_op 能力。

关联： changzherui1/hyper-parallel#6 TP 专项 · mindspore/hyper-parallel#269 loss_parallel

参考实现（PyTorch 侧）：

组件	路径
TP style async_op=True	torch/distributed/tensor/parallel/style.py
Redistribute autograd	torch/distributed/tensor/_redistribute.py
ACT / wait_tensor	torch/distributed/_functional_collectives.py

Hyper 侧现状：

组件	路径	现状
Col/Row/Seq hook	hyper_parallel/core/tensor_parallel/style.py	redistribute() 无 async_op
redistribute 实现	hyper_parallel/core/dtensor/tensor_redistribution.py	同步 collective
platform 原语	hyper_parallel/platform/torch/platform.py	all_gather_single(..., async_op=True) 已有；differentiable_all_gather_concat 同步
已有 ACT 用法	differentiable_all_to_all_single_async、FSDP async AG	PP/MoE/FSDP 路径，非 TP redistribute

本 Issue 范围： Torch 后端 + AsyncCollectiveTensor。MindSpore 走自有 AsyncCollectiveTensor（platform/mindspore/platform.py），另开子任务，不在本 Issue 首版交付。

---

二、PyTorch 异步机制摘要

redistribute(async_op=True)
  → funcol.all_gather_single / all_reduce / ...
  → _wrap_tensor_autograd(elem) → AsyncCollectiveTensor
  → 返回 DTensor(_local_tensor=ACT, ...)

下游 Linear 读 input：
  → ACT.__torch_dispatch__ → trigger_wait() → wait_tensor(elem)
  → 与已提交的 collective 重叠

async_op=False（默认）：
  → redistribute_local_tensor 末尾 new_local_tensor.wait()

关键约束：

• ACT 通过 __torch_dispatch__ 延迟 wait；view 类 op 可能不 wait（PyTorch 对 view 有特殊处理）。
• backward 需同样传递 async_op，否则只有 forward 重叠。
• to_local() 若直接返回 _local_tensor 而不经 autograd，ACT 会原样传给 nn.Linear——这正是期望行为。

---

二点五、Hyper 与 __torch_dispatch__ 的关系

结论： TP 异步不需要给 Hyper DTensorBase 实现全套 __torch_dispatch__；复用 PyTorch AsyncCollectiveTensor.__torch_dispatch__ 即可。

机制	Hyper	PyTorch DTensor / ACT
DTensor 算子拦截	__torch_function__ → _OP_DISPATCHER（platform/torch/dtensor.py）	DTensor.__torch_dispatch__
TP 边界 pending 张量	本 Issue：redistribute → ACT	AsyncCollectiveTensor
延迟 wait	ACT __torch_dispatch__ → trigger_wait()	同左
MindSpore	__ms_dispatch__（platform/mindspore/...）	—

路径 A（推荐，本 Issue）：
  redistribute(async_op=True) → DTensor._local_tensor = ACT
  use_local_output=True: to_local() → ACT → nn.Linear → ACT.__torch_dispatch__ wait

路径 B（不推荐）：
  给 Hyper DTensorBase 加 __torch_dispatch__ → 与 OpDispatcher 架构冲突

路径 C（需验证）：
  use_local_output=False → DTensor op → _OP_DISPATCHER._unwrap_value → to_local() 取 ACT
  → 已注册 op 调用时 ACT dispatch 应仍生效；full-gather 回退路径需避免 eager wait

与 #269 关系： loss_parallel CE 走融合 kernel + _OP_DISPATCHER，不依赖本 Issue 的 redistribute 重叠；两者可并行推进。

---

三、Hyper 当前差异

#	差异	影响
1	DTensor.redistribute() 无 async_op 参数	style 层无法开启异步
2	TensorRedistribution 全部同步 wait	无重叠窗口
3	differentiable_all_gather_concat 用同步 all_gather	TP 热路径 all_concat 阻塞
4	DTensor.to_local() 直接返回 _local_tensor	ACT 可透传（✅ 有利），但 op dispatch 路径需验证 ACT
5	backward 无 async redistribute	反向无重叠
6	Col/Row/Seq 未传 async_op=True	与 PyTorch TP style 行为不一致

---

四、实现方案

4.1 总体架构

parallelize_module
  → Col/Row/Seq._prepare_input_fn / _prepare_output_fn
      → DTensor.redistribute(..., async_op=True)    # 新增参数，默认 False
          → TensorRedistribution.redistribution(..., async_op)
              → all_concat / all_reduce / all_to_all（platform 层）
                  → async_op=True: 返回 ACT，不 wait
                  → async_op=False: work.wait()（保持现状）

原则： async_op=False 为默认，不改变现有行为；仅 TP style 显式传 True。

4.2 分步任务

Step 1 — Platform：异步 all_gather / all_reduce 包装 ACT

文件： hyper_parallel/platform/torch/platform.py

新增（或扩展）：

def differentiable_all_gather_concat_async(data, group, concat_size, concat_dim, rank_list=None):
    # all_gather_into_tensor(..., async_op=True)
    # 用 funcol 路径返回 ACT（参考 differentiable_all_to_all_single_async）
    ...

def differentiable_all_reduce_async(data, op, group):
    # all_reduce(..., async_op=True) → ACT
    ...

复用现有：

• wait_async_tensor() → wait_tensor()
• _wrap_tensor_autograd / AsyncCollectiveTensor（PyTorch 内置）

注意： differentiable_all_gather_concat 当前用 list(dist.all_gather)；异步版建议改为 all_gather_into_tensor 单 buffer，与 FSDP 路径一致。

Step 2 — TensorRedistribution 传递 async_op

文件： hyper_parallel/core/dtensor/tensor_redistribution.py

• redistribution(self, input_x, to_layout, *, async_op=False)
• _construct_all_concat / _construct_all_concat_new：根据 async_op 分支同步/异步
• all_reduce（partial → replicate）路径同理
• all_to_all：可复用 differentiable_all_to_all_single_async 或新增 async 变体

与 PyTorch 对齐的逻辑（单步 collective 后）：

if not async_op and isinstance(result, AsyncCollectiveTensor):
    result = result.wait()

Step 3 — DTensor.redistribute(async_op=False)

文件： hyper_parallel/core/dtensor/dtensor.py

def redistribute(self, device_mesh, placements, *, async_op=False) -> DTensor:
    ...
    out = _tensor_redistribution.redistribution(self, dst_layout, async_op=async_op)

Step 4 — Autograd：backward 支持 async

方案 A（推荐，与 PyTorch 一致）： 新增 RedistributeFunction(torch.autograd.Function)，forward/backward 保存 async_op，backward 调用反向 redistribution 时同样 async_op=True。

方案 B（最小 PoC）： 首版仅 forward async，backward 同步——有重叠收益但 backward 无重叠。

首版建议 方案 A，至少覆盖 Col/Row 主路径。

Step 5 — TP style 开启 async_op=True

文件： hyper_parallel/core/tensor_parallel/style.py

与 PyTorch 相同位置传参：

Style	位置
ColwiseParallel	_prepare_input_fn、_prepare_output_fn 的 redistribute
RowwiseParallel	同上
SequenceParallel	_prepare_input_fn 的 redistribute
PrepareModuleInput/Output	首版保持同步（PyTorch 亦未默认 async）

Step 6 — Op dispatch / to_local 验证

文件： hyper_parallel/core/shard/_op_dispatch.py、platform/torch/dtensor.py

Hyper OpDispatcher._unwrap_value 对 DTensor 调用 to_local()，不会对 ACT 提前 wait()——有利于 ACT 透传。需验证：

• use_local_output=True：to_local() → ACT → nn.Linear → ACT __torch_dispatch__ wait ✅
• use_local_output=False：DTensor 包裹 ACT → _OP_DISPATCHER → unwrap 后 local 仍为 ACT → 已注册 op 应触发 ACT wait
• SkipDTensorDispatch backward：plain tensor 路径仍靠 ACT 自身 dispatch
• 未注册 op full-gather 回退：可能对 ACT 不当处理，导致重叠失效或错误——TP 热路径 op 须走注册表

若 op dispatch 对 ACT 提前 wait，需在 dispatch 入口 透传 ACT 而非 eager wait（仅当实测重叠失效时再改）。

---

五、TP 边界 collective 映射

Style	边界	layout 变化	collective	async 收益
Colwise	input	→ Replicate	all_gather	高
Colwise	output	调整 Shard	视配置	中
Rowwise	input	→ Shard(-1)	all_to_all / slice	中
Rowwise	output	Partial → Replicate	all_reduce	高
SequenceParallel	input	→ Shard(seq)	all_gather	中

---

六、测试与验收

类型	内容
数值	async on/off 结果一致（Col+Row 单层、Llama3 block）
梯度	async on/off grad 一致
ACT	isinstance(local, AsyncCollectiveTensor) 在 hook 后、Linear 前为 True
Profiler	NPU/CUDA timeline：all_gather 与 mm 有时间重叠（定性）
回归	现有 tensor_parallel ST 全过（默认 async_op=False）

测试文件建议：

• tests/torch/tensor_parallel/test_tp_redistribute_async.py（新建）
• 扩展 examples/torch/llama3/tensor_parallel_example.py 可选 --async-redistribute

---

七、里程碑

阶段	交付	工作量
M1 PoC	platform async all_gather + redistribute(async_op) forward only + Colwise	中
M2 完整	Rowwise/Seq + backward async + autograd Function	中～大
M3 验收	数值/梯度 parity + profiler + 文档	小

---

八、风险与不在范围

风险	缓解
ACT 在 op dispatch 被提前 wait	dispatch 层 ACT 透传策略
NPU wait_tensor 与 stream 行为	在目标 CANN 版本上 profiler 验证
autograd 二阶	首版不保证；与 PyTorch NestedRedistribute scope 一致

不在本 Issue：

• MindSpore AsyncCollectiveTensor 路径（另开）
• torch.compile / AsyncTP（Inductor micro-pipeline）
• PrepareModuleInput async（PyTorch 亦未默认开启）

---

附录：PyTorch 参考代码位置

# style.py — ColwiseParallel
input_tensor.redistribute(placements=desired_input_layouts, async_op=True)
outputs.redistribute(placements=output_layouts, async_op=True)

# _redistribute.py
if not async_op and isinstance(new_local_tensor, funcol.AsyncCollectiveTensor):
    new_local_tensor = new_local_tensor.wait()

# _functional_collectives.py — AsyncCollectiveTensor.__torch_dispatch__
# 非 view op → trigger_wait()

---

关联：changzherui1/hyper-parallel#6 · mindspore/hyper-parallel#269

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

Start with hyper_parallel/platform/torch/platform.py and the existing differentiable_all_to_all_single_async path, then trace tensor_redistribution.py, dtensor.py, and tensor_parallel/style.py. Add tests in tests/torch/tensor_parallel/test_tp_redistribute_async.py covering Colwise, Rowwise, and SequenceParallel forward and backward behavior. Done means async and synchronous results and gradients agree, ACT reaches the boundary before Linear, existing tests still pass, and profiler evidence shows overlap.