已关闭

FSDP patch消除 #1788

zhenyu创建于 4月30日关闭于 5月27日

4月30日创建

FSDP2 foreach copy 原生对齐改动

评审对象：

torch-npu：28ee39f623d89ed260c5d316d1112ec3db63dc31，重点文件 torch_npu/distributed/fsdp/_add_fsdp_patch.py
op-plugin：8dca0d33f8ca9d520b34ffab51e6560c910d548a，重点文件 op_plugin/ops/opapi/ForeachCopyKernelOpApi.cpp

评审结论建议：允许删除 torch-npu FSDP2 中仅为 foreach copy non_blocking=True 服务的 Python collectives patch，将同设备 D2D foreach copy 的设备侧异步语义收敛到 op-plugin 算子实现中，从而复用 PyTorch 原生 FSDP2 collectives 路径，并对齐 torch+cuda 的实际执行语义。

一、背景描述

1.1 当前问题

torch-npu 为适配 FSDP2 fully-shard，历史上在 _add_fsdp_patch.py 中覆写过 PyTorch 原生 FSDP collectives 的部分逻辑，主要包括：

_patched_get_param_all_gather_inputs
_patched_all_gather_copy_in
_patched_finalize_backward
_apply_fsdp_patch

其中本次要消除的关键差异集中在两个 foreach copy 调用点：

# _patched_get_param_all_gather_inputs 中的历史 NPU patch
if splits[0].device == foreach_copy_inputs[0].device:
    torch._foreach_copy_(splits, foreach_copy_inputs, non_blocking=True)
else:
    torch._foreach_copy_(splits, foreach_copy_inputs)

# _patched_all_gather_copy_in 中的历史 NPU patch
if foreach_copy_dsts[0].device == all_gather_inputs[0].device:
    torch._foreach_copy_(foreach_copy_dsts, all_gather_inputs, non_blocking=True)
else:
    torch._foreach_copy_(foreach_copy_dsts, all_gather_inputs)

上游 PyTorch 原生 FSDP2 在对应位置直接调用：

torch._foreach_copy_(foreach_copy_dsts, all_gather_inputs)

从 Python API 形态看，torch._foreach_copy_ 默认参数是 non_blocking=False。但是在 torch+cuda 的同设备 D2D fast path 中，设备侧 copy 按当前 stream 异步排队执行，fast path 并不依赖 non_blocking 参数来决定是否阻塞 host。因此，torch-npu 若继续在 FSDP Python 层显式传 non_blocking=True，虽然可以补齐 NPU 行为，但会带来两个维护问题：

torch-npu 需要持续覆写 PyTorch FSDP2 内部私有函数，跟随上游变化成本高。
FSDP2 collectives 路径与原生 torch+cuda 不一致，不利于后续删除补丁和复用上游逻辑。

1.2 目标

本次改动目标是：

删除 torch-npu 中针对 FSDP2 collectives 的 Python 覆写，让 _get_param_all_gather_inputs 与 all_gather_copy_in 回到 PyTorch 原生实现。
将 NPU 同设备 D2D foreach copy 的异步语义放到 op-plugin 的 _foreach_copy_ 算子实现中处理。
尽量保持 op-plugin 既有 fast path / memcpyBatch / slow path 分支结构不变，只在同设备 fast path 上对齐 torch+cuda 的设备侧异步行为。
不改变 Python API 签名，不改变 Host<->Device、跨设备、fallback 等场景的原有 non_blocking 语义。

1.3 torch+cuda 原生 foreach copy 逻辑梳理

这一节是本次评审的关键：从 Python 层看，上游 FSDP2 的确没有显式传 non_blocking=True；但从 CUDA 后端实现看，同设备 D2D 命中 foreach fast path 后，并不是按 non_blocking=False 去做 host 同步阻塞 copy，而是走设备侧 fast path。

1.3.1 FSDP2 原生调用点：不显式传第三参

在 PyTorch 原生 torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py 中，两个对应调用点都是直接调用 torch._foreach_copy_：

@torch.library.impl(lib, "all_gather_copy_in", "CUDA")
@torch.library.impl(lib, "all_gather_copy_in", "XPU")
@torch.library.impl(lib, "all_gather_copy_in", "HPU")
@torch.library.impl(lib, "all_gather_copy_in", "CPU")
@torch.library.impl(lib, "all_gather_copy_in", "MTIA")
@torch.library.impl(lib, "all_gather_copy_in", "PrivateUse1")
def all_gather_copy_in_cuda(
    all_gather_inputs: list[torch.Tensor],
    all_gather_output: torch.Tensor,
    inp_split_sizes: list[int],
    all_gather_input_numel: int,
    rank: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    all_gather_input = all_gather_output.narrow(
        0, all_gather_input_numel * rank, all_gather_input_numel
    )
    foreach_copy_dsts = torch.split(all_gather_input, inp_split_sizes)
    with torch.no_grad():
        torch._foreach_copy_(foreach_copy_dsts, all_gather_inputs)
    return all_gather_input, all_gather_output

# _get_param_all_gather_inputs(...) 中
if foreach_copy_inputs:
    fsdp_param_0 = fsdp_params[foreach_copy_indices[0]]
    param_dtype, device = fsdp_param_0.param_dtype, fsdp_param_0.device
    flat_foreach_copy_input = torch.empty(
        (sum(foreach_copy_input_numels),), device=device, dtype=param_dtype
    )
    splits = torch.split(flat_foreach_copy_input, foreach_copy_input_numels)
    torch._foreach_copy_(splits, foreach_copy_inputs)
    for i, split in zip(foreach_copy_indices, splits):
        param_all_gather_inputs[i] = [split]

对应 schema 中，non_blocking 默认值确实是 False：

- func: _foreach_copy_(Tensor(a!)[] self, Tensor[] src, bool non_blocking=False) -> ()
  device_check: NoCheck
  variants: function
  dispatch:
    CompositeExplicitAutograd: foreach_tensor_copy_list_kernel_slow_
    CUDA: foreach_tensor_copy_list_kernel_cuda_
    MTIA: foreach_tensor_copy_list_kernel_mtia_
  autogen: _foreach_copy.out

1.3.2 CUDA foreach fast path：命中后不使用 `non_blocking` 控制是否同步

CUDA 后端实现位于 aten/src/ATen/native/cuda/ForeachBinaryOpList.cu。关键逻辑如下：

void foreach_tensor_copy_list_kernel_cuda_(
    TensorList self,
    TensorList src,
    const bool non_blocking) {
  check_foreach_api_restrictions(self, src);
  if (!(_check_tensors_share_device_and_dtype(
            {self, src}, /* skip_dtype_check */ true) &&
        std::all_of(
            src.cbegin(),
            src.cend(),
            [&src](const auto& t) -> bool {
              return t.dtype() == src[0].dtype();
            }) &&
        std::all_of(
            self.cbegin(),
            self.cend(),
            [&self](const auto& t) -> bool {
              return t.dtype() == self[0].dtype();
            }) &&
        _check_tensors_share_sizes_and_strides({self, src}))) {
    return at::native::foreach_tensor_copy_list_kernel_slow_(
        self, src, non_blocking);
  }

  std::vector<std::vector<at::Tensor>> tensor_lists{src.vec(), self.vec()};

  AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND7(
      ScalarType::Half,
      ScalarType::BFloat16,
      ScalarType::Bool,
      ScalarType::Float8_e4m3fn,
      ScalarType::Float8_e4m3fnuz,
      ScalarType::Float8_e5m2,
      ScalarType::Float8_e5m2fnuz,
      self[0].scalar_type(),
      "foreach_tensor_copy",
      [&]() {
        AT_DISPATCH_SOURCE_TYPES(src[0].scalar_type(), "foreach_tensor_copy", [&] {
          if constexpr (std::is_same_v<scalar_t, src_t>) {
            multi_tensor_apply<2>(
                tensor_lists,
                UnaryOpFunctor<
                    scalar_t,
                    /* depth */ 2,
                    /* r_args_depth */ 1,
                    /* res_arg_index */ 1>(),
                Copy<scalar_t, scalar_t>());
          } else {
            multi_tensor_apply<2>(
                tensor_lists,
                CopyFunctor<
                    scalar_t,
                    src_t,
                    /* depth */ 2,
                    /* r_args_depth */ 1,
                    /* res_arg_index */ 1>(),
                Copy<scalar_t, src_t>());
          }
        });
      });
  increment_version(self);
}

这里要注意两点：

non_blocking 只在未命中 fast path 时传给 slow path。
命中 fast path 后，主体执行是 multi_tensor_apply + copy functor，没有使用 non_blocking 分支去选择同步/异步。

也就是说，torch+cuda 的 FSDP2 原生路径虽然在 Python 层使用默认 False，但同设备 D2D foreach fast path 的实际行为不是“host 同步等待每个 copy 完成”。

1.3.3 CUDA slow path：才会逐个把 `non_blocking` 传给 `copy_`

slow path 位于 aten/src/ATen/native/ForeachOpsKernels.cpp：

void foreach_tensor_copy_list_kernel_slow_(
    TensorList self,
    TensorList src,
    const bool non_blocking) {
  check_foreach_api_restrictions(self, src);

  for (const auto i : c10::irange(self.size())) {
    self[i].copy_(src[i], non_blocking);
  }
}

这说明 non_blocking 是 fallback 逐 tensor copy_ 路径的重要参数；本次 NPU 改动也保留了这一点，不会把 fallback 路径统一强制成 true。

1.3.4 CUDA `copy_`：D2D 与 CPU<->CUDA 对 `non_blocking` 的角色不同

aten/src/ATen/native/cuda/Copy.cu 中，CUDA D2D 会进入 copy_device_to_device(...)：

// Copy on GPU (or between GPUs)
if (dst_device.is_cuda() && src_device.is_cuda()) {
  copy_device_to_device(iter, non_blocking, p2p_enabled);
  return;
}

copy_device_to_device 的核心逻辑是使用当前 CUDA stream 做设备侧 copy；同设备时不会走 CPU<->CUDA 那种 non_blocking ? cudaMemcpyAsync : memcpy_and_sync 分支：

// device-to-device copy, does type conversion
void copy_device_to_device(TensorIterator& iter,
                           bool non_blocking,
                           bool p2p_enabled) {
  int64_t numel = iter.numel();

  bool same_type = iter.dtype(0) == iter.dtype(1);
  bool same_conj = iter.tensor(0).is_conj() == iter.tensor(1).is_conj();
  bool same_neg = iter.tensor(0).is_neg() == iter.tensor(1).is_neg();
  bool memcpy_eligible = same_type && same_conj && same_neg && iter.is_contiguous();

  Device dst_device = iter.device(0);
  Device src_device = iter.device(1);

  CUDAGuard device_guard(src_device);
  CUDAStream copy_stream = getCurrentCUDAStream(src_device.index());

  if (memcpy_eligible) {
    void *dst = iter.data_ptr(0);
    void *src = iter.data_ptr(1);
    size_t size = numel * iter.element_size(0);
    if (src != dst || src_device != dst_device) {
      AT_CUDA_CHECK(CUDACachingAllocator::memcpyAsync(
        dst, dst_device.index(),
        src, src_device.index(),
        size, copy_stream, p2p_enabled));
    }
  } else {
    // non-contiguous / type-conversion cases run CUDA copy kernels
    direct_copy_kernel_cuda(iter);
  }

  AT_CUDA_CHECK(cudaGetLastError());
}

而 CPU<->CUDA copy 才显式按 non_blocking 选择异步 memcpy 还是同步 memcpy：

// Copy between CPU and GPU
CUDAStream stream = getCurrentCUDAStream();

if (non_blocking) {
  AT_CUDA_CHECK(cudaMemcpyAsync(dst, src, nbytes, kind, stream));
  at::getHostAllocator(at::kCUDA)->record_event(ptr, ctx, stream.unwrap());
} else {
  at::cuda::memcpy_and_sync(dst, src, nbytes, kind, stream);
}

因此，torch+cuda 的区别可以概括为：

路径	`non_blocking` 的作用
foreach CUDA fast path，同设备 D2D	fast path 主体不使用 `non_blocking` 控制 host 同步；按设备 stream 排队执行
foreach slow path	逐 tensor 调 `copy_(..., non_blocking)`
`copy_` 的 CPU<->CUDA	`true` 走 `cudaMemcpyAsync`，`false` 走 `memcpy_and_sync`
`copy_` 的 CUDA D2D	设备侧 copy / kernel 按 stream 执行，`non_blocking` 不是 host 同步开关

本次 NPU 改动要对齐的是第一行：foreach 同设备 D2D fast path 的设备侧异步执行语义。

二、方案设计

2.1 总体方案

采用“上层复用原生，下层补齐后端语义”的设计：

torch-npu 删除 FSDP2 collectives patch：
- 不再覆写 _get_param_all_gather_inputs
- 不再覆写 torch.ops.fsdp.all_gather_copy_in
- 不再覆写 FSDPParamGroup.finalize_backward
- fully_shard() 入口只保留 NPU 侧增强 patch，例如内存缓存和 recompute/prefetch 状态管理
op-plugin 在 foreach copy 算子中对齐同设备 D2D fast path：
- fast path 判定显式使用 PyTorch foreach 工具函数：
  - at::native::can_use_fast_route(self, src)
  - check_tensor_dtype_support_base(src)
- can_use_fast_route(self, src) 本身会检查 self/src dtype 一致，因此 dtype 不同不会进入 NPU fast path，保持当前 NPU 不支持跨 dtype fast path 的既有语义。
- 命中 fast path 时调用：

split_and_exec_npu_cmd_copy(self, src, /*non_blocking=*/true);

fallback 路径保持调用方传入语义：
- 未命中 fast path 时，memcpyBatch(self, src, non_blocking) 不变
- slow path foreach_tensor_copy_list_kernel_slow_(self, src, non_blocking) 不变
- DO_COMPATIBILITY 回退仍使用原始 non_blocking

2.2 op-plugin 关键设计

op-plugin 当前实现的核心改动如下：

if (!is_support_nd_out || !at::native::can_use_fast_route(self, src) ||
        !check_tensor_dtype_support_base(src)) {
    if (is_support_batch && ... && check_tensor_device_dtype_base(self, src, non_blocking)) {
        return memcpyBatch(self, src, non_blocking);
    }
    return at::native::foreach_tensor_copy_list_kernel_slow_(self, src, non_blocking);
}

split_and_exec_npu_cmd_copy(self, src, /*non_blocking=*/true);

该设计的关键点是：

命中 fast path 后，传给 split_and_exec_npu_cmd_copy 的 non_blocking 使用 true；由于 can_use_fast_route(self, src) 已经保证 self/src 同 device，这里不需要再单独判断 same_device。
没有把所有 _foreach_copy_ 调用都强制改成异步。
没有在 FSDP Python 层新增特判，避免继续耦合 PyTorch 内部 FSDP2 实现。
fast path 条件复用原有 can_use_fast_route(self, src)，只覆盖 NPU 当前支持的同 dtype D2D 场景；跨 dtype copy 继续走 fallback，避免把未支持场景提前放入 aclnnForeachCopy fast path。

can_use_fast_route(self, src) 的源码展开如下：

inline bool can_use_fast_route(
    TensorList tensors1,
    TensorList tensors2,
    bool does_op_promote_integer_inputs_to_float = false) {
  return can_use_fast_route(
      {tensors1, tensors2}, {}, does_op_promote_integer_inputs_to_float);
}

inline bool can_use_fast_route(
    ArrayRef<TensorList> tensorLists,
    ArrayRef<Scalar> scalarList = {},
    bool does_op_promote_integer_inputs_to_float = false) {
  return check_fast_path_restrictions(
      tensorLists, scalarList, does_op_promote_integer_inputs_to_float);
}

inline bool check_fast_path_restrictions(
    ArrayRef<TensorList> tensorLists,
    ArrayRef<Scalar> scalarList = {},
    bool does_op_promote_integer_inputs_to_float = false,
    bool skip_cross_list_dtype_check = false) {
  return _check_tensors_share_device_and_dtype(
             tensorLists, skip_cross_list_dtype_check) &&
      _check_tensors_share_sizes_and_strides(tensorLists) &&
      _check_tensors_do_type_promotion_with_scalars(
             tensorLists[0],
             scalarList,
             does_op_promote_integer_inputs_to_float);
}

这里 can_use_fast_route(self, src) 没有传入 skip_cross_list_dtype_check=true，所以 skip_cross_list_dtype_check 使用默认值 false。继续展开 _check_tensors_share_device_and_dtype：

inline bool _check_tensors_share_device_and_dtype(
    ArrayRef<TensorList> tensorLists,
    const bool skip_cross_list_dtype_check = false) {
  const auto expected_dtype = tensorLists[0][0].dtype();
  const auto expected_device = tensorLists[0][0].device();

  return std::all_of(
      tensorLists.cbegin(),
      tensorLists.cend(),
      [&](const TensorList& tensorList) {
        const auto list_dtype = tensorList[0].dtype();
        return std::all_of(
            tensorList.cbegin(), tensorList.cend(), [&](const Tensor& tensor) {
              return tensor.device() == expected_device &&
                  tensor.layout() == at::kStrided &&
                  tensor.is_non_overlapping_and_dense() &&
                  tensor.dtype() == list_dtype &&
                  (skip_cross_list_dtype_check ||
                   tensor.dtype() == expected_dtype);
            });
      });
}

因此 can_use_fast_route(self, src) 命中时已经能推出：

self 和 src 两个 tensor list 中的所有 tensor 都等于同一个 expected_device，所以 fast path 的 device 一定相同。
skip_cross_list_dtype_check=false，因此所有 tensor dtype 也必须等于同一个 expected_dtype，跨 dtype 不会进入 fast path。
_check_tensors_share_sizes_and_strides 已经保证对应 tensor 的 size/stride 一致。

所以对本次目标而言，保留原来的 can_use_fast_route(self, src) 是最小且最稳的做法；不需要改成 _check_tensors_share_device_and_dtype({self, src}, true)，也不需要额外写 same_device 判断。

2.3 与 torch+cuda foreach copy 的区别和对齐点

本次评审需要明确：NPU 不是逐字复制 CUDA 实现，而是对齐 CUDA 在 FSDP2 依赖场景中的行为语义。

对比项	torch+cuda 原生	op-plugin NPU 本次方案
FSDP2 Python 调用	直接 `torch._foreach_copy_(...)`，不传第三参	删除 torch-npu Python 特判后同样直接复用原生调用
API 默认值	`non_blocking=False`	不改 API 默认值
fast path 判定	`_check_tensors_share_device_and_dtype({self, src}, true)`、两侧 list 内 dtype 一致、size/stride 一致；CUDA fast path 内部支持跨 dtype copy functor	使用原有 `can_use_fast_route(self, src)`，要求 `self/src` 同设备、同 dtype、同 size/stride；NPU 当前不支持跨 dtype fast path
fast path 执行	`multi_tensor_apply` copy functor；函数主体不使用 `non_blocking` 控制同步	`split_and_exec_npu_cmd_copy(self, src, /non_blocking=/true)`；同 device fast path 使用 `true` 表达设备侧异步 copy 语义
fallback 执行	slow path 逐 tensor `copy_(..., non_blocking)`	`memcpyBatch(..., non_blocking)` 或 slow path，继续使用原始 `non_blocking`
本质语义	同设备 D2D fast path 按设备 stream 排队	对齐同设备 D2D fast path 的设备侧异步排队

这里最重要的区别是实现方式：

CUDA fast path 没有显式写 non_blocking=true，因为它进入的是 multi_tensor_apply 设备侧 fast path，non_blocking 本身不是这个 fast path 的控制项。
NPU op-plugin fast path 需要把 true 传给 split_and_exec_npu_cmd_copy，用于表达同设备 D2D fast path 应走设备侧异步 copy 语义。

所以，本次 NPU 代码里出现 /*non_blocking=*/true 并不表示改变 Python API 默认语义，而是把 CUDA fast path 已经具备的“同 device fast path 设备侧异步排队”行为显式补齐到 NPU fast path。

需要特别说明的差异是：CUDA foreach copy fast path 可以通过 AT_DISPATCH_SOURCE_TYPES 和 copy functor 支持跨 dtype D2D copy；NPU aclnnForeachCopy fast path 当前不支持跨 dtype。因此 NPU 侧不能完全复刻 CUDA 的跨 dtype fast path 判定，而是继续使用原有 can_use_fast_route(self, src) 保持同 dtype 约束。这样既对齐 FSDP2 依赖的同 dtype D2D fast path 异步语义，也不扩大 NPU fast path 的输入域。

2.4 torch-npu 关键设计

torch-npu commit 删除了 FSDP2 collectives 覆写后，_add_fsdp_patch.py 只保留 NPU 相关增强：

FSDPMemCache
DefaultAllGather.allocate / DefaultReduceScatter.allocate 内存缓存
foreach_reduce 后释放 reduce-scatter 输入缓存
FSDPState._post_forward
FSDPParamGroup.post_forward
FSDPParamGroup.post_backward
fully_shard() 首次调用时应用 _apply_fsdp_enhance_patch()

删除后效果：

FSDP2 的 all-gather input 构造回到上游 _get_param_all_gather_inputs
FSDP2 的 all-gather copy-in 回到上游 torch.ops.fsdp.all_gather_copy_in
torch-npu 不再维护仅为 foreach copy non_blocking=True 存在的 Python monkey patch

2.5 为什么不影响之前逻辑

影响边界限定如下：

场景	本次行为	是否改变旧逻辑
NPU 同设备 D2D，命中 foreach fast path	后端内部按 `non_blocking=true` 执行，贴近 CUDA fast path 设备流异步语义	这是原 Python patch 想补齐的语义，下沉到算子层
Host<->Device	仍使用调用方传入的 `non_blocking`	不改变
跨设备或不满足 fast path	走 `memcpyBatch` 或 slow path，仍使用调用方传入的 `non_blocking`	不改变
aclnn 不兼容回退	`DO_COMPATIBILITY` 仍回退到 slow kernel，传原始 `non_blocking`	不改变
Python API	`torch._foreach_copy_(..., non_blocking=False)` 默认签名不变	不改变
FSDP2 collectives	使用 PyTorch 原生实现	对齐 torch+cuda，减少 torch-npu 私有 patch

从语义上看，本次不是把公开 API 默认值从 False 改为 True，而是让 NPU 同设备 D2D fast path 具备与 CUDA fast path 一致的设备侧异步排队行为。旧的 torch-npu Python patch 只在 FSDP2 内部两个调用点补 non_blocking=True；新方案把同一语义收敛到 foreach copy 算子实现，覆盖更准确，也更接近后端职责边界。

三、使用说明

3.1 合入方式

建议按以下顺序合入和验证：

先合入 op-plugin foreach copy 改动：
- commit：8dca0d33f8ca9d520b34ffab51e6560c910d548a
- 重点确认 ForeachCopyKernelOpApi.cpp 中 fast path 强制使用 non_blocking=true
- 重点确认 fallback 仍保留原始 non_blocking
再合入 torch-npu FSDP patch 删除：
- commit：28ee39f623d89ed260c5d316d1112ec3db63dc31
- 重点确认 _patched_get_param_all_gather_inputs 和 _patched_all_gather_copy_in 已删除
- 重点确认 _apply_fsdp_patch 不再应用 collectives monkey patch
- 重点确认 _apply_fsdp_enhance_patch 仍保留 NPU 内存缓存和调度增强
torch-npu 与 op-plugin 必须配套发布：
- 若只删除 torch-npu Python patch，但 op-plugin 未具备同设备 D2D fast path 异步语义，FSDP2 all-gather copy-in 可能退回到不符合 CUDA 原生表现的执行行为。
- 因此评审结论应明确：删除 torch-npu patch 的前提是 op-plugin foreach copy fast path 对齐改动已合入并进入同一版本包。

3.2 开发注意事项

后续不要在 torch-npu FSDP2 Python 层重新添加只为 foreach copy non_blocking=True 服务的 collectives 覆写。
若后续 PyTorch FSDP2 上游逻辑变化，应优先复用上游实现，只在 NPU 后端算子或 NPU 专属增强点处理差异。
op-plugin 中应保持 fast path 与 fallback 的边界：
- fast path：同设备 D2D，按设备流异步语义执行。
- fallback：继续尊重调用方 non_blocking 参数。
测试中读取 NPU 同设备 D2D copy 结果前，应显式 torch.npu.synchronize()，避免异步语义下立即断言引入误判。

3.3 对专家的评审请求

请专家重点评审并确认：

是否认可将 FSDP2 foreach copy 的 non_blocking=True 特判从 torch-npu Python patch 下沉到 op-plugin 算子实现。
是否认可该改动是对齐 torch+cuda fast path 的设备侧异步语义，而不是改变 Python API 默认参数。
是否认可当前影响范围限定在同设备 D2D fast path，fallback 路径不受影响。
是否认可在 op-plugin 改动配套合入的前提下，torch-npu 删除 _patched_get_param_all_gather_inputs 与 _patched_all_gather_copy_in。

四、测试设计

4.1 op-plugin 单算子测试

目标：验证 _foreach_copy_ 在 NPU 同设备 D2D、跨 dtype、不同 tensor 数量下结果正确，并且异步语义下测试读取结果前有同步。

建议执行：

cd <op-plugin-root>
python test/test_v2r1_ops/test_foreach_copy.py
python test/test_v2r2_ops/test_foreach_copy.py

测试覆盖点：

同 dtype NPU D2D：
- float16
- float32
- bfloat16
- int32
- int64
- double
- bool
跨 dtype NPU D2D：
- 例如 int32 -> double
tensor 数量：
- 20
- 50
每个 NPU D2D 断言前显式：

torch.npu.synchronize()

4.2 fallback 路径测试

目标：确认未命中 fast path 的路径仍沿用调用方 non_blocking 参数。

建议覆盖：

Host -> NPU
NPU -> Host
不满足 fast path 的 shape / stride 组合
CANN 或 SoC 不支持 fast route 时的 slow kernel 回退

检查点：

memcpyBatch(self, src, non_blocking) 仍收到调用方传入值。
foreach_tensor_copy_list_kernel_slow_(self, src, non_blocking) 仍收到调用方传入值。
DO_COMPATIBILITY 回退仍不被 fast path 的 true 影响。

4.3 torch-npu FSDP2 主路径测试

目标：验证删除 Python collectives patch 后，FSDP2 fully-shard 主训练路径仍可用，并确认不再依赖旧 patch 函数。

建议执行或选取对应 CI：

cd <torch-npu-root>
python test/distributed/fsdp2/test_fully_shard_comm.py
python test/distributed/fsdp2/test_fully_shard_state.py
python test/distributed/fsdp2/test_fully_shard_autograd.py
python test/distributed/fsdp2/test_fully_shard_clip_grad_norm_.py
python test/distributed/fsdp2/test_fully_shard_extensions.py

建议增加一组 smoke：

import torch
import torch_npu

src = [torch.arange(1024, device="npu", dtype=torch.int32) for _ in range(8)]
dst = [torch.empty(1024, device="npu", dtype=torch.float32) for _ in src]
torch._foreach_copy_(dst, src)
torch.npu.synchronize()
assert torch.equal(dst[0].cpu(), src[0].cpu().to(torch.float32))

4.4 回归判断标准

通过标准：

op-plugin foreach copy 单算子测试全部通过。
FSDP2 fully-shard 关键分布式测试通过。
删除 torch-npu collectives patch 后，无 _patched_get_param_all_gather_inputs、_patched_all_gather_copy_in 等旧函数依赖。
同设备 D2D fast path 与 torch+cuda 原生行为一致：设备侧按当前 stream 异步排队，测试读取结果前需要同步。
Host<->Device、跨设备、fallback 路径的 non_blocking 行为保持原样。

建议评审通过条件：

在 op-plugin foreach copy fast path 对齐改动随版本合入的前提下，允许 torch-npu 删除 FSDP2 collectives Python patch，后续 FSDP2 collectives 直接复用 PyTorch 原生实现。

五、评审遗留问题补充与风险排除计划

5.1 `split_and_exec_npu_cmd_copy` 是否对齐 CUDA

split_and_exec_npu_cmd_copy 的职责只有两个：

void exec_npu_cmd_copy(const at::TensorList dst, at::TensorList src, bool non_blocking)
{
    EXEC_NPU_CMD(aclnnForeachCopy, src, dst);
}

void split_and_exec_npu_cmd_copy(const at::TensorList dst, at::TensorList src, bool non_blocking)
{
    if (tensor_count <= max_tensor_count) {
        exec_npu_cmd_copy(dst, src, non_blocking);
        return;
    }
    // 按 SINGLE_FOREACH_OP_TENSOR_COUNT 切分后多次调用 exec_npu_cmd_copy
}

结论：

CUDA fast path 用 multi_tensor_apply 对 tensor list 做批量设备侧 copy；NPU fast path 用 split_and_exec_npu_cmd_copy 按 SINGLE_FOREACH_OP_TENSOR_COUNT 切块后调用 aclnnForeachCopy。二者实现不同，但语义都是同设备 D2D fast path 在当前设备 stream 上排队执行。
split_and_exec_npu_cmd_copy 当前没有默认参数，fast path 调用处显式传 /*non_blocking=*/true。这里不再额外判断 same_device，因为原有 can_use_fast_route(self, src) 命中时已保证 self/src 同 device。
函数内部只把 non_blocking 继续传给 exec_npu_cmd_copy；而 exec_npu_cmd_copy 内部调用 EXEC_NPU_CMD(aclnnForeachCopy, src, dst)，并没有把 non_blocking 作为 ACLNN 参数使用。因此当前把 fast path 调用处设为 true，不会改变 split_and_exec_npu_cmd_copy 内部执行路径、切分逻辑或同步逻辑。
需要控制的风险是未来语义漂移：如果后续 exec_npu_cmd_copy 或 aclnnForeachCopy 包装层开始消费 non_blocking 参数，则 fast path 固定传 true 会变成真实行为变化。建议保留注释或后续去掉未使用参数，避免误读。

5.2 dtype fast path 与 `memcpyBatch` 影响

评审后 fast path 判定保持原逻辑，只调整 fast path 执行处：

if (!is_support_nd_out || !at::native::can_use_fast_route(self, src) ||
        !check_tensor_dtype_support_base(src)) {
    ...
    return at::native::foreach_tensor_copy_list_kernel_slow_(self, src, non_blocking);
}

split_and_exec_npu_cmd_copy(self, src, /*non_blocking=*/true);

该调整的语义是：

can_use_fast_route(self, src) 复用原有 fastpath 判断，要求 self/src 同设备、同 dtype、同 size/stride。dtype 不同不会进入 fast path，符合当前 NPU fast path 不支持跨 dtype copy 的实际能力。
CUDA fast path 支持跨 dtype 是 CUDA 后端 copy functor 的能力，不能直接作为 NPU fast path 放开的依据。
命中 fast path 后，传入 split_and_exec_npu_cmd_copy 的 non_blocking 使用 true；同 device 条件已由 can_use_fast_route(self, src) 保证。

对 memcpyBatch 的影响：

场景	`non_blocking=False`	`non_blocking=True`	本次调整影响
NPU D2D，同 dtype	命中 fast path，调用 `split_and_exec_npu_cmd_copy(self, src, /non_blocking=/true)`；当前 `non_blocking` 不被 `exec_npu_cmd_copy` 消费	命中 fast path，调用同左	不影响 `memcpyBatch`，因为 D2D 不走 `check_tensor_device_dtype_base`
NPU D2D，跨 dtype	不进 fast path；`check_tensor_device_dtype_base` 对同设备类型返回 false；走 slow path 并保留原始 `non_blocking=False`	不进 fast path；走 slow path 并保留原始 `non_blocking=True`	恢复既有语义，避免不支持场景进入 fast path
Host -> NPU / NPU -> Host	不进 fast path；满足条件时走 `AclrtMemcpyBatch`，调用前同步 stream	不进 fast path；满足条件时走 `LaunchBatchAsyncCopyTask`	不受 dtype fast path 修改影响，继续由调用方 `non_blocking` 控制
shape/stride/device 不满足 fast path	按原始 `non_blocking` 进入 `memcpyBatch` 或 slow path	按原始 `non_blocking` 进入 `memcpyBatch` 或 slow path	fallback 语义不变

因此，本次 dtype 判定收紧不会让原来不应该走 memcpyBatch 的 D2D 数据进入 memcpyBatch，也不会让原本应该执行 memcpyBatch 的 Host<->Device 场景被 fast path 截走。

5.3 op-plugin 内部调用范围

在 op-plugin 仓内按 split_and_exec_npu_cmd_copy|process_tensor_list_batch|exec_npu_cmd_copy|_foreach_copy_ 检索，直接调用关系只存在于 op_plugin/ops/opapi/ForeachCopyKernelOpApi.cpp：

exec_npu_cmd_copy 只被 split_and_exec_npu_cmd_copy 调用。
split_and_exec_npu_cmd_copy 只被 process_tensor_list_batch 的 fast path 调用。
process_tensor_list_batch 只被本文件 _foreach_copy_ 调用，分别覆盖未分组路径和 dtype 分组后的路径。
未发现其它 op-plugin 算子或公共接口直接调用这两个 helper 的路径。

结论：split_and_exec_npu_cmd_copy、process_tensor_list_batch 的行为变化只会通过 _foreach_copy_ 暴露，不会直接影响其它 op-plugin API。

5.4 `torch._foreach_copy_` 调用点梳理

foreach_copy 不是只有 FSDP2 调用。当前需要纳入评估的调用点如下。

PyTorch 原生 FSDP2：

torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py:281：all_gather_copy_in_cuda 中 torch._foreach_copy_(foreach_copy_dsts, all_gather_inputs)。
torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py:435：_get_param_all_gather_inputs 中 torch._foreach_copy_(splits, foreach_copy_inputs)。

torch-npu 代码仓：

torch_npu/npu/_graph_tree.py:1017：图树 replay 前的非 static 输入批量 copy，直接调用 torch._foreach_copy_(dst_tensors, src_tensors)。
torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/inductor_patch/lowering.py:6977：register_foreach_pointwise(aten._foreach_copy, copy)，属于编译 lowering 注册入口。

torch-npu 测试用例：

test/test_npu.py:301、:306：test_foreach_copy_d2h 覆盖 NPU -> CPU，non_blocking=True。
test/test_npu.py:311、:316：test_foreach_copy_h2d 覆盖 CPU -> NPU，non_blocking=True。
test/test_npu.py:320、:325：test_foreach_copy_h2d_sync 覆盖 CPU -> NPU，non_blocking=False。
test/npu/test_graph_tree.py:741：mock torch._foreach_copy_，验证 graph tree 输入 copy 使用批量 foreach copy。

op-plugin 测试用例：

test/test_v2r1_ops/test_foreach_copy.py、test/test_v2r2_ops/test_foreach_copy.py：覆盖同 dtype NPU D2D、多 dtype 类型、Host -> NPU，以及 test_foreach_copy_different_dtype 跨 dtype D2D fallback。
test/test_base_ops/test_foreach_copy.py：基础 _foreach_copy_ out/bfloat16 路径。
op_plugin/config/op_plugin_functions.yaml 和 test/core_tests/torch_npu_OpApi_schema_all.json 仅是 schema/接口声明，不是额外运行时调用点。

逐项影响分析如下：

调用点	是否命中本次 fast path 改动	影响判断	需要覆盖的回归
FSDP2 `all_gather_copy_in_cuda`	预期命中。`foreach_copy_dsts` 来自 `all_gather_input` split，`all_gather_inputs` 与 output 位于同一 device，size/stride 对齐且 dtype 已按 all-gather metadata 归一	这是本次目标场景。命中后 fast path 传 `non_blocking=true`，对齐 CUDA fast path 不依赖 API 默认 `non_blocking=False` 的设备侧排队语义	FSDP2 all-gather / fully-shard 分布式用例
FSDP2 `_get_param_all_gather_inputs`	同 dtype、同 device、size/stride 对齐时命中；如果出现跨 dtype，则 `can_use_fast_route(self, src)` 不通过，仍走 fallback	同 dtype 场景是目标影响；跨 dtype 不进入 NPU fast path，避免当前不支持场景被扩大。需要关注 mixed precision 配置是否产生跨 dtype fallback，但这不属于 fast path 改动面	FSDP2 mixed precision、all-gather input 构造 smoke
`torch_npu/npu/_graph_tree.py::_copy_inputs_and_remove_from_src`	动态输入 replay copy 通常是 NPU D2D、同 dtype、同 shape，可能命中 fast path	这是公共调用点，会受到 fast path 异步化影响。调用后紧接 `run_graph()`，如果 copy 和 graph replay 在同一 stream 顺序执行，语义不变；如果后续改成跨 stream，需要显式 wait/sync 兜底	`test/npu/test_graph_tree.py -k batched_copy`，再补一个真实 graph replay smoke 校验输出
inductor lowering `register_foreach_pointwise(aten._foreach_copy, copy)`	不直接调用 eager op-plugin `_foreach_copy_` fast path；它注册的是编译 lowering，内部走 lowering.py 的 `copy(self, src, non_blocking=False)`	直接影响较低。编译图中如降低到该 lowering，不会因为 op-plugin fast path 传 `true` 改变 lowering 逻辑；但如果编译失败回退 eager，仍会进入公共 `_foreach_copy_`	包含 `aten._foreach_copy` 的 inductor 编译 smoke
`test/test_npu.py` H2D/D2H 用例	不命中 fast path。device 不同导致 `can_use_fast_route(self, src)` 为 false	不受 fast path 末尾 `true` 影响，仍按原始 `non_blocking` 进入 `memcpyBatch` 或 slow path。`test_foreach_copy_h2d_sync` 的 `non_blocking=False` 同步语义应保持	`pytest test/test_npu.py -k foreach_copy`
`test/npu/test_graph_tree.py` mock 用例	mock 掉 `torch._foreach_copy_`，不执行 op-plugin fast path	不验证 fast path 语义，只验证 graph tree 确实合并为一次 foreach copy 调用；仍可作为调用点未丢失的结构性测试	保留 mock 测试，另加真实 replay smoke
op-plugin 同 dtype NPU D2D 单算子用例	命中 fast path	会受到本次改动影响：copy 设备侧排队后，断言前需要 `torch.npu.synchronize()` 或等价 stream 同步，避免测试读取早于 copy 完成	`test_v2r1_ops/test_foreach_copy.py`、`test_v2r2_ops/test_foreach_copy.py`
op-plugin 跨 dtype D2D 用例	不命中 fast path，`can_use_fast_route(self, src)` 因 dtype 不同返回 false	不受本次 fast path `true` 影响，继续验证 fallback 正确性。该用例是防止跨 dtype 误入 fast path 的关键保护	`test_foreach_copy_different_dtype`

结论：删除 FSDP2 Python patch 的主要收益和目标影响在 FSDP2 同 device D2D fast path。公共调用中，H2D/D2H、跨 dtype、inductor lowering 不会直接受 fast path 末尾 true 影响；graph tree 和 op-plugin 同 dtype D2D 单算子会进入同一 fast path，需要通过同步点和 graph replay 回归确认可见性。