分布式张量 (Distributed Tensors)

Monarch 内置一个张量引擎（tensor engine），提供分布式 torch 张量：让单个 Actor 直接用普通 torch API 计算跨进程分片的张量。这是 Monarch 与 PyTorch 深度集成的核心体现。

1. activate：激活网格

使用分布式特性的入口是激活一个 ProcMesh：

import torch
from monarch.actor import this_host

mesh = this_host().spawn_procs({"gpu": 8})

# 未激活时，torch 在本地运行
local_t = torch.rand(3, 4)

# 激活后，torch 在网格的每个设备上运行
with mesh.activate():
    t = torch.rand(3, 4, device="cuda")

激活后创建的 t 是一个分布式张量：网格中每个进程各持有自己的一份。控制器上执行的任何 torch 命令，都会作用到集合中的所有张量上。

2. 检查与调试：inspect / show / fetch_shard

调试分布式张量有三件套：

import monarch
from monarch import fetch_shard

# inspect：把 rank-0 的副本拉回控制器检查（调试用）
monarch.inspect(t)

# show：在 notebook 中显示
monarch.show(t)

# 指定坐标，检查其他 rank 的副本
monarch.show(t, gpu=1)

# fetch_shard：按坐标取回某个分片
print("one", fetch_shard(t, hosts=0, gpus=0).get())
print("two", fetch_shard(t, hosts=0, gpus=1).get())

3. 直接计算

激活状态下，任何 torch 函数都能直接用于分布式张量，引擎自动生成必要的通信原语：

with mesh.activate():
    dist_compute = t @ t.T

print("dist", fetch_shard(dist_compute, hosts=0, gpus=0).get())

错误不致命：如果某条命令失败（比如张量太大），worker 依然存活，可以继续执行后续命令——这让你可以边试边搜可行配置：

N = 1
while True:
    try:
        with mesh.activate():
            batch = torch.rand(N, 1024 * 1024 * 1024, device="cuda")
        monarch.inspect(batch.sum())
        N = 2 * N
    except Exception:
        print(f"max is 2^{N} elements")
        break

4. 集合通信：沿维度归约

每个设备持有自己的张量副本（类似 torch.distributed）。跨设备的计算通过通信算子完成——以维度名为参数：

with mesh.activate():
    a = torch.rand(3, 4, device="cuda")
    r = a.reduce("gpu", "sum")   # 沿 "gpu" 维度求和

reduce(维度名, 归约方式) 与 NCCL 集合通信类似，但直接按网格维度表达，例如 "gpu"（机内 NVLink）或 "host"（跨机网络）。

5. 重命名与重塑：把物理拓扑变成并行策略

rename 和 split 可以把物理网格重整成你想要的并行形状：

# 远程网格：4 主机 × 4 GPU
remote_mesh = this_host().spawn_procs({"host": 4, "gpu": 4})

# 重命名：host 维度叫 dp（数据并行），gpu 维度叫 tp（张量并行）
mesh_2d_parallel = remote_mesh.rename(host="dp", gpu="tp")

# 拆分：把 host 维度拆成 dp + pp 两个维度
mesh_3d_parallel = remote_mesh.split(host=("dp", "pp"), gpu=("tp",), pp=2)

print(mesh_2d_parallel.to_table())
print(mesh_3d_parallel.to_table())

6. 流水线并行：切片 + 跨网格搬运

流水线并行的实现思路：把网格按 pp 维度切片成多个 stage，张量用 to_mesh 在网格之间搬运：

import torch.nn as nn

pipeline_mesh = remote_mesh.rename(host="pp")
meshes = [pipeline_mesh.slice(pp=i) for i in range(pipeline_mesh.size("pp"))]

# 每个 stage 上初始化自己的模型层
layers_per_stage = 2
stages = []
for stage_mesh in meshes:
    with stage_mesh.activate():
        layers = []
        for _ in range(layers_per_stage):
            layers.extend([nn.Linear(4, 4), nn.ReLU()])
        stages.append(nn.Sequential(*layers))

def forward_pipeline(x):
    with torch.no_grad():
        for stage_mesh, stage in zip(meshes, stages):
            x = x.to_mesh(stage_mesh)      # 张量搬到下一级网格
            with stage_mesh.activate():
                x = stage(x)
        return x

7. 传给 Actor：rref 副词

分布式张量可以传给同一网格上的任何 ActorMesh，每个 Actor 自动收到与自己同进程的那个分片：

print(dtrainers.step.call(dist_compute[0]).get())

Actor 也可以通过 rref 副词返回分布式张量。@endpoint(propagate=...) 里的传播函数类似类型签名——根据输入形状推算输出形状，让本地引用无需同步即可立即使用：

class LinearActor(Actor):
    def __init__(self):
        self.w = torch.rand(3, 3)

    @endpoint(propagate=lambda x: x)
    def forward(self, input):
        return input @ self.w

linear = trainer_procs.spawn("linear", LinearActor)

o = torch.relu(linear.forward.rref(t))  # 返回分布式张量引用

8. 进阶：梯度生成器与多流并行

官方 examples 还演示了两个高级技巧：

• monarch.grad_generator(loss, params)：增量执行反向传播，返回逐个产出参数梯度的迭代器——可以边反传边归约边更新（DDP 的核心套路）；
• monarch.Stream("comms") + borrow API：把通信（allreduce）放到独立流上，与计算重叠。borrow 的独占语义消除了 torch.cuda.stream 的竞态条件，保证确定的内存占用。

现状提示

官方明确说明：分布式张量的消息性能尚未优化到可支撑完整训练框架的程度，目前非常适合交互式调试——在控制器里像操作本地张量一样试探集群的行为。生产级训练请结合 Actor 消息 + RDMA（见高级教程）。

关键点

• with mesh.activate(): 是分布式张量的总开关；
• reduce(维度名, 方式) 按拓扑维度做集合通信；
• rename / split 把物理网格重整为 dp/tp/pp 等并行策略；
• inspect / show / fetch_shard 三件套负责调试取证；
• rref 副词让 Actor 的返回值直接成为分布式张量。