Buckets:

hf-doc-build
/

doc

Files

xet

hf-doc-build/doc / diffusers /v0.43.0 /zh /optimization /tensor_parallelism.md

mishig

about 3 hours ago

preview code

download

raw

10.4 kB

# 张量并行

[张量并行](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) 同时使用了两种打包样式。

Xet Storage Details

Size:: 10.4 kB
Xet hash:: f6cd17dae55931dfa10f94b574ca35023a1fe9a9f0b44aac640b1fd25d34ceee

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	# 张量并行

	[张量并行](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) 同时使用了两种打包样式。