已开启
[接口梳理1] DTensor / DeviceMesh 全量接口对标总表(Hyper · PyTorch · TorchTitan) #7
changzherui创建于  7月1日
changzherui
changzherui成员
7月1日 创建

一、文档说明

📎 上级文档: HyperParallel / PyTorch / TorchTitan 用户接口对标 #1

📎 姊妹专项(差距与路线图): DTensor / DeviceMesh 对标 PyTorch #2

本 Issue 定位: 在 #1 模块划分下,汇总 §1 DTensor / 设备 Mesh 子模块的全量对外接口对标总表(表头:接口名 · 参数 · 功能 · hyper · pytorch · titan),作为 #2 差距梳理的接口索引附录。

分析依据(源码,2026-07):

侧 路径
HyperParallel hyper_parallel/__init__.py、hyper_parallel/core/dtensor/
PyTorch torch.distributed.tensor、torch.distributed.device_mesh
TorchTitan torchtitan/distributed/parallel_dims.py、protocols/module.py、distributed/utils.py

图例: ✅ 有对标 · ⚠️ 部分对标 · ❌ 无 · 🔶 Hyper 扩展

列说明:

  • hyper:HyperParallel 导出路径或实现位置
  • pytorch:PyTorch 对标 API
  • titan:TorchTitan 实际调用的接口(多为 PyTorch 原生,非 Titan 自有封装)

二、全量接口对标总表

接口名 参数 功能 hyper pytorch titan
顶层函数
init_device_mesh device_type, mesh_shape, *, mesh_dim_names=None, rank_list=None, init_backend=True 按形状创建并缓存 DeviceMesh,可选初始化进程组 hyper_parallel.init_device_mesh torch.distributed.device_mesh.init_device_mesh(无 rank_list/init_backend) init_device_mesh(device_type, (world_size,), mesh_dim_names=("world",))
get_current_mesh 无 取当前线程 with mesh: 栈顶 mesh hyper_parallel.get_current_mesh _mesh_resources.get_current_mesh()(内部) 不直接调用
distribute_module module, device_mesh=None, partition_fn=None, input_fn=None, output_fn=None 将模块参数/缓冲区变为 DTensor;可选 I/O hook hyper_parallel.distribute_module torch.distributed.tensor.distribute_module 少用;主路径 distribute_tensor + DTensor.from_local
distribute_tensor tensor, device_mesh, placements, *, src_data_rank=0 将逻辑全局张量按 placement 分布到各 rank core.dtensor.dtensor.distribute_tensor(未进 __all__) torch.distributed.tensor.distribute_tensor(无 src_data_rank) torch.distributed.tensor.distribute_tensor
init_parameters module, stage_index=0 按 DTensor layout 分片初始化参数 hyper_parallel.init_parameters ❌ ❌;走 distribute_tensor + 权重加载
init_empty_weights include_buffers=False meta 设备上下文,延迟分配权重 hyper_parallel.init_empty_weights ❌(类似 accelerate) ❌
init_on_device device, include_buffers=False 在指定设备上创建参数/缓冲区 hyper_parallel.init_on_device ❌ ❌;model.to(device)
manual_seed seed, device_mesh DTensor 随机算子 SPMD 一致 RNG hyper_parallel.manual_seed torch.distributed.tensor._random.manual_seed(seed, mesh) torch.distributed.tensor._random.manual_seed(seed, world_mesh)
is_rng_supported_mesh device_mesh=None 检查 mesh 是否支持 DTensor RNG core.dtensor.random(未导出) torch.distributed.tensor._random.is_rng_supported_mesh ❌
ones / zeros / empty / full size, device_mesh, placements;full 另有 fill_value 按 placement 创建 DTensor dtensor.py(未导出) torch.distributed.tensor.ones/zeros/empty/full ❌ 少用
local_shard_size_and_offset curr_local_size, num_chunks, rank 计算不均匀分片 local size/offset 子模块内部 _utils 内部 ❌
Placement
Placement — placement 抽象基类 placement_types.Placement placement_types.Placement 间接使用
Shard dim 沿张量维切分 ✅ ✅ ✅
StridedShard dim, split_factor 同维跨多 mesh 轴切分 ✅(distribute_tensor 通信路径暂不支持) ✅(较新版本) 视版本
Replicate 无 全复制 ✅ ✅ ✅
Partial reduce_op="sum" partial 待归约状态 ✅ ✅ ✅
is_shard / is_replicate / is_partial 可选 dim / reduce_op 判断 placement 类型 ✅ ✅ ✅
Layout(Hyper 扩展)
Layout mesh_shape, alias_name, rank_list=None mesh 与张量维映射(MindSpore 血统) hyper_parallel.Layout ❌ ❌
Layout.from_device_mesh device_mesh 从 DeviceMesh 建 Layout ✅ ❌ ❌
layout(*alias) / layout(placements) 别名或 Placement 列表 声明张量各维 sharding ✅ 用 placements 替代 ❌
placement_to_tensor_map / tensor_map_to_placement dim 等 Placement ↔ tensor_map ✅ 内部 内部 dim_map ❌
set_partial_by_dev_axis / reset_partial / is_partial axis, op Layout 级 partial 管理 ✅ Partial placement 部分覆盖 ❌
get_global_shape / get_split_id / get_dim_split_num slice_shape / tensor_dim 全局 shape 与分片位置 ✅ ❌ ❌
DTensor 类
DTensor(...) local_tensor, device_mesh, placements 构造分布式张量 hyper_parallel.DTensor 内部;用户多用 from_local 少用
DTensor.from_local local_tensor, device_mesh, placements 从 local shard 包装 ✅(支持别名字符串 🔶) ✅ ✅ protocols/module.py
.device_mesh / .placements — mesh 与 placement 元数据 ✅ 属性 ✅ ✅
.layout — 内部 Layout ✅ 属性 ❌;用户面向 .spec ❌
.shape / .size() / .numel() dim(可选) 全局 shape / 元素数 ✅ ✅ ✅
.local_shape — local shard shape ✅ 属性 to_local().shape —
to_local() — 取 local shard ✅ ✅ ✅
redistribute device_mesh, placements 重分布到目标 layout ✅(mesh 必填;无 async_op) ✅(placements, device_mesh=None, async_op=False) 经 TP style / op 间接使用
full_tensor() — all-gather 得完整张量 ✅ ✅ 调试/测试
reduce_partial() — 归约 partial 状态 ✅ 近似 redistribute 去 Partial —
copy_ src, non_blocking=False 同 mesh/placement in-place 拷贝 ✅ 部分对齐 —
zero_ / fill_ value(fill_) in-place 填零/填值 ✅ local 语义 —
to / float *args, **kwargs dtype/device 转换,保留分布 ✅ ✅ ✅
grad / requires_grad / backward / detach 等 autograd 相关 转发到 local tensor ✅(DTensorBase) ✅ ✅
.data — 读写字段 ✅ ✅ FSDP 加载相关
.spec / DTensorSpec — mesh+placements 元数据 ❌ ✅ checkpoint/工具
is_distributed() — 是否为分布式张量 ❌ ✅ —
is_partial()(api_reference 记载) — 是否 partial ❌ 未实现(Layout.is_partial 有) 查 placements —
from_local(..., run_check=) run_check 校验 global/local 一致 ❌ ✅ —
rand / randn / randint 等 size, mesh, placements 分布式随机张量工厂 ❌ ✅ —
check_tensor_meta — 元数据校验 ❌ ✅ —
算子分发 — ATen 拦截 __torch_function__ + OpDispatcher __torch_dispatch__ 依赖 PyTorch 原生
DeviceMesh 类
DeviceMesh(...) device_type, mesh, *, mesh_dim_names=None 构造多维设备拓扑 hyper_parallel.DeviceMesh torch.distributed.device_mesh.DeviceMesh 经 init_device_mesh
.ndim / .mesh_shape / .shape — 维数与形状 ✅ ✅ ✅
.mesh_dim_names — 维度名称 ✅ ✅ ✅ 必须有名轴
.device_type / .rank / .rank_list — 设备类型与 rank 信息 ✅ ✅ ✅
.root_mesh / .sub_mesh — 子 mesh 树 ✅ ✅ 类似 —
__getitem__ sub_mesh_dim_names 按名切子 mesh ✅ ✅ ✅ 广泛使用
size mesh_dim=None 某维或总设备数 ✅ ✅ ✅
get_coordinate — 当前 rank 坐标 ✅ ✅ ✅
get_local_rank mesh_dim=None 沿某维 local rank ✅ ✅ ✅
get_group / get_all_groups mesh_dim 取通信进程组 ✅ ✅ ✅
flatten mesh_dim_name=None 压平为 1D mesh ✅ ✅ ._flatten("loss_mesh")
_unflatten dim, mesh_sizes, mesh_dim_names, backend_override=None 一维展开为多维 ✅ ✅ _unflatten ✅ parallel_dims.unflatten_mesh
_flatten mesh_dim_name, backend_override 带 backend 的 flatten ✅ ✅ ✅ backend_override="fake"
DeviceMesh.concatenate meshes 合并子 mesh 视图 ✅ ✅ —
DeviceMesh.from_group group, device_type, mesh, mesh_dim_names 从已有 PG 建 mesh ✅ ✅ —
__enter__ / __exit__ — current mesh 上下文栈 ✅ ✅ 隐式
get_device_num_along_axis axis(名) 按名取轴设备数 ✅ 🔶 ❌ ❌
get_rank_list_along_axis mesh_dim(名) 同列 rank 列表 ✅ 🔶 ❌ ❌
get_comm_group_by_axis mesh_dim 按名/索引取 PG ✅ 🔶 get_group 部分覆盖 —
get_devices_for_axis mesh_dim, rank 指定 rank 的 peer 列表 ✅ 🔶 ❌ ❌
get_global_shape slice_shape, tensor_map layout 级 global shape ✅ 🔶 ❌ ❌
axis_id / axis_index axis 轴名→内部 ID/索引 ✅ 🔶 ❌ ❌
assert_axis axis, operate_name 校验轴名 ✅ 🔶 ❌ ❌
to_hash — mesh 配置哈希 ✅ 内部类似 —
XLA device mesh — XLA 后端 mesh ❌ 不支持 ✅ —
其他
SkipDTensorDispatch no_skip=None 上下文内绕过 DTensor dispatch hyper_parallel.SkipDTensorDispatch ❌ 无公开同名 API ❌
parallel_dims.unflatten_mesh world_mesh, dim_names, dim_degrees 封装 _unflatten + fake backend ❌ 内部 _unflatten ✅ Titan 自有辅助
spmd_distribute_tensor tensor, mesh, layout SPMD 类型后端分布 ❌ ❌ ✅ torchtitan.distributed.spmd_types

三、差距速查(与 #2 联动)

本节从 #7 总表 中筛出 hyper 列为 ❌ 或 ⚠️ 的项,按来源拆为三块。详细 P0/P1 补齐路线见 #2。

3.1 仅 PyTorch 有 · Hyper 无

接口 / 参数 PyTorch 路径 功能 影响
DTensor.spec / DTensorSpec torch.distributed.tensor 封装 mesh + placements 元数据 DCP/checkpoint、工具链、类型检查
DTensor.is_distributed() 同上 判断是否为分布式张量 框架与测试中的类型分支
DTensor.from_local(..., run_check=...) 同上 创建时校验 global/local 一致性 调试与严格 SPMD 校验
check_tensor_meta 同上 元数据合法性校验 分布式张量调试
rand / randn / randint / … torch.distributed.tensor 按 placement 创建随机 DTensor 分布式随机初始化(Hyper 仅有 manual_seed)
redistribute(..., async_op=True) 同上 异步 collective + AsyncCollectiveTensor TP 边界通信-计算重叠(#270)
全量 DTensor.__torch_dispatch__ ATen 全量拦截 算子经 __torch_dispatch__ 分发 ATen 覆盖范围与 PyTorch 原生行为差异;Hyper 用 __torch_function__ + OpDispatcher
XLA / torch_xla DeviceMesh torch.distributed.device_mesh XLA 设备拓扑 Hyper 文档明确不支持
DTensor.is_partial() 生态/文档 判断 partial 状态 Hyper 未实现(api_reference.md 有记载);可用 Layout.is_partial() 或查 placements

3.2 仅 TorchTitan 有 · Hyper 无

Titan 在 DTensor/Mesh 底层多数直接调用 PyTorch(init_device_mesh、distribute_tensor、DTensor.from_local、manual_seed 等 Hyper 均有对标)。下表为 Titan 编排/实验层 特有、Hyper 无对等公开 API 的项。

接口 / 能力 Titan 路径 功能 影响
parallel_dims.unflatten_mesh torchtitan/distributed/parallel_dims.py 封装 world_mesh._unflatten(0, …) + backend_override="fake" 多维并行 mesh 一键展开;Hyper 需手写 _unflatten + backend 配置
ParallelDims 多维 mesh 视图 同上 构建 dataloading / loss / full_dense / full_sparse 等命名 mesh 训练脚本级 mesh 编排;Hyper 无对等 trainer.parallel_dims 体系
mesh._flatten("loss_mesh") 等命名 flatten parallel_dims.py 为 loss / dataloader 等场景压平子 mesh Hyper 有 flatten(),但无 Titan 这套约定命名与组合
spmd_distribute_tensor torchtitan/distributed/spmd_types.py SPMD 类型后端下的张量分布 spmd_types 实验后端;Hyper 无此路径
ShardingPlan / local_map / Module.parallelize() torchtitan/protocols/module.py、protocols/sharding.py 声明式模块 SPMD 协议(非 parallelize_module) Titan 新编排栈;Hyper 主路径为 parallelize_module + distribute_module
distinct_seed_mesh_dims 多种子策略 torchtitan/distributed/utils.py 按 PP 等维叠加 seed offset Hyper 有 manual_seed,无 Titan 级训练种子编排辅助

3.3 PyTorch 有 · Hyper 有但语义偏弱(⚠️)

接口 / 参数 PyTorch Hyper 现状 差距说明
redistribute(placements, device_mesh=None) mesh 可省略(沿用当前) mesh 必填 API 签名不等价;调用方须显式传 mesh
redistribute(async_op=…) 支持异步 仅同步 见 3.1;TP Col/Row/Seq 边界无重叠
distribute_tensor + StridedShard 完整支持 类型有;src_data_rank 通信路径不支持 多维同轴切分 + rank0 scatter 场景受限
distribute_tensor + Partial 完整 partial 语义 仅 broadcast + warn partial 初始分布不完整
ones / zeros / empty / full 官方导出 dtensor.py 有实现、未进 __all__ 功能可用但非顶层公开 API
is_rng_supported_mesh _random 子模块 有实现、未导出 需在子模块导入
DTensor.copy_ 与 Tensor 语义对齐 约束更严(mesh/placement/shape) 部分 in-place 场景行为不完全一致
DeviceMesh._unflatten + backend_override 一等支持 有 _unflatten,无 Titan 式 unflatten_mesh 封装 可用但编排成本高于 Titan
get_current_mesh _mesh_resources 内部 已导出 功能对齐;PyTorch 侧非公开顶层 API

3.4 Hyper 有 · PyTorch 用户 API 无直接对标

接口 说明
Layout 全套 MindSpore 半自动并行
placements 别名字符串 ("dp","None") 声明式可读性
init_parameters / init_empty_weights / init_on_device 大模型初始化编排
distribute_tensor(..., src_data_rank=None) 无通信本地 slice 遗留路径
SkipDTensorDispatch 梯度 hook 绕过 dispatch
DeviceMesh.get_device_num_along_axis 等 按名查询 mesh 几何
interleaved_parallel mesh 维 MindSpore 虚拟切分

四、顶层导出清单(hyper_parallel.__all__,本模块相关)

"DTensor", "Layout", "DeviceMesh", "init_device_mesh", "get_current_mesh",
"distribute_module", "init_parameters", "init_empty_weights", "init_on_device",
"manual_seed", "SkipDTensorDispatch"

未导出但常用: distribute_tensor,ones/zeros/empty/full,Shard/Replicate/Partial/StridedShard。


五、维护说明

  • 本表为 §1 DTensor / 设备 Mesh 的接口索引;差距分析与 P0/P1 补齐路线见 #2。
  • 后续模块(HSDP/FSDP、TP、PP…)将按 #1 划分另开 Issue 附录。
  • 源码变更导致接口漂移时,请同步更新本 Issue 总表。

父文档:HyperParallel / PyTorch / TorchTitan 用户接口对标 #1 · 姊妹专项:DTensor / DeviceMesh #2

likedislike
changzheruichangzherui成员
7月1日 修改标题为 “[接口梳理] DTensor / DeviceMesh 全量接口对标总表(Hyper · PyTorch · TorchTitan)”,原标题为“[基础能力1] DTensor / DeviceMesh 全量接口对标总表(Hyper · PyTorch · TorchTitan)”
changzheruichangzherui成员
7月1日 修改了issue 的描述
changzheruichangzherui成员
7月1日 修改标题为 “[接口梳理1] DTensor / DeviceMesh 全量接口对标总表(Hyper · PyTorch · TorchTitan)”,原标题为“[接口梳理] DTensor / DeviceMesh 全量接口对标总表(Hyper · PyTorch · TorchTitan)”