revert the head_dim assignment from PR 23862 (#28571)

This commit is contained in:
cctry
2026-06-17 15:46:28 -07:00
committed by GitHub
parent 732b81d5b9
commit 5d6b35eabb
2 changed files with 20 additions and 92 deletions
+20 -21
View File
@@ -595,28 +595,27 @@ class ModelConfig:
def _derive_model_shapes(self): def _derive_model_shapes(self):
# Unify the config keys for hf_text_config # Unify the config keys for hf_text_config
self.head_dim = getattr(self.hf_text_config, "head_dim", None) self.head_dim = getattr(
if self.head_dim is None: self.hf_text_config,
self.head_dim = ( "head_dim",
self.hf_text_config.hidden_size self.hf_text_config.hidden_size // self.hf_text_config.num_attention_heads,
// self.hf_text_config.num_attention_heads )
) self.v_head_dim = getattr(
setattr(self.hf_text_config, "head_dim", self.head_dim) self.hf_text_config,
"v_head_dim",
self.head_dim,
)
self.v_head_dim = getattr(self.hf_text_config, "v_head_dim", None) self.swa_head_dim = getattr(
if self.v_head_dim is None: self.hf_text_config,
self.v_head_dim = self.head_dim "swa_head_dim",
setattr(self.hf_text_config, "v_head_dim", self.v_head_dim) self.head_dim,
)
self.swa_head_dim = getattr(self.hf_text_config, "swa_head_dim", None) self.swa_v_head_dim = getattr(
if self.swa_head_dim is None: self.hf_text_config,
self.swa_head_dim = self.head_dim "swa_v_head_dim",
setattr(self.hf_text_config, "swa_head_dim", self.swa_head_dim) self.swa_head_dim,
)
self.swa_v_head_dim = getattr(self.hf_text_config, "swa_v_head_dim", None)
if self.swa_v_head_dim is None:
self.swa_v_head_dim = self.swa_head_dim
setattr(self.hf_text_config, "swa_v_head_dim", self.swa_v_head_dim)
# FIXME: temporary special judge for MLA architecture # FIXME: temporary special judge for MLA architecture
if ( if (
"DeepseekV2ForCausalLM" in self.hf_config.architectures "DeepseekV2ForCausalLM" in self.hf_config.architectures
@@ -1,71 +0,0 @@
"""Unit tests for ModelConfig shape normalization."""
import unittest
from types import SimpleNamespace
from sglang.srt.configs.model_config import ModelConfig
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")
def _make_text_config(**overrides):
defaults = dict(
architectures=["MixtralForCausalLM"],
model_type="mixtral",
hidden_size=4096,
num_attention_heads=32,
num_hidden_layers=2,
vocab_size=32000,
num_key_value_heads=8,
)
defaults.update(overrides)
return SimpleNamespace(**defaults)
class TestModelConfigShapes(CustomTestCase):
def _derive_shapes(self, text_config):
model_config = ModelConfig.__new__(ModelConfig)
model_config.hf_config = text_config
model_config.hf_text_config = text_config
model_config._derive_model_shapes()
return model_config
def test_optional_head_dims_default_when_none(self):
text_config = _make_text_config(
head_dim=None,
v_head_dim=None,
swa_head_dim=None,
swa_v_head_dim=None,
)
model_config = self._derive_shapes(text_config)
self.assertEqual(model_config.head_dim, 128)
self.assertEqual(model_config.v_head_dim, 128)
self.assertEqual(model_config.swa_head_dim, 128)
self.assertEqual(model_config.swa_v_head_dim, 128)
self.assertEqual(text_config.head_dim, 128)
self.assertEqual(text_config.v_head_dim, 128)
self.assertEqual(text_config.swa_head_dim, 128)
self.assertEqual(text_config.swa_v_head_dim, 128)
def test_explicit_head_dims_are_preserved(self):
text_config = _make_text_config(
head_dim=128,
v_head_dim=96,
swa_head_dim=64,
swa_v_head_dim=48,
)
model_config = self._derive_shapes(text_config)
self.assertEqual(model_config.head_dim, 128)
self.assertEqual(model_config.v_head_dim, 96)
self.assertEqual(model_config.swa_head_dim, 64)
self.assertEqual(model_config.swa_v_head_dim, 48)
if __name__ == "__main__":
unittest.main()