并行加速指南¶
本指南介绍了如何在 vLLM-Omni 中使用并行方法来加速扩散模型推理，以及减少每个设备的内存需求。
概述¶
vLLM-Omni 目前支持以下并行方法：
- DeepSpeed Ulysses 序列并行 (Ulysses-SP) (论文): Ulysses-SP 沿序列维度分割输入，并使用 all-to-all 通信，让每个设备只计算一部分注意力头。
下表显示了并行方法目前支持的模型
| 模型 | 模型标识符 | Ulysses-SP |
|---|---|---|
| Qwen-Image | Qwen/Qwen-Image | ✅ |
| Z-Image | Tongyi-MAI/Z-Image-Turbo | ❌ |
| Qwen-Image-Edit | Qwen/Qwen-Image-Edit | ✅ |
序列并行¶
Ulysses-SP¶
快速入门¶
下面展示了一个使用 Ulysses-SP 的示例
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
ulysses_degree = 2
omni = Omni(
model="Qwen/Qwen-Image",
parallel_config=DiffusionParallelConfig(ulysses_degree=2)
)
outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
请参阅 examples/offline_inference/text_to_image/text_to_image.py 获取完整的可运行示例。
在线服务¶
您可以通过 --usp 在线服务中启用 Ulysses-SP 来进行扩散模型推理。
基准测试¶
基准测试免责声明
这些基准测试仅供**一般参考**。所示配置使用了默认或常用的参数设置，并未经过详尽优化以达到最佳性能。实际性能可能因以下因素而异：
- 特定的模型和用例
- 硬件配置
- 仔细的参数调优
- 不同的推理设置（例如，步数、图像分辨率）
为了衡量并行方法的效果，我们使用 **Qwen/Qwen-Image** 模型生成图像（**2048x2048** 作为长序列输入），进行 50 次推理步。硬件设备为 NVIDIA H800 GPU。sdpa 是注意力后端。
| 配置 | Ulysses 度 | 生成时间 | 加速比 |
|---|---|---|---|
| 基线 (diffusers) | - | 112.5s | 1.0x |
| Ulysses-SP | 2 | 65.2s | 1.73x |
| Ulysses-SP | 4 | 39.6s | 2.84x |
| Ulysses-SP | 8 | 30.8s | 3.65x |
如何并行化新模型¶
如果某个扩散模型已部署在 vLLM-Omni 中并支持单卡推理，您可以参考以下说明，了解如何使用 Ulysses-SP 对该模型进行并行化。
首先，请编辑 TransformerModel 的 forward 函数，位于 xxx_model_transformer.py 文件中，使输入（图像隐藏状态、位置嵌入等）成为沿序列维度分割的块。以 qwen_image_transformer.py 为例：
class QwenImageTransformer2DModel(nn.Module):
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
...
):
+ if self.parallel_config.sequence_parallel_size > 1:
+ hidden_states = torch.chunk(hidden_states, get_sequence_parallel_world_size(), dim=-2)[
+ get_sequence_parallel_rank()
+ ]
hidden_states = self.img_in(hidden_states)
...
image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
+ def get_rotary_emb_chunk(freqs):
+ freqs = torch.chunk(freqs, get_sequence_parallel_world_size(), dim=0)[get_sequence_parallel_rank()]
+ return freqs
+ if self.parallel_config.sequence_parallel_size > 1:
+ img_freqs, txt_freqs = image_rotary_emb
+ img_freqs = get_rotary_emb_chunk(img_freqs)
+ image_rotary_emb = (img_freqs, txt_freqs)
接下来，在 forward 函数的末尾，请调用 get_sp_group().all_gather 来聚合跨设备的块状输出，并在序列维度上将它们连接起来。
class QwenImageTransformer2DModel(nn.Module):
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
...
):
# Use only the image part (hidden_states) from the dual-stream blocks
hidden_states = self.norm_out(hidden_states, temb)
output = self.proj_out(hidden_states)
+ if self.parallel_config.sequence_parallel_size > 1:
+ output = get_sp_group().all_gather(output, dim=-2)
return Transformer2DModelOutput(sample=output)
最后，您可以设置并行配置并将其传递给 Omni，然后开始并行推理：