sgl-project--sglang/python/sglang/multimodal_gen/runtime/models/encoders/ideogram.py

# SPDX-License-Identifier: Apache-2.0

from collections.abc import Iterable
from typing import Tuple

import torch
from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig

from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
from sglang.multimodal_gen.configs.models.encoders.ideogram import (
    Ideogram4TextEncoderConfig,
)
from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
    BitsAndBytesConfig,
    attach_bitsandbytes_4bit_quant_states,
    build_bitsandbytes_4bit_quant_states,
    is_bitsandbytes_4bit_state_name,
)
from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder
from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLTextModel


class IdeogramQwen3VLTextEncoder(TextEncoder):
    """Language-only Qwen3-VL text encoder stored inside Ideogram checkpoints."""

    _activation_layers = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)

    def __init__(self, config: Ideogram4TextEncoderConfig) -> None:
        super().__init__(config)
        arch_config = config.arch_config
        text_config = getattr(arch_config, "text_config")
        if isinstance(text_config, dict):
            text_config = Qwen3VLTextConfig(**text_config)
        self._uses_bitsandbytes_4bit = getattr(
            arch_config, "ideogram_bnb_4bit_weight_only", False
        )
        self._uses_weight_only_fp8 = getattr(
            arch_config, "ideogram_fp8_weight_only", False
        )
        quant_config = None
        if self._uses_bitsandbytes_4bit:
            source_quant_config = getattr(arch_config, "quantization_config")
            if isinstance(source_quant_config, dict):
                quant_config_dict = source_quant_config
            else:
                quant_config_dict = source_quant_config.to_dict()
            quant_config = BitsAndBytesConfig.from_config(quant_config_dict)
        self.language_model = Qwen3VLTextModel(
            text_config,
            quant_config=quant_config,
            use_weight_only_fp8=self._uses_weight_only_fp8,
            use_tensor_parallel=True,
        )

    @torch.no_grad()
    def forward(
        self,
        input_ids: torch.Tensor | None,
        position_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        output_hidden_states: bool | None = None,
        **kwargs,
    ) -> BaseEncoderOutput:
        outputs = self.language_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            use_cache=False,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )
        return BaseEncoderOutput(
            last_hidden_state=outputs.last_hidden_state,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def encode_ideogram_features(
        self,
        token_ids: torch.Tensor,
        text_position_ids: torch.Tensor,
        indicator: torch.Tensor,
        llm_token_indicator: int,
    ) -> torch.Tensor:
        batch_size, seq_len = token_ids.shape
        hidden_size = self.language_model.config.hidden_size
        out_dim = hidden_size * len(self._activation_layers)
        features = torch.zeros(
            batch_size,
            seq_len,
            out_dim,
            dtype=torch.float32,
            device=token_ids.device,
        )
        for batch_idx in range(batch_size):
            text_mask = indicator[batch_idx] == llm_token_indicator
            cur_token_ids = token_ids[batch_idx, text_mask].unsqueeze(0)
            if cur_token_ids.numel() == 0:
                continue
            pos_2d = text_position_ids[batch_idx, text_mask, 0].unsqueeze(0)
            position_ids = pos_2d[None, ...].expand(4, 1, -1)
            attention_mask = torch.ones_like(cur_token_ids)
            with set_forward_context(current_timestep=0, attn_metadata=None):
                outputs = self.forward(
                    input_ids=cur_token_ids,
                    position_ids=position_ids,
                    attention_mask=attention_mask,
                    output_hidden_states=True,
                )
            assert outputs.hidden_states is not None
            selected = [outputs.hidden_states[i] for i in self._activation_layers]
            stacked = torch.stack(selected, dim=0).permute(1, 2, 3, 0)
            features[batch_idx, text_mask] = stacked.reshape(
                1, cur_token_ids.shape[1], -1
            )[0].to(torch.float32)
        return features

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        if self._uses_bitsandbytes_4bit:
            return self._load_bitsandbytes_4bit_weights(weights)

        loaded_params: set[str] = set()
        params_dict = dict(self.named_parameters(remove_duplicate=False))
        for name, loaded_weight in weights:
            if name.startswith("visual."):
                continue
            if "rotary_emb.inv_freq" in name:
                continue
            param = params_dict.get(name)
            if param is None:
                raise KeyError(
                    f"Unexpected weight name while loading Ideogram text encoder: {name}"
                )
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
            weight_loader(param, loaded_weight.to(param.dtype))
            loaded_params.add(name)
        return loaded_params

    def _load_bitsandbytes_4bit_weights(
        self, weights: Iterable[Tuple[str, torch.Tensor]]
    ):
        params_dict = dict(self.named_parameters(remove_duplicate=False))
        raw_quant_state: dict[str, torch.Tensor] = {}
        normal_weight_names: list[str] = []
        loaded_params: set[str] = set()
        for name, loaded_weight in weights:
            if is_bitsandbytes_4bit_state_name(name):
                if "quant_state.bitsandbytes" in name:
                    loaded_weight = loaded_weight.cpu().data
                raw_quant_state[name] = loaded_weight
                continue
            if name.startswith("visual."):
                continue
            if "rotary_emb.inv_freq" in name:
                continue
            param = params_dict.get(name)
            if param is None:
                raise KeyError(
                    f"Unexpected weight name while loading Ideogram text encoder: {name}"
                )
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
            weight_loader(param, loaded_weight.to(param.dtype))
            normal_weight_names.append(name)
            loaded_params.add(name)

        quant_states = build_bitsandbytes_4bit_quant_states(
            normal_weight_names,
            raw_quant_state,
            next(self.parameters()).device,
        )
        attach_bitsandbytes_4bit_quant_states(params_dict, quant_states)
        quantized_params_missing_state = [
            name
            for name, param in params_dict.items()
            if getattr(param, "use_bitsandbytes_4bit", False)
            and name not in quant_states
        ]
        if quantized_params_missing_state:
            raise ValueError(
                "Missing bitsandbytes quant_state for Ideogram text encoder weights: "
                f"{quantized_params_missing_state[:8]}"
            )
        return loaded_params


EntryClass = IdeogramQwen3VLTextEncoder