[Bug]: fully_shard 在 reduce_dtype 与梯度 dtype 相同时释放了调用方的梯度 storage，导致 ConcatD 访问非法 GM 地址-hyper-parallel-AtomGit

已开启

[Bug]: fully_shard 在 reduce_dtype 与梯度 dtype 相同时释放了调用方的梯度 storage，导致 ConcatD 访问非法 GM 地址 #355

Java不加糖创建于 29 天前

29 天前创建

Checklist

1. I have searched the existing issues.
2. I have read the relevant documentation.
3. I have created a minimal reproduction case that clearly demonstrates the issue.

🐛 Describe the bug

fully_shard 在 reduce-scatter 完成后释放通信输入的 storage，但在 reduce_dtype
与梯度 dtype 相同时，该"通信输入"与调用方的真实梯度共享同一块 storage，导致
释放了 fully_shard 并不拥有的内存。

后果：该内存块被分配器回收再分配给后续算子，训练首步即在 Ascend 上崩溃：

fault kernel_name=ConcatD_36655e11f3e4558425f72c5a4171ecc6_high_performance_4000001
aivec error, error code = 0x800000
errorStr: MTE accesses an invalid GM address or the cross-device memory access times out

崩溃在下一次 host 同步处暴露为 RuntimeError: SyncCopy failed for Tensor(shape=[], dtype=Float32)。

缺陷位置

hyper_parallel/platform/mindspore/fully_shard/param.py

def reduce_scatter_grad(self, ...):
    ...
    self._grad = grad.to(self.reduce_comm_dtype(grad))   # (1) 同 dtype 时不拷贝
    ...
    self._grad = self._grad.reshape(-1)                  # (2) 仍共享 storage

def reduce_scatter_output(self):
    if self.reduce_scatter_comm_ctx.reduce_scatter_handle is not None:
        self.reduce_scatter_comm_ctx.reduce_scatter_handle.wait()
        self._grad.untyped_storage().resize_(0)          # (3) 释放调用方的梯度

reduce_comm_dtype 在 reduce_dtype 已设置时直接返回它。当
MixedPrecisionPolicy(reduce_dtype=float32) 且梯度本身为 fp32 时，(1) 是同 dtype
转换——MindSpore 返回新的 Tensor 对象但共享源 storage，(2) 继续共享，于是 (3)
把真实梯度缓冲的 storage 置零。

仅当 reduce_dtype 与梯度 dtype 不同时才会产生独立副本，因此该缺陷是配置相关的。

最小复现（纯 CPU，无需 NPU）

import numpy as np
import mindspore as ms

def check(make_comm_buffer):
    grad = ms.Tensor(np.arange(16, dtype=np.float32).reshape(4, 4))
    comm_buf = make_comm_buffer(grad)
    comm_buf.untyped_storage().resize_(0)          # reduce_scatter_output 所做的事
    return grad.untyped_storage().size() == 0      # 调用方的梯度是否被释放

print(check(lambda g: g.to(ms.float32)))               # True  <- 缺陷
print(check(lambda g: g.to(ms.float32).reshape(-1)))   # True  <- 缺陷
print(check(lambda g: g.to(ms.bfloat16)))              # False <- 真拷贝，安全

注意：不要用 untyped_storage().data_ptr() 判断是否共享——reshape 之后该指针会
变化，具有误导性；resize_ 的传播才是可靠判据。

定位过程

在 mindspore/mindformers 门禁上做三点二分，同一份 mindformers 代码，仅更换 hyper_parallel：

hyper_parallel 构建	含 `e85a106a`	含 `54793eb2`	门禁
`r1.0.0_20260715000008`	✗	✗	全部通过
`master_20260824010007_454d844f`	✓	✗	全部通过
`master_20260827010008_095c6ab6`	✓	✓	UT 失败

绿→红区间内仅 54793eb2「refactor(fully_shard): align MindSpore FSDP lifecycle with
Torch」(MR !1245) 一个提交触及 fully_shard。该提交删除了原先带 .contiguous() 的
pack_utils.py 打包路径，改为在 reduce_scatter_grad 内联 chunk/cat，并新增了
上述 storage 释放。

已排除的因素：

mindspore/hyper_parallel 版本错配：将 mindspore 由 master_20260630010018
升至配套的 master_20260827010013 后重跑，fault kernel 名称与前次逐字节相同。
调用方缺少同步：在优化器读取梯度前补 hsdp_sync_stream() 无效——该接口内部
执行的正是这段释放逻辑，只会让释放更早发生，而非避免。

触发配置

deepseek3，2 卡，data_parallel_shard: -1、data_parallel_shard_strategy: optim_grads_params
MixedPrecisionPolicy(reduce_dtype=float32, apply_grad_on_fp32_main_grad=True)
模型 params_dtype: float32（梯度为 fp32，与 reduce_dtype 相同）
comm_fusion=False

建议修复方向

在释放前确保 self._grad 是 fully_shard 自有的缓冲，例如在 (1) 处当转换未产生新
storage 时显式 clone()，或在 (3) 处仅释放确认由本模块分配的 storage。

Expected behavior

fully_shard 只应释放自身分配的通信缓冲。当 grad.to(reduce_dtype) 未产生新
storage 时（reduce_dtype 与梯度 dtype 相同），reduce_scatter_output() 不应对该
storage 调用 resize_(0)，否则会破坏调用方仍在使用的梯度缓冲。

Environment info

hyper_parallel: master_20260827010008_095c6ab / master_20260829010008_15a0afec（均复现）
mindspore: master_20260827010013_b007b0db / master_20260829010012_c5d5d454
硬件: Ascend 910B, 2 卡
未复现版本: r1.0.0_20260715000008, master_20260824010007_454d844f
最小复现脚本为纯 CPU，任意 MindSpore 安装均可运行

Java不加糖

29 天前评论：

补充一个候选补丁与验证结果，供参考。

候选补丁

hyper_parallel/platform/mindspore/fully_shard/param.py::reduce_scatter_grad

         if self.unsharded_accumulated_grad is not None:
             grad = self.unsharded_accumulated_grad_data
         else:
             grad = self.unsharded_grad_data
-        self._grad = grad.to(self.reduce_comm_dtype(grad))
+        # ``Tensor.to`` is a no-op when the dtype already matches, returning a
+        # tensor that SHARES the caller's storage. ``reduce_scatter_output`` later
+        # resizes ``self._grad``'s storage to 0, so without an explicit copy that
+        # release would free the live gradient buffer this module does not own.
+        comm_grad = grad.to(self.reduce_comm_dtype(grad))
+        if comm_grad is grad or comm_grad.untyped_storage().data_ptr() == grad.untyped_storage().data_ptr():
+            comm_grad = comm_grad.clone()
+        self._grad = comm_grad

思路是让 self._grad 始终为 fully_shard 自有的缓冲，从而 reduce_scatter_output()
中的 resize_(0) 只作用于本模块分配的 storage。

另一种等价方向是不改此处、改为在 reduce_scatter_output() 中仅释放确认自有的
storage；两者取其一即可。

验证边界（重要）

在机理层面验证了该守卫的有效性：

修复前:  同 dtype -> 释放调用方梯度 = True    异 dtype -> False
修复后:  同 dtype -> 释放调用方梯度 = False   异 dtype -> False

但未经由真实的 reduce_scatter_grad 路径端到端验证（需分布式环境，提单方无
NPU 环境）。补丁中 data_ptr() 的比较仅在此处（未经 reshape 的即时结果）可靠；
如放到 reshape 之后判断会失真，详见正文中的说明。

请以贵方的 FSDP 用例为准做最终验证。

影响面

mindspore/mindformers 门禁在升级 hyper-parallel 至 master 线时被此问题阻塞，目前
将 pin 暂时固定在 master_20260824010007_454d844f（引入该缺陷之前的最新 master
构建）。相关 PR：mindspore/mindformers!8758。

Java不加糖

29 天前评论：

更正上一条评论：其中基于 untyped_storage().data_ptr() 比较的写法不可用，请勿采纳。

为何不可用

该比较不是一次无副作用的观测——读取源张量的 storage 会改变后续 reshape 的别名
关系。确定性复现（同一环境重复三次结果一致）：

def fresh(): return ms.Tensor(np.arange(4, dtype=np.float32))

# 不做观测
t = fresh(); c = t.to(ms.float32).reshape(-1)
c.untyped_storage().resize_(0)
t.untyped_storage().size()          # -> 0   （共享，被释放）

# 在 .to() 与 .reshape() 之间读一次 data_ptr
t = fresh(); a = t.to(ms.float32)
_ = a.untyped_storage().data_ptr() == t.untyped_storage().data_ptr()
c = a.reshape(-1)
c.untyped_storage().resize_(0)
t.untyped_storage().size()          # -> 16  （未被释放）

即：探测报告"共享"，但探测之后强制释放却伤不到源张量。以它作为"是否可释放"的判据
并不成立。上一条评论里"data_ptr() 在此处可靠"的说法是错的，一并更正。

改用不依赖运行时探测的写法

所有权可以直接从代码路径推出，无需观测 storage：Tensor.to 仅在 dtype 真正变化时
分配，reshape 从不分配，mint.cat（非零维分片路径与 dim-0 补齐路径）总是分配。

+    #: Whether ``_grad`` is a communication base this object allocated, and may
+    #: therefore reclaim in :meth:`reduce_scatter_output`.
+    _grad_owns_storage = False
@@ reduce_scatter_grad
-        self._grad = grad.to(self.reduce_comm_dtype(grad))
+        comm_dtype = self.reduce_comm_dtype(grad)
+        self._grad_owns_storage = comm_dtype is not None and grad.dtype != comm_dtype
+        self._grad = grad.to(comm_dtype)
         shard_dim = self.hsdp_placement.dim
         if self.shard_world_size <= 1:
             self._grad = self._grad.reshape(-1)
         elif shard_dim != 0:
             grad_chunks = self._grad.chunk(self.shard_world_size, dim=shard_dim)
             self._grad = ms.mint.cat(grad_chunks, dim=0).reshape(-1)
+            self._grad_owns_storage = True
         else:
             padded_unsharded_dim0 = self.padded_sharded_param_size[0] * self.shard_world_size
             if self._grad.shape[0] != padded_unsharded_dim0:
                 self._grad = _pad_dim0_for_communication(self._grad, padded_unsharded_dim0)
+                self._grad_owns_storage = True
             self._grad = self._grad.reshape(-1)
@@ reduce_scatter_output
             self.reduce_scatter_comm_ctx.reduce_scatter_handle.wait()
-            self._grad.untyped_storage().resize_(0)
+            if self._grad_owns_storage:
+                self._grad.untyped_storage().resize_(0)
             self._grad = None
+            self._grad_owns_storage = False

这样保留了 test_reduce_scatter_scales_gradient_view_without_materializing 所固化的
设计意图（偶数 dim-0 路径刻意复用调用方梯度、不materialize 副本），只是不再回收这块
并非自有的 storage。

验证

tests/ut/platform/mindspore/ 全量 283 passed, 11 skipped（含上述既有用例）。
另补了两个用例覆盖两侧：别名路径下调用方梯度存活、补齐路径下自有 buffer 正常回收。

仍未在 NPU 上端到端验证（提单方无该环境），请以贵方 FSDP 用例为准。

J Java不加糖

29 天前关联了pull request：[bugfix] fully_shard 仅回收自有的 reduce-scatter 通信缓冲，修复释放调用方梯度导致的 ConcatD 非法 GM 访问