尧图网站建设 
vllm-omni 扩散模型张量并行Tensor Parallel接入指南以 Z-Image 为参考的完整实现【免费下载链接】vllm-omniA framework for efficient model inference with omni-modality models项目地址: https://gitcode.com/GitHub_Trending/vl/vllm-omni本指南面向希望在 vllm-omni 中为扩散 TransformerDiT模型添加张量并行Tensor ParallelTP支持的研究与工程人员完整梳理从识别可切分线性层、替换为 vLLM 并行层、校验维度约束到端到端测试与问题排查的全过程。读完本文你将掌握以 Z-Image 为参考的标准化 TP 改造范式并能复用到 FLUX、Qwen-Image 等其他扩散模型上。一、背景为什么扩散 Transformer 需要 Tensor ParallelTensor ParallelTP是一种模型并行技术其核心思想是将模型权重按张量维度切分到多张 GPU 上每张 GPU 只持有模型参数的一部分也只计算每一层输出的一个分片。对于包含大规模 Attention 与 MLP 层的扩散 Transformer 而言沿模型维度dim切分可以带来两个直接收益显存可承载更大模型得以在单卡放不下的情况下分布式运行接近线性的加速切分后的每张卡计算量约降为原来的 1/N配合通信即可逼近线性加速比。vllm-omni 的张量并行实现建立在 vLLM 的 Parallel Layers 之上这些并行层位于vllm.model_executor.layers.linear其类型与职责如下层类型用途权重切分方式ColumnParallelLinearFFN 第一层、独立的 QKV 投影按列切分输出维度RowParallelLinearFFN 第二层、Attention 输出投影按行切分输入维度QKVParallelLinear多头 / 分组查询注意力GQA的 QKV 投影自动处理 head 复制按 head 切分ReplicatedLinear不应切分的层如时间步嵌入、最终输出投影不切分全卡复制说明本文所依据的设计文档为 docs/design/feature/tensor_parallel.md其中 Z-Image 作为参考实现本节关于 vLLM Parallel Layers 的对应关系来自 vLLM 上游的模型接入文档本文不再赘述外部链接。二、改造前的第一步识别需要切分的 Linear 层在动手改代码之前需要先盘点模型中所有nn.Linear并回答两个关键问题哪些层应该按列并行权重按列切分—— 典型答案是 FFN 的第一个投影、Attention 的 QKV 投影哪些层应该按行并行权重按行切分—— 典型答案是 FFN 的第二个投影、Attention 的输出投影。一个实用的判断基准一个列并行 → 激活函数/注意力 → 行并行的配对构成一个完整的切分闭环。列并行层输出的分片张量shape 为[batch, seq, hidden/N]经过无需跨卡通信的激活或注意力计算后由行并行层通过 All-Reduce 归约为完整维度从而保证下一层拿到的是与单卡完全一致的输入。此外模型中还有一些不能切分的层例如时间步嵌入 MLP输出参与 adaLN 调制对精度敏感patch embedding、caption embedding 等入口层最终输出投影层直接产生输出 latent。这些层在 Z-Image 参考实现中统一使用ReplicatedLinear且很多刻意保持全精度quant_configNone详见源码注释见 z_image_transformer.py 中 TimestepEmbedder 的说明。三、改造核心用并行层替换nn.Linear3.1 MLP 块Up-Down 模式扩散模型 FFN 的标准形态是升维再降维。参考实现见 FeedForward中Z-Image 使用MergedColumnParallelLinear将 w1/w3 两个门控投影合并配合SiluAndMul激活class FeedForward(nn.Module): def __init__(self, dim: int, hidden_dim: int): super().__init__() # 列并行权重按列切分输出为 [hidden_dim/N] self.w13 MergedColumnParallelLinear( dim, [hidden_dim] * 2, # w1 与 w3 合并各输出 hidden_dim biasFalse, return_biasFalse, ) self.act SiluAndMul() # 作用于切分后的张量无需通信 # 行并行输入来自列并行层的输出已切分 self.w2 RowParallelLinear( hidden_dim, dim, biasFalse, input_is_parallelTrue, # 关键输入已被 w13 切分 return_biasFalse, ) def forward(self, x): # x: [batch, seq, dim]每张卡上完整复制 # w13 输出切分为 [batch, seq, hidden_dim/N] hidden_states self.w13(x) hidden_states self.act(hidden_states) # w2 通过 All-Reduce 恢复完整 dim hidden_states self.w2(hidden_states) return hidden_states对照设计文档中的示意代码ColumnParallelLinear→RowParallelLinear实际源码采用了MergedColumnParallelLinear合并 w1/w3 的写法两者在 TP 语义上完全一致合并写法还减少了 Kernel 调用次数。注意 Z-Image 的 FFN 隐藏维度取int(dim / 3 * 8)在dim3840时为10240。3.2 Attention 块QKV-Out 模式Attention 的 QKV 投影使用QKVParallelLinear它会把总 head 数平均切分到各卡每卡持有num_heads/N个头并自动处理 GQA 下的 KV head 复制。参考实现见 ZImageAttentionfrom vllm_omni.diffusion.attention.layer import Attention class ZImageAttention(nn.Module): def __init__(self, dim: int, num_heads: int, num_kv_heads: int): super().__init__() self.head_dim dim // num_heads # 列并行QKV 权重按 head 切分每卡获得 num_heads/N 个头 self.to_qkv QKVParallelLinear( hidden_sizedim, head_sizeself.head_dim, total_num_headsnum_heads, # 模型总 head 数 total_num_kv_headsnum_kv_heads, # 模型总 KV head 数 biasFalse, return_biasFalse, ) # 行并行Attention 输出按行切分All-Reduce 恢复完整维度 self.to_out RowParallelLinear( dim, dim, biasFalse, input_is_parallelTrue, # 输入来自注意力输出按 head 切分 return_biasFalse, ) # 关键Attention 使用“本地 head 数”而非模型总 head 数 self.attn Attention( num_headsself.to_qkv.num_heads, # 本卡实际持有的 head 数 head_sizeself.head_dim, softmax_scale1.0 / (self.head_dim**0.5), causalFalse, num_kv_headsself.to_qkv.num_kv_heads, # 本卡持有的 KV head 数 ) def forward(self, x): qkv, _ self.to_qkv(x) # [batch, seq, (qkv)*head_dim/N] q_size self.to_qkv.num_heads * self.head_dim # 本地 q 尺寸 kv_size self.to_qkv.num_kv_heads * self.head_dim # 本地 kv 尺寸 q, k, v qkv.split([q_size, kv_size, kv_size], dim-1) # 每张卡独立计算注意力无需跨卡通信 out self.attn(q, k, v) out self.to_out(out) # All-Reduce 回完整 dim return out三个必须遵守的要点ColumnParallelLinear/QKVParallelLinear与RowParallelLinear成对出现这是 TP 的标准配对RowParallelLinear必须设置input_is_parallelTrue因为其输入来自列并行层的切分输出Attention 使用本地 head 数self.to_qkv.num_heads不能用模型总 head 数去切分 QKV否则会出现维度不匹配。3.3 权重加载与切分并行层替换完成后还需要让权重加载逻辑适配合并参数的形态。Z-Image 在load_weights中声明了stacked_params_mapping见 z_image_transformer.pystacked_params_mapping [ (.to_qkv., .to_q., q), # to_q/to_k/to_v 合并进 to_qkv (.to_qkv., .to_k., k), (.to_qkv., .to_v., v), (.w13, .w1, 0), # w1/w3 合并进 w13 (.w13, .w3, 1), ]加载权重时切分后的每个并行层参数通过其weight_loader将完整权重按 TP 维度切片到对应卡上。类属性packed_modules_mapping同时服务于量化 checkpoint 适配器与 LoRA 对融合投影的处理。这也解释了为什么只替换nn.Linear、不动权重加载会导致权重形状对不上——两者必须同步改造。四、第三步校验 TP 约束可整除性TP 正确运行的硬性前提是所有会被切分的维度必须能被tensor_parallel_size整除。设计文档给出的约束表如下维度原因错误示例num_headshead 数由 QKVParallelLinear 按列切分num_heads30, tp4❌30 % 4 ≠ 0num_kv_headsKV head 数由 QKVParallelLinear 切分num_kv_heads30, tp4❌30 % 4 ≠ 0Z-Image 在此基础上将校验落到了代码中validate_zimage_tp_constraints见 z_image_transformer.py会在模型初始化时对dim、n_heads、n_kv_heads、ffn_hidden_dimdim/3*8、最终输出维度final_out_dimspatch_size² × f_patch_size × in_channels逐一检查可整除性不满足时抛出带支持的 TP 候选值提示的ValueError。TP size 从 forward context 中读取_get_tensor_parallel_size_from_context这意味着同一个模型代码在单卡推理时tp1与多卡 TP 推理时都能复用。对应的单元测试test_zimage_tp_constraints.py验证了典型场景dim3840, n_heads30, n_kv_heads30, tp2→ 通过ffn_hidden_dim10240final_out_dims[64]支持的 TP 候选为[1, 2]tp4→ 因n_heads % tp ! 0抛错30 无法被 4 整除tp3→ 因ffn_hidden_dim % tp ! 0抛错10240 无法被 3 整除。在规划 TP 规模时应先用这些约束倒推候选值对 Z-Image 而言30 个 head 与 10240 的 FFN 隐藏维度共同决定了实际可用 TP 规模只有 1 和 2。五、端到端验证如何测试 TP 是否生效5.1 Python API 方式按设计文档TP 通过DiffusionParallelConfig(tensor_parallel_sizeN)开启from vllm_omni import Omni from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams parallel_config DiffusionParallelConfig(tensor_parallel_size2) omni Omni(modelyour-model-name, parallel_configparallel_config) output omni.generate( a cup of coffee on the table, OmniDiffusionSamplingParams(num_inference_steps50), )DiffusionParallelConfig见 vllm_omni/diffusion/data.py是扩散模型分布式执行的统一配置除tensor_parallel_size默认 1外还包含pipeline_parallel_size、data_parallel_size、sequence_parallel_sizeUlysses/Ring/AllGather-KV、cfg_parallel_size、vae_patch_parallel_size、text_encoder_tp_size等可按需与 TP 组合使用。5.2 命令行方式官方离线推理示例examples/offline_inference/text_to_image提供了现成的--tensor-parallel-size参数见 text_to_image.pycd examples/offline_inference/text_to_image python text_to_image.py \ --model Your-org/your-model \ --prompt a cup of coffee on the table \ --negative-prompt ugly, unclear \ --cfg-scale 4.0 \ --num-inference-steps 50 \ --output tp_enabled.png \ --tensor-parallel-size 2启动日志中会打印并行配置例如Parallel configuration: tensor_parallel_size2, ...见 text_to_image.py可用于确认 TP 确实被传入。5.3 仓库自带的自动化验证仓库在 tests/diffusion/distributed/test_tensor_parallel.py 中提供了完整的 TP 端到端E2E回归测试以Tongyi-MAI/Z-Image-Turbo为基准模型对比TP1与TP2正确性同一 seed 下分别用 TP1 与 TP2 生成 512×512 图像计算两图的 mean 绝对差与 P99 绝对差断言mean_abs_diff 3e-2且p99_abs_diff 2.5e-1从像素级验证切分不改变生成语义性能使用DeviceMemoryMonitor以 20ms 间隔采样显存断言 TP2 的每请求中位耗时低于 TP1ROCm 平台除外显存断言 TP2 的峰值显存低于 TP1直接印证权重切分降低单卡显存的目标。该测试标记为hardware_test需要至少 2 张 CUDA 卡或 2 张 ROCm 卡NPU 平台当前会跳过测试文件中注明 TP parity 测试目前仅支持 CUDA 与 ROCm。5.4 人工验证清单设计文档给出的手工验证步骤查看日志中的e2e_time_ms确认是否有加速对比 TP 关闭与开启时生成图像的质量确认显存占用按比例下降将对比结果记录在 PR 中含正确性、速度、显存三项数据。六、常见问题排查问题 1TP 未生效仍跑在单卡上无显存节省与加速症状模型在单张 GPU 上运行日志显示tensor_parallel_size1。根因与解法仍然使用nn.Linear并行层未被替换。解法替换为并行等价层# ❌ 错误普通线性层无法切分 self.proj nn.Linear(dim, dim) # ✅ 正确行并行层 self.proj RowParallelLinear(dim, dim, input_is_parallelTrue)未正确传入并行配置确认DiffusionParallelConfig(tensor_parallel_sizeN)已传入且启动日志中的并行配置显示tensor_parallel_sizeN而非默认的deploy/default。问题 2前向过程报RuntimeError: shape mismatch症状forward 中张量形状对不上。根因与解法缺少input_is_parallelTrueRowParallelLinear期望接收已切分的输入却收到了完整张量。解法当输入来自列并行层时务必显式开启# ✅ 正确配对 self.w1 ColumnParallelLinear(dim, hidden_dim, return_biasFalse) self.w2 RowParallelLinear( hidden_dim, dim, input_is_parallelTrue, # 输入已被 w1 切分 return_biasFalse, )QKV 切分尺寸用了总 head 数切分后的 QKV 必须按本地localhead 数计算# ❌ 错误用模型总 head 数 q_size self.total_num_heads * self.head_dim # ✅ 正确用并行层暴露的本地 head 数 q_size self.to_qkv.num_heads * self.head_dim问题 3初始化报维度不可整除症状加载模型时抛出ValueError: ... requires n_heads % tensor_parallel_size 0之类错误并附带Supported tp candidates列表。根因与解法所选 TP 规模与模型结构不匹配。解法按错误提示中给出的候选值调整tensor_parallel_size或修改模型结构使相关维度可被目标 TP 规模整除后者一般不建议属于架构级改动。七、仓库内参考实现一览以下为改造时的对照范本与测试锚点模型 / 组件路径模式备注Z-Imagez_image_transformer.py标准 TP完整实现含约束校验validate_zimage_tp_constraints与序列并行_sp_planFLUXflux_transformer.py双流Dual-stream图像流与文本流分别切分Qwen-Imageqwen_image_transformer.py标准 TP RoPE带旋转位置编码的切分处理TP E2E 测试test_tensor_parallel.py端到端TP1 vs TP2 的正确性、性能、显存三方对比约束单元测试test_zimage_tp_constraints.py单元测试验证可整除性校验逻辑Z-Image 除 TP 外还叠加了序列并行SP对应 diffusers 中的上下文并行 CP_sp_plan指定了在unified_prepare模块边界对 unified 序列、RoPE cos/sin、注意力掩码按序列维切分并在all_final_layer处聚拢输出见 z_image_transformer.py。实际落地时TP 解决权重放不下SP 解决序列太长两者可组合使用。八、总结为扩散 Transformer 接入 Tensor Parallel 支持遵循四步流程即可✅识别线性层——判断哪些层该切分列并行/行并行哪些层必须复制ReplicatedLinear✅替换为并行层——Attention 用QKVParallelLinearRowParallelLinearFFN 用ColumnParallelLinear/MergedColumnParallelLinearRowParallelLinear并同步改造权重加载的合并映射✅校验 TP 约束——确保num_heads、num_kv_heads、dim、FFN 隐藏维度、最终输出维度均可被tensor_parallel_size整除✅测试——以tensor_parallel_sizeN运行核对e2e_time_ms、峰值显存与生成图像质量并以 TP1 为基线做正确性对比。以 Z-Image 为参考的这一套识别 → 替换 → 校验 → 测试方法论可以直接迁移到仓库中的其他扩散模型FLUX、Qwen-Image 等是 vllm-omni 中为任意 DiT 模型落地多卡推理的标准路径。【免费下载链接】vllm-omniA framework for efficient model inference with omni-modality models项目地址: https://gitcode.com/GitHub_Trending/vl/vllm-omni创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 vllm-omni 扩散模型张量并行（Tensor Parallel）接入指南：以 Z-Image 为参考的完整实现
vllm-omni 扩散模型张量并行Tensor Parallel接入指南以 Z-Image 为参考的完整实现【免费下载链接】vllm-omniA framework for efficient model inference with omni-modality models项目地址: https://gitcode.com/GitHub_Trending/vl/vllm-omni本指南面向希望在 vllm-omni 中为扩散 TransformerDiT模型添加张量并行Tensor ParallelTP支持的研究与工程人员完整梳理从识别可切分线性层、替换为 vLLM 并行层、校验维度约束到端到端测试与问题排查的全过程。读完本文你将掌握以 Z-Image 为参考的标准化 TP 改造范式并能复用到 FLUX、Qwen-Image 等其他扩散模型上。一、背景为什么扩散 Transformer 需要 Tensor ParallelTensor ParallelTP是一种模型并行技术其核心思想是将模型权重按张量维度切分到多张 GPU 上每张 GPU 只持有模型参数的一部分也只计算每一层输出的一个分片。对于包含大规模 Attention 与 MLP 层的扩散 Transformer 而言沿模型维度dim切分可以带来两个直接收益显存可承载更大模型得以在单卡放不下的情况下分布式运行接近线性的加速切分后的每张卡计算量约降为原来的 1/N配合通信即可逼近线性加速比。vllm-omni 的张量并行实现建立在 vLLM 的 Parallel Layers 之上这些并行层位于vllm.model_executor.layers.linear其类型与职责如下层类型用途权重切分方式ColumnParallelLinearFFN 第一层、独立的 QKV 投影按列切分输出维度RowParallelLinearFFN 第二层、Attention 输出投影按行切分输入维度QKVParallelLinear多头 / 分组查询注意力GQA的 QKV 投影自动处理 head 复制按 head 切分ReplicatedLinear不应切分的层如时间步嵌入、最终输出投影不切分全卡复制说明本文所依据的设计文档为 docs/design/feature/tensor_parallel.md其中 Z-Image 作为参考实现本节关于 vLLM Parallel Layers 的对应关系来自 vLLM 上游的模型接入文档本文不再赘述外部链接。二、改造前的第一步识别需要切分的 Linear 层在动手改代码之前需要先盘点模型中所有nn.Linear并回答两个关键问题哪些层应该按列并行权重按列切分—— 典型答案是 FFN 的第一个投影、Attention 的 QKV 投影哪些层应该按行并行权重按行切分—— 典型答案是 FFN 的第二个投影、Attention 的输出投影。一个实用的判断基准一个列并行 → 激活函数/注意力 → 行并行的配对构成一个完整的切分闭环。列并行层输出的分片张量shape 为[batch, seq, hidden/N]经过无需跨卡通信的激活或注意力计算后由行并行层通过 All-Reduce 归约为完整维度从而保证下一层拿到的是与单卡完全一致的输入。此外模型中还有一些不能切分的层例如时间步嵌入 MLP输出参与 adaLN 调制对精度敏感patch embedding、caption embedding 等入口层最终输出投影层直接产生输出 latent。这些层在 Z-Image 参考实现中统一使用ReplicatedLinear且很多刻意保持全精度quant_configNone详见源码注释见 z_image_transformer.py 中 TimestepEmbedder 的说明。三、改造核心用并行层替换nn.Linear3.1 MLP 块Up-Down 模式扩散模型 FFN 的标准形态是升维再降维。参考实现见 FeedForward中Z-Image 使用MergedColumnParallelLinear将 w1/w3 两个门控投影合并配合SiluAndMul激活class FeedForward(nn.Module): def __init__(self, dim: int, hidden_dim: int): super().__init__() # 列并行权重按列切分输出为 [hidden_dim/N] self.w13 MergedColumnParallelLinear( dim, [hidden_dim] * 2, # w1 与 w3 合并各输出 hidden_dim biasFalse, return_biasFalse, ) self.act SiluAndMul() # 作用于切分后的张量无需通信 # 行并行输入来自列并行层的输出已切分 self.w2 RowParallelLinear( hidden_dim, dim, biasFalse, input_is_parallelTrue, # 关键输入已被 w13 切分 return_biasFalse, ) def forward(self, x): # x: [batch, seq, dim]每张卡上完整复制 # w13 输出切分为 [batch, seq, hidden_dim/N] hidden_states self.w13(x) hidden_states self.act(hidden_states) # w2 通过 All-Reduce 恢复完整 dim hidden_states self.w2(hidden_states) return hidden_states对照设计文档中的示意代码ColumnParallelLinear→RowParallelLinear实际源码采用了MergedColumnParallelLinear合并 w1/w3 的写法两者在 TP 语义上完全一致合并写法还减少了 Kernel 调用次数。注意 Z-Image 的 FFN 隐藏维度取int(dim / 3 * 8)在dim3840时为10240。3.2 Attention 块QKV-Out 模式Attention 的 QKV 投影使用QKVParallelLinear它会把总 head 数平均切分到各卡每卡持有num_heads/N个头并自动处理 GQA 下的 KV head 复制。参考实现见 ZImageAttentionfrom vllm_omni.diffusion.attention.layer import Attention class ZImageAttention(nn.Module): def __init__(self, dim: int, num_heads: int, num_kv_heads: int): super().__init__() self.head_dim dim // num_heads # 列并行QKV 权重按 head 切分每卡获得 num_heads/N 个头 self.to_qkv QKVParallelLinear( hidden_sizedim, head_sizeself.head_dim, total_num_headsnum_heads, # 模型总 head 数 total_num_kv_headsnum_kv_heads, # 模型总 KV head 数 biasFalse, return_biasFalse, ) # 行并行Attention 输出按行切分All-Reduce 恢复完整维度 self.to_out RowParallelLinear( dim, dim, biasFalse, input_is_parallelTrue, # 输入来自注意力输出按 head 切分 return_biasFalse, ) # 关键Attention 使用“本地 head 数”而非模型总 head 数 self.attn Attention( num_headsself.to_qkv.num_heads, # 本卡实际持有的 head 数 head_sizeself.head_dim, softmax_scale1.0 / (self.head_dim**0.5), causalFalse, num_kv_headsself.to_qkv.num_kv_heads, # 本卡持有的 KV head 数 ) def forward(self, x): qkv, _ self.to_qkv(x) # [batch, seq, (qkv)*head_dim/N] q_size self.to_qkv.num_heads * self.head_dim # 本地 q 尺寸 kv_size self.to_qkv.num_kv_heads * self.head_dim # 本地 kv 尺寸 q, k, v qkv.split([q_size, kv_size, kv_size], dim-1) # 每张卡独立计算注意力无需跨卡通信 out self.attn(q, k, v) out self.to_out(out) # All-Reduce 回完整 dim return out三个必须遵守的要点ColumnParallelLinear/QKVParallelLinear与RowParallelLinear成对出现这是 TP 的标准配对RowParallelLinear必须设置input_is_parallelTrue因为其输入来自列并行层的切分输出Attention 使用本地 head 数self.to_qkv.num_heads不能用模型总 head 数去切分 QKV否则会出现维度不匹配。3.3 权重加载与切分并行层替换完成后还需要让权重加载逻辑适配合并参数的形态。Z-Image 在load_weights中声明了stacked_params_mapping见 z_image_transformer.pystacked_params_mapping [ (.to_qkv., .to_q., q), # to_q/to_k/to_v 合并进 to_qkv (.to_qkv., .to_k., k), (.to_qkv., .to_v., v), (.w13, .w1, 0), # w1/w3 合并进 w13 (.w13, .w3, 1), ]加载权重时切分后的每个并行层参数通过其weight_loader将完整权重按 TP 维度切片到对应卡上。类属性packed_modules_mapping同时服务于量化 checkpoint 适配器与 LoRA 对融合投影的处理。这也解释了为什么只替换nn.Linear、不动权重加载会导致权重形状对不上——两者必须同步改造。四、第三步校验 TP 约束可整除性TP 正确运行的硬性前提是所有会被切分的维度必须能被tensor_parallel_size整除。设计文档给出的约束表如下维度原因错误示例num_headshead 数由 QKVParallelLinear 按列切分num_heads30, tp4❌30 % 4 ≠ 0num_kv_headsKV head 数由 QKVParallelLinear 切分num_kv_heads30, tp4❌30 % 4 ≠ 0Z-Image 在此基础上将校验落到了代码中validate_zimage_tp_constraints见 z_image_transformer.py会在模型初始化时对dim、n_heads、n_kv_heads、ffn_hidden_dimdim/3*8、最终输出维度final_out_dimspatch_size² × f_patch_size × in_channels逐一检查可整除性不满足时抛出带支持的 TP 候选值提示的ValueError。TP size 从 forward context 中读取_get_tensor_parallel_size_from_context这意味着同一个模型代码在单卡推理时tp1与多卡 TP 推理时都能复用。对应的单元测试test_zimage_tp_constraints.py验证了典型场景dim3840, n_heads30, n_kv_heads30, tp2→ 通过ffn_hidden_dim10240final_out_dims[64]支持的 TP 候选为[1, 2]tp4→ 因n_heads % tp ! 0抛错30 无法被 4 整除tp3→ 因ffn_hidden_dim % tp ! 0抛错10240 无法被 3 整除。在规划 TP 规模时应先用这些约束倒推候选值对 Z-Image 而言30 个 head 与 10240 的 FFN 隐藏维度共同决定了实际可用 TP 规模只有 1 和 2。五、端到端验证如何测试 TP 是否生效5.1 Python API 方式按设计文档TP 通过DiffusionParallelConfig(tensor_parallel_sizeN)开启from vllm_omni import Omni from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams parallel_config DiffusionParallelConfig(tensor_parallel_size2) omni Omni(modelyour-model-name, parallel_configparallel_config) output omni.generate( a cup of coffee on the table, OmniDiffusionSamplingParams(num_inference_steps50), )DiffusionParallelConfig见 vllm_omni/diffusion/data.py是扩散模型分布式执行的统一配置除tensor_parallel_size默认 1外还包含pipeline_parallel_size、data_parallel_size、sequence_parallel_sizeUlysses/Ring/AllGather-KV、cfg_parallel_size、vae_patch_parallel_size、text_encoder_tp_size等可按需与 TP 组合使用。5.2 命令行方式官方离线推理示例examples/offline_inference/text_to_image提供了现成的--tensor-parallel-size参数见 text_to_image.pycd examples/offline_inference/text_to_image python text_to_image.py \ --model Your-org/your-model \ --prompt a cup of coffee on the table \ --negative-prompt ugly, unclear \ --cfg-scale 4.0 \ --num-inference-steps 50 \ --output tp_enabled.png \ --tensor-parallel-size 2启动日志中会打印并行配置例如Parallel configuration: tensor_parallel_size2, ...见 text_to_image.py可用于确认 TP 确实被传入。5.3 仓库自带的自动化验证仓库在 tests/diffusion/distributed/test_tensor_parallel.py 中提供了完整的 TP 端到端E2E回归测试以Tongyi-MAI/Z-Image-Turbo为基准模型对比TP1与TP2正确性同一 seed 下分别用 TP1 与 TP2 生成 512×512 图像计算两图的 mean 绝对差与 P99 绝对差断言mean_abs_diff 3e-2且p99_abs_diff 2.5e-1从像素级验证切分不改变生成语义性能使用DeviceMemoryMonitor以 20ms 间隔采样显存断言 TP2 的每请求中位耗时低于 TP1ROCm 平台除外显存断言 TP2 的峰值显存低于 TP1直接印证权重切分降低单卡显存的目标。该测试标记为hardware_test需要至少 2 张 CUDA 卡或 2 张 ROCm 卡NPU 平台当前会跳过测试文件中注明 TP parity 测试目前仅支持 CUDA 与 ROCm。5.4 人工验证清单设计文档给出的手工验证步骤查看日志中的e2e_time_ms确认是否有加速对比 TP 关闭与开启时生成图像的质量确认显存占用按比例下降将对比结果记录在 PR 中含正确性、速度、显存三项数据。六、常见问题排查问题 1TP 未生效仍跑在单卡上无显存节省与加速症状模型在单张 GPU 上运行日志显示tensor_parallel_size1。根因与解法仍然使用nn.Linear并行层未被替换。解法替换为并行等价层# ❌ 错误普通线性层无法切分 self.proj nn.Linear(dim, dim) # ✅ 正确行并行层 self.proj RowParallelLinear(dim, dim, input_is_parallelTrue)未正确传入并行配置确认DiffusionParallelConfig(tensor_parallel_sizeN)已传入且启动日志中的并行配置显示tensor_parallel_sizeN而非默认的deploy/default。问题 2前向过程报RuntimeError: shape mismatch症状forward 中张量形状对不上。根因与解法缺少input_is_parallelTrueRowParallelLinear期望接收已切分的输入却收到了完整张量。解法当输入来自列并行层时务必显式开启# ✅ 正确配对 self.w1 ColumnParallelLinear(dim, hidden_dim, return_biasFalse) self.w2 RowParallelLinear( hidden_dim, dim, input_is_parallelTrue, # 输入已被 w1 切分 return_biasFalse, )QKV 切分尺寸用了总 head 数切分后的 QKV 必须按本地localhead 数计算# ❌ 错误用模型总 head 数 q_size self.total_num_heads * self.head_dim # ✅ 正确用并行层暴露的本地 head 数 q_size self.to_qkv.num_heads * self.head_dim问题 3初始化报维度不可整除症状加载模型时抛出ValueError: ... requires n_heads % tensor_parallel_size 0之类错误并附带Supported tp candidates列表。根因与解法所选 TP 规模与模型结构不匹配。解法按错误提示中给出的候选值调整tensor_parallel_size或修改模型结构使相关维度可被目标 TP 规模整除后者一般不建议属于架构级改动。七、仓库内参考实现一览以下为改造时的对照范本与测试锚点模型 / 组件路径模式备注Z-Imagez_image_transformer.py标准 TP完整实现含约束校验validate_zimage_tp_constraints与序列并行_sp_planFLUXflux_transformer.py双流Dual-stream图像流与文本流分别切分Qwen-Imageqwen_image_transformer.py标准 TP RoPE带旋转位置编码的切分处理TP E2E 测试test_tensor_parallel.py端到端TP1 vs TP2 的正确性、性能、显存三方对比约束单元测试test_zimage_tp_constraints.py单元测试验证可整除性校验逻辑Z-Image 除 TP 外还叠加了序列并行SP对应 diffusers 中的上下文并行 CP_sp_plan指定了在unified_prepare模块边界对 unified 序列、RoPE cos/sin、注意力掩码按序列维切分并在all_final_layer处聚拢输出见 z_image_transformer.py。实际落地时TP 解决权重放不下SP 解决序列太长两者可组合使用。八、总结为扩散 Transformer 接入 Tensor Parallel 支持遵循四步流程即可✅识别线性层——判断哪些层该切分列并行/行并行哪些层必须复制ReplicatedLinear✅替换为并行层——Attention 用QKVParallelLinearRowParallelLinearFFN 用ColumnParallelLinear/MergedColumnParallelLinearRowParallelLinear并同步改造权重加载的合并映射✅校验 TP 约束——确保num_heads、num_kv_heads、dim、FFN 隐藏维度、最终输出维度均可被tensor_parallel_size整除✅测试——以tensor_parallel_sizeN运行核对e2e_time_ms、峰值显存与生成图像质量并以 TP1 为基线做正确性对比。以 Z-Image 为参考的这一套识别 → 替换 → 校验 → 测试方法论可以直接迁移到仓库中的其他扩散模型FLUX、Qwen-Image 等是 vllm-omni 中为任意 DiT 模型落地多卡推理的标准路径。【免费下载链接】vllm-omniA framework for efficient model inference with omni-modality models项目地址: https://gitcode.com/GitHub_Trending/vl/vllm-omni创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考