[Feature] Stronger transformers modeling backend with TP, PP, MoE, VLMs, and torch compile (#19163)
This commit is contained in:
@@ -158,7 +158,10 @@ class ModelConfig:
|
||||
"Llama4ForConditionalGeneration",
|
||||
"Step3VLForConditionalGeneration",
|
||||
]
|
||||
if self.hf_config.architectures[0] in mm_disabled_models:
|
||||
if (
|
||||
self.hf_config.architectures[0] in mm_disabled_models
|
||||
and self.model_impl != ModelImpl.TRANSFORMERS
|
||||
):
|
||||
enable_multimodal = False
|
||||
logger.info(
|
||||
f"Multimodal is disabled for {self.hf_config.model_type}. To enable it, set --enable-multimodal."
|
||||
@@ -177,8 +180,14 @@ class ModelConfig:
|
||||
self.is_generation = is_generation_model(
|
||||
self.hf_config.architectures, is_embedding
|
||||
)
|
||||
self.is_multimodal = enable_multimodal and is_multimodal_model(
|
||||
self.hf_config.architectures
|
||||
has_multimodal_subconfig = (
|
||||
self.hf_config is not self.hf_text_config
|
||||
or hasattr(self.hf_config, "vision_config")
|
||||
or hasattr(self.hf_config, "audio_config")
|
||||
)
|
||||
self.is_multimodal = enable_multimodal and (
|
||||
is_multimodal_model(self.hf_config.architectures)
|
||||
or has_multimodal_subconfig
|
||||
)
|
||||
self.is_audio_model = enable_multimodal and is_audio_model(
|
||||
self.hf_config.architectures
|
||||
|
||||
@@ -672,6 +672,11 @@ class MMReceiverBase(ABC):
|
||||
server_args,
|
||||
_processor,
|
||||
transport_mode,
|
||||
model_config=(
|
||||
getattr(self.scheduler, "model_config", None)
|
||||
if self.scheduler is not None
|
||||
else None
|
||||
),
|
||||
skip_mm_pool=not enable_adaptive_dispatch_to_encoder,
|
||||
)
|
||||
|
||||
|
||||
@@ -743,6 +743,8 @@ class TokenizedGenerateReqInput(BaseReq):
|
||||
# Whether to return entropy
|
||||
return_entropy: bool = False
|
||||
|
||||
token_type_ids: Optional[List[int]] = None
|
||||
|
||||
need_wait_for_mm_inputs: bool = False
|
||||
num_items_assigned: Optional[Dict[Modality, List[int]]] = None
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@ import inspect
|
||||
import logging
|
||||
import pkgutil
|
||||
|
||||
from sglang.srt.configs.model_config import ModelImpl
|
||||
from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
|
||||
@@ -41,14 +42,41 @@ def import_processors(package_name: str, overwrite: bool = False):
|
||||
|
||||
|
||||
def get_mm_processor(
|
||||
hf_config, server_args: ServerArgs, processor, transport_mode, **kwargs
|
||||
hf_config,
|
||||
server_args: ServerArgs,
|
||||
processor,
|
||||
transport_mode,
|
||||
model_config=None,
|
||||
**kwargs,
|
||||
) -> BaseMultimodalProcessor:
|
||||
model_impl = str(getattr(server_args, "model_impl", "auto")).lower()
|
||||
uses_transformers_backend = model_impl == "transformers"
|
||||
if model_impl == "auto" and model_config is not None:
|
||||
from sglang.srt.model_loader.utils import get_resolved_model_impl
|
||||
|
||||
uses_transformers_backend = (
|
||||
get_resolved_model_impl(model_config) == ModelImpl.TRANSFORMERS
|
||||
)
|
||||
|
||||
for model_cls, processor_cls in PROCESSOR_MAPPING.items():
|
||||
if model_cls.__name__ in hf_config.architectures:
|
||||
if model_cls.__name__ not in hf_config.architectures:
|
||||
continue
|
||||
if not uses_transformers_backend or getattr(
|
||||
processor_cls, "supports_transformers_backend", False
|
||||
):
|
||||
return processor_cls(
|
||||
hf_config, server_args, processor, transport_mode, **kwargs
|
||||
)
|
||||
|
||||
if uses_transformers_backend:
|
||||
from sglang.srt.multimodal.processors.transformers_auto import (
|
||||
TransformersAutoMultimodalProcessor,
|
||||
)
|
||||
|
||||
return TransformersAutoMultimodalProcessor(
|
||||
hf_config, server_args, processor, transport_mode, **kwargs
|
||||
)
|
||||
|
||||
raise ValueError(
|
||||
f"No processor registered for architecture: {hf_config.architectures}.\n"
|
||||
f"Registered architectures: {[model_cls.__name__ for model_cls in PROCESSOR_MAPPING.keys()]}"
|
||||
|
||||
@@ -38,7 +38,7 @@ from torch.cuda import Stream as CudaStream
|
||||
from torch.distributed import barrier
|
||||
|
||||
from sglang.jit_kernel.ngram_embedding import update_token_table
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.configs.model_config import ModelConfig, ModelImpl
|
||||
from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
|
||||
from sglang.srt.constrained.grammar_manager import GrammarManager
|
||||
from sglang.srt.disaggregation.decode import (
|
||||
@@ -185,6 +185,7 @@ from sglang.srt.mem_cache.common import release_kv_cache
|
||||
from sglang.srt.mem_cache.radix_cache import RadixCache
|
||||
from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardMode, PPProxyTensors
|
||||
from sglang.srt.model_loader.utils import get_resolved_model_impl
|
||||
from sglang.srt.multiplex.multiplexing_mixin import SchedulerMultiplexMixin
|
||||
from sglang.srt.observability.req_time_stats import (
|
||||
real_time,
|
||||
@@ -699,6 +700,9 @@ class Scheduler(
|
||||
|
||||
def init_cache_with_memory_pool(self):
|
||||
server_args = self.server_args
|
||||
uses_transformers_backend = (
|
||||
get_resolved_model_impl(self.model_config) == ModelImpl.TRANSFORMERS
|
||||
)
|
||||
|
||||
# Hybrid memory pool
|
||||
self.is_hybrid_swa = self.tp_worker.is_hybrid_swa
|
||||
@@ -718,9 +722,21 @@ class Scheduler(
|
||||
self.tp_worker.get_memory_pool()
|
||||
)
|
||||
|
||||
# Create cache
|
||||
self.disable_radix_cache = server_args.disable_radix_cache or (
|
||||
self.model_config.is_multimodal and uses_transformers_backend
|
||||
)
|
||||
if self.disable_radix_cache and not server_args.disable_radix_cache:
|
||||
logger.warning(
|
||||
"Radix cache is disabled for multimodal models with the "
|
||||
"Transformers backend to avoid multimodal prefix-cache mismatches."
|
||||
)
|
||||
|
||||
effective_chunked_prefill_size = server_args.chunked_prefill_size
|
||||
if self.model_config.is_multimodal and uses_transformers_backend:
|
||||
effective_chunked_prefill_size = None
|
||||
|
||||
params = CacheInitParams(
|
||||
disable=server_args.disable_radix_cache,
|
||||
disable=self.disable_radix_cache,
|
||||
req_to_token_pool=self.req_to_token_pool,
|
||||
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
|
||||
page_size=self.page_size,
|
||||
@@ -736,14 +752,11 @@ class Scheduler(
|
||||
enable_mamba_extra_buffer=server_args.enable_mamba_extra_buffer(),
|
||||
pp_rank=self.pp_rank,
|
||||
pp_size=self.pp_size,
|
||||
chunked_prefill_size=server_args.chunked_prefill_size,
|
||||
chunked_prefill_size=effective_chunked_prefill_size,
|
||||
sliding_window_size=self.sliding_window_size,
|
||||
)
|
||||
|
||||
if (
|
||||
server_args.chunked_prefill_size is not None
|
||||
and server_args.disable_radix_cache
|
||||
):
|
||||
if effective_chunked_prefill_size is not None and self.disable_radix_cache:
|
||||
if not self.is_hybrid_swa:
|
||||
from sglang.srt.mem_cache.chunk_cache import ChunkCache
|
||||
|
||||
@@ -844,9 +857,22 @@ class Scheduler(
|
||||
self._engine_paused = False
|
||||
|
||||
def init_chunked_prefill(self):
|
||||
# Init chunked prefill
|
||||
self.chunked_prefill_size = self.server_args.chunked_prefill_size
|
||||
if self.chunked_prefill_size <= 0: # -1 means disable
|
||||
uses_transformers_backend = (
|
||||
get_resolved_model_impl(self.model_config) == ModelImpl.TRANSFORMERS
|
||||
)
|
||||
if (
|
||||
self.chunked_prefill_size is not None
|
||||
and self.chunked_prefill_size > 0
|
||||
and self.model_config.is_multimodal
|
||||
and uses_transformers_backend
|
||||
):
|
||||
logger.warning(
|
||||
"Chunked prefill is disabled for multimodal models with the "
|
||||
"Transformers backend to avoid partial multimodal chunk mismatches."
|
||||
)
|
||||
self.chunked_prefill_size = None
|
||||
elif self.chunked_prefill_size is not None and self.chunked_prefill_size <= 0:
|
||||
self.chunked_prefill_size = None
|
||||
self.chunked_req = None
|
||||
self.is_mixed_chunk = (
|
||||
@@ -1724,6 +1750,7 @@ class Scheduler(
|
||||
stream=recv_req.stream,
|
||||
lora_id=recv_req.lora_id,
|
||||
input_embeds=recv_req.input_embeds,
|
||||
token_type_ids=recv_req.token_type_ids,
|
||||
custom_logit_processor=recv_req.custom_logit_processor,
|
||||
require_reasoning=recv_req.require_reasoning,
|
||||
return_hidden_states=recv_req.return_hidden_states,
|
||||
@@ -1806,10 +1833,12 @@ class Scheduler(
|
||||
SessionController.adjust_mm_offsets(recv_req, req, image_inputs)
|
||||
|
||||
# The following steps are already fast, execute locally on each rank.
|
||||
# Expand a single image token into multiple dummy tokens for receiving image embeddings
|
||||
req.origin_input_ids = self.pad_input_ids_func(
|
||||
req.origin_input_ids, image_inputs
|
||||
)
|
||||
# Expand a single image token into multiple dummy tokens for receiving image embeddings.
|
||||
# The pad function is model-specific and can be None for some backends.
|
||||
if self.pad_input_ids_func:
|
||||
req.origin_input_ids = self.pad_input_ids_func(
|
||||
req.origin_input_ids, image_inputs
|
||||
)
|
||||
req.extend_image_inputs(image_inputs)
|
||||
self._maybe_compute_mrope_positions(req)
|
||||
|
||||
|
||||
@@ -266,7 +266,11 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
||||
# We create mm_processor for any skip_tokenizer_init to make sure we still encode
|
||||
# images even with skip_tokenizer_init=False.
|
||||
self.mm_processor = get_mm_processor(
|
||||
self.model_config.hf_config, server_args, _processor, transport_mode
|
||||
self.model_config.hf_config,
|
||||
server_args,
|
||||
_processor,
|
||||
transport_mode,
|
||||
model_config=self.model_config,
|
||||
)
|
||||
|
||||
if server_args.skip_tokenizer_init:
|
||||
@@ -747,6 +751,10 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
||||
|
||||
if mm_inputs and "input_ids" in mm_inputs:
|
||||
input_ids = mm_inputs["input_ids"]
|
||||
if mm_inputs and "token_type_ids" in mm_inputs:
|
||||
token_type_ids = mm_inputs.pop("token_type_ids")
|
||||
if not isinstance(token_type_ids, list):
|
||||
token_type_ids = token_type_ids.flatten().tolist()
|
||||
if (
|
||||
envs.SGLANG_MM_PRECOMPUTE_HASH.get()
|
||||
and mm_inputs
|
||||
@@ -971,6 +979,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
||||
priority=obj.priority,
|
||||
extra_key=obj.extra_key,
|
||||
routing_key=obj.routing_key,
|
||||
token_type_ids=token_type_ids,
|
||||
need_wait_for_mm_inputs=obj.need_wait_for_mm_inputs,
|
||||
num_items_assigned=obj.num_items_assigned,
|
||||
)
|
||||
|
||||
@@ -2118,6 +2118,16 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
|
||||
if self.server_args.enable_torch_compile:
|
||||
set_torch_compile_config()
|
||||
should_disable_torch_compile = not getattr(
|
||||
self.model, "_can_torch_compile", True
|
||||
)
|
||||
if should_disable_torch_compile:
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
"Transformers backend model reports it is not torch.compile "
|
||||
"compatible (e.g. dynamic rope scaling). Disabling torch.compile.",
|
||||
)
|
||||
self.server_args.enable_torch_compile = False
|
||||
|
||||
if self.eagle_use_aux_hidden_state:
|
||||
self.model.set_eagle3_layers_to_capture()
|
||||
|
||||
@@ -27,9 +27,87 @@ def set_default_torch_dtype(dtype: torch.dtype):
|
||||
torch.set_default_dtype(old_dtype)
|
||||
|
||||
|
||||
def _is_moe_model(model_config: ModelConfig, architectures: list[str]) -> bool:
|
||||
lowered_arches = [arch.lower() for arch in architectures]
|
||||
if any("moe" in arch or "mixtral" in arch for arch in lowered_arches):
|
||||
return True
|
||||
|
||||
text_config = model_config.hf_text_config
|
||||
expert_attrs = (
|
||||
"num_local_experts",
|
||||
"num_experts",
|
||||
"num_experts_per_tok",
|
||||
"moe_intermediate_size",
|
||||
"n_routed_experts",
|
||||
)
|
||||
for attr in expert_attrs:
|
||||
value = getattr(text_config, attr, None)
|
||||
if value is None:
|
||||
continue
|
||||
if isinstance(value, bool):
|
||||
if value:
|
||||
return True
|
||||
continue
|
||||
if isinstance(value, (int, float)):
|
||||
threshold = 0 if attr == "moe_intermediate_size" else 1
|
||||
if value > threshold:
|
||||
return True
|
||||
continue
|
||||
if isinstance(value, (list, tuple, set, dict)):
|
||||
if len(value) > 0:
|
||||
return True
|
||||
continue
|
||||
if isinstance(value, str) and value == "":
|
||||
continue
|
||||
if value is not None:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _is_sequence_classification_model(architectures: list[str]) -> bool:
|
||||
return any(
|
||||
"sequenceclassification" in lowered or "rewardmodel" in lowered
|
||||
for lowered in (arch.lower() for arch in architectures)
|
||||
)
|
||||
|
||||
|
||||
def _get_transformers_backend_arch(
|
||||
model_config: ModelConfig, architectures: list[str]
|
||||
) -> str:
|
||||
is_pooling = not model_config.is_generation
|
||||
is_multimodal = model_config.is_multimodal or (
|
||||
model_config.hf_config is not model_config.hf_text_config
|
||||
)
|
||||
is_moe = _is_moe_model(model_config, architectures)
|
||||
base_arch = "ForCausalLM"
|
||||
if is_pooling:
|
||||
base_arch = (
|
||||
"ForSequenceClassification"
|
||||
if _is_sequence_classification_model(architectures)
|
||||
else "EmbeddingModel"
|
||||
)
|
||||
|
||||
arch = "Transformers"
|
||||
if is_multimodal:
|
||||
arch += "MultiModal"
|
||||
if is_moe:
|
||||
arch += "MoE"
|
||||
return arch + base_arch
|
||||
|
||||
|
||||
def _model_impl_from_architecture(architecture: str) -> ModelImpl:
|
||||
if architecture.startswith("Transformers"):
|
||||
return ModelImpl.TRANSFORMERS
|
||||
if architecture.startswith("MindSpore"):
|
||||
return ModelImpl.MINDSPORE
|
||||
return ModelImpl.SGLANG
|
||||
|
||||
|
||||
def resolve_transformers_arch(model_config: ModelConfig, architectures: list[str]):
|
||||
for i, arch in enumerate(architectures):
|
||||
if arch == "TransformersForCausalLM":
|
||||
backend_arch = _get_transformers_backend_arch(model_config, architectures)
|
||||
|
||||
for arch in architectures:
|
||||
if arch.startswith("Transformers"):
|
||||
continue
|
||||
auto_map: dict[str, str] = (
|
||||
getattr(model_config.hf_config, "auto_map", None) or dict()
|
||||
@@ -42,15 +120,33 @@ def resolve_transformers_arch(model_config: ModelConfig, architectures: list[str
|
||||
# "AutoModel": "<your-repo-name>--<config-name>",
|
||||
# "AutoModelFor<Task>": "<your-repo-name>--<config-name>",
|
||||
# },
|
||||
auto_modules = {
|
||||
name: get_class_from_dynamic_module(
|
||||
module, model_config.model_path, revision=model_config.revision
|
||||
auto_modules = {}
|
||||
try:
|
||||
auto_modules = {
|
||||
name: get_class_from_dynamic_module(
|
||||
module, model_config.model_path, revision=model_config.revision
|
||||
)
|
||||
for name, module in sorted(auto_map.items(), key=lambda x: x[0])
|
||||
}
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Failed to load dynamic modules from auto_map for '%s': %s. "
|
||||
"Skipping remote model compatibility checks.",
|
||||
arch,
|
||||
e,
|
||||
)
|
||||
for name, module in sorted(auto_map.items(), key=lambda x: x[0])
|
||||
}
|
||||
model_module = getattr(transformers, arch, None)
|
||||
if model_module is None:
|
||||
if "AutoModel" not in auto_map:
|
||||
has_auto_model = "AutoModel" in auto_modules
|
||||
if not has_auto_model and model_config.model_impl == ModelImpl.TRANSFORMERS:
|
||||
logger.warning(
|
||||
"Cannot resolve model class for '%s' and no auto_map.AutoModel "
|
||||
"is present. Skipping compatibility gate because "
|
||||
"--model-impl=transformers is explicitly requested.",
|
||||
arch,
|
||||
)
|
||||
continue
|
||||
if not has_auto_model and "AutoModel" not in auto_map:
|
||||
raise ValueError(
|
||||
f"Cannot find model module. '{arch}' is not a registered "
|
||||
"model in the Transformers library (only relevant if the "
|
||||
@@ -58,16 +154,25 @@ def resolve_transformers_arch(model_config: ModelConfig, architectures: list[str
|
||||
"not present in the model config's 'auto_map' (relevant "
|
||||
"if the model is custom)."
|
||||
)
|
||||
if not has_auto_model:
|
||||
raise ValueError(
|
||||
f"Cannot find model module. '{arch}' is not a registered "
|
||||
"model in the Transformers library and loading the custom "
|
||||
f"model from auto_map failed. The remote model code may be "
|
||||
f"incompatible with the installed transformers version."
|
||||
)
|
||||
model_module = auto_modules["AutoModel"]
|
||||
if model_config.model_impl == ModelImpl.TRANSFORMERS:
|
||||
if hasattr(model_module, "is_backend_compatible") and (
|
||||
not model_module.is_backend_compatible()
|
||||
):
|
||||
raise ValueError(
|
||||
f"The Transformers implementation of {arch} is not "
|
||||
"compatible with SGLang."
|
||||
logger.warning(
|
||||
"The Transformers implementation of %s reports it is not "
|
||||
"backend-compatible (_supports_attention_backend=False). "
|
||||
"Proceeding anyway because --model-impl=transformers was "
|
||||
"explicitly requested. The model may not work correctly.",
|
||||
arch,
|
||||
)
|
||||
architectures[i] = "TransformersForCausalLM"
|
||||
if model_config.model_impl == ModelImpl.AUTO:
|
||||
if hasattr(model_module, "is_backend_compatible") and (
|
||||
not model_module.is_backend_compatible()
|
||||
@@ -82,8 +187,7 @@ def resolve_transformers_arch(model_config: ModelConfig, architectures: list[str
|
||||
"performance may not be optimal.",
|
||||
arch,
|
||||
)
|
||||
architectures[i] = "TransformersForCausalLM"
|
||||
return architectures
|
||||
return [backend_arch]
|
||||
|
||||
|
||||
def get_model_architecture(model_config: ModelConfig) -> Tuple[Type[nn.Module], str]:
|
||||
@@ -114,7 +218,29 @@ def get_model_architecture(model_config: ModelConfig) -> Tuple[Type[nn.Module],
|
||||
architectures = ["MindSporeForCausalLM"]
|
||||
elif not is_native_supported or model_config.model_impl == ModelImpl.TRANSFORMERS:
|
||||
architectures = resolve_transformers_arch(model_config, architectures)
|
||||
return ModelRegistry.resolve_model_cls(architectures)
|
||||
model_cls, resolved_arch = ModelRegistry.resolve_model_cls(architectures)
|
||||
setattr(model_config, "_resolved_model_arch", resolved_arch)
|
||||
setattr(
|
||||
model_config,
|
||||
"_resolved_model_impl",
|
||||
_model_impl_from_architecture(resolved_arch),
|
||||
)
|
||||
return model_cls, resolved_arch
|
||||
|
||||
|
||||
def get_resolved_model_impl(model_config: ModelConfig) -> ModelImpl:
|
||||
resolved_model_impl = getattr(model_config, "_resolved_model_impl", None)
|
||||
if resolved_model_impl is not None:
|
||||
return resolved_model_impl
|
||||
|
||||
resolved_arch = getattr(model_config, "_resolved_model_arch", None)
|
||||
if resolved_arch is None:
|
||||
_, resolved_arch = get_model_architecture(model_config)
|
||||
|
||||
resolved_model_impl = _model_impl_from_architecture(resolved_arch)
|
||||
setattr(model_config, "_resolved_model_arch", resolved_arch)
|
||||
setattr(model_config, "_resolved_model_impl", resolved_model_impl)
|
||||
return resolved_model_impl
|
||||
|
||||
|
||||
def get_architecture_class_name(model_config: ModelConfig) -> str:
|
||||
|
||||
@@ -269,6 +269,7 @@ class Qwen2Model(nn.Module):
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.padding_idx = getattr(config, "pad_token_id", None)
|
||||
self.vocab_size = config.vocab_size
|
||||
self.pp_group = get_pp_group()
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -13,6 +13,7 @@
|
||||
# ==============================================================================
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
from collections.abc import Iterable, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from functools import lru_cache
|
||||
@@ -28,6 +29,7 @@ from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled
|
||||
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
|
||||
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.model_loader.weight_utils import default_weight_loader
|
||||
from sglang.srt.utils import get_current_device_stream_fast, is_cuda, is_hip
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
@@ -49,6 +51,13 @@ class WeightsMapper:
|
||||
orig_to_new_prefix: WeightsMapping = field(default_factory=dict)
|
||||
orig_to_new_suffix: WeightsMapping = field(default_factory=dict)
|
||||
|
||||
def __or__(self, other: "WeightsMapper") -> "WeightsMapper":
|
||||
return WeightsMapper(
|
||||
orig_to_new_substr={**self.orig_to_new_substr, **other.orig_to_new_substr},
|
||||
orig_to_new_prefix={**self.orig_to_new_prefix, **other.orig_to_new_prefix},
|
||||
orig_to_new_suffix={**self.orig_to_new_suffix, **other.orig_to_new_suffix},
|
||||
)
|
||||
|
||||
def _map_name(self, key: str) -> Optional[str]:
|
||||
for substr, new_key in sorted(
|
||||
self.orig_to_new_substr.items(), key=lambda i: len(i[0]), reverse=True
|
||||
@@ -106,6 +115,161 @@ class WeightsMapper:
|
||||
}
|
||||
|
||||
|
||||
class AutoWeightsLoader:
|
||||
ROTARY_EMBEDS_UNUSED_WEIGHTS = [
|
||||
"rotary_pos_emb.inv_freq",
|
||||
"rotary_emb.inv_freq",
|
||||
"rotary_emb.cos_cached",
|
||||
"rotary_emb.sin_cached",
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
module: torch.nn.Module,
|
||||
*,
|
||||
skip_prefixes: list[str] | None = None,
|
||||
skip_substrs: list[str] | None = None,
|
||||
ignore_unexpected_prefixes: list[str] | None = None,
|
||||
ignore_unexpected_suffixes: list[str] | None = None,
|
||||
) -> None:
|
||||
self.module = module
|
||||
self.skip_prefixes = list(skip_prefixes or [])
|
||||
self.skip_substrs = [
|
||||
*(skip_substrs or []),
|
||||
*self.ROTARY_EMBEDS_UNUSED_WEIGHTS,
|
||||
]
|
||||
self.ignore_unexpected_prefixes = list(ignore_unexpected_prefixes or [])
|
||||
self.ignore_unexpected_suffixes = list(ignore_unexpected_suffixes or [])
|
||||
|
||||
def _groupby_prefix(
|
||||
self,
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
) -> Iterable[tuple[str, Iterable[tuple[str, torch.Tensor]]]]:
|
||||
weights_by_parts = (
|
||||
(weight_name.split(".", 1), weight_data)
|
||||
for weight_name, weight_data in weights
|
||||
)
|
||||
for prefix, group in itertools.groupby(weights_by_parts, key=lambda x: x[0][0]):
|
||||
yield prefix, (
|
||||
("" if len(parts) == 1 else parts[1], weight_data)
|
||||
for parts, weight_data in group
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_qualname(prefix: str, rest: str) -> str:
|
||||
if prefix == "":
|
||||
return rest
|
||||
if rest == "":
|
||||
return prefix
|
||||
return f"{prefix}.{rest}"
|
||||
|
||||
def _can_skip(self, qualname: str) -> bool:
|
||||
return any(qualname.startswith(p) for p in self.skip_prefixes) or any(
|
||||
sub in qualname for sub in self.skip_substrs
|
||||
)
|
||||
|
||||
def _can_ignore_unexpected(self, qualname: str) -> bool:
|
||||
return any(
|
||||
qualname.startswith(p) for p in self.ignore_unexpected_prefixes
|
||||
) or any(qualname.endswith(s) for s in self.ignore_unexpected_suffixes)
|
||||
|
||||
def _load_param(
|
||||
self,
|
||||
base_prefix: str,
|
||||
param: torch.nn.Parameter,
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
) -> Iterable[str]:
|
||||
for weight_name, weight_data in weights:
|
||||
weight_qualname = self._get_qualname(base_prefix, weight_name)
|
||||
if self._can_skip(weight_qualname):
|
||||
continue
|
||||
if weight_name != "":
|
||||
if self._can_ignore_unexpected(weight_qualname):
|
||||
continue
|
||||
raise ValueError(
|
||||
f"Attempted to load nested weight {weight_qualname!r} "
|
||||
f"into parameter {base_prefix!r}"
|
||||
)
|
||||
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, weight_data)
|
||||
yield weight_qualname
|
||||
|
||||
def _load_module(
|
||||
self,
|
||||
base_prefix: str,
|
||||
module: torch.nn.Module,
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
) -> Iterable[str]:
|
||||
if module.__class__.__name__ == "PPMissingLayer":
|
||||
return
|
||||
|
||||
if module is not self.module:
|
||||
module_load_weights = getattr(module, "load_weights", None)
|
||||
if callable(module_load_weights):
|
||||
loaded = module_load_weights(weights)
|
||||
if loaded is not None:
|
||||
yield from (
|
||||
self._get_qualname(base_prefix, loaded_name)
|
||||
for loaded_name in loaded
|
||||
)
|
||||
return
|
||||
|
||||
child_modules = dict(module.named_children())
|
||||
child_params = dict(module.named_parameters(recurse=False))
|
||||
child_buffers = dict(module.named_buffers(recurse=False))
|
||||
for child_prefix, child_weights in self._groupby_prefix(weights):
|
||||
prefix = self._get_qualname(base_prefix, child_prefix)
|
||||
if child_prefix in child_modules:
|
||||
if self._can_skip(prefix + "."):
|
||||
continue
|
||||
yield from self._load_module(
|
||||
prefix,
|
||||
child_modules[child_prefix],
|
||||
child_weights,
|
||||
)
|
||||
continue
|
||||
|
||||
if child_prefix in child_params:
|
||||
if self._can_skip(prefix):
|
||||
continue
|
||||
yield from self._load_param(
|
||||
prefix, child_params[child_prefix], child_weights
|
||||
)
|
||||
continue
|
||||
|
||||
if child_prefix in child_buffers:
|
||||
if self._can_skip(prefix):
|
||||
continue
|
||||
yield from self._load_param(
|
||||
prefix, child_buffers[child_prefix], child_weights
|
||||
)
|
||||
continue
|
||||
|
||||
if self._can_skip(prefix) or self._can_skip(prefix + "."):
|
||||
continue
|
||||
if self._can_ignore_unexpected(prefix) or self._can_ignore_unexpected(
|
||||
prefix + "."
|
||||
):
|
||||
continue
|
||||
raise ValueError(
|
||||
f"No module or parameter named {prefix!r} in {self.module._get_name()}."
|
||||
)
|
||||
|
||||
def load_weights(
|
||||
self,
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
*,
|
||||
mapper: WeightsMapper | None = None,
|
||||
) -> set[str]:
|
||||
if mapper is not None:
|
||||
weights = mapper.apply(weights)
|
||||
weights = (
|
||||
(name, weight) for name, weight in weights if not self._can_skip(name)
|
||||
)
|
||||
return set(self._load_module("", self.module, weights))
|
||||
|
||||
|
||||
def enable_fused_set_kv_buffer(forward_batch: ForwardBatch):
|
||||
"""Enable fused set_kv_buffer only on CUDA with bfloat16 KV cache."""
|
||||
return (
|
||||
|
||||
@@ -234,6 +234,7 @@ async def preprocess_video(
|
||||
|
||||
# Compatible with Qwen-VL & Qwen-Omni Series
|
||||
class QwenVLImageProcessor(SGLangBaseProcessor):
|
||||
supports_transformers_backend = True
|
||||
models = [
|
||||
Qwen2VLForConditionalGeneration,
|
||||
Qwen2_5_VLForConditionalGeneration,
|
||||
|
||||
@@ -0,0 +1,215 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
|
||||
from sglang.srt.multimodal.processors.base_processor import (
|
||||
BaseMultimodalProcessor,
|
||||
MultimodalSpecialTokens,
|
||||
)
|
||||
from sglang.srt.utils import load_image
|
||||
|
||||
|
||||
def _first_attr(obj, names: tuple[str, ...], default=None):
|
||||
for name in names:
|
||||
value = getattr(obj, name, None)
|
||||
if value is not None:
|
||||
return value
|
||||
return default
|
||||
|
||||
|
||||
def _uses_mrope(hf_config) -> bool:
|
||||
text_config = getattr(hf_config, "text_config", hf_config)
|
||||
rope_scaling = getattr(text_config, "rope_scaling", None) or {}
|
||||
if isinstance(rope_scaling, dict) and "mrope_section" in rope_scaling:
|
||||
return True
|
||||
rope_type = str(getattr(text_config, "rope_type", "")).lower()
|
||||
return "mrope" in rope_type
|
||||
|
||||
|
||||
class TransformersAutoMultimodalProcessor(BaseMultimodalProcessor):
|
||||
"""Generic multimodal processor for the Transformers backend.
|
||||
|
||||
Unlike model-specific processors that rely on regex-based token matching
|
||||
in the raw prompt, this processor applies the HF processor directly to
|
||||
the prompt text + raw media. This handles models like Gemma3 where the
|
||||
chat template uses a marker (``<start_of_image>``) that the HF processor
|
||||
internally expands into placeholder tokens.
|
||||
"""
|
||||
|
||||
models = []
|
||||
|
||||
def __init__(self, hf_config, server_args, _processor, *args, **kwargs):
|
||||
super().__init__(hf_config, server_args, _processor, *args, **kwargs)
|
||||
self.mm_tokens = MultimodalSpecialTokens(
|
||||
image_token=getattr(_processor, "image_token", None),
|
||||
video_token=getattr(_processor, "video_token", None),
|
||||
audio_token=getattr(_processor, "audio_token", None),
|
||||
image_token_id=_first_attr(
|
||||
hf_config,
|
||||
("image_token_id", "image_token_index", "im_token_id"),
|
||||
),
|
||||
video_token_id=_first_attr(
|
||||
hf_config,
|
||||
("video_token_id",),
|
||||
),
|
||||
audio_token_id=_first_attr(
|
||||
hf_config,
|
||||
("audio_token_id",),
|
||||
),
|
||||
).build(_processor)
|
||||
|
||||
self._is_mrope = _uses_mrope(hf_config)
|
||||
if self._is_mrope:
|
||||
vision_config = getattr(hf_config, "vision_config", None)
|
||||
self._spatial_merge_size = getattr(vision_config, "spatial_merge_size", 2)
|
||||
self._tokens_per_second = getattr(vision_config, "tokens_per_second", None)
|
||||
self._vision_start_token_id = _first_attr(
|
||||
hf_config, ("vision_start_token_id",)
|
||||
)
|
||||
self._model_type = getattr(hf_config, "model_type", "")
|
||||
|
||||
def _compute_mrope_positions(
|
||||
self,
|
||||
input_ids: list[int],
|
||||
image_grid_thw: Optional[torch.Tensor] = None,
|
||||
video_grid_thw: Optional[torch.Tensor] = None,
|
||||
):
|
||||
from sglang.srt.layers.rotary_embedding import MRotaryEmbedding
|
||||
|
||||
input_ids_tensor = torch.tensor(input_ids, dtype=torch.long).unsqueeze(0)
|
||||
mrope_positions, mrope_position_delta = MRotaryEmbedding.get_rope_index(
|
||||
spatial_merge_size=self._spatial_merge_size,
|
||||
image_token_id=self.mm_tokens.image_token_id,
|
||||
video_token_id=self.mm_tokens.video_token_id or -1,
|
||||
vision_start_token_id=self._vision_start_token_id,
|
||||
model_type=self._model_type,
|
||||
input_ids=input_ids_tensor,
|
||||
image_grid_thw=image_grid_thw,
|
||||
video_grid_thw=video_grid_thw,
|
||||
tokens_per_second=self._tokens_per_second,
|
||||
)
|
||||
return mrope_positions.squeeze(1), mrope_position_delta
|
||||
|
||||
def _load_images(self, image_data) -> list:
|
||||
"""Download / decode images from URLs, file paths, or base64."""
|
||||
if not image_data:
|
||||
return []
|
||||
images = []
|
||||
for data in image_data:
|
||||
img, _ = load_image(data)
|
||||
if img.mode != "RGB":
|
||||
img = img.convert("RGB")
|
||||
images.append(img)
|
||||
return images
|
||||
|
||||
def _apply_hf_processor(self, text: str, images=None, videos=None):
|
||||
"""Run the HF processor on text + media and return the full output.
|
||||
|
||||
This is the key method that makes the generic processor work for
|
||||
models with non-trivial token expansion (Gemma3, PaliGemma, etc.).
|
||||
The HF processor handles chat-template expansion, image token
|
||||
insertion, and tokenization in one shot.
|
||||
"""
|
||||
kwargs = {}
|
||||
if images:
|
||||
kwargs["images"] = images
|
||||
if videos:
|
||||
kwargs["videos"] = videos
|
||||
return self._processor(text=text, return_tensors="pt", **kwargs)
|
||||
|
||||
def _build_mm_items(
|
||||
self, processor_output: dict, input_ids: torch.Tensor
|
||||
) -> list[MultimodalDataItem]:
|
||||
"""Extract MultimodalDataItem objects from the HF processor output."""
|
||||
items = self.collect_mm_items_from_processor_output(processor_output)
|
||||
|
||||
modality_to_token_id = {
|
||||
Modality.IMAGE: self.mm_tokens.image_token_id,
|
||||
Modality.MULTI_IMAGES: self.mm_tokens.image_token_id,
|
||||
Modality.VIDEO: self.mm_tokens.video_token_id,
|
||||
Modality.AUDIO: self.mm_tokens.audio_token_id,
|
||||
}
|
||||
|
||||
for item in items:
|
||||
token_id = modality_to_token_id.get(item.modality)
|
||||
if token_id is not None:
|
||||
item.offsets = self.get_mm_items_offset(input_ids, token_id)
|
||||
|
||||
return items
|
||||
|
||||
async def process_mm_data_async(
|
||||
self,
|
||||
image_data,
|
||||
audio_data,
|
||||
input_text,
|
||||
request_obj,
|
||||
**kwargs,
|
||||
):
|
||||
video_data = getattr(request_obj, "video_data", None)
|
||||
if video_data is not None and not isinstance(video_data, list):
|
||||
video_data = [video_data]
|
||||
|
||||
# Load raw media
|
||||
images = self._load_images(image_data)
|
||||
# TODO: video / audio loading when needed
|
||||
|
||||
# Apply HF processor — handles token expansion internally
|
||||
processor_output = self._apply_hf_processor(
|
||||
text=input_text,
|
||||
images=images or None,
|
||||
videos=video_data or None,
|
||||
)
|
||||
|
||||
input_ids = processor_output["input_ids"].flatten()
|
||||
|
||||
# Build mm_items from processor output
|
||||
mm_items = self._build_mm_items(processor_output, input_ids)
|
||||
|
||||
ret = {
|
||||
"input_ids": input_ids.tolist(),
|
||||
"mm_items": mm_items,
|
||||
}
|
||||
|
||||
# Propagate token_type_ids for models that need it (Gemma3, PaliGemma)
|
||||
token_type_key = (
|
||||
"mm_token_type_ids"
|
||||
if "mm_token_type_ids" in processor_output
|
||||
else "token_type_ids"
|
||||
)
|
||||
if token_type_key in processor_output:
|
||||
ret["token_type_ids"] = processor_output[token_type_key].flatten().tolist()
|
||||
|
||||
if self.mm_tokens.image_token_id is not None:
|
||||
ret["im_token_id"] = self.mm_tokens.image_token_id
|
||||
if self.mm_tokens.video_token_id is not None:
|
||||
ret["video_token_id"] = self.mm_tokens.video_token_id
|
||||
if self.mm_tokens.audio_token_id is not None:
|
||||
ret["audio_token_id"] = self.mm_tokens.audio_token_id
|
||||
|
||||
image_start_id = _first_attr(
|
||||
self.hf_config,
|
||||
("image_start_token_id", "vision_start_token_id", "im_start_id"),
|
||||
)
|
||||
image_end_id = _first_attr(
|
||||
self.hf_config,
|
||||
("image_end_token_id", "vision_end_token_id", "im_end_id"),
|
||||
)
|
||||
if image_start_id is not None:
|
||||
ret["im_start_id"] = image_start_id
|
||||
if image_end_id is not None:
|
||||
ret["im_end_id"] = image_end_id
|
||||
|
||||
# M-RoPE positions (Qwen2.5-VL, Qwen3-VL)
|
||||
if self._is_mrope:
|
||||
image_grid_thw = processor_output.get("image_grid_thw")
|
||||
video_grid_thw = processor_output.get("video_grid_thw")
|
||||
mrope_positions, mrope_position_delta = self._compute_mrope_positions(
|
||||
ret["input_ids"],
|
||||
image_grid_thw=image_grid_thw,
|
||||
video_grid_thw=video_grid_thw,
|
||||
)
|
||||
ret["mrope_positions"] = mrope_positions
|
||||
ret["mrope_position_delta"] = mrope_position_delta
|
||||
|
||||
return ret
|
||||
Reference in New Issue
Block a user