---
title: "MPI 执行"
description: "读取已有收缩序和切片，在多个 MPI 进程间分配计算并合并结果。"
eyebrow: "并行与分布式"
---

## 按切片分配计算 {#main-path}

`tnmpi` 读取已保存的收缩序和切片集合，将不同切片分给不同 MPI 进程。所有切片使用同一条收缩路径，各进程先累加自己的部分结果，再通过 MPI 的 Allreduce 将各进程的结果逐元素相加，让每个进程都得到完整结果。

收缩序优化与切片选择在执行前完成，可使用 Light、Heavy 或独立算法。各 MPI 进程沿保存的路径执行分配到的切片，在本地完成每片的矩阵计算，并保存完整输入张量和完整输出。

## 保存收缩序与切片 {#modes}

用 [arctn_plan](/docs/python-api#plan) 生成收缩序文件：

```python
from arctn import arctn_plan

plan = arctn_plan(
    inputs, output, size_dict,
    preset="heavy", target_size=1048576,
)
plan.save("plan.json")
```

文件保存网络结构、SSA 路径和切片 index，不保存数组数值。MPI 执行使用其中的最终路径和切片，不受生成时采用 fixed 或 dynamic 模式的影响。切片集合为空时，只有一个完整网络任务。

也可以用 [tnpath](/docs/cli#tnpath) 的 `--save-path` 保存独立算法的结果。第 2 版文件包含网络，无需再传 `--net`；第 1 版文件需要额外提供对应的网络 JSON。

## 构建与运行 {#run}

安装系统 MPI 后构建执行程序，见[安装说明](/docs/installation#mpi)。

```bash
cargo build --release --locked --features mpi --bin tnmpi

mpirun -n 4 target/release/tnmpi \
  --load-path plan.json --data tensors.bin
```

`--dtype` 选择输入和收缩计算的数据类型，默认 `f64`：

| `--dtype` | 每个元素的表示 | 每个元素的字节数 |
| --- | --- | --- |
| `f32` | 32 位实数 | 4 |
| `f64` | 64 位实数 | 8 |
| `complex64` | 实部和虚部各为一个 `f32` | 8 |
| `complex128` | 实部和虚部各为一个 `f64` | 16 |

`tensors.bin` 按输入张量顺序连续存放数组，各数组使用行优先、小端格式，不含文件头。复数的每个元素先存实部，再存虚部；所有输入必须使用与 `--dtype` 一致的数据类型。`complex128` 表示总共 128 位的复数，不是 128 位实数。例如：

```bash
mpirun -n 4 target/release/tnmpi \
  --load-path plan.json --data tensors-complex128.bin --dtype complex128
```

各进程必须能访问相同数据。未指定 `--data` 时，会按 `--seed` 生成所选类型的相同演示数组；这个 seed 不参与收缩序优化。最终 Allreduce 使用与所选类型对应的 MPI 数据类型逐元素求和。

执行后，编号为 0 的 MPI 进程（rank 0）输出 JSON 摘要，包括进程数、切片数、单片路径指标和时间：

- `load_wall_seconds_max`：各进程读取、校验、广播收缩序与准备输入所用时间的最大值。
- `execution_wall_seconds_max`：各进程执行时间的最大值，包括局部切片计算和最终 Allreduce，不包括输入准备或数值对照。

复用同一份收缩序文件和输入，就能比较不同进程数下的执行。若切片数小于进程数，多出的进程没有切片可算，但仍参与结果合并。

## 结果核对 {#correctness}

```bash
mpirun -n 2 target/release/tnmpi \
  --load-path plan.json --data tensors.bin --check
```

`--check` 在 rank 0 用单进程切片收缩核对结果。输出超过 $2^{22}$ 个元素或切片超过 $2^{12}$ 个时跳过对照，并在摘要中标明。浮点求和顺序可能随进程数和 MPI 实现变化，因此不同运行的结果不保证逐位相同。

`target_size` 限制每片产生的单个中间张量元素数，不是进程 RSS 上限。输入、最终输出和执行临时空间仍需各进程容纳，增加进程数不会使这些内存自动按比例缩小。

## 线程与作业设置 {#threads-and-time}

每个 MPI 进程的切片循环顺序执行。Rayon 默认每进程 1 个线程，可通过 `RAYON_NUM_THREADS` 设置。启用 `mt` feature 后，可通过 `MATMUL_NUM_THREADS` 设置进程内矩阵乘法的线程数。

CPU 绑定和作业时间上限由 MPI 启动器或作业调度器管理。