From 3b62604ceca8f27ccdaf94c004a010cb5e03a87a Mon Sep 17 00:00:00 2001 From: Makcum888e <79456407+Makcum888e@users.noreply.github.com> Date: Tue, 19 May 2026 17:27:21 +0300 Subject: [PATCH] [Diffusion] Support parallelism for GLM-Image (#25645) --- python/pyproject_npu.toml | 2 +- .../configs/pipeline_configs/glm_image.py | 7 +- .../runtime/pipelines/glm_image.py | 11 +- .../executors/parallel_executor.py | 18 + .../runtime/pipelines_core/stages/base.py | 2 + .../stages/model_specific_stages/glm_image.py | 371 +++++++++--------- 6 files changed, 227 insertions(+), 184 deletions(-) diff --git a/python/pyproject_npu.toml b/python/pyproject_npu.toml index 4a0c7e728..6fb3b95d6 100644 --- a/python/pyproject_npu.toml +++ b/python/pyproject_npu.toml @@ -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", diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py index 2801eeb37..c7c94df4d 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py @@ -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 { diff --git a/python/sglang/multimodal_gen/runtime/pipelines/glm_image.py b/python/sglang/multimodal_gen/runtime/pipelines/glm_image.py index df3e58e36..1d73372ea 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/glm_image.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/glm_image.py @@ -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", ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py index 07f765cbb..997981685 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py index 05b6b06de..1044e95fc 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py @@ -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): diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/glm_image.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/glm_image.py index 47246b419..c909619a6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/glm_image.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/glm_image.py @@ -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 H W shape specification - - Returns: - Tuple of (expanded_prompt, token_h, token_w, prev_token_h, prev_token_w) - """ - match = re.search(r"(\d+)\s+(\d+)", prompt) - if match is None: - raise ValueError( - f"Prompt must contain shape info in format 'H W', 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"{token_h} {token_w}" - new_shape = ( - f"{token_h} {token_w}{prev_token_h} {prev_token_w}" - ) - 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 H W shape specification + + Returns: + Tuple of (expanded_prompt, token_h, token_w, prev_token_h, prev_token_w) + """ + match = re.search(r"(\d+)\s+(\d+)", prompt) + if match is None: + raise ValueError( + f"Prompt must contain shape info in format 'H W', 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"{token_h} {token_w}" + new_shape = ( + f"{token_h} {token_w}{prev_token_h} {prev_token_w}" + ) + 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(