From 4dc7dc851880cd46ac18d104668a1363bff75ce2 Mon Sep 17 00:00:00 2001 From: Pavan Sivaram Girijala Date: Mon, 31 Aug 2026 08:35:50 +0530 Subject: [PATCH] [Fix] Transformers-fallback (GPT-NeoX) + KV pool config (DeepSeek-VL2) (#35244) Co-authored-by: YanbingJiang --- python/sglang/srt/configs/model_config.py | 2 +- python/sglang/srt/models/transformers.py | 6 +- .../unit/configs/test_model_config_shapes.py | 14 ++++ .../test_transformers_fallback.py | 77 +++++++++++++++++++ 4 files changed, 97 insertions(+), 2 deletions(-) create mode 100644 test/registered/unit/model_loader/test_transformers_fallback.py diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 1785a7c8d..e58734bf7 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -907,7 +907,7 @@ class ModelConfig: setattr(self.hf_text_config, "head_dim", self.head_dim) self.v_head_dim = getattr(self.hf_text_config, "v_head_dim", None) - if self.v_head_dim is None: + if self.v_head_dim is None or self.v_head_dim == 0: self.v_head_dim = self.head_dim setattr(self.hf_text_config, "v_head_dim", self.v_head_dim) diff --git a/python/sglang/srt/models/transformers.py b/python/sglang/srt/models/transformers.py index 20d600e90..343f601d9 100644 --- a/python/sglang/srt/models/transformers.py +++ b/python/sglang/srt/models/transformers.py @@ -544,6 +544,8 @@ class TransformersBase(nn.Module): "model.score.": "classifier.", "model.classifier.": "classifier.", "transformer.": "model.", + "gpt_neox.": "model.", + "embed_out.": "lm_head.", "model.": "model.", "lm_head.": "lm_head.", "score.": "classifier.", @@ -581,7 +583,9 @@ class TransformersBase(nn.Module): self.skip_substrs: list[str] = [] self.ignore_unexpected_prefixes: list[str] = [] self.ignore_unexpected_suffixes: list[str] = [] - self.skip_substrs.extend([".attn.bias", ".attn.masked_bias", ".masked_bias"]) + self.skip_substrs.extend( + [".attn.bias", ".attn.masked_bias", ".attention.bias", ".masked_bias"] + ) self.ignore_unexpected_prefixes.extend(["classifier.", "score."]) if self.quant_config is not None: diff --git a/test/registered/unit/configs/test_model_config_shapes.py b/test/registered/unit/configs/test_model_config_shapes.py index 8cdd4dc10..dc12db3ac 100644 --- a/test/registered/unit/configs/test_model_config_shapes.py +++ b/test/registered/unit/configs/test_model_config_shapes.py @@ -69,6 +69,20 @@ class TestModelConfigShapes(CustomTestCase): self.assertEqual(model_config.swa_head_dim, 64) self.assertEqual(model_config.swa_v_head_dim, 48) + def test_v_head_dim_zero_falls_back_to_head_dim(self): + # deepseek-vl2-tiny's language_config sets use_mla=False + v_head_dim=0 + # as an MLA-disabled sentinel. The non-MLA KV pool sizes its V buffer + # from v_head_dim directly, so a literal 0 collapses V-cache to a + # zero-width tensor and `v_cache[indices] = v` fails with a shape + # mismatch inside `_set_kv_buffer_impl`. v_head_dim=0 must be treated + # identically to `None` and fall back to head_dim. + text_config = _make_text_config(head_dim=128, v_head_dim=0) + + model_config = self._derive_shapes(text_config) + + self.assertEqual(model_config.v_head_dim, 128) + self.assertEqual(text_config.v_head_dim, 128) + def test_ling_mla_nope_shapes(self): text_config = _make_text_config( architectures=["BailingMoeV3ForCausalLM"], diff --git a/test/registered/unit/model_loader/test_transformers_fallback.py b/test/registered/unit/model_loader/test_transformers_fallback.py new file mode 100644 index 000000000..743403bd4 --- /dev/null +++ b/test/registered/unit/model_loader/test_transformers_fallback.py @@ -0,0 +1,77 @@ +"""Unit tests for the shared Transformers-fallback loader path in SGLang.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from sglang.srt.models.transformers import TransformersBase +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestTransformersFallbackWeightMapper(CustomTestCase): + def test_gpt_neox_prefix_rewrites(self): + # GPT-NeoX checkpoints (Pythia, gpt-neox-20b) wrap the backbone under + # `gpt_neox.` and expose the LM head as `embed_out.`. `AutoModel` + # returns the backbone as `self.model` and `CausalMixin` installs a + # `ParallelLMHead` at `self.lm_head`, so the fallback mapper must + # rewrite both prefixes. Without these rules, `load_weights` raises + # unexpected-key errors on every `gpt_neox.*` / `embed_out.*` entry. + mapped = TransformersBase.hf_to_sglang_mapper.apply_list( + [ + "gpt_neox.embed_in.weight", + "gpt_neox.layers.0.mlp.dense_h_to_4h.weight", + "gpt_neox.final_layer_norm.weight", + "embed_out.weight", + ] + ) + + self.assertEqual( + mapped, + [ + "model.embed_in.weight", + "model.layers.0.mlp.dense_h_to_4h.weight", + "model.final_layer_norm.weight", + "lm_head.weight", + ], + ) + + +class TestTransformersFallbackSkipSubstrs(CustomTestCase): + def test_init_registers_attention_bias_skip(self): + # Older GPT-NeoX checkpoints (pythia-1.4b, gpt-neox-20b) ship a + # persistent `attention.bias` causal-mask buffer. Newer transformers + # builds register it as `persistent=False`, so it is absent from the + # constructed module tree and `AutoWeightsLoader` raises on the + # unexpected key unless the fallback tells it to skip. The pre-existing + # `.attn.bias` covers GPT-2, not NeoX, so `.attention.bias` must be its + # own entry. + stub = TransformersBase.__new__(TransformersBase) + + class _Stop(RuntimeError): + pass + + with ( + patch( + "sglang.srt.models.transformers.get_pp_group", + return_value=SimpleNamespace(), + ), + patch( + "sglang.srt.models.transformers.get_hf_text_config", + return_value=SimpleNamespace(), + ), + patch( + "sglang.srt.models.transformers._resolve_attention_backend_model_cls", + side_effect=_Stop, + ), + self.assertRaises(_Stop), + ): + TransformersBase.__init__(stub, config=SimpleNamespace()) + + self.assertIn(".attention.bias", stub.skip_substrs) + + +if __name__ == "__main__": + unittest.main()