网格操作:切片与广播
拥有了 Mesh 之后,我们可以像操作 NumPy 数组一样对分布式集群进行精细控制。
1. 广播调用 (Broadcast)
当你在一个 Mesh 对象上调用方法时,默认行为是广播。
python
# 假设 actor_mesh 是 (2, 4) 的网格
# 这条指令会发送给全部 8 个 Actor
actor_mesh.init_weights.call(seed=42)2. 网格切片 (Slicing)
你可以使用标准的 Python 切片语法,或者基于维度的命名切片,来选中 Mesh 的一部分。
场景:混合并行策略
假设我们有 8 个 GPU。我们想让前 4 个作为 "Group A",后 4 个作为 "Group B"。
python
# 假设 mesh 是 1D 的,长度 8,维度名为 "gpus"
# 选中前 4 个
group_a = mesh.slice(gpus=slice(0, 4))
# 选中后 4 个
group_b = mesh.slice(gpus=slice(4, 8))
# 分别发送不同的指令
group_a.set_role.call("Encoder")
group_b.set_role.call("Decoder")3. 复杂索引
Monarch 支持更灵活的索引方式。例如,只选取每台机器的第 0 号 GPU(通常作为该机器的通讯协调者)。
python
# 假设 mesh 是 (Hosts=2, GPUs=4)
# [:, 0] 表示选取所有 Host 的第 0 列
leaders = mesh[:, 0]
# 此时 leaders 的逻辑形状是 (2,)
# 只有这 2 个进程会收到指令
leaders.sync_logs.call()这种能力使得实现复杂的调度逻辑(如 Ring-AllReduce 的手动调度、复杂的流水线并行分配)变得异常直观和简洁。
