[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/core/tensor_parallel/style.py, core/dtensor/dtensor.py, and core/dtensor/tensor_redistribution.py, then compare the async implementations in hyper_parallel/platform/torch/platform.py and PyTorch's referenced files. Run the existing tensor-parallel tests before adding tests/torch/tensor_parallel/test_tp_redistribute_async.py. Done means Col/Row/Seq async redistribution preserves values and gradients, retains ACT until consumption, and leaves synchronous behavior unchanged.