[AutoParallel] Enhance processmesh

[xuxinyi389]

PR Category

Auto Parallel

PR Types

Improvements

Description

card-73263

1. Processmesh 支持get_group方法，支持转换为对应的 Group
2. ProcessMesh 支持 get_submesh_with_dim方法，返回对应“dim”维度的通信SubMesh
3. ProcessMesh 索引方法新增支持“str”类型，本质上调用 get_submesh_with_dim

Example：

mesh_2d = dist.ProcessMesh([[0, 1, 2, 3], [4, 5, 6, 7]], dim_names=["dp", "tp"])
dp_mesh = mesh_2d["dp"]
tp_mesh = mesh_2d["tp"]

Calling mesh_2d["dp"] on rank 0, 4 returns a 1D submesh of DeviceMesh:([0, 4]).
Calling mesh_2d["dp"] on rank 1, 5 returns a 1D submesh of DeviceMesh:([1, 5]).
Calling mesh_2d["dp"] on rank 2, 6 returns a 1D submesh of DeviceMesh:([2, 6]).
Calling mesh_2d["dp"] on rank 3, 7 returns a 1D submesh of DeviceMesh:([3, 7]).
Calling mesh_2d["tp"] on rank 0, 1, 2, 3 returns a 1D submesh of DeviceMesh:([0, 1, 2, 3]).
Calling mesh_2d["tp"] on rank 4, 5, 6, 7 returns a 1D submesh of DeviceMesh:([4, 5, 6, 7]).

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[paddle-ci-bot]

Sorry to inform you that b589534's CIs have passed for more than 7 days. To prevent PR conflicts, you need to re-run all CIs manually.

[xuxinyi389]

/re-run approval

[From00]

这个方法和get_mesh_with_dim有什么区别？

[xuxinyi389]

get_mesh_with_dim 只是对 mesh 的一个简单重排，如mesh.get_mesh_with_dim(“dp”)只是把mesh的dp维放在最外维，并没有减少mesh内process_ids。mesh.get_submesh_with_dim("dp")则是获取包含当前rank的dp通信组的submesh。比如说：mesh_2d = dist.ProcessMesh([[0, 1, 2, 3], [4, 5, 6, 7]], dim_names=["dp", "tp"])
dp_mesh = mesh_2d.get_submesh_with_dim("dp")
on rank 0, 4 returns a 1D submesh of ProcessMesh:([0, 4]).
on rank 1, 5 returns a 1D submesh of ProcessMesh:([1, 5]).
on rank 2, 6 returns a 1D submesh of ProcessMesh:([2, 6]).
on rank 3, 7 returns a 1D submesh of ProcessMesh:([3, 7]).