[RFC] DeviceMesh：分布式设备拓扑抽象

状态（Status）： Draft
作者（Authors）： @lzy0920232
相关 PR： https://gitcode.com/mindspore/hyper-parallel/pull/94 （Expanded DeviceMesh Class Implementation）
对应代码： hyper_parallel/core/device_mesh.py、tests/mindspore/ut/test_device_mesh.py

[RFC] DeviceMesh：分布式设备拓扑抽象

1. 概述

1.1 简介

本文描述 Hyper-Parallel 中 DeviceMesh 的设计。DeviceMesh 将物理设备（global rank）组织成多维逻辑网格，为数据并行（DP）、张量并行（TP）、流水线并行（PP）、上下文并行（CP）等多维并行策略提供统一的拓扑抽象与通信组管理。上层并行模块（DTensor、ParallelStyle、ContextParallel 等）均以 DeviceMesh 作为唯一的拓扑入口。

1.2 动机

多维并行训练需要回答三类问题：

1. 拓扑：某个 rank 在 DP/TP/CP 各维度上的坐标是什么、同维度上的 peer 有哪些。
2. 通信域：沿某一并行维度做集合通信（all-reduce/all-gather/all-to-all）时，进程组由哪些 rank 组成。
3. 组合与切分：从全局网格中提取子网格（如只取 tp 维），或把多维网格扁平化为一维（如 dp_shard x cp -> fsdp）。

在没有统一抽象前，上述逻辑散落在各并行模块中，通信组重复创建、坐标计算易错、子网格语义不一致。DeviceMesh 将其收敛为单一抽象。

1.3 目标

1. 提供多维设备网格的构造与属性查询（mesh_shape/alias_name/rank_list/ndim）。
2. 支持按维度名/索引提取子网格（__getitem__），并维护 root_mesh 父子关系。
3. 支持按维度获取通信组（get_group）与本地排名（get_local_rank）。
4. 支持多维网格扁平化为一维（flatten），用于全局通信、checkpoint 等。
5. 提供通信组缓存与 DeviceMesh 对象缓存，避免重复创建。
6. 提供 init_device_mesh 工厂：自动生成顺序 rank_list、参数校验、复用缓存。

非目标：

1. 不实现具体集合通信算法（all-reduce 等），仅负责通信组的组织；算法由 DTensor/CP 等上层模块调用。
2. 不负责并行策略的语义（如哪个维度切权重），只提供拓扑与通信域。
3. 不处理弹性伸缩 / 动态 rank 变化。

2. 用例分析（功能场景）

2.1 构建多维并行拓扑

训练启动时按并行配置构造全局网格，例如 (dp, tp) = (2, 4)：

mesh = init_device_mesh(mesh_shape=(2, 4), alias_name=("dp", "tp"))
# mesh.rank_list == (0,1,2,3,4,5,6,7)

2.2 提取子网格做定向通信

TP 模块只需 tp 维通信域，CP 模块只需 cp 维：

tp_mesh = mesh["tp"]          # 子网格
tp_group = tp_mesh.get_group()  # 1D 子网格直接取组

2.3 跨维度组合（扁平化）

FSDP 需要把 dp_shard x cp 合并成单一 shard 维：

fsdp_mesh = mesh["dp_shard", "cp"].flatten()  # alias "dp_shard_cp"

2.4 本地坐标定位

确定当前 rank 在某维的本地分片下标（选择本地权重/数据分片）：

local_tp = mesh.get_local_rank("tp")

3. 方案设计

3.1 总体方案

DeviceMesh 持有「设备布局张量 mesh + 维度别名 alias_name」两个核心字段，所有拓扑查询都从二者派生。对象与通信组均带缓存。

flowchart TD
    A[init_device_mesh mesh_shape alias_name] --> B[参数校验]
    B --> C[生成顺序 rank_list 并 reshape]
    C --> D[_create_device_mesh 查 _DEVICE_MESH_MAP 缓存]
    D --> E[命中则复用 否则 new DeviceMesh]
    E --> F[DeviceMesh 实例]
    F --> G[getitem 取子网格]
    F --> H[get_group 取通信组]
    F --> I[get_local_rank 取本地坐标]
    F --> J[flatten 扁平为一维]
    G --> K[set root_mesh 与 _sub_mesh_cache]
    H --> L[get_comm_group_by_axis 查 _group_map 缓存]
    J --> M[注册到 root flatten_mapping]

核心要素：

要素	代码位置	责任
DeviceMesh	core/device_mesh.py	拓扑抽象、子网格、通信组、坐标、扁平化
init_device_mesh	同上	工厂：自动 rank_list + 校验 + 缓存
_create_device_mesh / _DEVICE_MESH_MAP	同上	DeviceMesh 对象缓存
get_comm_group_by_axis / _group_map	同上	通信组缓存（按 axis, rank）
platform.create_group	platform/*	真正创建进程组（集合操作）

3.1.1 类图（静态结构）

classDiagram
    class DeviceMesh {
        -_mesh : Tensor
        -_mesh_shape : tuple
        -_alias_name : tuple
        -_rank_list : tuple
        -_root_mesh : DeviceMesh
        -_sub_mesh_cache : dict
        -_flatten_mapping : dict
        +mesh
        +mesh_shape
        +alias_name
        +rank_list
        +ndim
        +root_mesh
        +__getitem__(sub_alias_name) DeviceMesh
        +get_group(mesh_dim) Group
        +get_local_rank(mesh_dim) int
        +flatten() DeviceMesh
        +get_comm_group_by_axis(axis, rank) Group
        +get_rank_list_along_axis(axis) list
        +to_hash() tuple
    }
    class init_device_mesh {
        <<factory>>
    }
    class GroupCache {
        <<module global _group_map>>
    }
    class MeshCache {
        <<module global _DEVICE_MESH_MAP>>
    }
    class platform {
        <<adapter>>
        +get_rank()
        +create_group(rank_list)
    }

    DeviceMesh --> DeviceMesh : root_mesh and sub_mesh
    DeviceMesh ..> GroupCache : get_comm_group_by_axis 读写
    DeviceMesh ..> platform : create_group and get_rank
    init_device_mesh ..> MeshCache : 复用对象
    init_device_mesh ..> DeviceMesh : 构造

3.1.2 时序图（get_group 解析通信组）

get_group 需处理三种情况：1D 自动取维、按别名/索引取维、以及扁平化别名经 root 的 flatten_mapping 解析。

sequenceDiagram
    autonumber
    participant U as Caller
    participant M as DeviceMesh
    participant R as root_mesh
    participant Cache as _group_map
    participant P as platform

    U->>M: get_group mesh_dim
    alt ndim gt 1 and mesh_dim is None
        M-->>U: raise RuntimeError need mesh_dim
    end
    M->>R: _get_root_mesh and check flatten_mapping
    alt mesh_dim is flattened name
        R-->>M: flattened sub mesh
        M->>M: get_comm_group_by_axis on flattened mesh
    else normal axis name or index
        M->>M: resolve axis from alias_name
        M->>Cache: lookup axis and rank
        alt cache hit
            Cache-->>M: cached group
        else cache miss
            M->>M: get_rank_list_along_axis
            M->>P: create_group rank_list
            P-->>M: new group
            M->>Cache: store axis and rank
        end
    end
    M-->>U: process group

3.1.3 进程/通信模型（集合一致性）

platform.create_group(rank_list) 是集合操作：参与该组的所有 rank 必须以相同顺序、相同 rank_list 调用，否则部分 rank 阻塞/死锁。设计据此保证一致性：

• 确定性 rank_list：get_rank_list_along_axis(axis) 由 mesh_shape + alias_name 纯函数推导，同一 axis 在所有 rank 上得到一致的 rank_list。
• 缓存幂等：_group_map[(axis, rank)] 保证同一 (axis, rank) 只创建一次组，重复 get_group 不再触发集合调用。
• 构造期无通信：DeviceMesh.__init__、__getitem__、flatten 仅做本地元数据计算（reshape/坐标/别名），不创建组；通信组创建被惰性延迟到首次 get_group，由各 rank 对称调用。

sequenceDiagram
    participant R0 as rank 0
    participant R1 as rank 1
    participant CC as Collective create_group
    Note over R0,R1: 所有 rank 对称调用 同序同 rank_list
    R0->>CC: create_group along tp
    R1->>CC: create_group along tp
    CC-->>R0: tp group
    CC-->>R1: tp group
    Note over R0,R1: 二次 get_group 命中 _group_map 不再进入集合

3.1.4 缓存设计

缓存	作用域	key	用途
_DEVICE_MESH_MAP	模块全局	hash(mesh_shape, alias_name, (first, last rank))	复用相同拓扑的 DeviceMesh 对象
_group_map	模块全局	(axis, rank)	复用通信组，保证集合幂等
_sub_mesh_cache	实例	子维度名组合	复用子网格对象
_cache_rank_list_along_axis	实例	axis	复用沿轴 rank_list
_global_shape_map	实例	slice_shape+map	复用 DTensor 全局 shape 推导
_flatten_mapping	实例（root）	扁平别名	支持 get_group 解析扁平维

3.2 技术选型

方案	优点	缺点	结论
各模块自行维护 rank 组	局部直观	通信组重复创建、坐标逻辑分散易错	不采用
仅用一维 rank list + 手算坐标	无新抽象	多维并行坐标/子组语义难维护	不采用
统一 DeviceMesh 多维抽象 + 缓存	拓扑/通信域单点收敛，子网格/扁平化语义一致	需维护别名与缓存规则	采用
构造期即建所有通信组	调用方零延迟	启动开销大、易因不对称构造死锁	不采用
通信组惰性创建 + 缓存	启动轻、按需建组、幂等	首次 get_group 含集合开销	采用

3.3 安全隐私与 DFX 设计

兼容性：

• mesh 接受 Tensor / list / np.ndarray，统一转 int32。
• get_group / get_local_rank 对 1D 网格允许 mesh_dim=None 自动取维；多维强制显式指定。

可维护性：

• 拓扑与通信域集中在单一类，别名到坐标的映射（_dev_name_to_index）一处定义。
• 子网格通过 root_mesh 链回溯（_get_root_mesh 递归），扁平化别名统一注册到 root 的 flatten_mapping。

可测试性：

• 纯拓扑计算（rank_list、坐标、子网格 shape）可在 mock platform 下单测，无需真实多卡。
• 通信组创建通过 mock_platform 打桩验证调用与缓存命中。

可靠性：

• 完整输入校验（见 §3.4 异常）：维度数与别名数一致、别名非空/唯一/非 "None"、interleaved_parallel 必须末维。
• 通信组确定性 + 幂等，避免不对称建组导致死锁。
• 当前 rank 不在 rank_list 时 get_local_rank 显式报错，避免错误坐标。

3.4 接口定义

3.4.1 init_device_mesh

def init_device_mesh(mesh_shape: Tuple[int, ...], alias_name: Tuple[str, ...]) -> DeviceMesh

参数	类型	描述
mesh_shape	tuple[int]	各维大小，如 (2, 4)
alias_name	tuple[str]	各维别名，长度需与 mesh_shape 一致、唯一

行为：自动生成顺序 rank_list = range(prod(mesh_shape))，reshape 成 mesh_shape，经 _create_device_mesh 走对象缓存返回。
异常：mesh_shape/alias_name 非 tuple → TypeError；长度不一致或别名重复/空 → ValueError。

3.4.2 DeviceMesh(mesh, alias_name)

参数	类型	描述
mesh	Tensor/list/np.ndarray	global rank ID 的多维布局
alias_name	tuple[str]	各维别名

属性：mesh mesh_shape alias_name rank_list ndim root_mesh sub_mesh。
异常：mesh 为 0 维 → ValueError；alias_name 非 tuple/含非 str/空串/"None"/重复，或 interleaved_parallel 非末维 → TypeError/ValueError。

3.4.3 __getitem__(sub_alias_name)

def __getitem__(self, sub_alias_name: Union[str, Tuple[str, ...]]) -> DeviceMesh

行为：按一个或多个维度名提取子网格，设置子网格 root_mesh，并写入 _sub_mesh_cache。
异常：类型非 str/tuple → TypeError；空 → ValueError；维度名不存在 → KeyError。

3.4.4 get_group(mesh_dim=None)

行为：返回某维通信组；支持别名或索引；扁平化别名经 root flatten_mapping 解析；1D 网格可省略 mesh_dim。组经 get_comm_group_by_axis 缓存。
异常：多维且 mesh_dim=None → RuntimeError；别名不存在或索引越界 → ValueError。

3.4.5 get_local_rank(mesh_dim=None)

行为：计算当前 rank 在指定维的本地下标（由 rank_list 中的 index 反推多维坐标）。
异常：多维且 mesh_dim=None → RuntimeError；维度非法或当前 rank 不在 rank_list → ValueError。

3.4.6 flatten()

行为：把当前（子）网格扁平化为 1D 网格，别名为各维别名以 _ 连接（如 dp_shard_cp），并注册到 root 的 flatten_mapping，供后续 get_group 解析。
异常：生成别名与 root 已有别名冲突 → ValueError。

3.4.7 to_hash()

返回 (mesh_shape, alias_name, (rank_list[0], rank_list[-1]))，用于在 DTensor Layout 等处做网格等价判断与缓存键。

3.5 使用约束

1. alias_name 各元素必须唯一、非空、非 "None"。
2. interleaved_parallel 若出现必须是最后一维。
3. 多维网格上的 get_group/get_local_rank 必须显式给 mesh_dim。
4. get_group 必须由参与该组的所有 rank 对称调用（同序同 rank_list）。
5. flatten 生成的别名不得与 root 现有别名冲突。

4. 测试设计

4.1 单元测试（对应 test_device_mesh.py）

1. test_init_device_mesh_basic：自动 rank_list 生成与基本属性。
2. test_init_device_mesh 缓存：相同 (mesh_shape, alias_name) 返回同一对象。
3. test_device_mesh_direct_construction_with_tensor/list/numpy：三种输入构造，自定义布局 rank_list 正确扁平。
4. test_device_mesh_getitem_valid：2D/3D 子网格提取、root_mesh 引用、子网格 rank_list。
5. test_device_mesh_get_group_valid：按别名/索引取组，_group_map 缓存命中。
6. test_device_mesh_get_local_rank：各维本地坐标计算。
7. test_device_mesh_flatten：3D 扁平为 1D、别名拼接、rank_list 与 root_mesh。
8. test_device_mesh_mesh_property：mesh 属性返回张量。

4.2 异常用例

• 别名数与维度数不一致、别名重复/空/"None"、interleaved_parallel 非末维。
• 多维网格 get_group(None) / get_local_rank(None)。
• __getitem__ 不存在的维度名。
• 当前 rank 不在 rank_list。

4.3 门禁测试

1. pylint / 语法检查。
2. mock platform 下全量单测（无需真实多卡）。

5. 缺点与风险

1. 通信组创建依赖各 rank 对称调用；若上层在不同分支条件下创建不同组，可能引入死锁，需在调用方保证对称性。
2. 缓存键 to_hash 仅取 rank_list 首尾两个 rank，理论上存在「首尾相同但中间不同」的极端布局碰撞风险（顺序 rank_list 场景下不会发生）。
3. 别名约定（None/interleaved_parallel）为隐式契约，跨模块需文档对齐。

6. 现有技术

1. PyTorch DeviceMesh / init_device_mesh 提供类似的多维拓扑与子网格 API，本设计在接口与语义上对齐并适配 MindSpore/Hyper-Parallel platform 抽象。
2. Megatron / TorchTitan 以并行维度组合 shard group，本设计的子网格与扁平化为其提供统一拓扑底座。

7. 未解决问题

1. 弹性训练 / 动态 rank 变化下网格与缓存的失效与重建策略。
2. 多机异构拓扑（非规则 mesh）的支持。
3. to_hash 碰撞在自定义非顺序布局下的健壮性增强。

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

Start with the related PR 94, then read hyper_parallel/core/device_mesh.py and tests/mindspore/ut/test_device_mesh.py. Review the listed DeviceMesh APIs, cache behavior, validation rules, and mock-platform tests; the work is done when the documented topology, subgroup, local-rank, flattening, and caching cases pass the unit and lint checks.