338 lines
13 KiB
Python
338 lines
13 KiB
Python
# Copyright 2023-2024 SGLang Team
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Config loading utilities."""
|
|
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from transformers import PretrainedConfig
|
|
from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
|
|
|
|
from sglang.srt.configs.model_config_parser_registry import (
|
|
ModelConfigParserBase,
|
|
get_model_config_parser,
|
|
register_model_config_parser,
|
|
)
|
|
from sglang.srt.connector import create_remote_connector
|
|
from sglang.srt.utils import is_remote_url, lru_cache_frozenset
|
|
|
|
from ..hf_transformers_patches import _ensure_gguf_version
|
|
from .common import (
|
|
_CONFIG_REGISTRY,
|
|
AutoConfig,
|
|
DeepseekVLV2Config,
|
|
_is_deepseek_ocr2_model,
|
|
_is_deepseek_ocr_model,
|
|
_override_v_head_dim_if_zero,
|
|
check_gguf_file,
|
|
get_hf_text_config,
|
|
gguf_sidecar_dir,
|
|
resolve_runai_obj_uri,
|
|
)
|
|
from .gguf_native import build_gguf_config, has_native_gguf_support
|
|
from .mistral_utils import is_mistral_model, load_mistral_config
|
|
|
|
|
|
def _set_architectures(config, arch_name):
|
|
config.update({"architectures": [arch_name]})
|
|
|
|
|
|
def _apply_deepseek_ocr_overrides(config, model):
|
|
_override_v_head_dim_if_zero(config)
|
|
_set_architectures(config, "DeepseekOCRForCausalLM")
|
|
config._name_or_path = model
|
|
|
|
|
|
_LONGCAT_ARCHS = {
|
|
"LongcatCausalLM",
|
|
"LongcatFlashForCausalLM",
|
|
"LongcatFlashNgramForCausalLM",
|
|
}
|
|
|
|
|
|
def _try_load_longcat_config(model, revision: Optional[str], **kwargs):
|
|
config_dict, _ = PretrainedConfig.get_config_dict(
|
|
model, revision=revision, **kwargs
|
|
)
|
|
architectures = config_dict.get("architectures") or []
|
|
if not any(arch in _LONGCAT_ARCHS for arch in architectures):
|
|
return None
|
|
|
|
return _CONFIG_REGISTRY["longcat_flash"].from_pretrained(
|
|
model, revision=revision, **kwargs
|
|
)
|
|
|
|
|
|
def _try_load_raw_mamba_config(model, revision: Optional[str], **kwargs):
|
|
"""Recognize the original state-spaces Mamba-1 checkpoints.
|
|
|
|
The raw `state-spaces/mamba-*` repos (e.g. mamba-130m/790m/2.8b, as opposed
|
|
to the `-hf` conversions) ship a minimal `config.json` with `d_model` /
|
|
`n_layer` / `ssm_cfg` and NO `model_type` / `architectures`, so
|
|
`AutoConfig.from_pretrained` rejects them with "Unrecognized model ...".
|
|
Detect that shape and build our `MambaConfig` (model_type `mamba`, arch
|
|
`MambaForCausalLM`) with the field-name mapping the SGLang Mamba model
|
|
expects. Uses `get_config_dict` (which does not require a model_type) so
|
|
this runs before the failing `AutoConfig` path.
|
|
"""
|
|
config_dict, _ = PretrainedConfig.get_config_dict(
|
|
model, revision=revision, **kwargs
|
|
)
|
|
# Raw state-spaces Mamba: has d_model + ssm_cfg, and no model_type/arch.
|
|
if config_dict.get("model_type") or config_dict.get("architectures"):
|
|
return None
|
|
if "d_model" not in config_dict or "ssm_cfg" not in config_dict:
|
|
return None
|
|
|
|
from sglang.srt.configs.mamba import MambaConfig
|
|
|
|
d_model = config_dict["d_model"]
|
|
# The embedding is padded up to a multiple of pad_vocab_size_multiple; match
|
|
# the checkpoint (e.g. 50277 -> 50280) so weight shapes line up.
|
|
pad = config_dict.get("pad_vocab_size_multiple", 1)
|
|
vocab_size = config_dict.get("vocab_size", 50280)
|
|
if pad > 1:
|
|
vocab_size = ((vocab_size + pad - 1) // pad) * pad
|
|
return MambaConfig(
|
|
vocab_size=vocab_size,
|
|
hidden_size=d_model,
|
|
num_hidden_layers=config_dict["n_layer"],
|
|
state_size=config_dict.get("ssm_cfg", {}).get("d_state", 16),
|
|
layer_norm_epsilon=config_dict.get("layer_norm_epsilon", 1e-5),
|
|
residual_in_fp32=config_dict.get("residual_in_fp32", True),
|
|
architectures=["MambaForCausalLM"],
|
|
)
|
|
|
|
|
|
@register_model_config_parser("hf")
|
|
class HfModelConfigParser(ModelConfigParserBase):
|
|
def parse(
|
|
self,
|
|
model,
|
|
trust_remote_code: bool,
|
|
revision: Optional[str] = None,
|
|
**kwargs,
|
|
):
|
|
config = _try_load_longcat_config(model, revision, **kwargs)
|
|
if config is None:
|
|
config = _try_load_raw_mamba_config(model, revision, **kwargs)
|
|
if config is None:
|
|
config = AutoConfig.from_pretrained(
|
|
model,
|
|
trust_remote_code=trust_remote_code,
|
|
revision=revision,
|
|
**kwargs,
|
|
)
|
|
|
|
if (
|
|
config.architectures is not None
|
|
and config.architectures[0] == "Phi4MMForCausalLM"
|
|
):
|
|
from transformers import SiglipVisionConfig
|
|
|
|
config.vision_config = SiglipVisionConfig(
|
|
hidden_size=1152,
|
|
image_size=448,
|
|
intermediate_size=4304,
|
|
model_type="siglip_vision_model",
|
|
num_attention_heads=16,
|
|
num_hidden_layers=26,
|
|
patch_size=14,
|
|
)
|
|
|
|
if config.architectures in [
|
|
["LongcatCausalLM"],
|
|
["LongcatFlashForCausalLM"],
|
|
["LongcatFlashNgramForCausalLM"],
|
|
]:
|
|
config.model_type = "longcat_flash"
|
|
|
|
text_config = get_hf_text_config(config=config)
|
|
|
|
if isinstance(model, str) and text_config is not None:
|
|
items = (
|
|
text_config.items()
|
|
if hasattr(text_config, "items")
|
|
else vars(text_config).items()
|
|
)
|
|
for key, val in items:
|
|
if not hasattr(config, key) and val is not None:
|
|
setattr(config, key, val)
|
|
|
|
is_ocr = _is_deepseek_ocr_model(config)
|
|
is_ocr2 = _is_deepseek_ocr2_model(config)
|
|
|
|
if is_ocr2:
|
|
_override_v_head_dim_if_zero(config)
|
|
config.model_type = "deepseek-ocr"
|
|
_set_architectures(config, "DeepseekOCRForCausalLM")
|
|
config = DeepseekVLV2Config.from_pretrained(model, revision=revision)
|
|
_apply_deepseek_ocr_overrides(config, model)
|
|
elif config.model_type in _CONFIG_REGISTRY:
|
|
model_type = config.model_type
|
|
if model_type == "deepseek_vl_v2" and is_ocr:
|
|
model_type = "deepseek-ocr"
|
|
# Raw state-spaces Mamba configs are built by
|
|
# _try_load_raw_mamba_config with architectures injected; reloading
|
|
# from the checkpoint would drop them, so skip it when the config is
|
|
# already one of our classes.
|
|
from sglang.srt.configs.mamba import FalconMambaConfig, MambaConfig
|
|
from sglang.srt.configs.mamba2 import Mamba2Config
|
|
|
|
if not isinstance(config, (Mamba2Config, MambaConfig, FalconMambaConfig)):
|
|
config = _CONFIG_REGISTRY[model_type].from_pretrained(
|
|
model, revision=revision
|
|
)
|
|
|
|
# Re-check after reloading config from registry
|
|
if _is_deepseek_ocr_model(config) or _is_deepseek_ocr2_model(config):
|
|
_apply_deepseek_ocr_overrides(config, model)
|
|
else:
|
|
config._name_or_path = model
|
|
|
|
if isinstance(model, str) and config.model_type == "internvl_chat":
|
|
for key, val in config.llm_config.__dict__.items():
|
|
if not hasattr(config, key):
|
|
setattr(config, key, val)
|
|
|
|
if config.model_type == "multi_modality":
|
|
_set_architectures(config, "MultiModalityCausalLM")
|
|
|
|
if config.model_type in (
|
|
"gemma4",
|
|
"gemma4_assistant",
|
|
"gemma4_unified",
|
|
"gemma4_unified_assistant",
|
|
):
|
|
# Gemma4 configs use base attributes for SWA layers and `global_*`
|
|
# variants for full-attention layers. SGLang expects the opposite:
|
|
# base = full-attention, `swa_*` = sliding-window overrides.
|
|
text_config = config.text_config
|
|
global_head_dim = getattr(text_config, "global_head_dim", None)
|
|
global_kv_heads = getattr(text_config, "num_global_key_value_heads", None)
|
|
|
|
swa_head_dim = text_config.head_dim
|
|
swa_kv_heads = text_config.num_key_value_heads
|
|
|
|
text_config.swa_head_dim = swa_head_dim
|
|
text_config.swa_v_head_dim = swa_head_dim
|
|
text_config.swa_num_key_value_heads = swa_kv_heads
|
|
|
|
if global_head_dim is not None:
|
|
text_config.head_dim = global_head_dim
|
|
if global_kv_heads is not None:
|
|
text_config.num_key_value_heads = global_kv_heads
|
|
|
|
if not hasattr(text_config, "v_head_dim"):
|
|
text_config.v_head_dim = text_config.head_dim
|
|
if not hasattr(text_config, "swa_v_head_dim"):
|
|
text_config.swa_v_head_dim = text_config.swa_head_dim
|
|
|
|
# Unified Gemma4 names the end-of-audio token `eoa_token_index`,
|
|
# but the multimodal processor expects `eoa_token_id`.
|
|
if not hasattr(config, "eoa_token_id") and hasattr(
|
|
config, "eoa_token_index"
|
|
):
|
|
config.eoa_token_id = config.eoa_token_index
|
|
|
|
if config.model_type == "longcat_flash":
|
|
_set_architectures(config, "LongcatFlashForCausalLM")
|
|
|
|
return config
|
|
|
|
|
|
@register_model_config_parser("mistral")
|
|
class MistralModelConfigParser(ModelConfigParserBase):
|
|
def parse(
|
|
self,
|
|
model,
|
|
trust_remote_code: bool,
|
|
revision: Optional[str] = None,
|
|
**kwargs,
|
|
):
|
|
del kwargs
|
|
return load_mistral_config(
|
|
model, trust_remote_code=trust_remote_code, revision=revision
|
|
)
|
|
|
|
|
|
@lru_cache_frozenset(maxsize=32)
|
|
def get_config(
|
|
model: str,
|
|
trust_remote_code: bool,
|
|
revision: Optional[str] = None,
|
|
model_override_args: Optional[dict] = None,
|
|
model_config_parser: str = "auto",
|
|
**kwargs,
|
|
):
|
|
is_gguf = check_gguf_file(model)
|
|
gguf_has_sidecar_config = False
|
|
if is_gguf:
|
|
if model_config_parser not in ("auto", "hf"):
|
|
raise ValueError(
|
|
f"model_config_parser={model_config_parser!r} is incompatible "
|
|
"with GGUF inputs; only 'hf' (or 'auto') is supported."
|
|
)
|
|
_ensure_gguf_version()
|
|
gguf_has_sidecar_config = gguf_sidecar_dir(model, "config.json") is not None
|
|
if not gguf_has_sidecar_config and has_native_gguf_support(model):
|
|
config = build_gguf_config(model)
|
|
if model_override_args:
|
|
config.update(model_override_args)
|
|
return config
|
|
if not gguf_has_sidecar_config:
|
|
kwargs["gguf_file"] = model
|
|
model = Path(model).parent
|
|
# Skip auto-resolution for GGUF: the name-based Mistral heuristic
|
|
# would misfire on the rewritten parent dir.
|
|
model_config_parser = "hf"
|
|
|
|
model = resolve_runai_obj_uri(model)
|
|
|
|
if is_remote_url(model):
|
|
client = create_remote_connector(model)
|
|
client.pull_files(ignore_pattern=["*.pt", "*.safetensors", "*.bin"])
|
|
model = client.get_local_dir()
|
|
|
|
if model_config_parser == "auto":
|
|
# `model` is post-rewrite (gguf parent / runai uri / remote pull).
|
|
model_config_parser = "mistral" if is_mistral_model(model) else "hf"
|
|
|
|
parser = get_model_config_parser(model_config_parser)
|
|
config = parser.parse(
|
|
model, trust_remote_code=trust_remote_code, revision=revision, **kwargs
|
|
)
|
|
|
|
if model_override_args:
|
|
# A plain update() setattrs a dict-valued override straight onto the
|
|
# config, so '{"text_config": {...}}' on a VLM would replace the whole
|
|
# sub-config with a dict and break attribute access downstream.
|
|
for key, value in model_override_args.items():
|
|
current = getattr(config, key, None)
|
|
if isinstance(value, dict) and isinstance(current, PretrainedConfig):
|
|
current.update(value)
|
|
else:
|
|
setattr(config, key, value)
|
|
|
|
if is_gguf and not gguf_has_sidecar_config:
|
|
if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
|
|
raise RuntimeError(
|
|
f"Can't get gguf config for {config.model_type}. Place a "
|
|
"config.json next to the .gguf file to load the config from "
|
|
"there instead."
|
|
)
|
|
_set_architectures(config, MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type])
|
|
|
|
return config
|