分布式有限元方法¶

当问题在单设备上无法容纳——无论是内存还是时间方面——TensorMesh 可以对网格进行分区，在每个子域上分别于各自的 GPU（或 CPU 进程）进行装配，然后将结果作为单一逻辑系统交给 torch-sla 的分布式稀疏求解器。整个流程保持熟悉的形态：

Mesh → DistributedMesh → @distributed Assembler → DSparseMatrix → Condenser → torch_sla.solve

所有功能均位于 tensormesh.distributed（参见 API 参考），求解器端则是 torch-sla 的 DSparseTensor 机制——基于「拥有行 + halo」分区的域分解，上游基准测试已将其扩展至超过 4 亿自由度（DOF）的多 GPU 节点。

小心

分布式层已完成功能性端到端验证（装配 → 边界条件 → 多 GPU 共轭梯度法 → 与解析解对比验证），但部分接口为临时性的，可能发生变化：Condenser 桥接目前需通过汇聚到单设备矩阵再返回，且分布式右端项（RHS）装配仍在每个进程上返回完整长度的向量。该装饰器仅覆盖单场装配器——MixedElementAssembler 尚未支持（其块自由度布局需要跨分区的全局编号层；详见 ROADMAP 中的计划）。多 GPU 基准测试的完整示例将在多卡集群上验证后添加至示例库。

网格分区¶

DistributedMesh 封装了网格分区器：它将全局 Mesh 分割为子网格，并将每个子网格分配到一个设备。每个子网格在其 point_data 中保留 orig_nid 条目——即装配层用于散射到全局坐标的局部到全局节点映射。

from tensormesh import Mesh
from tensormesh.distributed import DistributedMesh

mesh = Mesh.gen_rectangle(chara_length=0.02)
dmesh = DistributedMesh(mesh, num_partitions=4)   # defaults: one per GPU
print(dmesh)          # partition sizes + devices
sub0 = dmesh[0]       # a plain Mesh on its assigned device

可通过 method= 选择三种分区方法："coordinate"（递归坐标二分法——快速，默认）、"spectral"（Fiedler 向量二分法）和 "metis"（图分区，需要 pymetis）。先在 CPU 上构建并分区，然后让子网格携带设备分配——分区器在 CPU 上建立索引。

使任意装配器支持分布式¶

无需为分布式路径编写第二个装配器。distributed() 类装饰器可将任何现有的 ElementAssembler 或 NodeAssembler 子类——内置或自定义——转换为其分布式对应版本，通过交换两个入口点实现：from_mesh 接受 DistributedMesh，__call__ 执行分布式装配：

from tensormesh.assemble import LaplaceElementAssembler
from tensormesh.distributed import distributed

DistLaplace = distributed(LaplaceElementAssembler)
K_dist = DistLaplace.from_mesh(dmesh)()     # -> DSparseMatrix

弱形式、求积设置和单元核函数均原样继承——装饰器仅重新路由单元循环在何处运行以及结果如何拼接。构造参数仍通过 from_mesh(dmesh, **kwargs) 传递，调用时数据（point_data 等）通过调用传递，与单设备约定完全一致。

底层有两种执行模式覆盖常见部署场景：

单进程，多设备。 在未初始化 torch.distributed 进程组的情况下，子网格以线程方式并行装配，每个设备一个线程，然后合并。这是低门槛模式：无需启动器，可在 notebook 中运行，且在仅 CPU 机器上它仅模拟分区——适用于在租用 GPU 之前测试依赖分区的代码。
每 rank 一个进程（SPMD）。 在 torchrun 下，每个 rank 仅装配自己的子网格；生成的 COO 三元组通过单次 all_to_all 交换路由到拥有其行的 rank（CUDA 上使用 NCCL，CPU 上使用 gloo），halo 列自动发现，每个 rank 最终获得精确的拥有行加 halo 的切片。这是可扩展的模式。

装饰器背后的自由函数——distributed_element_assemble()、distributed_element_assemble_per_rank()、distributed_node_assemble()——也是公开的，外加 distributed_element_assemble_to_sparse，它将所有内容合并回普通的全局 SparseMatrix（便于将分布式结果与单设备结果对比验证）。

分布式矩阵¶

分布式装配返回 DSparseMatrix——它是 torch_sla.distributed.DSparseTensor 的子类，正如 SparseMatrix 是 torch_sla.SparseTensor 的子类。因为它就是 DSparseTensor，所有 torch-sla API 均可直接接受：@ 矩阵向量乘法、算术运算、I/O，以及——最重要的——分布式求解器。两个 FEM 特定的扩展：

一个分区 UUID 嵌入 layout_signature——在 rank 0 上生成并广播，因此同一分区上的重复装配共享凝聚缓存，而两个独立分区的矩阵永远不会共享（即使它们的局部布局恰好重合）；
to_single()——将分布式矩阵全收集（allgather）回单个 SparseMatrix，用于验证和小规模调试。

边界条件与分布式求解¶

未作改动的 Condenser 根据矩阵类型进行分发：传给它 DSparseMatrix 即通过分布式桥接路由。当前（临时）约定中的两条实用规则：

在 CPU 上构建狄利克雷掩码（mesh.boundary_mask.cpu()）——凝聚缓冲区随后保留在 CPU 上，桥接将凝聚输出移回 GPU；
求解器消耗每个 rank 的 RHS 拥有切片，并返回解的拥有切片——先切片，后收集。

import torch_sla
from torch_sla import SolverConfig
from torch_sla.distributed import gather_owned_to_global

cond = Condenser(mesh.boundary_mask.cpu())
K_inner, b_inner = cond(K_dist, b)

owned = K_inner.partition.owned_nodes.long().to(device)
with SolverConfig(method="cg", atol=1e-12, rtol=1e-10, maxiter=1000,
                  verbose=(rank == 0)):
    u_owned = torch_sla.solve(K_inner, b_inner.to(device)[owned])

u_inner = gather_owned_to_global(owned, u_owned, b_inner.shape[0]).cpu()
u = cond.recover(u_inner)

注意求解通过 torch_sla.solve 自由函数进行——DSparseTensor 故意不设实例 .solve()——正因为子类关系，它才能在 DSparseMatrix 上工作。分布式求解为迭代式（FEM 产生的 SPD 系统采用预处理 CG）；参见稀疏求解器了解其在后端矩阵中的位置。

保持各 rank 一致¶

在 SPMD 下每个 rank 执行整个脚本，因此任何随机性——采样系数场、随机化载荷——都会在各 rank 间静默发散，求解器将收敛到随 rank 而异的答案。broadcast_from_rank0() 封装了该抽取：

from tensormesh.distributed import broadcast_from_rank0

a = broadcast_from_rank0(lambda: torch.empty((K, K)).uniform_(-1, 1))

rank 0 求值闭包；其他每个 rank 接收广播的张量。在单进程运行中它短路为普通调用，因此同一脚本在两种方式下均有效。

完整的 `torchrun` 骨架¶

examples/distributed/poisson_distributed_cuda.py 是参考端到端脚本（兼作集成测试）。骨架如下：

import os, torch, torch.distributed as dist
from tensormesh import Mesh, Condenser
from tensormesh.assemble import LaplaceElementAssembler
from tensormesh.distributed import DistributedMesh, distributed

def main():
    rank = int(os.environ["LOCAL_RANK"])
    world = int(os.environ["WORLD_SIZE"])
    dist.init_process_group(backend="nccl", rank=rank, world_size=world)
    torch.cuda.set_device(rank)
    device = torch.device(f"cuda:{rank}")
    try:
        mesh = Mesh.gen_rectangle(chara_length=0.02)        # CPU first
        dmesh = DistributedMesh(mesh, num_partitions=world,
                                devices=[device] * world)
        mesh = mesh.to(device=device)

        K = distributed(LaplaceElementAssembler).from_mesh(dmesh)()
        # ...assemble b, condense, torch_sla.solve, recover (above)...
    finally:
        dist.destroy_process_group()

if __name__ == "__main__":
    main()

启动方式：

torchrun --nproc-per-node=2 examples/distributed/poisson_distributed_cuda.py

使用 backend="gloo" 替代 "nccl" 可在 CPU rank 上运行相同的 SPMD 流程（例如用于 CI 或笔记本电脑）。

后续内容¶

稀疏求解器——完整的求解器后端矩阵，包括分布式路径从单设备求解接管的位置。
批量化工作流——适用于单设备装配的 batch_size 分块。
torch-sla 文档——DSparseTensor、Partition 以及分布式求解器内部机制。