diffusers 如何用 TensorParallelConfig 分片注意力权重以跨 GPU 运行超大模型
【免费下载链接】diffusers🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers
当扩散模型的某个组件（通常是 transformer）权重放不进单卡显存时，diffusers 的张量并行（Tensor parallelism）可以把权重矩阵按行/列切分到多张 GPU 上：每个设备持有每层的列切分（"colwise"）或行切分（"rowwise"），各算部分结果，在层边界通过AllReduce/AllGather还原完整输出。与上下文并行不同，张量并行直接降低单卡的权重显存占用，这正是模型放不下单卡时的解法。入口是TensorParallelConfig配合~ModelMixin.enable_parallelism。
需要先了解三点前提：
- 该能力是实验性功能，API 可能随时破坏性变更（
enable_parallelism调用时会打印警告，并行模块源码中也标注了 Experimental）； tp_degree（分片到多少张卡）必须能整除模型的注意力头数num_attention_heads，这一点在~ModelMixin.enable_parallelism中做了校验；- 模型类必须定义
_tp_plan属性（模块名 glob 到分片方式的扁平映射），否则调用enable_parallelism时直接抛错。仓库中已内置 plan 的模型包括 QwenImageTransformer2DModel、FluxTransformer2DModel、Flux2Transformer2DModel。
准备环境
张量并行运行在 PyTorch 分布式环境下，支持的设备类型为"cuda"和"neuron"（见TensorParallelConfig的说明）。脚本通过torchrun启动，--nproc-per-node指定参与的 GPU 数。
下面以 4 卡为例，完整脚本可直接保存为tensor_parallel_flux.py：
import torch from torch import distributed as dist from diffusers import DiffusionPipeline, TensorParallelConfig def setup_distributed(): if not dist.is_initialized(): dist.init_process_group(backend="nccl") rank = dist.get_rank() device = torch.device(f"cuda:{rank}") torch.cuda.set_device(device) return device def main(): device = setup_distributed() world_size = dist.get_world_size() pipeline = DiffusionPipeline.from_pretrained( "black-forest-labs/FLUX.2-dev", torch_dtype=torch.bfloat16 ) # weights stay on CPU # Shard the transformer first, then move only each rank's slice onto the accelerator. pipeline.transformer.enable_parallelism(config=TensorParallelConfig(tp_degree=world_size)) pipeline.transformer.to(device) # Move the remaining, non-sharded components onto the accelerator individually. 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 dist.get_rank() == 0: image.save("output.png") if dist.is_initialized(): dist.destroy_process_group() if __name__ == "__main__": main()操作步骤与其中的关键取舍：
DiffusionPipeline.from_pretrained先让权重停留在 CPU 上（文档明确注释 "weights stay on CPU"），避免先整卡加载再分片造成显存峰值；pipeline.transformer.enable_parallelism(config=TensorParallelConfig(tp_degree=world_size))只对 transformer 做分片——分片必须先于.to(device)，这样每张卡只搬运属于自己的那一份切片；- text encoder 和 VAE 未参与分片，逐个
.to(device)挪到加速器上； - 固定
generator种子保证各 rank 从相同 latents 起步，只有 rank 0 负责保存output.png。
启动命令（会启动 4 个进程，每个进程占一张卡）：
torchrun --nproc-per-node 4 tensor_parallel_flux.pytp_degree取自world_size，所以--nproc-per-node 4就是把 transformer 切到 4 张卡上。如果你的模型注意力头数不能被 4 整除，需要把进程数改为能整除头数的值，否则enable_parallelism会抛出`tp_degree` (4) must divide the number of attention heads (...)。
如何写一个_tp_plan
使用自定义模型或想理解分片细节时，需要读懂 plan 的配对规则。每条规则的核心是把扩张隐藏维度的投影和把它收缩回来的投影配对：
| Style | 分片对象 | 每个 rank | 适用于 |
|---|---|---|---|
"colwise" | 输出特征（weightdim 0） | 计算输出的一片 | to_q、to_k、to_v、FFN 输入投影 |
"rowwise" | 输入特征（weightdim 1） | 计算部分和 | to_out.0、FFN 输出投影 |
必须按这个顺序配对："colwise"投影的输出处在分片状态，紧随其后的"rowwise"投影直接消费该分片，块边界的一次AllReduce即可还原结果。其它切法会在中间引入额外 gather，通信量大得多。对注意力而言，这意味着每个 rank 持有部分注意力头，这正是tp_degree必须整除头数的原因。编码器流（add_q_proj、to_add_out、ff_context）遵循与其图像流对应模块相同的模式。
plan 的 key 语法限制：
- 每个 key 最多含一个
*，且*前的前缀必须解析到nn.ModuleList，这样一条规则能覆盖所有 block； - 不含
*的 key 作用于模型本身； - 路径相对于模型定义。
融合投影：当一个Linear把多个逻辑张量沿被分片维度打包（如 QKV 融合、SwiGLU gate+up 融合），直接用"colwise"/"rowwise"会横切过拼接导致错位，应改用PackedColwiseParallel/PackedRowwiseParallel：
# in src/diffusers/models/transformers/your_model.py from ...hooks.tensor_parallel import PackedColwiseParallel, PackedRowwiseParallelblocks是一组比例整数，其和必须能整除被打包的维度——例如等半 SwiGLU gate+up 投影用[1, 1]，mlp_ratio=3的 Q+K+V+gate+up 融合投影用[1, 1, 1, 3, 3]：
"transformer_blocks.*.ff.linear_in": PackedColwiseParallel([1, 1]),如果 block 大小只能从 config 推算（如依赖mlp_ratio），可以省略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 融合投影因依赖mlp_ratio和mlp_mult_factor而把尺寸存在Flux2ParallelSelfAttention.__init__里。写新 plan 时建议从相似架构的现有 plan 起步——QwenImageTransformer2DModel完全未融合、条目全是纯"colwise"/"rowwise"，FluxTransformer2DModel多一条 packed rowwise 投影，Flux2Transformer2DModel两种 packed 风格都有。
应该留空的部分：不在 plan 里出现的模块会在每个 rank 上保持完整复制，这对归一化层、AdaLN 调制（img_mod/txt_mod）、patch/text embedding 以及末尾的norm_out/proj_out是正确选择——它们体积小，分片省不了多少显存反而增加通信。
验证与已知限制
数值验证。文档给的方法很直接：用固定种子在单卡上生成一次，再用张量并行生成一次，对比输出。因为colwise/rowwise放错位置时代码通常仍然能跑，只会产出看似合理但错误的图像，肉眼看 plan 是不可靠的。
约束条件（在~ModelMixin.enable_parallelism与文档中明确列出）：
tp_degree必须整除config.num_attention_heads；- 每个 packed block 必须单独能被
tp_degree整除，仅总和可整除不够。
策略选择。diffusers 的 分布式推理指南 在"Choosing a strategy"一节给出了四条路线的分工：数据并行（Accelerate/DDP）解决吞吐、device_map按组件搬卡、上下文并行切输入序列降激活显存、张量并行切权重矩阵降权重显存。文档对张量并行的适用判断是：当单个组件的权重放不下且互连带宽高（如 NVLink）时优先用它，因为它同时降低显存和延迟；它每个块边界都要通信，走 PCIe 时层间流量可能抵消收益，此时device_map是更好的选择。若权重和序列都过大，可用mesh参数让张量并行与上下文并行共享一个设备网格，但文档明确标注该组合尚未验证，属于实验性质。
【免费下载链接】diffusers🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
