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

[Diffusion] Support parallelism for GLM-Image (#25645)

Browse Source
This commit is contained in:
Makcum888e
2026-05-19 17:27:21 +03:00
committed by GitHub
parent aad00b0ed8
commit 3b62604cec
6 changed files with 227 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/pyproject_npu.toml
+1 -1
View File
@@ -77,7 +77,7 @@ diffusion = [
"moviepy>=2.0.0",
"opencv-python==4.10.0.84",
"remote-pdb",
"cache-dit==1.2.1",
"cache-dit==1.3.5",
"addict",
"scikit-image==0.25.2",
"trimesh>=4.0.0",
python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py
+5 -2
View File
@@ -11,6 +11,7 @@ from sglang.multimodal_gen.configs.models.vaes.glmimage import GlmImageVAEConfig
from sglang.multimodal_gen.configs.pipeline_configs.base import (
ModelTaskType,
SpatialImagePipelineConfig,
shard_rotary_emb_for_sp,
)
@@ -51,8 +52,10 @@ class GlmImagePipelineConfig(SpatialImagePipelineConfig):
height = batch.height // self.vae_scale_factor
width = batch.width // self.vae_scale_factor
hidden_states = torch.empty(1, 1, height, width, device=device, dtype=dtype)
freqs_cis = rotary_emb(hidden_states)
return freqs_cis
cos, sin = rotary_emb(hidden_states)
cos = shard_rotary_emb_for_sp(cos)
sin = shard_rotary_emb_for_sp(sin)
return cos, sin
def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
return {
python/sglang/multimodal_gen/runtime/pipelines/glm_image.py
+9 -2
View File
@@ -4,6 +4,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import
)
from sglang.multimodal_gen.runtime.pipelines_core.stages import DenoisingStage
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.glm_image import (
GlmImageAR,
GlmImageBeforeDenoisingStage,
)
from sglang.multimodal_gen.runtime.server_args import ServerArgs
@@ -26,15 +27,21 @@ class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase):
]
def create_pipeline_stages(self, server_args: ServerArgs):
self.add_stage(
GlmImageAR(
processor=self.get_module("processor"),
vision_language_encoder=self.get_module("vision_language_encoder"),
),
"glm_image_ar",
)
self.add_stage(
GlmImageBeforeDenoisingStage(
vae=self.get_module("vae"),
text_encoder=self.get_module("text_encoder"),
tokenizer=self.get_module("tokenizer"),
processor=self.get_module("processor"),
transformer=self.get_module("transformer"),
scheduler=self.get_module("scheduler"),
vision_language_encoder=self.get_module("vision_language_encoder"),
),
"glm_image_before_denoising_stage",
)
python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py
+18
View File
@@ -8,6 +8,7 @@ from sglang.multimodal_gen.runtime.distributed import get_sp_group
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_cfg_group,
get_classifier_free_guidance_rank,
get_world_group,
get_world_rank,
)
from sglang.multimodal_gen.runtime.pipelines_core import Req
@@ -65,6 +66,7 @@ class ParallelExecutor(PipelineExecutor):
else:
rank = get_world_rank()
cfg_group = get_cfg_group()
group = get_world_group()
self.begin_component_residency_request(stages, batch, server_args)
try:
@@ -101,6 +103,22 @@ class ParallelExecutor(PipelineExecutor):
self.before_stage(stage, stage_index, batch, server_args)
batch = stage(batch, server_args)
self.after_stage(stage_index)
elif paradigm == StageParallelismType.MAIN_RANK_ONLY_AND_SEND_TO_OTHERS:
if rank == 0:
# Only main rank executes, others just wait
self.before_stage(stage, stage_index, batch, server_args)
batch = stage(batch, server_args)
self.after_stage(stage_index)
torch.distributed.barrier()
# Send batch to other ranks
obj_list = [batch] if rank == 0 else []
broadcasted_list = broadcast_pyobj(
obj_list, rank=rank, dist_group=group.cpu_group, src=0
)
if rank != 0:
batch = broadcasted_list[0]
torch.distributed.barrier()
finally:
self.finish_component_residency_request()
return batch
python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py
+2
View File
@@ -40,6 +40,8 @@ class StageParallelismType(Enum):
MAIN_RANK_ONLY = auto()
# this stage requires a cfg-parallel
CFG_PARALLEL = auto()
# executed on main rank only and send result to other ranks
MAIN_RANK_ONLY_AND_SEND_TO_OTHERS = auto()
class StageVerificationError(Exception):
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/glm_image.py
+192 -179
View File
@@ -15,7 +15,10 @@ from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_c
from sglang.multimodal_gen.runtime.models.dits.glm_image import GlmImageKVCache
from sglang.multimodal_gen.runtime.models.vision_utils import load_image
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
PipelineStage,
StageParallelismType,
)
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -103,194 +106,31 @@ def retrieve_latents(
raise AttributeError("Could not access latents of provided encoder_output")
class GlmImageBeforeDenoisingStage(PipelineStage):
class GlmImageAR(PipelineStage):
r"""
Pipeline for text-to-image generation using GLM-Image.
This pipeline integrates both the AR (autoregressive) model for token generation and the DiT (diffusion
transformer) model for image decoding.
This stage for the AR (autoregressive) model for token generation.
Args:
vae ([`AutoencoderKL`]):
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
text_encoder ([`T5EncoderModel`]):
Frozen text-encoder for glyph embeddings.
tokenizer (`PreTrainedTokenizer`):
Tokenizer for the text encoder.
processor (`AutoProcessor`):
Processor for the AR model to handle chat templates and tokenization.
vision_language_encoder ([`GlmImageForConditionalGeneration`]):
The AR model that generates image tokens from text prompts.
transformer ([`GlmImageTransformer2DModel`]):
A text conditioned transformer to denoise the encoded image latents (DiT).
scheduler ([`SchedulerMixin`]):
A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
"""
def __init__(
self,
tokenizer,
processor,
text_encoder,
vision_language_encoder,
vae,
transformer,
scheduler,
):
super().__init__()
self.tokenizer = tokenizer
self.processor = processor
self.text_encoder = text_encoder
self.vision_language_encoder = vision_language_encoder
self.vae = vae
self.transformer = transformer
self.scheduler = scheduler
self.vae_scale_factor = (
2 ** (len(self.vae.config.block_out_channels) - 1)
if getattr(self, "vae", None)
else 8
)
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
self.default_sample_size = (
self.transformer.config.sample_size
if hasattr(self, "transformer")
and self.transformer is not None
and hasattr(self.transformer.config, "sample_size")
else 128
)
def _parse_and_expand_shape_info(
self, prompt: str
) -> Tuple[str, int, int, int, int]:
"""
Parse the shape info from prompt and expand it for AR model.
Args:
prompt: The prompt containing <sop>H W<eop> shape specification
Returns:
Tuple of (expanded_prompt, token_h, token_w, prev_token_h, prev_token_w)
"""
match = re.search(r"<sop>(\d+)\s+(\d+)<eop>", prompt)
if match is None:
raise ValueError(
f"Prompt must contain shape info in format '<sop>H W<eop>', got: {prompt}"
)
token_h, token_w = int(match.group(1)), int(match.group(2))
ratio = token_h / token_w
prev_token_h = int(sqrt(ratio) * 16)
prev_token_w = int(sqrt(1 / ratio) * 16)
old_shape = f"<sop>{token_h} {token_w}<eop>"
new_shape = (
f"<sop>{token_h} {token_w}<eop><sop>{prev_token_h} {prev_token_w}<eop>"
)
expanded_prompt = prompt.replace(old_shape, new_shape)
return expanded_prompt, token_h, token_w, prev_token_h, prev_token_w
def _build_image_grid_thw(
self,
token_h: int,
token_w: int,
prev_token_h: int,
prev_token_w: int,
existing_grid: Optional[torch.Tensor] = None,
device: Optional[torch.device] = None,
) -> torch.Tensor:
"""
Build image grid tensor for AR model.
For text-to-image: creates grid for large image + small image For image-to-image: appends new image to existing
grid
"""
if existing_grid is None or existing_grid.numel() == 0:
# Text-to-image: large image + small image
return torch.tensor(
[
[1, token_h, token_w],
[1, prev_token_h, prev_token_w],
],
device=device,
)
else:
# Image-to-image: append to existing
return torch.cat(
[existing_grid, torch.tensor([[1, token_h, token_w]], device=device)],
dim=0,
)
def _calculate_ar_generation_params(
self,
token_h: int,
token_w: int,
prev_token_h: int,
prev_token_w: int,
is_text_to_image: bool,
) -> Tuple[int, int]:
"""
Calculate max_new_tokens and large_image_start_offset for AR generation.
"""
large_image_tokens = token_h * token_w
small_image_tokens = prev_token_h * prev_token_w
if is_text_to_image:
max_new_tokens = small_image_tokens + large_image_tokens + 1
large_image_start_offset = small_image_tokens
else:
max_new_tokens = large_image_tokens + 1
large_image_start_offset = 0
return max_new_tokens, large_image_start_offset
def _extract_large_image_tokens(
self,
outputs: torch.Tensor,
input_length: int,
large_image_start_offset: int,
large_image_tokens: int,
) -> torch.Tensor:
"""
Extract the large image tokens from AR model output.
"""
generated_tokens = outputs[0][input_length:]
large_image_start = large_image_start_offset
large_image_end = large_image_start + large_image_tokens
return generated_tokens[large_image_start:large_image_end]
def _upsample_d32_to_d16(
self, token_ids: torch.Tensor, token_h: int, token_w: int
) -> torch.Tensor:
"""
Upsample token IDs from d32 format to d16 format.
AR model generates tokens at d32 resolution (each token = 32x32 pixels). DiT expects tokens at d16 resolution
(each token = 16x16 pixels). This function performs 2x nearest-neighbor upsampling.
Args:
token_ids: Token IDs of shape [N] where N = token_h * token_w
token_h: Height in d32 token units
token_w: Width in d32 token units
Returns:
Upsampled token IDs of shape [1, N*4] where N*4 = (token_h*2) * (token_w*2)
"""
# Reshape to spatial format: [1, 1, H, W]
token_ids = token_ids.view(1, 1, token_h, token_w)
# 2x nearest-neighbor upsampling
token_ids = torch.nn.functional.interpolate(
token_ids.float(), scale_factor=2, mode="nearest"
).to(dtype=torch.long)
# Flatten back to [1, H*W*4]
token_ids = token_ids.view(1, -1)
return token_ids
@property
def parallelism_type(self) -> StageParallelismType:
return StageParallelismType.MAIN_RANK_ONLY_AND_SEND_TO_OTHERS
@staticmethod
def _compute_generation_params(
@@ -407,6 +247,185 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
return prior_token_ids, prior_token_image_ids
@torch.no_grad()
def forward(
self,
batch: Req,
server_args: ServerArgs,
) -> Req:
prompt = batch.prompt
height = batch.height
width = batch.width
device = get_local_torch_device()
time_start = time.time()
prior_token_id, prior_token_image_ids = self.generate_prior_tokens(
prompt=prompt,
height=height,
width=width,
)
prior_token_id = prior_token_id.to(device=device)
time_end = time.time()
logger.info(f"generate_prior_tokens time: {time_end - time_start}")
batch.prior_token_id = prior_token_id
batch.prior_token_image_ids = prior_token_image_ids
return batch
def _extract_large_image_tokens(
self,
outputs: torch.Tensor,
input_length: int,
large_image_start_offset: int,
large_image_tokens: int,
) -> torch.Tensor:
"""
Extract the large image tokens from AR model output.
"""
generated_tokens = outputs[0][input_length:]
large_image_start = large_image_start_offset
large_image_end = large_image_start + large_image_tokens
return generated_tokens[large_image_start:large_image_end]
class GlmImageBeforeDenoisingStage(PipelineStage):
r"""
Pipeline for text-to-image generation using GLM-Image.
This stage for preparations before denoising stage like encoding, latents, timesteps
Args:
vae ([`AutoencoderKL`]):
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
text_encoder ([`T5EncoderModel`]):
Frozen text-encoder for glyph embeddings.
tokenizer (`PreTrainedTokenizer`):
Tokenizer for the text encoder.
transformer ([`GlmImageTransformer2DModel`]):
A text conditioned transformer to denoise the encoded image latents (DiT).
scheduler ([`SchedulerMixin`]):
A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
"""
def __init__(
self,
tokenizer,
text_encoder,
vae,
transformer,
scheduler,
):
super().__init__()
self.tokenizer = tokenizer
self.text_encoder = text_encoder
self.vae = vae
self.transformer = transformer
self.scheduler = scheduler
self.vae_scale_factor = (
2 ** (len(self.vae.config.block_out_channels) - 1)
if getattr(self, "vae", None)
else 8
)
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
self.default_sample_size = (
self.transformer.config.sample_size
if hasattr(self, "transformer")
and self.transformer is not None
and hasattr(self.transformer.config, "sample_size")
else 128
)
def _parse_and_expand_shape_info(
self, prompt: str
) -> Tuple[str, int, int, int, int]:
"""
Parse the shape info from prompt and expand it for AR model.
Args:
prompt: The prompt containing <sop>H W<eop> shape specification
Returns:
Tuple of (expanded_prompt, token_h, token_w, prev_token_h, prev_token_w)
"""
match = re.search(r"<sop>(\d+)\s+(\d+)<eop>", prompt)
if match is None:
raise ValueError(
f"Prompt must contain shape info in format '<sop>H W<eop>', got: {prompt}"
)
token_h, token_w = int(match.group(1)), int(match.group(2))
ratio = token_h / token_w
prev_token_h = int(sqrt(ratio) * 16)
prev_token_w = int(sqrt(1 / ratio) * 16)
old_shape = f"<sop>{token_h} {token_w}<eop>"
new_shape = (
f"<sop>{token_h} {token_w}<eop><sop>{prev_token_h} {prev_token_w}<eop>"
)
expanded_prompt = prompt.replace(old_shape, new_shape)
return expanded_prompt, token_h, token_w, prev_token_h, prev_token_w
def _build_image_grid_thw(
self,
token_h: int,
token_w: int,
prev_token_h: int,
prev_token_w: int,
existing_grid: Optional[torch.Tensor] = None,
device: Optional[torch.device] = None,
) -> torch.Tensor:
"""
Build image grid tensor for AR model.
For text-to-image: creates grid for large image + small image For image-to-image: appends new image to existing
grid
"""
if existing_grid is None or existing_grid.numel() == 0:
# Text-to-image: large image + small image
return torch.tensor(
[
[1, token_h, token_w],
[1, prev_token_h, prev_token_w],
],
device=device,
)
else:
# Image-to-image: append to existing
return torch.cat(
[existing_grid, torch.tensor([[1, token_h, token_w]], device=device)],
dim=0,
)
def _calculate_ar_generation_params(
self,
token_h: int,
token_w: int,
prev_token_h: int,
prev_token_w: int,
is_text_to_image: bool,
) -> Tuple[int, int]:
"""
Calculate max_new_tokens and large_image_start_offset for AR generation.
"""
large_image_tokens = token_h * token_w
small_image_tokens = prev_token_h * prev_token_w
if is_text_to_image:
max_new_tokens = small_image_tokens + large_image_tokens + 1
large_image_start_offset = small_image_tokens
else:
max_new_tokens = large_image_tokens + 1
large_image_start_offset = 0
return max_new_tokens, large_image_start_offset
def get_glyph_texts(self, prompt):
prompt = prompt[0] if isinstance(prompt, list) else prompt
ocr_texts = (
@@ -664,16 +683,10 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
if ar_condition_images is not None:
height = height or ar_condition_images[0].height
width = width or ar_condition_images[0].width
time_start = time.time()
prior_token_id, prior_token_image_ids = self.generate_prior_tokens(
prompt=prompt,
image=ar_condition_images,
height=height,
width=width,
)
prior_token_id = prior_token_id.to(device=device)
time_end = time.time()
logger.info(f"generate_prior_tokens time: {time_end - time_start}")
prior_token_id = batch.prior_token_id
prior_token_image_ids = batch.prior_token_image_ids
prior_token_id = prior_token_id.to(device)
# 3. Encode input prompt
prompt_embeds, negative_prompt_embeds = self.encode_prompt(