diff --git a/docs_new/cookbook/diffusion/Ideogram/Ideogram4.mdx b/docs_new/cookbook/diffusion/Ideogram/Ideogram4.mdx index 1432c4a77..392686c44 100644 --- a/docs_new/cookbook/diffusion/Ideogram/Ideogram4.mdx +++ b/docs_new/cookbook/diffusion/Ideogram/Ideogram4.mdx @@ -6,7 +6,7 @@ metatags: ## 1. Model introduction -[Ideogram 4](https://huggingface.co/ideogram-ai/ideogram-4-nf4) is Ideogram's text-to-image diffusion model. SGLang Diffusion supports the official NF4 checkpoint, the official FP8 checkpoint, and the Comfy-Org NVFP4 transformer checkpoint. +[Ideogram 4](https://huggingface.co/ideogram-ai/ideogram-4-nf4) is Ideogram's text-to-image diffusion model. SGLang Diffusion supports the official NF4 and FP8 checkpoints, the Comfy-Org NVFP4 transformer checkpoint, and fal's single-branch Fast and Instant variants. Compared with previous open-source image models, Ideogram 4 provides a significant aesthetic lift, with stronger composition, more polished visual style, and better typography-aware generation. @@ -15,12 +15,16 @@ Compared with previous open-source image models, Ideogram 4 provides a significa | NF4 | `ideogram-ai/ideogram-4-nf4` | Official bitsandbytes NF4 checkpoint. Use this path first for low-memory deployment. | | FP8 | `ideogram-ai/ideogram-4-fp8` | Official FP8 checkpoint. | | NVFP4 | `Comfy-Org/Ideogram-4` | Comfy-Org NVFP4 transformer weights. SGLang loads non-transformer components from `ideogram-ai/ideogram-4-fp8`. | +| Fast | `fal/ideogram-v4-fast` | 20-step, CFG-distilled, FP4-targeted floating checkpoint. Defaults to `V4_FAST_20`. | +| Instant | `fal/ideogram-v4-instant` | 8-step, CFG- and timestep-distilled BF16 checkpoint. Defaults to `V4_INSTANT_8`. | + +The fal repositories contain only the transformer component. When either model ID is passed directly to `--model-path`, SGLang loads the text encoder, tokenizer, VAE, and scheduler from the `ideogram-ai/ideogram-4-nf4-diffusers` revision referenced by fal's model cards, then runs the distilled transformer without the unconditional branch. ## 2. Prerequisites - NVIDIA CUDA GPU. - SGLang installed with diffusion dependencies. -- `bitsandbytes>=0.46.1` for the NF4 checkpoint. +- `bitsandbytes>=0.46.1` and `accelerate>=1.1.0` for the NF4 checkpoint and the fal variants' shared NF4 text encoder. - `HF_TOKEN` with access to the Ideogram 4 gated repositories. ## 3. Serve the model @@ -57,6 +61,69 @@ HF_TOKEN=$HF_TOKEN sglang serve \ Use B200 or another Blackwell GPU for NVFP4. +fal Fast: + +```bash Command +HF_TOKEN=$HF_TOKEN sglang serve \ + --model-path fal/ideogram-v4-fast \ + --num-gpus 1 \ + --performance-mode auto \ + --port 30010 +``` + +fal Instant: + +```bash Command +HF_TOKEN=$HF_TOKEN sglang serve \ + --model-path fal/ideogram-v4-instant \ + --num-gpus 1 \ + --performance-mode auto \ + --port 30010 +``` + +### Tensor and sequence parallelism + +Both fal variants support native DiT tensor parallelism and Ulysses sequence parallelism. The following two-GPU layouts were validated at 1024×1024 on B200 GPUs for both `fal/ideogram-v4-fast` and `fal/ideogram-v4-instant`. + +TP2 shards the distilled DiT weights. The shared bitsandbytes NF4 text encoder is replicated because 4-bit row-parallel quantization states cannot be safely sharded: + +```bash Command +HF_TOKEN=$HF_TOKEN sglang serve \ + --model-path fal/ideogram-v4-instant \ + --num-gpus 2 \ + --tp-size 2 \ + --port 30010 +``` + +Ulysses/SP2 shards the image-token sequence while keeping the DiT weights replicated: + +```bash Command +HF_TOKEN=$HF_TOKEN sglang serve \ + --model-path fal/ideogram-v4-instant \ + --num-gpus 2 \ + --ulysses-degree 2 \ + --port 30010 +``` + +Replace the model ID with `fal/ideogram-v4-fast` to use the same layouts for Fast. TP2 and Ulysses/SP2 are the validated configurations; Ring SP is not yet validated for this pipeline. Ideogram 4 has 18 attention heads, so any other TP degree must divide 18. + +### Layerwise offload + +If the distilled DiT does not fit in GPU memory, enable transformer layerwise offload. This keeps the shared NF4 text encoder on the GPU while streaming the DiT blocks from pinned CPU memory. It reduces peak VRAM at the cost of additional latency: + +```bash Command +HF_TOKEN=$HF_TOKEN sglang serve \ + --model-path fal/ideogram-v4-instant \ + --num-gpus 1 \ + --dit-layerwise-offload \ + --layerwise-offload-components transformer \ + --port 30010 +``` + +The same option applies to Fast. Layerwise offload was validated with a real Instant generation; TP2 and Ulysses/SP2 were validated separately without offload. + +The released Fast weights are stored in a floating-point pre-pack form and can be loaded directly, but fal trained them with QAD for an NVFP4 execution path. SGLang does not currently quantize dense NVIDIA DiTs to NVFP4 at load time, so direct Fast inference bypasses the intended quantization path and may be visibly degraded. Treat direct floating-point Fast inference as compatibility/testing support, not the production-quality path. The released Instant checkpoint is BF16 and is the recommended local checkpoint today. + ## 4. Generate an image ```python Example @@ -79,4 +146,40 @@ with open("ideogram4.png", "wb") as f: f.write(image_bytes) ``` -Ideogram 4 presets are `V4_DEFAULT_20`, `V4_QUALITY_48`, and `V4_TURBO_12`. The preset controls both `num_inference_steps` and guidance, so do not set those fields directly. +fal's hosted endpoint expands natural-language prompts, but these local checkpoints expect Ideogram 4's structured JSON caption format. Pass a complete caption as the API `prompt`, for example: + +```python Example +import json + +prompt = json.dumps( + { + "high_level_description": ( + "A bold typographic poster centered on the exact words INSTANT BY FAL, " + "printed in black and electric orange on warm white paper." + ), + "compositional_deconstruction": { + "background": "Warm white textured paper with generous negative space.", + "elements": [ + { + "type": "text", + "text": "INSTANT BY FAL", + "desc": "Large uppercase geometric sans-serif lettering, precisely centered.", + } + ], + }, + }, + ensure_ascii=False, + separators=(",", ":"), +) + +response = client.images.generate( + model="fal/ideogram-v4-instant", + prompt=prompt, + size="1024x1024", + n=1, + response_format="b64_json", + extra_body={"seed": 42}, +) +``` + +Base Ideogram 4 presets are `V4_DEFAULT_20`, `V4_QUALITY_48`, and `V4_TURBO_12`. The fal variants automatically select `V4_FAST_20` and `V4_INSTANT_8`, respectively. A preset controls both `num_inference_steps` and guidance, so do not set those fields directly. diff --git a/docs_new/docs/sglang-diffusion/compatibility_matrix.mdx b/docs_new/docs/sglang-diffusion/compatibility_matrix.mdx index b54622796..a582c14c7 100644 --- a/docs_new/docs/sglang-diffusion/compatibility_matrix.mdx +++ b/docs_new/docs/sglang-diffusion/compatibility_matrix.mdx @@ -55,7 +55,7 @@ Rows are grouped when a family shares the same runtime path or optimization supp Other image pipelines -
zai-org/GLM-Imagetencent/Hunyuan3D-2baidu/ERNIE-Imagebaidu/ERNIE-Image-Turboideogram-ai/ideogram-4-fp8ideogram-ai/ideogram-4-nf4Comfy-Org/Ideogram-4
+
zai-org/GLM-Imagetencent/Hunyuan3D-2baidu/ERNIE-Imagebaidu/ERNIE-Image-Turboideogram-ai/ideogram-4-fp8ideogram-ai/ideogram-4-nf4Comfy-Org/Ideogram-4fal/ideogram-v4-fastfal/ideogram-v4-instant
diff --git a/python/sglang/multimodal_gen/configs/models/dits/__init__.py b/python/sglang/multimodal_gen/configs/models/dits/__init__.py index 94e7b98ac..83ef2327c 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/dits/__init__.py @@ -4,7 +4,10 @@ from sglang.multimodal_gen.configs.models.dits.cosmos3video import Cosmos3VideoC from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig from sglang.multimodal_gen.configs.models.dits.hunyuan3d import Hunyuan3DDiTConfig from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig -from sglang.multimodal_gen.configs.models.dits.ideogram import Ideogram4DiTConfig +from sglang.multimodal_gen.configs.models.dits.ideogram import ( + Ideogram4DistilledDiTConfig, + Ideogram4DiTConfig, +) from sglang.multimodal_gen.configs.models.dits.lingbot_world import ( LingBotWorldVideoConfig, ) @@ -21,6 +24,7 @@ __all__ = [ "HeliosConfig", "HunyuanVideoConfig", "Ideogram4DiTConfig", + "Ideogram4DistilledDiTConfig", "LingBotWorldVideoConfig", "LongLive2VideoConfig", "WanVideoConfig", diff --git a/python/sglang/multimodal_gen/configs/models/dits/ideogram.py b/python/sglang/multimodal_gen/configs/models/dits/ideogram.py index 16da3fec0..a23bb600f 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/ideogram.py +++ b/python/sglang/multimodal_gen/configs/models/dits/ideogram.py @@ -19,6 +19,26 @@ class Ideogram4DiTArchConfig(DiTArchConfig): num_attention_heads: int = 18 num_layers: int = 34 rope_theta: int = 5_000_000 + param_names_mapping: dict = field( + default_factory=lambda: { + r"^(layers\.\d+\.attention)\.to_q\.(.*)$": ( + r"\1.qkv.\2", + 0, + 3, + ), + r"^(layers\.\d+\.attention)\.to_k\.(.*)$": ( + r"\1.qkv.\2", + 1, + 3, + ), + r"^(layers\.\d+\.attention)\.to_v\.(.*)$": ( + r"\1.qkv.\2", + 2, + 3, + ), + r"^(layers\.\d+\.attention)\.to_out\.0\.(.*)$": r"\1.o.\2", + } + ) _fsdp_shard_conditions: list = field(default_factory=lambda: [is_layer]) _supported_attention_backends: set[AttentionBackendEnum] = field( default_factory=lambda: { @@ -37,3 +57,13 @@ class Ideogram4DiTArchConfig(DiTArchConfig): class Ideogram4DiTConfig(DiTConfig): arch_config: DiTArchConfig = field(default_factory=Ideogram4DiTArchConfig) prefix: str = "ideogram4" + # The official FP8 checkpoint stores row-wise FP8 weights without a + # quantization_config, so its native loader intentionally defaults to the + # dedicated weight-only FP8 linears. Distilled fal checkpoints instead + # store floating-point weights and must use ordinary TP-aware linears. + use_weight_only_fp8_linears: bool = True + + +@dataclass +class Ideogram4DistilledDiTConfig(Ideogram4DiTConfig): + use_weight_only_fp8_linears: bool = False diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py index 2801e2e4c..765dea30a 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -29,6 +29,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( Hunyuan3D2PipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.ideogram import ( + Ideogram4DistilledPipelineConfig, Ideogram4PipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.lingbot_world import ( @@ -64,6 +65,7 @@ __all__ = [ "FastHunyuanConfig", "Hunyuan3D2PipelineConfig", "Ideogram4PipelineConfig", + "Ideogram4DistilledPipelineConfig", "FluxPipelineConfig", "Flux2PipelineConfig", "Flux2KleinPipelineConfig", diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/ideogram.py b/python/sglang/multimodal_gen/configs/pipeline_configs/ideogram.py index d83ce21b8..946c2b87c 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/ideogram.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/ideogram.py @@ -3,7 +3,10 @@ from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig -from sglang.multimodal_gen.configs.models.dits.ideogram import Ideogram4DiTConfig +from sglang.multimodal_gen.configs.models.dits.ideogram import ( + Ideogram4DistilledDiTConfig, + Ideogram4DiTConfig, +) from sglang.multimodal_gen.configs.models.encoders.ideogram import ( Ideogram4TextEncoderConfig, ) @@ -300,3 +303,8 @@ class Ideogram4PipelineConfig(ImagePipelineConfig): grid_h = batch.height // patch grid_w = batch.width // patch return (batch_size, grid_h * grid_w, self.dit_config.arch_config.in_channels) + + +@dataclass +class Ideogram4DistilledPipelineConfig(Ideogram4PipelineConfig): + dit_config: DiTConfig = field(default_factory=Ideogram4DistilledDiTConfig) diff --git a/python/sglang/multimodal_gen/configs/sample/ideogram.py b/python/sglang/multimodal_gen/configs/sample/ideogram.py index 09e653027..62afbc5f1 100644 --- a/python/sglang/multimodal_gen/configs/sample/ideogram.py +++ b/python/sglang/multimodal_gen/configs/sample/ideogram.py @@ -26,6 +26,18 @@ IDEOGRAM4_PRESETS: dict[str, dict[str, object]] = { "mu": 0.5, "std": 1.75, }, + "V4_FAST_20": { + "num_steps": 20, + "guidance_schedule": (1.0,) * 20, + "mu": 0.0, + "std": 1.75, + }, + "V4_INSTANT_8": { + "num_steps": 8, + "guidance_schedule": (1.0,) * 8, + "mu": 0.0, + "std": 1.75, + }, } @@ -76,3 +88,13 @@ class Ideogram4SamplingParams(SamplingParams): self.num_inference_steps = preset_steps self.guidance_scale = float(preset_cfg["guidance_schedule"][-1]) super().__post_init__() + + +@dataclass +class Ideogram4FastSamplingParams(Ideogram4SamplingParams): + preset: str = "V4_FAST_20" + + +@dataclass +class Ideogram4InstantSamplingParams(Ideogram4SamplingParams): + preset: str = "V4_INSTANT_8" diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index f2c8a089b..381b3445e 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -59,6 +59,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( Hunyuan3D2PipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.ideogram import ( + Ideogram4DistilledPipelineConfig, Ideogram4PipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.joy_echo import ( @@ -119,7 +120,11 @@ from sglang.multimodal_gen.configs.sample.hunyuan import ( HunyuanSamplingParams, ) from sglang.multimodal_gen.configs.sample.hunyuan3d import Hunyuan3DSamplingParams -from sglang.multimodal_gen.configs.sample.ideogram import Ideogram4SamplingParams +from sglang.multimodal_gen.configs.sample.ideogram import ( + Ideogram4FastSamplingParams, + Ideogram4InstantSamplingParams, + Ideogram4SamplingParams, +) from sglang.multimodal_gen.configs.sample.joy_echo import JoyEchoSamplingParams from sglang.multimodal_gen.configs.sample.joy_image import ( JoyImageEditSamplingParams, @@ -1112,6 +1117,16 @@ def _register_configs(): ) # Ideogram 4 + register_configs( + sampling_param_cls=Ideogram4FastSamplingParams, + pipeline_config_cls=Ideogram4DistilledPipelineConfig, + hf_model_paths=["fal/ideogram-v4-fast"], + ) + register_configs( + sampling_param_cls=Ideogram4InstantSamplingParams, + pipeline_config_cls=Ideogram4DistilledPipelineConfig, + hf_model_paths=["fal/ideogram-v4-instant"], + ) register_configs( sampling_param_cls=Ideogram4SamplingParams, pipeline_config_cls=Ideogram4PipelineConfig, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py b/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py index 805c2a486..2b13484a4 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py @@ -36,6 +36,9 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( Qwen3VLTextRotaryEmbedding, qwen3_apply_rotary_pos_emb, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + LayerwiseOffloadableModuleMixin, +) from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT OUTPUT_IMAGE_INDICATOR = 2 @@ -83,10 +86,11 @@ def _linear( quant_config: QuantizationConfig | None = None, prefix: str = "", gather_output: bool = True, + use_weight_only_fp8_linears: bool = True, ): tp_size = _tp_size() use_column_parallel = tp_size > 1 and out_features % tp_size == 0 - if quant_config is None: + if quant_config is None and use_weight_only_fp8_linears: if use_column_parallel: return WeightOnlyFP8ColumnParallelLinear( in_features, @@ -119,13 +123,14 @@ def _merged_column_linear( bias: bool = True, quant_config: QuantizationConfig | None = None, prefix: str = "", + use_weight_only_fp8_linears: bool = True, ): tp_size = _tp_size() use_column_parallel = tp_size > 1 and all( output_size % tp_size == 0 for output_size in output_sizes ) out_features = sum(output_sizes) - if quant_config is None: + if quant_config is None and use_weight_only_fp8_linears: if use_column_parallel: return WeightOnlyFP8MergedColumnParallelLinear( in_features, @@ -158,10 +163,11 @@ def _row_linear( bias: bool = True, quant_config: QuantizationConfig | None = None, prefix: str = "", + use_weight_only_fp8_linears: bool = True, ): tp_size = _tp_size() use_row_parallel = tp_size > 1 and in_features % tp_size == 0 - if quant_config is None: + if quant_config is None and use_weight_only_fp8_linears: if use_row_parallel: return WeightOnlyFP8RowParallelLinear( in_features, @@ -197,6 +203,7 @@ class Ideogram4Attention(nn.Module): supported_attention_backends, quant_config: QuantizationConfig | None = None, prefix: str = "", + use_weight_only_fp8_linears: bool = True, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -211,6 +218,7 @@ class Ideogram4Attention(nn.Module): bias=False, quant_config=quant_config, prefix=f"{prefix}.qkv", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) self.norm_q = Ideogram4RMSNorm(self.head_dim, eps=eps) self.norm_k = Ideogram4RMSNorm(self.head_dim, eps=eps) @@ -228,6 +236,7 @@ class Ideogram4Attention(nn.Module): bias=False, quant_config=quant_config, prefix=f"{prefix}.o", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) def forward(self, x, cos, sin, attn_mask, attn_mask_meta): @@ -251,6 +260,7 @@ class Ideogram4MLP(nn.Module): hidden_dim: int, quant_config: QuantizationConfig | None = None, prefix: str = "", + use_weight_only_fp8_linears: bool = True, ) -> None: super().__init__() self.w1 = _linear( @@ -260,6 +270,7 @@ class Ideogram4MLP(nn.Module): quant_config=quant_config, prefix=f"{prefix}.w1", gather_output=False, + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) self.w2 = _row_linear( hidden_dim, @@ -267,6 +278,7 @@ class Ideogram4MLP(nn.Module): bias=False, quant_config=quant_config, prefix=f"{prefix}.w2", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) self.w3 = _linear( dim, @@ -275,6 +287,7 @@ class Ideogram4MLP(nn.Module): quant_config=quant_config, prefix=f"{prefix}.w3", gather_output=False, + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) def forward(self, x): @@ -292,6 +305,7 @@ class Ideogram4TransformerBlock(nn.Module): supported_attention_backends, quant_config: QuantizationConfig | None = None, prefix: str = "", + use_weight_only_fp8_linears: bool = True, ): super().__init__() self.attention = Ideogram4Attention( @@ -301,12 +315,14 @@ class Ideogram4TransformerBlock(nn.Module): supported_attention_backends=supported_attention_backends, quant_config=quant_config, prefix=f"{prefix}.attention", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) self.feed_forward = Ideogram4MLP( hidden_size, intermediate_size, quant_config=quant_config, prefix=f"{prefix}.feed_forward", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) self.attention_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps) self.ffn_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps) @@ -318,6 +334,7 @@ class Ideogram4TransformerBlock(nn.Module): bias=True, quant_config=quant_config, prefix=f"{prefix}.adaln_modulation", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) def forward(self, x, cos, sin, adaln_input, attn_mask, attn_mask_meta): @@ -359,6 +376,7 @@ class Ideogram4EmbedScalar(nn.Module): input_range: tuple[float, float], quant_config: QuantizationConfig | None = None, prefix: str = "", + use_weight_only_fp8_linears: bool = True, ) -> None: super().__init__() self.dim = dim @@ -369,6 +387,7 @@ class Ideogram4EmbedScalar(nn.Module): bias=True, quant_config=quant_config, prefix=f"{prefix}.mlp_in", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) self.mlp_out = _linear( dim, @@ -376,6 +395,7 @@ class Ideogram4EmbedScalar(nn.Module): bias=True, quant_config=quant_config, prefix=f"{prefix}.mlp_out", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) def forward(self, x): @@ -394,6 +414,7 @@ class Ideogram4FinalLayer(nn.Module): adaln_dim: int, quant_config: QuantizationConfig | None = None, prefix: str = "", + use_weight_only_fp8_linears: bool = True, ) -> None: super().__init__() self.norm_final = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False) @@ -403,6 +424,7 @@ class Ideogram4FinalLayer(nn.Module): bias=True, quant_config=quant_config, prefix=f"{prefix}.linear", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) self.adaln_modulation = _linear( adaln_dim, @@ -410,6 +432,7 @@ class Ideogram4FinalLayer(nn.Module): bias=True, quant_config=quant_config, prefix=f"{prefix}.adaln_modulation", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) def forward(self, x, c): @@ -417,14 +440,15 @@ class Ideogram4FinalLayer(nn.Module): return self.linear(self.norm_final(x) * scale) -class Ideogram4Transformer2DModel(BaseDiT): +class Ideogram4Transformer2DModel(BaseDiT, LayerwiseOffloadableModuleMixin): _repeated_blocks = ["Ideogram4TransformerBlock"] + layer_names = ["layers"] _fsdp_shard_conditions = Ideogram4DiTConfig().arch_config._fsdp_shard_conditions _compile_conditions = Ideogram4DiTConfig().arch_config._compile_conditions _supported_attention_backends = ( Ideogram4DiTConfig().arch_config._supported_attention_backends ) - param_names_mapping = {} + param_names_mapping = Ideogram4DiTConfig().arch_config.param_names_mapping reverse_param_names_mapping = {} def __init__( @@ -436,6 +460,7 @@ class Ideogram4Transformer2DModel(BaseDiT): ) -> None: super().__init__(config, hf_config, **kwargs) cfg = config.arch_config + use_weight_only_fp8_linears = config.use_weight_only_fp8_linears self._supported_attention_backends = cfg._supported_attention_backends hidden_size = cfg.num_attention_heads * cfg.attention_head_dim self.hidden_size = hidden_size @@ -447,6 +472,7 @@ class Ideogram4Transformer2DModel(BaseDiT): bias=True, quant_config=quant_config, prefix="input_proj", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) self.llm_cond_norm = Ideogram4RMSNorm(cfg.llm_features_dim, eps=1e-6) self.llm_cond_proj = _linear( @@ -455,12 +481,14 @@ class Ideogram4Transformer2DModel(BaseDiT): bias=True, quant_config=quant_config, prefix="llm_cond_proj", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) self.t_embedding = Ideogram4EmbedScalar( hidden_size, input_range=(0.0, 1.0), quant_config=quant_config, prefix="t_embedding", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) self.adaln_proj = _linear( hidden_size, @@ -468,6 +496,7 @@ class Ideogram4Transformer2DModel(BaseDiT): bias=True, quant_config=quant_config, prefix="adaln_proj", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) self.embed_image_indicator = nn.Embedding(2, hidden_size) self.rotary_emb = Qwen3VLTextRotaryEmbedding( @@ -486,6 +515,7 @@ class Ideogram4Transformer2DModel(BaseDiT): supported_attention_backends=self._supported_attention_backends, quant_config=quant_config, prefix=f"layers.{i}", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) for i in range(cfg.num_layers) ] @@ -496,6 +526,7 @@ class Ideogram4Transformer2DModel(BaseDiT): adaln_dim=cfg.adaln_dim, quant_config=quant_config, prefix="final_layer", + use_weight_only_fp8_linears=use_weight_only_fp8_linears, ) def post_load_weights(self) -> None: diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/ideogram.py b/python/sglang/multimodal_gen/runtime/models/encoders/ideogram.py index 49399397a..54dc1fca1 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/ideogram.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/ideogram.py @@ -51,7 +51,11 @@ class IdeogramQwen3VLTextEncoder(TextEncoder): text_config, quant_config=quant_config, use_weight_only_fp8=self._uses_weight_only_fp8, - use_tensor_parallel=True, + # bitsandbytes 4-bit quant states can be sliced safely for output + # (column-parallel) shards, but not for the row-parallel input + # shards used by attention/MLP output projections. Replicate the + # relatively small NF4 text encoder while keeping DiT TP enabled. + use_tensor_parallel=not self._uses_bitsandbytes_4bit, ) @torch.no_grad() diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ideogram.py b/python/sglang/multimodal_gen/runtime/pipelines/ideogram.py index d1f8e1a7b..4613f777c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ideogram.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ideogram.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 +import json import os from dataclasses import dataclass from functools import lru_cache from typing import Any, cast +from huggingface_hub import hf_hub_download, snapshot_download + from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( @@ -34,6 +37,15 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) _IDEOGRAM4_BASE_MODEL = "ideogram-ai/ideogram-4-fp8" +_IDEOGRAM4_DISTILLED_COMPONENTS_MODEL = "ideogram-ai/ideogram-4-nf4-diffusers" +_IDEOGRAM4_DISTILLED_COMPONENTS_REVISION = "1874bc70267ba2c823a7239e1d70dd308c8d64dc" +_IDEOGRAM4_DISTILLED_COMPONENT_PATTERNS = [ + "model_index.json", + "scheduler/*", + "text_encoder/*", + "tokenizer/*", + "vae/*", +] _IDEOGRAM4_NVFP4_COND_FILE = "diffusion_models/ideogram4_nvfp4_mixed.safetensors" _IDEOGRAM4_NVFP4_UNCOND_FILE = ( "diffusion_models/ideogram4_unconditional_nvfp4_mixed.safetensors" @@ -53,6 +65,20 @@ def _resolve_ideogram4_base_model_path() -> str: return maybe_download_model(_IDEOGRAM4_BASE_MODEL, force_diffusers_model=True) +@lru_cache(maxsize=1) +def _resolve_ideogram4_distilled_components_path() -> str: + # fal's model cards explicitly reference the NF4 Diffusers repository for + # these shared components. Download only those components: its conditional + # and unconditional base transformers are unused by distilled variants. + return snapshot_download( + repo_id=_IDEOGRAM4_DISTILLED_COMPONENTS_MODEL, + revision=_IDEOGRAM4_DISTILLED_COMPONENTS_REVISION, + allow_patterns=_IDEOGRAM4_DISTILLED_COMPONENT_PATTERNS, + ignore_patterns=["*.onnx", "*.msgpack"], + max_workers=8, + ) + + def _resolve_ideogram4_unconditional_transformer_weights_path( transformer_weights_path: str, ) -> str | None: @@ -241,4 +267,86 @@ class Ideogram4Nvfp4Pipeline(Ideogram4Pipeline): return super().load_modules(server_args, loaded_modules) -EntryClass = [Ideogram4Pipeline, Ideogram4Nvfp4Pipeline] +class Ideogram4DistilledPipeline(Ideogram4Pipeline): + _required_config_modules = [ + "text_encoder", + "tokenizer", + "vae", + "transformer", + "scheduler", + ] + _distilled_transformer_path: str | None = None + + def _load_config(self) -> dict[str, Any]: + logger.info( + "Using '%s' for distilled config and non-transformer components", + _IDEOGRAM4_DISTILLED_COMPONENTS_MODEL, + ) + model_index_path = hf_hub_download( + repo_id=_IDEOGRAM4_DISTILLED_COMPONENTS_MODEL, + filename="model_index.json", + revision=_IDEOGRAM4_DISTILLED_COMPONENTS_REVISION, + ) + with open(model_index_path, encoding="utf-8") as model_index_file: + return cast(dict[str, Any], json.load(model_index_file)) + + def _resolve_distilled_transformer_path(self) -> str: + if self._distilled_transformer_path is None: + model_path = ( + self.model_path + if os.path.exists(self.model_path) + else snapshot_download( + repo_id=self.model_path, + allow_patterns=["transformer/*"], + ignore_patterns=["*.onnx", "*.msgpack"], + max_workers=8, + ) + ) + self._distilled_transformer_path = os.path.join(model_path, "transformer") + return self._distilled_transformer_path + + def _resolve_component_path( + self, + server_args: ServerArgs, + module_name: str, + load_module_name: str, + ) -> str: + override_path = server_args.component_paths.get(module_name) + if override_path is not None: + return maybe_download_model(override_path) + if module_name == "transformer": + return self._resolve_distilled_transformer_path() + return os.path.join( + _resolve_ideogram4_distilled_components_path(), load_module_name + ) + + def _create_denoising_stage(self): + transformer = self.get_module("transformer") + return ProgressiveDenoisingStageRouter( + standard_stage=Ideogram4DenoisingStage( + transformer=transformer, + unconditional_transformer=None, + pipeline=self, + ), + progressive_stage_factory=lambda: Ideogram4ProgressiveDenoisingStage( + transformer=transformer, + unconditional_transformer=None, + pipeline=self, + ), + ) + + +class Ideogram4FastPipeline(Ideogram4DistilledPipeline): + pipeline_name = "Ideogram4FastPipeline" + + +class Ideogram4InstantPipeline(Ideogram4DistilledPipeline): + pipeline_name = "Ideogram4InstantPipeline" + + +EntryClass = [ + Ideogram4Pipeline, + Ideogram4Nvfp4Pipeline, + Ideogram4FastPipeline, + Ideogram4InstantPipeline, +] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ideogram.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ideogram.py index 375572a4d..bddce267b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ideogram.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ideogram.py @@ -285,6 +285,8 @@ class Ideogram4DenoisingStage(DenoisingStage): def _dual_transformer_execution_mode( self, ) -> DualTransformerExecutionMode | None: + if self.unconditional_transformer is None: + return None return DualTransformerExecutionMode.PAIRED_PER_STEP def _cache_dit_secondary_uses_primary_config(self) -> bool: @@ -457,30 +459,34 @@ class Ideogram4DenoisingStage(DenoisingStage): ) pos_v = pos_out[:, max_text_tokens : max_text_tokens + num_image_tokens] - self._manage_unconditional_transformer_use_site(batch) - with set_forward_context( - current_timestep=i, - attn_metadata=step.attn_metadata, - forward_batch=batch, - ): - neg_v = self._run_ideogram_transformer( - self.unconditional_transformer, - dict( - llm_features=ctx.extra["ideogram4_neg_llm_features"], - x=z, - t=t, - position_ids=ctx.extra["ideogram4_neg_position_ids"], - segment_ids=ctx.extra["ideogram4_neg_segment_ids"], - indicator=ctx.extra["ideogram4_neg_indicator"], - attn_mask=ctx.extra["ideogram4_neg_attn_mask"], - attn_mask_meta=ctx.extra["ideogram4_neg_attn_mask_meta"], - ), - ) + neg_v = None + if self.unconditional_transformer is not None: + self._manage_unconditional_transformer_use_site(batch) + with set_forward_context( + current_timestep=i, + attn_metadata=step.attn_metadata, + forward_batch=batch, + ): + neg_v = self._run_ideogram_transformer( + self.unconditional_transformer, + dict( + llm_features=ctx.extra["ideogram4_neg_llm_features"], + x=z, + t=t, + position_ids=ctx.extra["ideogram4_neg_position_ids"], + segment_ids=ctx.extra["ideogram4_neg_segment_ids"], + indicator=ctx.extra["ideogram4_neg_indicator"], + attn_mask=ctx.extra["ideogram4_neg_attn_mask"], + attn_mask_meta=ctx.extra["ideogram4_neg_attn_mask_meta"], + ), + ) with maybe_nvtx_range("scheduler_step", use_nvtx): - velocity = ( - guidance_schedule[i] * pos_v + (1.0 - guidance_schedule[i]) * neg_v - ) + velocity = pos_v + if neg_v is not None: + velocity = ( + guidance_schedule[i] * pos_v + (1.0 - guidance_schedule[i]) * neg_v + ) ctx.latents = z + velocity * schedule_deltas[i] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/ideogram.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/ideogram.py index f73c7fc8f..aede267c4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/ideogram.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/ideogram.py @@ -153,9 +153,10 @@ class Ideogram4ProgressiveDenoisingStage( """Progressive-resolution denoising stage for Ideogram 4. Inherits the progressive loop from ProgressiveDenoisingStage and the - Ideogram-specific dual-transformer forward pass from Ideogram4DenoisingStage - via MRO. __init__ calls DenoisingStage directly to avoid cooperative-init - incompatibility between the two parent signatures. + Ideogram-specific forward pass from Ideogram4DenoisingStage via MRO. The + base checkpoint uses two transformers, while distilled checkpoints pass no + unconditional transformer. __init__ calls DenoisingStage directly to avoid + cooperative-init incompatibility between the two parent signatures. MRO for method resolution: Ideogram4ProgressiveDenoisingStage diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index eb86fcdbd..0d382aca6 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -129,10 +129,14 @@ DEFAULT_BCG_TEXT_BUCKETS = (64, 128, 256, 512, 1024) BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS = frozenset( { "comfy-org/ideogram-4", + "fal/ideogram-v4-fast", + "fal/ideogram-v4-instant", "glm-image", "ideogram-4", "ideogram-4-fp8", "ideogram-4-nf4", + "ideogram-v4-fast", + "ideogram-v4-instant", "ideogram-ai/ideogram-4-fp8", "ideogram-ai/ideogram-4-nf4", "qwen/qwen-image", diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py index f4849dd24..f83f18197 100644 --- a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py +++ b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py @@ -251,6 +251,14 @@ class TestDiffusionBCGPadding(unittest.TestCase): "comfy-org/ideogram-4", BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS, ) + self.assertIn( + "fal/ideogram-v4-fast", + BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS, + ) + self.assertIn( + "fal/ideogram-v4-instant", + BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS, + ) self.assertIn( "Ideogram4PipelineConfig", BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS, diff --git a/python/sglang/multimodal_gen/test/unit/test_ideogram4.py b/python/sglang/multimodal_gen/test/unit/test_ideogram4.py index 88be3d994..ad569cd67 100644 --- a/python/sglang/multimodal_gen/test/unit/test_ideogram4.py +++ b/python/sglang/multimodal_gen/test/unit/test_ideogram4.py @@ -9,15 +9,21 @@ import torch import torch.nn.functional as F from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig -from sglang.multimodal_gen.configs.models.dits.ideogram import Ideogram4DiTConfig +from sglang.multimodal_gen.configs.models.dits.ideogram import ( + Ideogram4DistilledDiTConfig, + Ideogram4DiTConfig, +) from sglang.multimodal_gen.configs.models.encoders.ideogram import ( Ideogram4TextEncoderConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.ideogram import ( + Ideogram4DistilledPipelineConfig, Ideogram4PipelineConfig, ) from sglang.multimodal_gen.configs.sample.ideogram import ( IDEOGRAM4_PRESETS, + Ideogram4FastSamplingParams, + Ideogram4InstantSamplingParams, Ideogram4SamplingParams, ) from sglang.multimodal_gen.registry import _get_config_info, get_model_info @@ -51,8 +57,17 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader i from sglang.multimodal_gen.runtime.loader.fsdp_load import ( load_model_from_full_model_state_dict, ) +from sglang.multimodal_gen.runtime.loader.utils import ( + get_param_names_mapping, + hf_to_custom_state_dict, +) from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + LayerwiseOffloadableModuleMixin, +) from sglang.multimodal_gen.runtime.models.dits.ideogram import ( + Ideogram4ColumnParallelLinear, + Ideogram4MergedColumnParallelLinear, Ideogram4RowParallelLinear, Ideogram4Transformer2DModel, ) @@ -60,6 +75,8 @@ from sglang.multimodal_gen.runtime.models.encoders.ideogram import ( IdeogramQwen3VLTextEncoder, ) from sglang.multimodal_gen.runtime.pipelines.ideogram import ( + Ideogram4FastPipeline, + _resolve_ideogram4_distilled_components_path, _resolve_ideogram4_unconditional_transformer_weights_path, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req @@ -173,6 +190,12 @@ def _fake_ideogram_pipeline(transformer, unconditional_transformer): class TestIdeogram4(unittest.TestCase): + def test_ideogram_dit_supports_layerwise_offload(self): + self.assertTrue( + issubclass(Ideogram4Transformer2DModel, LayerwiseOffloadableModuleMixin) + ) + self.assertEqual(Ideogram4Transformer2DModel.layer_names, ["layers"]) + def test_registry_resolves_model_index_class_name(self): get_model_info.cache_clear() _get_config_info.cache_clear() @@ -206,6 +229,77 @@ class TestIdeogram4(unittest.TestCase): self.assertIs(info.pipeline_config_cls, Ideogram4PipelineConfig) self.assertIs(info.sampling_param_cls, Ideogram4SamplingParams) + def test_registry_resolves_fal_distilled_repos_to_native_pipelines(self): + get_model_info.cache_clear() + _get_config_info.cache_clear() + + fast = get_model_info("fal/ideogram-v4-fast", backend="sglang") + instant = get_model_info("fal/ideogram-v4-instant", backend="sglang") + + self.assertEqual(fast.pipeline_cls.__name__, "Ideogram4FastPipeline") + self.assertIs(fast.pipeline_config_cls, Ideogram4DistilledPipelineConfig) + self.assertIs(fast.sampling_param_cls, Ideogram4FastSamplingParams) + self.assertEqual(instant.pipeline_cls.__name__, "Ideogram4InstantPipeline") + self.assertIs(instant.pipeline_config_cls, Ideogram4DistilledPipelineConfig) + self.assertIs(instant.sampling_param_cls, Ideogram4InstantSamplingParams) + + def test_fal_distilled_pipeline_resolves_component_only_repo(self): + pipeline = object.__new__(Ideogram4FastPipeline) + pipeline.model_path = "fal/ideogram-v4-fast" + pipeline._distilled_transformer_path = None + server_args = SimpleNamespace(component_paths={}) + + with ( + patch( + "sglang.multimodal_gen.runtime.pipelines.ideogram.snapshot_download", + return_value="/cache/fast", + ) as download, + patch( + "sglang.multimodal_gen.runtime.pipelines.ideogram._resolve_ideogram4_distilled_components_path", + return_value="/cache/components", + ), + ): + transformer = pipeline._resolve_component_path( + server_args, "transformer", "transformer" + ) + vae = pipeline._resolve_component_path(server_args, "vae", "vae") + + self.assertEqual(transformer, "/cache/fast/transformer") + self.assertEqual(vae, "/cache/components/vae") + self.assertNotIn("unconditional_transformer", pipeline._required_config_modules) + download.assert_called_once_with( + repo_id="fal/ideogram-v4-fast", + allow_patterns=["transformer/*"], + ignore_patterns=["*.onnx", "*.msgpack"], + max_workers=8, + ) + + def test_fal_distilled_pipeline_downloads_pinned_shared_components(self): + _resolve_ideogram4_distilled_components_path.cache_clear() + try: + with patch( + "sglang.multimodal_gen.runtime.pipelines.ideogram.snapshot_download", + return_value="/cache/components", + ) as download: + path = _resolve_ideogram4_distilled_components_path() + finally: + _resolve_ideogram4_distilled_components_path.cache_clear() + + self.assertEqual(path, "/cache/components") + download.assert_called_once_with( + repo_id="ideogram-ai/ideogram-4-nf4-diffusers", + revision="1874bc70267ba2c823a7239e1d70dd308c8d64dc", + allow_patterns=[ + "model_index.json", + "scheduler/*", + "text_encoder/*", + "tokenizer/*", + "vae/*", + ], + ignore_patterns=["*.onnx", "*.msgpack"], + max_workers=8, + ) + def test_registry_resolves_official_nf4_repo_to_native_pipeline(self): get_model_info.cache_clear() _get_config_info.cache_clear() @@ -348,6 +442,17 @@ class TestIdeogram4(unittest.TestCase): with self.assertRaisesRegex(ValueError, "Unknown Ideogram 4 preset"): Ideogram4SamplingParams(preset="V4_FAST") + def test_ideogram_distilled_sampling_defaults(self): + fast = Ideogram4FastSamplingParams() + instant = Ideogram4InstantSamplingParams() + + self.assertEqual(fast.preset, "V4_FAST_20") + self.assertEqual(fast.num_inference_steps, 20) + self.assertEqual(fast.guidance_scale, 1.0) + self.assertEqual(instant.preset, "V4_INSTANT_8") + self.assertEqual(instant.num_inference_steps, 8) + self.assertEqual(instant.guidance_scale, 1.0) + def test_ideogram_sampling_params_merge_recomputes_preset_fields(self): target = Ideogram4SamplingParams() user = Ideogram4SamplingParams( @@ -450,6 +555,25 @@ class TestIdeogram4(unittest.TestCase): ) self.assertTrue(all(use.target_dtype is None for use in uses)) + def test_ideogram_distilled_denoiser_uses_one_transformer(self): + import sglang.multimodal_gen.runtime.server_args as server_args_module + + prev_args = server_args_module._global_server_args + try: + set_global_server_args(_fake_server_args()) + transformer = FakeIdeogramTransformer() + stage = Ideogram4DenoisingStage( + transformer=transformer, + unconditional_transformer=None, + pipeline=_fake_ideogram_pipeline(transformer, None), + ) + uses = stage.component_uses(_fake_server_args(), "stage") + finally: + set_global_server_args(prev_args) + + self.assertEqual([use.component_name for use in uses], ["transformer"]) + self.assertIsNone(stage._dual_transformer_execution_mode()) + def test_ideogram_stages_inherit_common_stage_bases(self): import sglang.multimodal_gen.runtime.server_args as server_args_module @@ -664,6 +788,99 @@ class TestIdeogram4(unittest.TestCase): ) self.assertEqual(state["layers.0.attention.qkv.weight"].dtype, FP8_WEIGHT_DTYPE) + def test_distilled_ideogram_dit_uses_unquantized_linears(self): + import sglang.multimodal_gen.runtime.server_args as server_args_module + + prev_args = server_args_module._global_server_args + try: + set_global_server_args( + SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False) + ) + with patch( + "sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size", + return_value=1, + ): + with torch.device("meta"): + model = Ideogram4Transformer2DModel( + Ideogram4DistilledDiTConfig(), {} + ) + finally: + set_global_server_args(prev_args) + + self.assertIsInstance(model.input_proj.quant_method, UnquantizedLinearMethod) + self.assertIsInstance( + model.layers[0].attention.qkv.quant_method, UnquantizedLinearMethod + ) + state = model.state_dict() + self.assertEqual(tuple(state["input_proj.weight"].shape), (4608, 128)) + self.assertNotIn("input_proj.weight_scale", state) + self.assertNotIn("layers.0.attention.qkv.weight_scale", state) + + def test_distilled_ideogram_maps_diffusers_attention_weights(self): + mapping = get_param_names_mapping( + Ideogram4DistilledDiTConfig().arch_config.param_names_mapping + ) + weights = [ + ("layers.0.attention.to_q.weight", torch.full((2, 2), 1.0)), + ("layers.0.attention.to_k.weight", torch.full((2, 2), 2.0)), + ("layers.0.attention.to_v.weight", torch.full((2, 2), 3.0)), + ("layers.0.attention.to_out.0.weight", torch.full((2, 2), 4.0)), + ] + + mapped, _ = hf_to_custom_state_dict(iter(weights), mapping) + + torch.testing.assert_close( + mapped["layers.0.attention.qkv.weight"], + torch.cat([weight for _, weight in weights[:3]], dim=0), + ) + torch.testing.assert_close(mapped["layers.0.attention.o.weight"], weights[3][1]) + + def test_distilled_ideogram_dit_uses_tp_unquantized_linears(self): + import sglang.multimodal_gen.runtime.server_args as server_args_module + + fake_tp_group = SimpleNamespace(world_size=2, rank_in_group=1) + prev_args = server_args_module._global_server_args + try: + set_global_server_args( + SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False) + ) + with ( + patch( + "sglang.multimodal_gen.runtime.models.dits.ideogram.model_parallel_is_initialized", + return_value=True, + ), + patch( + "sglang.multimodal_gen.runtime.models.dits.ideogram.get_tp_world_size", + return_value=2, + ), + patch( + "sglang.multimodal_gen.runtime.layers.linear.get_tp_group", + return_value=fake_tp_group, + ), + patch( + "sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size", + return_value=1, + ), + ): + with torch.device("meta"): + model = Ideogram4Transformer2DModel( + Ideogram4DistilledDiTConfig(), {} + ) + finally: + set_global_server_args(prev_args) + + self.assertIsInstance(model.input_proj, Ideogram4ColumnParallelLinear) + self.assertIsInstance( + model.layers[0].attention.qkv, Ideogram4MergedColumnParallelLinear + ) + self.assertIsInstance(model.layers[0].attention.o, Ideogram4RowParallelLinear) + self.assertIsInstance(model.input_proj.quant_method, UnquantizedLinearMethod) + self.assertEqual(tuple(model.input_proj.weight.shape), (2304, 128)) + self.assertEqual( + tuple(model.layers[0].attention.qkv.weight.shape), (6912, 4608) + ) + self.assertEqual(tuple(model.layers[0].attention.o.weight.shape), (4608, 2304)) + def test_ideogram_dit_uses_tp_fp8_linears_when_tp_is_initialized(self): import sglang.multimodal_gen.runtime.server_args as server_args_module @@ -1070,6 +1287,43 @@ class TestIdeogram4(unittest.TestCase): self.assertTrue(layer.mlp.down_proj.input_is_parallel) self.assertTrue(layer.mlp.down_proj.reduce_results) + def test_ideogram_nf4_text_encoder_is_replicated_under_tp(self): + config = Ideogram4TextEncoderConfig() + config.update_model_arch( + { + "quantization_config": { + "quant_method": "bitsandbytes", + "load_in_4bit": True, + "bnb_4bit_quant_type": "nf4", + } + } + ) + config.arch_config.text_config = Qwen3VLTextConfig( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=8, + max_position_embeddings=64, + pad_token_id=0, + ) + + with ( + patch( + "sglang.multimodal_gen.runtime.models.encoders.ideogram.BitsAndBytesConfig.from_config", + return_value=object(), + ), + patch( + "sglang.multimodal_gen.runtime.models.encoders.ideogram.Qwen3VLTextModel" + ) as text_model, + ): + text_model.return_value = torch.nn.Identity() + IdeogramQwen3VLTextEncoder(config) + + self.assertFalse(text_model.call_args.kwargs["use_tensor_parallel"]) + def test_denoise_and_decode_shape_check(self): import sglang.multimodal_gen.runtime.server_args as server_args_module diff --git a/python/sglang/utils.py b/python/sglang/utils.py index a04d184cd..1bf5eaa8c 100644 --- a/python/sglang/utils.py +++ b/python/sglang/utils.py @@ -38,6 +38,10 @@ KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS: dict[str, str] = { "pi0.5": "Pi05Pipeline", "hunyuan3d": "Hunyuan3D2Pipeline", "flux.2-dev-nvfp4": "Flux2NvfpPipeline", + "fal/ideogram-v4-fast": "Ideogram4FastPipeline", + "fal--ideogram-v4-fast": "Ideogram4FastPipeline", + "fal/ideogram-v4-instant": "Ideogram4InstantPipeline", + "fal--ideogram-v4-instant": "Ideogram4InstantPipeline", "comfy-org/ideogram-4": "Ideogram4Nvfp4Pipeline", "comfy-org--ideogram-4": "Ideogram4Nvfp4Pipeline", }