anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Feature] add Tensor Parallelism to Omnigen2 (#2065)

Browse Source 
Signed-off-by: zhou zhuoxin <zhouzhuoxin1508@outlook.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
This commit is contained in:
zhou zhuoxinand汪志鹏 authored and GitHub committed 2026-03-23 11:46:58 +08:00
1 parent 0833470d41
commit b147553d67
2 files changed
+29 -83
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/user_guide/diffusion/parallelism_acceleration.md
+1
View File
@@ -40,6 +40,7 @@ The following table shows which models are currently supported by parallelism me
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | ✅ |
| **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ |
| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ✅ | ❌ | ❌ | N/A | ❌ |
| **OmniGen2** | `OmniGen2/OmniGen2` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | ❌ |
!!! note "TP Limitations for Diffusion Models"
We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py
+28 -83
View File
@@ -5,14 +5,16 @@ from types import SimpleNamespace
import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.models.activations import get_activation
from diffusers.models.embeddings import Timesteps, get_1d_rotary_pos_embed
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from einops import rearrange, repeat
from torch.nn import RMSNorm
from vllm.model_executor.layers.activation import get_act_and_mul_fn
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
QKVParallelLinear,
RowParallelLinear,
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -21,10 +23,6 @@ from vllm_omni.diffusion.attention.layer import Attention
logger = logging.getLogger(__name__)
def swiglu(x, y):
return F.silu(x.float(), inplace=False).to(x.dtype) * y
class OmniGen2Attention(nn.Module):
def __init__(
self,
@@ -53,7 +51,6 @@ class OmniGen2Attention(nn.Module):
self.norm_k = RMSNorm(self.head_dim, eps=eps)
self.to_out = nn.ModuleList([nn.Linear(dim, dim, bias=False)])
self.attn = Attention(
num_heads=num_heads,
head_size=self.head_dim,
@@ -78,22 +75,19 @@ class OmniGen2Attention(nn.Module):
Returns:
torch.Tensor: Processed hidden states after attention computation
"""
batch_size, sequence_length, _ = hidden_states.shape
batch_size = hidden_states.shape[0]
# Get Query-Key-Value Pair
qkv, _ = self.to_qkv(hidden_states)
q_dim = self.num_heads * self.head_dim
kv_dim = self.num_kv_heads * self.head_dim
query = qkv[..., :q_dim]
key = qkv[..., q_dim : q_dim + kv_dim]
value = qkv[..., q_dim + kv_dim : q_dim + 2 * kv_dim]
q_size = self.num_heads * self.head_dim
kv_size = self.num_kv_heads * self.head_dim
query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
# Reshape tensors for attention computation
query = query.view(batch_size, sequence_length, self.num_heads, self.head_dim)
key = key.view(batch_size, sequence_length, self.num_kv_heads, self.head_dim)
value = value.view(batch_size, sequence_length, self.num_kv_heads, self.head_dim)
query = query.unflatten(-1, (self.num_heads, -1))
key = key.unflatten(-1, (self.num_kv_heads, -1))
value = value.unflatten(-1, (self.num_kv_heads, -1))
# Apply Query-Key normalization
if self.norm_q is not None:
@@ -164,14 +158,6 @@ class TimestepEmbedding(nn.Module):
else:
self.post_act = get_activation(post_act_fn)
self.initialize_weights()
def initialize_weights(self):
nn.init.normal_(self.linear_1.weight, std=0.02)
nn.init.zeros_(self.linear_1.bias)
nn.init.normal_(self.linear_2.weight, std=0.02)
nn.init.zeros_(self.linear_2.bias)
def forward(self, sample, condition=None):
if condition is not None:
sample = sample + self.cond_proj(condition)
@@ -293,7 +279,7 @@ class LuminaLayerNormContinuous(nn.Module):
if norm_type == "layer_norm":
self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
elif norm_type == "rms_norm":
self.norm = RMSNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine)
self.norm = RMSNorm(embedding_dim, eps=eps)
else:
raise ValueError(f"unknown norm_type {norm_type}")
@@ -340,32 +326,31 @@ class LuminaFeedForward(nn.Module):
ffn_dim_multiplier: float | None = None,
):
super().__init__()
self.swiglu = swiglu
# custom hidden_size factor multiplier
if ffn_dim_multiplier is not None:
inner_dim = int(ffn_dim_multiplier * inner_dim)
inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of)
self.linear_1 = nn.Linear(
self.gate_up_proj = MergedColumnParallelLinear(
dim,
inner_dim,
[inner_dim, inner_dim],
bias=False,
return_bias=False,
)
self.linear_2 = nn.Linear(
self.act_fn = get_act_and_mul_fn("silu")
self.down_proj = RowParallelLinear(
inner_dim,
dim,
bias=False,
)
self.linear_3 = nn.Linear(
dim,
inner_dim,
bias=False,
input_is_parallel=True,
return_bias=False,
)
def forward(self, x):
h1, h2 = self.linear_1(x), self.linear_3(x)
return self.linear_2(self.swiglu(h1, h2))
x = self.gate_up_proj(x)
x = self.act_fn(x)
return self.down_proj(x)
class Lumina2CombinedTimestepCaptionEmbedding(nn.Module):
@@ -395,12 +380,6 @@ class Lumina2CombinedTimestepCaptionEmbedding(nn.Module):
nn.Linear(text_feat_dim, hidden_size, bias=True),
)
self._initialize_weights()
def _initialize_weights(self):
nn.init.trunc_normal_(self.caption_embedder[1].weight, std=0.02)
nn.init.zeros_(self.caption_embedder[1].bias)
def forward(
self,
timestep: torch.Tensor,
@@ -642,25 +621,6 @@ class OmniGen2TransformerBlock(nn.Module):
self.norm2 = RMSNorm(dim, eps=norm_eps)
self.ffn_norm2 = RMSNorm(dim, eps=norm_eps)
def initialize_weights(self) -> None:
"""
Initialize the weights of the transformer block.
Uses Xavier uniform initialization for linear layers and zero initialization for biases.
"""
nn.init.xavier_uniform_(self.attn.to_q.weight)
nn.init.xavier_uniform_(self.attn.to_k.weight)
nn.init.xavier_uniform_(self.attn.to_v.weight)
nn.init.xavier_uniform_(self.attn.to_out[0].weight)
nn.init.xavier_uniform_(self.feed_forward.linear_1.weight)
nn.init.xavier_uniform_(self.feed_forward.linear_2.weight)
nn.init.xavier_uniform_(self.feed_forward.linear_3.weight)
if self.modulation:
nn.init.zeros_(self.norm1.linear.weight)
nn.init.zeros_(self.norm1.linear.bias)
def forward(
self,
hidden_states: torch.Tensor,
@@ -873,27 +833,6 @@ class OmniGen2Transformer2DModel(nn.Module):
# Add learnable embeddings to distinguish different images
self.image_index_embedding = nn.Parameter(torch.randn(5, hidden_size)) # support max 5 ref images
self.initialize_weights()
def initialize_weights(self) -> None:
"""
Initialize the weights of the model.
Uses Xavier uniform initialization for linear layers.
"""
nn.init.xavier_uniform_(self.x_embedder.weight)
nn.init.constant_(self.x_embedder.bias, 0.0)
nn.init.xavier_uniform_(self.ref_image_patch_embedder.weight)
nn.init.constant_(self.ref_image_patch_embedder.bias, 0.0)
nn.init.zeros_(self.norm_out.linear_1.weight)
nn.init.zeros_(self.norm_out.linear_1.bias)
nn.init.zeros_(self.norm_out.linear_2.weight)
nn.init.zeros_(self.norm_out.linear_2.bias)
nn.init.normal_(self.image_index_embedding, std=0.02)
def img_patch_embed_and_refine(
self,
hidden_states,
@@ -1182,6 +1121,9 @@ class OmniGen2Transformer2DModel(nn.Module):
(".to_qkv", ".to_q", "q"),
(".to_qkv", ".to_k", "k"),
(".to_qkv", ".to_v", "v"),
# feed-forward
(".feed_forward.gate_up_proj", ".feed_forward.linear_1", 0),
(".feed_forward.gate_up_proj", ".feed_forward.linear_3", 1),
]
params_dict = dict(self.named_parameters())
@@ -1197,6 +1139,9 @@ class OmniGen2Transformer2DModel(nn.Module):
weight_loader(param, loaded_weight, shard_id)
break
else:
# feed_forward.linear_2 -> feed_forward.down_proj rename
if ".feed_forward.linear_2." in name:
name = name.replace(".feed_forward.linear_2.", ".feed_forward.down_proj.")
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)