[DTensor] 开发任务：distribute_tensor 合入、from_local(run_check)、rand/randn 工厂、顶层 API 导出

背景

本 Issue 是 HSDP/FSDP 用户接口对标（changzherui1/hyper-parallel#2） 在 DTensor 基础能力 方向的开发任务清单。

下文描述对外行为、对用户的影响、不做时的缺口，以及对应的 PR 实现进度。

---

PR 与任务映射

任务	内容	PR	状态
任务 1	distribute_tensor + src_data_rank scatter/broadcast	!987	🔄 Review 中
任务 2	from_local(run_check=True)	!988	🔄 Review 中
任务 3	rand / randn DTensor 工厂	!988	🔄 Review 中
任务 4	顶层 API 导出	!987 导出 distribute_tensor；!988 导出工厂 + Placement	🔄 Review 中

合入顺序建议：先合 !987（通信路径），再合 !988（run_check / rand / 完整导出）；!988 rebase 后可复用 !987 的 _collective_utils.mesh_broadcast。

---

任务 1：合入 PR !987 — distribute_tensor 支持 src_data_rank

关联 PR： mindspore/hyper-parallel!987

PR 实现要点：

• 新增 _collective_utils.py：mesh_scatter / mesh_broadcast（group-local group_src）
• distribute_tensor(..., src_data_rank=0) 默认走通信路径；src_data_rank=None 保留本地 slice
• Platform 双栈 scatter / broadcast 支持 group_src；TP 透传 src_data_rank
• UT + 双卡 ST（rank0_only / 非零 src / Replicate / 2D placement）

对外行为（合入后用户得到什么）

• 调用 distribute_tensor(global_tensor, mesh, placements) 时，只有源 rank 持有完整 global tensor 即可，框架自动 scatter/broadcast 到各 rank 的 local shard。
• 默认与 PyTorch distribute_tensor(..., src_data_rank=0) 一致；src_data_rank=None 时保持旧行为。

对用户的影响

• checkpoint 加载：rank0 读权重即可，其余 rank 自动收到正确分片。
• distribute_module 初始化：与 PyTorch「单点 global → 自动分布」一致。
• PyTorch 迁移：distribute_tensor 调用可直接照搬。

不做的话

缺失能力	用户影响
rank0_only scatter/broadcast	仅 rank0 有 checkpoint 时无法正确灌权重
与 PyTorch 语义对齐	迁移脚本行为不一致，易 silent wrong shard
TP 透传 src_data_rank	单点初始化场景与 PyTorch 不符

---

任务 2：DTensor.from_local(run_check=True)

关联 PR： mindspore/hyper-parallel!988

PR 实现要点：

• 新增 _from_local_utils.py：check_tensor_meta + mesh_broadcast
• DTensor.from_local(..., run_check=True) 跨 rank 元数据校验 + Replicate 维 broadcast
• 默认 run_check=False，热路径行为不变

对外行为（合入后用户得到什么）

• run_check=True 时：dtype/shape 不一致立即报错；Replicate 维自动对齐。
• 默认关闭，与现有行为相同。

对用户的影响

• 调试并行代码时可尽早发现分片/shape/Replicate 不一致。
• 与 PyTorch from_local(run_check=True) 调试开关对齐。

不做的话

缺失能力	用户影响
跨 rank 分片校验	分片错误可能 silent 到 loss/精度才暴露
Replicate 维自动对齐	须用户自行保证各 rank 一致
PyTorch 对等 API	迁移体验不完整

---

任务 3：DTensor.rand / DTensor.randn 工厂

关联 PR： mindspore/hyper-parallel!988

PR 实现要点：

• _dtensor_init_helper(..., rng_tracked=True) + OffsetBasedRNGTracker._distribute_region
• 新增 rand() / randn() 模块级工厂；Torch / MindSpore platform 补齐

对外行为（合入后用户得到什么）

w = DTensor.randn(1024, 4096, device_mesh=mesh, placements=[Shard(0)])

• 等价于 global randn 再按 placement 切分；Replicate 维各 rank 相同；配合 manual_seed 可复现。

对用户的影响

• 一行代码按 mesh 正确随机初始化 DTensor 参数。
• 可替换 PyTorch torch.distributed.tensor.randn / rand。

不做的话

缺失能力	用户影响
随机 DTensor 工厂	须绕路（Module + shard / checkpoint / 手动 RNG）
PyTorch API 对齐	迁移脚本不可用 DTensor.randn(...)
正确分布式 RNG 语义	每 rank 独立 randn 可能语义错误

---

任务 4：顶层 API 导出

关联 PR：

• !987 — distribute_tensor（若 !988 先合，由 !988 一并导出）
• !988 — ones/zeros/empty/full/rand/randn + Shard/Replicate/Partial/Placement

PR 实现要点（!988）：

from hyper_parallel import (
    DTensor, distribute_tensor, distribute_module,
    ones, zeros, empty, full, rand, randn,
    Shard, Replicate,
)

对外行为（合入后用户得到什么）

• 与 PyTorch from torch.distributed.tensor import ... 习惯一致。
• __all__ 明确公开 API 边界。

对用户的影响

• 降低上手成本，避免深路径 hyper_parallel.core.dtensor.*。
• 按 PyTorch 习惯 import 不再 ImportError。

不做的话

缺失能力	用户影响
顶层导出	用户以为功能不存在
文档一致	示例混用内部路径，易误用

---

里程碑进度

里程碑	内容	PR	状态
M1	distribute_tensor 通信 + 顶层导出	!987	🔄
M2	from_local(run_check)	!988	🔄
M3	rand/randn + 完整顶层导出	!988	🔄

关联

• 上级对标：changzherui1/hyper-parallel#2
• PR !987：feat(dtensor): distribute_tensor scatter/broadcast via src_data_rank
• PR !988：feat(dtensor): from_local run_check, rand/randn factories, and public API exports

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

Start by reviewing linked PRs !987 and !988, then inspect the mentioned _collective_utils.py and _from_local_utils.py modules and their UT and two-card ST coverage. Done means the listed distribute_tensor, from_local(run_check), rand/randn, and top-level export tasks are reviewed and merged in the stated order.