张量并行

张量并行将模型的权重矩阵分片到多个设备上。每个设备持有每一层的按列（"colwise"）或按行（"rowwise"）切片，并计算部分结果。在层边界处，通过 AllReduce 或 AllGather 将部分结果合并为完整输出。与 上下文并行 不同，它减少了每个设备上的 权重 内存占用，这在模型无法装入单个设备时非常有用。请参阅 分布式推理 以将其与其他策略进行比较。

加载分片模型

Pass a TensorParallelConfig to the parallel_config argument of the model’s 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.

要将计划添加到另一个模型，请参阅 为模型添加张量并行。

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

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 启动它。

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（仅 transformer，32B 参数，bf16）在 4x A10G（23GB）上的结果。

局限性

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

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

为模型添加张量并行

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

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

每个键最多可以包含一个 *，且 * 之前的部分必须解析为一个 nn.ModuleList，因此一个条目可以覆盖所有块。不包含 * 的键应用于模型本身。路径是相对于模型的。

列向和行向

请务必按顺序将 "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，它们会独立地对每个打包块进行分片。

# 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]。

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

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

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

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

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() 和 enable_parallelism() 会在加载或分片任何权重之前同时检查这两项。

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

从一个类似架构的现有计划开始。

• transformer_qwenimage.py 没有融合投影，因此每个条目都是普通的 "colwise" 或 "rowwise"。
• transformer_flux.py 添加了一个打包的行投影。
• transformer_flux2.py 同时使用了两种打包样式。