# 张量并行

[张量并行](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism)将模型的权重矩阵分片到多个设备上。每个设备持有每一层的按列（`"colwise"`）或按行（`"rowwise"`）切片，并计算部分结果。在层边界处，通过 `AllReduce` 或 `AllGather` 将部分结果合并为完整输出。与 [上下文并行](./context_parallelism) 不同，它减少了每个设备上的 *权重* 内存占用，这在模型无法装入单个设备时非常有用。请参阅 [分布式推理](../training/distributed_inference) 以将其与其他策略进行比较。

## 加载分片模型

Pass a [TensorParallelConfig](/docs/diffusers/v0.43.0/zh/api/parallel#diffusers.TensorParallelConfig) to the `parallel_config` argument of the model's [from_pretrained()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.from_pretrained). `tp_degree` is the number of devices to shard across and must divide the model's number of attention heads. The model's class must define a `_tp_plan`, which maps module names to a sharding style. To check whether a model supports tensor parallelism, look for this attribute on its class, such as [Flux2Transformer2DModel](/docs/diffusers/v0.43.0/zh/api/models/flux2_transformer#diffusers.Flux2Transformer2DModel).

要将计划添加到另一个模型，请参阅 [为模型添加张量并行](#add-tensor-parallelism-to-a-model)。

以这种方式加载时，检查点在读取过程中会被分片。每个 rank 仅读取每个权重中属于自己的切片，并将其直接放置到其 GPU 上，因此没有任何 rank 会持有完整大小的权重。随着 `tp_degree` 的增加，每个 rank 的内存占用会下降。

```py
import os
import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, Flux2Transformer2DModel, TensorParallelConfig

def main():
    dist.init_process_group(backend="nccl")
    rank, world_size = dist.get_rank(), dist.get_world_size()
    # torchrun sets LOCAL_RANK to this process's GPU index on its node
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)

    # Each rank reads only its own shard of every planned weight, straight onto its own GPU.
    transformer = Flux2Transformer2DModel.from_pretrained(
        "black-forest-labs/FLUX.2-dev",
        subfolder="transformer",
        dtype=torch.bfloat16,
        parallel_config=TensorParallelConfig(tp_degree=world_size),
    )

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-dev", transformer=transformer, dtype=torch.bfloat16
    )
    # The transformer is already on its device; move the remaining components individually. Do not call
    # `pipeline.to(device)` — that would move every rank's shards onto the same device.
    pipeline.text_encoder.to(device)
    pipeline.vae.to(device)

    generator = torch.Generator().manual_seed(42)
    image = pipeline(prompt="a cat holding a sign that says hello", generator=generator).images[0]
    if rank == 0:
        image.save("output.png")
    dist.destroy_process_group()

if __name__ == "__main__":
    main()
```

将脚本保存为 `tensor_parallel_flux2.py`，并使用 `torchrun` 启动它。

```shell
torchrun --nproc-per-node 4 tensor_parallel_flux2.py
```

该脚本将 `tp_degree` 设置为 `world_size`，因此 `--nproc-per-node 4` 会将 transformer 分片到 4 个 GPU 上。

仅对 transformer 进行分片。每个 rank 都持有文本编码器的完整副本，因此每个 GPU 需要足够的内存来容纳文本编码器及其所分配的 transformer 部分。

## 加载性能

使用 `parallel_config` 进行加载比加载完整模型后再调用 `enable_parallelism()` 速度更快，且每个 rank 占用的 CPU 内存更少。下表展示了 [black-forest-labs/FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev)（仅 transformer，32B 参数，bf16）在 4x A10G（23GB）上的结果。

| tp_degree | 方法 | 加载时间 | 每个 rank 的峰值 CPU |
|---|---|---|---|
| 4 | `from_pretrained(parallel_config=...)` | 12.5s | 6.8GB |
| 4 | `from_pretrained` + `enable_parallelism` | 30.4s | 64.1GB |

## 局限性

Tensor parallelism 与某些加载选项和功能不兼容，并且无法保存分片模型。

- 使用张量并行的 `parallel_config` 加载时，与 `device_map`、`quantization_config`、`low_cpu_mem_usage=False`、`use_flashpack=True` 或非 safetensors 权重不兼容。上述任一情况都会引发错误，而不是静默加载完整的检查点。
- 将张量并行与量化、卸载或 LoRA 适配器结合使用时，无论模型如何分片，都会引发错误。
- 要对已在内存中的模型进行分片，请使用相同的配置调用 `enable_parallelism()`。它会先加载完整模型，然后再进行分片，因此每个 rank 都需要容纳完整检查点的内存。使用 `parallel_config` 进行加载可以避免这一开销。
- 保存张量并行模型时，[save_pretrained()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.save_pretrained) 会引发错误。请在分片之前保存模型。

## 为模型添加张量并行

本节面向希望为 Diffusers 模型添加张量并行支持的贡献者。

张量并行仅适用于定义了 `_tp_plan` 的模型，这是一个将模块名称映射到分片样式的类属性。编写该属性通常意味着将每个 *扩展*隐藏维度的投影与将其 *收缩*回原状的投影配对。

每个键最多可以包含一个 `*`，且 `*` 之前的部分必须解析为一个 [`nn.ModuleList`](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html)，因此一个条目可以覆盖所有块。不包含 `*` 的键应用于模型本身。路径是相对于模型的。

### 列向和行向

| 风格 | 分片 | 每个 rank | 用于 |
|---|---|---|---|
| `"colwise"` | 输出特征（`weight` 的第 0 维） | 计算输出的一个切片 | `to_q`、`to_k`、`to_v`、FFN 输入投影 |
| `"rowwise"` | 输入特征（`weight` 维度 1） | 计算部分和 | `to_out.0`，FFN 输出投影 |
| `"rowwise_split_input"` | 输入特征（`weight` 维度 1） | 从其复制输入的切片中计算部分和 | 一个独立的投影层，其输入是复制的，例如读取时间步嵌入的 AdaLN 调制投影 |

请务必按顺序将 `"colwise"` 与 `"rowwise"` 配对使用。`"colwise"` 投影会使其输出保持分片状态，紧随其后的 `"rowwise"` 投影会直接消费该分片，并在块边界处执行一次 `AllReduce` 来重构结果。若以其他方式对该组合进行分片，则会在中间强制触发一次 gather 操作，从而导致通信量大幅增加。

在注意力机制中，按列划分的 Q、K 和 V 投影为每个 rank 分配一组头，因此 `tp_degree` 必须能整除头的数量。编码器流的重复部分（`add_q_proj`、`to_add_out`、`ff_context`）遵循与其图像流对应部分相同的模式。

### 融合投影

当一个 `Linear` 沿分片维度打包了多个逻辑张量时，普通的 `"colwise"`/`"rowwise"` 会直接跨拼接进行切片，导致各部分错位。请改用 `PackedColwiseParallel`/`PackedRowwiseParallel`，它们会独立地对每个打包块进行分片。

```py
# in src/diffusers/models/transformers/your_model.py
from ...hooks.tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel
```

传入 `blocks`，这是一个成比例整数列表，其总和必须能整除打包维度。例如，对于等分的 SwiGLU 门控和上投影，使用 `[1, 1]`；对于融合的 Q、K、V、门控和上投影且 `mlp_ratio=3`，使用 `[1, 1, 1, 3, 3]`。

```py
"transformer_blocks.*.ff.linear_in": PackedColwiseParallel([1, 1]),
```

当块大小仅能从配置中获知时，请省略 `blocks`，并改为在 `__init__` 期间将绝对大小存储在 `Linear` 上，作为 `_tp_packed_col_blocks` 或 `_tp_packed_row_blocks`。

```py
# in the attention module's __init__
self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim]
```

```py
# in the model's _tp_plan
"single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(),
```

[`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) 中的计划同时使用了这两种形式。双流 SwiGLU 对直接传递 `blocks`。融合的单流 QKV 和 MLP 投影将其尺寸存储在 `Flux2ParallelSelfAttention.__init__` 中的 `Linear` 上，因为它们依赖于 `mlp_ratio` 和 `mlp_mult_factor`。

### 需要排除的内容

未包含在计划中的模块会在每个 rank 上复制。请排除归一化层、patch 和文本嵌入，以及最终的 `norm_out` 和 `proj_out`。这些层很小，对它们进行分片节省的内存很少，却会增加通信开销。AdaLN 调制（`img_mod` 和 `txt_mod`）通常也足够小，可以排除。当它很大时，例如 MiniMax-H3 的 `adaln_proj.linear`，请使用 `"rowwise_split_input"` 对其进行分片。

### 约束与验证

- `tp_degree` 必须能整除 `config.num_attention_heads`。
- 每个打包块都必须*单独*能被`tp_degree`整除，而不仅仅是它们的总和。

[from_pretrained()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.from_pretrained) 和 `enable_parallelism()` 会在加载或分片任何权重之前同时检查这两项。

通过数值验证而非目视检查来验证新方案：在单个设备上使用固定种子生成，然后在张量并行下再次生成，并比较输出。错误放置的 `"colwise"`/`"rowwise"` 通常仍能运行，并生成看似合理但实际错误的图像。

> [!TIP]
> 从一个类似架构的现有计划开始。
>
> - [`transformer_qwenimage.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_qwenimage.py) 没有融合投影，因此每个条目都是普通的 `"colwise"` 或 `"rowwise"`。 - [`transformer_flux.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py) 增加了一个打包的行式投影。 - [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) 同时使用了两种打包样式。