chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Diffusion] Fix Ulysses/Ring process group construction under TP to enable correct Wan2.2 tensor parallelism (#16532)

Browse Source 
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Xiaoyu Zhang
2026-01-07 20:52:36 +08:00
committed by GitHub
parent 62d0280f62
commit 32a6540afc
8 changed files with 144 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/sglang/multimodal_gen/runtime/distributed/parallel_state.py
View File

@@ -406,9 +406,11 @@ def initialize_model_parallel(
try:
from .yunchang import PROCESS_GROUP as _YC_PROCESS_GROUP
from .yunchang import set_seq_parallel_pg as _set_seq_parallel_pg
from .yunchang import (
set_seq_parallel_pg_by_sp_groups as _set_seq_parallel_pg_by_sp_groups,
)
except ImportError:
_set_seq_parallel_pg = None
_set_seq_parallel_pg_by_sp_groups = None
class _DummyProcessGroup:
ULYSSES_PG = torch.distributed.group.WORLD
@@ -416,11 +418,15 @@ def initialize_model_parallel(
PROCESS_GROUP = _DummyProcessGroup()
else:
_set_seq_parallel_pg(
# Build yunchang SP sub-groups based on the true SP groups. This is
# critical when TP>1, because SP groups may be strided in global ranks
# (e.g., tp-sp order).
sp_groups = rank_generator.get_ranks("sp")
_set_seq_parallel_pg_by_sp_groups(
sp_ulysses_degree=ulysses_degree,
sp_ring_degree=ring_degree,
rank=get_world_group().rank_in_group,
world_size=dit_parallel_size,
rank=get_world_group().rank,
sp_groups=sp_groups,
)
PROCESS_GROUP = _YC_PROCESS_GROUP

69
python/sglang/multimodal_gen/runtime/distributed/yunchang.py
View File

@@ -82,3 +82,72 @@ def set_seq_parallel_pg(
PROCESS_GROUP.ULYSSES_PG = ulyssess_pg
PROCESS_GROUP.RING_PG = ring_pg
def set_seq_parallel_pg_by_sp_groups(
sp_ulysses_degree,
sp_ring_degree,
rank: int,
sp_groups: list[list[int]],
use_ulysses_low: bool = True,
):
"""Create Ulysses/Ring process groups inside each SP group.
This is required when TP>1, because SP groups are not necessarily made of
consecutive global ranks (e.g., tp-sp order makes SP ranks strided).
Args:
sp_ulysses_degree: ulysses degree inside SP.
sp_ring_degree: ring degree inside SP.
rank: global rank of current process.
sp_groups: list of global-rank lists for each SP group.
use_ulysses_low: keep the same semantics as the original function.
"""
sp_degree = sp_ring_degree * sp_ulysses_degree
assert sp_degree > 0
assert all(
len(g) == sp_degree for g in sp_groups
), f"Each SP group must have size {sp_degree}, got sizes {[len(g) for g in sp_groups]}"
ulyssess_pg = None
ring_pg = None
num_ulysses_pgs = sp_ring_degree
num_ring_pgs = sp_ulysses_degree
def _map_indices_to_ranks(ranks: list[int], indices: list[int]) -> list[int]:
return [ranks[i] for i in indices]
# Important: call torch.distributed.new_group in the same order on all ranks.
for sp_ranks in sp_groups:
if use_ulysses_low:
for i in range(num_ulysses_pgs):
idx = list(range(i * sp_ulysses_degree, (i + 1) * sp_ulysses_degree))
ulysses_ranks = _map_indices_to_ranks(sp_ranks, idx)
group = torch.distributed.new_group(ulysses_ranks)
if rank in ulysses_ranks:
ulyssess_pg = group
for i in range(num_ring_pgs):
idx = list(range(i, sp_degree, num_ring_pgs))
ring_ranks = _map_indices_to_ranks(sp_ranks, idx)
group = torch.distributed.new_group(ring_ranks)
if rank in ring_ranks:
ring_pg = group
else:
for i in range(num_ring_pgs):
idx = list(range(i * sp_ring_degree, (i + 1) * sp_ring_degree))
ring_ranks = _map_indices_to_ranks(sp_ranks, idx)
group = torch.distributed.new_group(ring_ranks)
if rank in ring_ranks:
ring_pg = group
for i in range(num_ulysses_pgs):
idx = list(range(i, sp_degree, num_ulysses_pgs))
ulysses_ranks = _map_indices_to_ranks(sp_ranks, idx)
group = torch.distributed.new_group(ulysses_ranks)
if rank in ulysses_ranks:
ulyssess_pg = group
PROCESS_GROUP.ULYSSES_PG = ulyssess_pg
PROCESS_GROUP.RING_PG = ring_pg

10
python/sglang/multimodal_gen/runtime/layers/visual_embedding.py
View File

@@ -21,7 +21,7 @@ except Exception as _e:
pass
from sglang.multimodal_gen.runtime.layers.activation import get_act_fn
from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear
from sglang.multimodal_gen.runtime.layers.mlp import MLP
@@ -219,8 +219,12 @@ class ModulateProjection(nn.Module):
super().__init__()
self.factor = factor
self.hidden_size = hidden_size
self.linear = ReplicatedLinear(
hidden_size, hidden_size * factor, bias=True, params_dtype=dtype
self.linear = ColumnParallelLinear(
hidden_size,
hidden_size * factor,
bias=True,
gather_output=True,
params_dtype=dtype,
)
self.act = get_act_fn(act_layer)

1
python/sglang/multimodal_gen/runtime/managers/scheduler.py
View File

@@ -170,6 +170,7 @@ class Scheduler:
if isinstance(req, Req):
warmup_req = deepcopy(req)
warmup_req.is_warmup = True
warmup_req.extra["cache_dit_num_inference_steps"] = req.num_inference_steps
warmup_req.num_inference_steps = 1
recv_reqs.insert(0, (identity, warmup_req))
logger.info("Server warming up....")

34
python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
View File

@@ -25,7 +25,7 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
ScaleResidual,
ScaleResidualLayerNormScaleShift,
)
from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear
from sglang.multimodal_gen.runtime.layers.mlp import MLP
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
NDRotaryEmbedding,
@@ -132,10 +132,10 @@ class WanSelfAttention(nn.Module):
self.parallel_attention = parallel_attention
# layers
self.to_q = ReplicatedLinear(dim, dim)
self.to_k = ReplicatedLinear(dim, dim)
self.to_v = ReplicatedLinear(dim, dim)
self.to_out = ReplicatedLinear(dim, dim)
self.to_q = ColumnParallelLinear(dim, dim, gather_output=True)
self.to_k = ColumnParallelLinear(dim, dim, gather_output=True)
self.to_v = ColumnParallelLinear(dim, dim, gather_output=True)
self.to_out = ColumnParallelLinear(dim, dim, gather_output=True)
self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
@@ -218,8 +218,8 @@ class WanI2VCrossAttention(WanSelfAttention):
supported_attention_backends=supported_attention_backends,
)
self.add_k_proj = ReplicatedLinear(dim, dim)
self.add_v_proj = ReplicatedLinear(dim, dim)
self.add_k_proj = ColumnParallelLinear(dim, dim, gather_output=True)
self.add_v_proj = ColumnParallelLinear(dim, dim, gather_output=True)
self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
self.norm_added_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
@@ -272,11 +272,11 @@ class WanTransformerBlock(nn.Module):
# 1. Self-attention
self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
self.to_q = ReplicatedLinear(dim, dim, bias=True)
self.to_k = ReplicatedLinear(dim, dim, bias=True)
self.to_v = ReplicatedLinear(dim, dim, bias=True)
self.to_q = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
self.to_k = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
self.to_v = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
self.to_out = ReplicatedLinear(dim, dim, bias=True)
self.to_out = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
if attention_type == "sla":
self.attn1 = MinimalA2AAttnOp(
SparseLinearAttention(
@@ -454,12 +454,14 @@ class WanTransformerBlock_VSA(nn.Module):
# 1. Self-attention
self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
self.to_q = ReplicatedLinear(dim, dim, bias=True)
self.to_k = ReplicatedLinear(dim, dim, bias=True)
self.to_v = ReplicatedLinear(dim, dim, bias=True)
self.to_gate_compress = ReplicatedLinear(dim, dim, bias=True)
self.to_q = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
self.to_k = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
self.to_v = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
self.to_gate_compress = ColumnParallelLinear(
dim, dim, bias=True, gather_output=True
)
self.to_out = ReplicatedLinear(dim, dim, bias=True)
self.to_out = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
self.attn1 = UlyssesAttention_VSA(
num_heads=num_heads,
head_size=dim // num_heads,

10
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
View File

@@ -487,6 +487,12 @@ class DenoisingStage(PipelineStage):
"""
assert self.transformer is not None
pipeline = self.pipeline() if self.pipeline else None
# NOTE: In warmup requests we may override req.num_inference_steps (e.g. set to 1)
# for latency amortization, but cache-dit needs the *original* total steps to
# initialize/refresh its context correctly.
cache_dit_num_inference_steps = batch.extra.get(
"cache_dit_num_inference_steps", batch.num_inference_steps
)
if not server_args.model_loaded["transformer"]:
# FIXME: reuse more code
loader = TransformerLoader()
@@ -494,13 +500,13 @@ class DenoisingStage(PipelineStage):
server_args.model_paths["transformer"], server_args, "transformer"
)
# enable cache-dit before torch.compile (delayed mounting)
self._maybe_enable_cache_dit(batch.num_inference_steps)
self._maybe_enable_cache_dit(cache_dit_num_inference_steps)
self.compile_module_with_torch_compile(self.transformer)
if pipeline:
pipeline.add_module("transformer", self.transformer)
server_args.model_loaded["transformer"] = True
else:
self._maybe_enable_cache_dit(batch.num_inference_steps)
self._maybe_enable_cache_dit(cache_dit_num_inference_steps)
# Prepare extra step kwargs for scheduler
extra_step_kwargs = self.prepare_extra_func_kwargs(

18
python/sglang/multimodal_gen/runtime/server_args.py
View File

@@ -366,10 +366,14 @@ class ServerArgs:
if self.vae_cpu_offload is None:
self.vae_cpu_offload = False
else:
self.dit_cpu_offload = True
self.text_encoder_cpu_offload = True
self.image_encoder_cpu_offload = True
self.vae_cpu_offload = True
if self.dit_cpu_offload is None:
self.dit_cpu_offload = True
if self.text_encoder_cpu_offload is None:
self.text_encoder_cpu_offload = True
if self.image_encoder_cpu_offload is None:
self.image_encoder_cpu_offload = True
if self.vae_cpu_offload is None:
self.vae_cpu_offload = True
def __post_init__(self):
# configure logger before use
@@ -1005,9 +1009,9 @@ class ServerArgs:
has_sp = self.sp_degree > 1
has_tp = self.tp_size > 1
if has_sp and has_tp:
raise ValueError(
"cache-dit does not support hybrid parallelism (SP + TP). "
"Please use either sequence parallelism or tensor parallelism, not both."
logger.warning(
"cache-dit is enabled with hybrid parallelism (SP + TP). "
"Proceeding anyway (SGLang integration may support this mode)."
)
def _set_default_attention_backend(self) -> None:

20
python/sglang/multimodal_gen/runtime/utils/cache_dit_integration.py
View File

@@ -28,6 +28,8 @@ from cache_dit import (
from cache_dit.caching.block_adapters import BlockAdapterRegister
from cache_dit.parallelism import ParallelismBackend, ParallelismConfig
from sglang.multimodal_gen.runtime.distributed.parallel_state import get_dit_group
_original_similarity = None
@@ -53,7 +55,8 @@ def _patch_cache_dit_similarity():
sp_group = getattr(self, "_sglang_sp_group", None)
tp_group = getattr(self, "_sglang_tp_group", None)
target_group = sp_group or tp_group
tp_sp_group = getattr(self, "_sglang_tp_sp_group", None)
target_group = tp_sp_group or sp_group or tp_group
if target_group is None:
return _original_similarity(
@@ -316,6 +319,14 @@ def enable_cache_on_transformer(
if context_manager is not None:
context_manager._sglang_sp_group = sp_group
context_manager._sglang_tp_group = tp_group
# In mixed TP + SP (Ulysses/Ring) mode, cache-dit decisions must be consistent
# across the full TP×SP model-parallel slice. Prefer using SGLang's DIT group
# as a conservative superset group; fallback to None.
tp_sp_group = None
if sp_group is not None and tp_group is not None:
tp_sp_group = get_dit_group()
context_manager._sglang_tp_sp_group = tp_sp_group
return transformer
@@ -488,5 +499,12 @@ def enable_cache_on_dual_transformer(
if context_manager is not None:
context_manager._sglang_sp_group = sp_group
context_manager._sglang_tp_group = tp_group
tp_sp_group = None
if sp_group is not None and tp_group is not None:
try:
tp_sp_group = get_dit_group()
except Exception:
tp_sp_group = None
context_manager._sglang_tp_sp_group = tp_sp_group
return transformer, transformer_2