已开启
[接口梳理1] DTensor / DeviceMesh 全量接口对标总表(Hyper · PyTorch · TorchTitan) #7
changzherui创建于 7月1日
7月1日 修改标题为 “[接口梳理] DTensor / DeviceMesh 全量接口对标总表(Hyper · PyTorch · TorchTitan)”,原标题为“[基础能力1] DTensor / DeviceMesh 全量接口对标总表(Hyper · PyTorch · TorchTitan)”
7月1日 修改标题为 “[接口梳理] DTensor / DeviceMesh 全量接口对标总表(Hyper · PyTorch · TorchTitan)”,原标题为“[基础能力1] DTensor / DeviceMesh 全量接口对标总表(Hyper · PyTorch · TorchTitan)”
7月1日 修改了issue 的描述
7月1日 修改标题为 “[接口梳理1] DTensor / DeviceMesh 全量接口对标总表(Hyper · PyTorch · TorchTitan)”,原标题为“[接口梳理] DTensor / DeviceMesh 全量接口对标总表(Hyper · PyTorch · TorchTitan)”
7月1日 修改标题为 “[接口梳理1] DTensor / DeviceMesh 全量接口对标总表(Hyper · PyTorch · TorchTitan)”,原标题为“[接口梳理] DTensor / DeviceMesh 全量接口对标总表(Hyper · PyTorch · TorchTitan)”
一、文档说明
本 Issue 定位: 在 #1 模块划分下,汇总 §1 DTensor / 设备 Mesh 子模块的全量对外接口对标总表(表头:接口名 · 参数 · 功能 · hyper · pytorch · titan),作为 #2 差距梳理的接口索引附录。
分析依据(源码,2026-07):
hyper_parallel/__init__.py、hyper_parallel/core/dtensor/torch.distributed.tensor、torch.distributed.device_meshtorchtitan/distributed/parallel_dims.py、protocols/module.py、distributed/utils.py图例: ✅ 有对标 · ⚠️ 部分对标 · ❌ 无 · 🔶 Hyper 扩展
列说明:
二、全量接口对标总表
init_device_meshdevice_type,mesh_shape,*,mesh_dim_names=None,rank_list=None,init_backend=Truehyper_parallel.init_device_meshtorch.distributed.device_mesh.init_device_mesh(无rank_list/init_backend)init_device_mesh(device_type, (world_size,), mesh_dim_names=("world",))get_current_meshwith mesh:栈顶 meshhyper_parallel.get_current_mesh_mesh_resources.get_current_mesh()(内部)distribute_modulemodule,device_mesh=None,partition_fn=None,input_fn=None,output_fn=Nonehyper_parallel.distribute_moduletorch.distributed.tensor.distribute_moduledistribute_tensor+DTensor.from_localdistribute_tensortensor,device_mesh,placements,*,src_data_rank=0core.dtensor.dtensor.distribute_tensor(未进__all__)torch.distributed.tensor.distribute_tensor(无src_data_rank)torch.distributed.tensor.distribute_tensorinit_parametersmodule,stage_index=0hyper_parallel.init_parametersdistribute_tensor+ 权重加载init_empty_weightsinclude_buffers=Falsehyper_parallel.init_empty_weightsinit_on_devicedevice,include_buffers=Falsehyper_parallel.init_on_devicemodel.to(device)manual_seedseed,device_meshhyper_parallel.manual_seedtorch.distributed.tensor._random.manual_seed(seed, mesh)torch.distributed.tensor._random.manual_seed(seed, world_mesh)is_rng_supported_meshdevice_mesh=Nonecore.dtensor.random(未导出)torch.distributed.tensor._random.is_rng_supported_meshones/zeros/empty/fullsize,device_mesh,placements;full另有fill_valuedtensor.py(未导出)torch.distributed.tensor.ones/zeros/empty/fulllocal_shard_size_and_offsetcurr_local_size,num_chunks,rank_utils内部Placementplacement_types.Placementplacement_types.PlacementSharddimStridedSharddim,split_factordistribute_tensor通信路径暂不支持)ReplicatePartialreduce_op="sum"is_shard/is_replicate/is_partialdim/reduce_opLayoutmesh_shape,alias_name,rank_list=Nonehyper_parallel.LayoutLayout.from_device_meshdevice_meshlayout(*alias)/layout(placements)placements替代placement_to_tensor_map/tensor_map_to_placementdim等set_partial_by_dev_axis/reset_partial/is_partialaxis,opget_global_shape/get_split_id/get_dim_split_numslice_shape/tensor_dimDTensor(...)local_tensor,device_mesh,placementshyper_parallel.DTensorfrom_localDTensor.from_locallocal_tensor,device_mesh,placementsprotocols/module.py.device_mesh/.placements.layout.spec.shape/.size()/.numel()dim(可选).local_shapeto_local().shapeto_local()redistributedevice_mesh,placementsasync_op)placements,device_mesh=None,async_op=False)full_tensor()reduce_partial()copy_src,non_blocking=Falsezero_/fill_value(fill_)to/float*args,**kwargsgrad/requires_grad/backward/detach等DTensorBase).data.spec/DTensorSpecis_distributed()is_partial()(api_reference 记载)Layout.is_partial有)from_local(..., run_check=)run_checkrand/randn/randint等size,mesh,placementscheck_tensor_meta__torch_function__+ OpDispatcher__torch_dispatch__DeviceMesh(...)device_type,mesh,*,mesh_dim_names=Nonehyper_parallel.DeviceMeshtorch.distributed.device_mesh.DeviceMeshinit_device_mesh.ndim/.mesh_shape/.shape.mesh_dim_names.device_type/.rank/.rank_list.root_mesh/.sub_mesh__getitem__sub_mesh_dim_namessizemesh_dim=Noneget_coordinateget_local_rankmesh_dim=Noneget_group/get_all_groupsmesh_dimflattenmesh_dim_name=None._flatten("loss_mesh")_unflattendim,mesh_sizes,mesh_dim_names,backend_override=None_unflattenparallel_dims.unflatten_mesh_flattenmesh_dim_name,backend_overridebackend_override="fake"DeviceMesh.concatenatemeshesDeviceMesh.from_groupgroup,device_type,mesh,mesh_dim_names__enter__/__exit__get_device_num_along_axisaxis(名)get_rank_list_along_axismesh_dim(名)get_comm_group_by_axismesh_dimget_group部分覆盖get_devices_for_axismesh_dim,rankget_global_shapeslice_shape,tensor_mapaxis_id/axis_indexaxisassert_axisaxis,operate_nameto_hashSkipDTensorDispatchno_skip=Nonehyper_parallel.SkipDTensorDispatchparallel_dims.unflatten_meshworld_mesh,dim_names,dim_degrees_unflatten+ fake backend_unflattenspmd_distribute_tensortensor,mesh,layouttorchtitan.distributed.spmd_types三、差距速查(与 #2 联动)
3.1 仅 PyTorch 有 · Hyper 无
DTensor.spec/DTensorSpectorch.distributed.tensorDTensor.is_distributed()DTensor.from_local(..., run_check=...)check_tensor_metarand/randn/randint/ …torch.distributed.tensormanual_seed)redistribute(..., async_op=True)AsyncCollectiveTensorDTensor.__torch_dispatch____torch_dispatch__分发__torch_function__+OpDispatchertorch_xlaDeviceMeshtorch.distributed.device_meshDTensor.is_partial()api_reference.md有记载);可用Layout.is_partial()或查placements3.2 仅 TorchTitan 有 · Hyper 无
parallel_dims.unflatten_meshtorchtitan/distributed/parallel_dims.pyworld_mesh._unflatten(0, …)+backend_override="fake"_unflatten+ backend 配置ParallelDims多维 mesh 视图dataloading/loss/full_dense/full_sparse等命名 meshtrainer.parallel_dims体系mesh._flatten("loss_mesh")等命名 flattenparallel_dims.pyflatten(),但无 Titan 这套约定命名与组合spmd_distribute_tensortorchtitan/distributed/spmd_types.pyspmd_types实验后端;Hyper 无此路径ShardingPlan/local_map/Module.parallelize()torchtitan/protocols/module.py、protocols/sharding.pyparallelize_module)parallelize_module+distribute_moduledistinct_seed_mesh_dims多种子策略torchtitan/distributed/utils.pymanual_seed,无 Titan 级训练种子编排辅助3.3 PyTorch 有 · Hyper 有但语义偏弱(⚠️)
redistribute(placements, device_mesh=None)redistribute(async_op=…)distribute_tensor+StridedShardsrc_data_rank通信路径不支持distribute_tensor+Partialones/zeros/empty/fulldtensor.py有实现、未进__all__is_rng_supported_mesh_random子模块DTensor.copy_DeviceMesh._unflatten+backend_override_unflatten,无 Titan 式unflatten_mesh封装get_current_mesh_mesh_resources内部3.4 Hyper 有 · PyTorch 用户 API 无直接对标
Layout全套("dp","None")init_parameters/init_empty_weights/init_on_devicedistribute_tensor(..., src_data_rank=None)SkipDTensorDispatchDeviceMesh.get_device_num_along_axis等interleaved_parallelmesh 维四、顶层导出清单(
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。五、维护说明
父文档:HyperParallel / PyTorch / TorchTitan 用户接口对标 #1 · 姊妹专项:DTensor / DeviceMesh #2