Revert "revert the head_dim assignment from PR 23862" (#28583)

This commit is contained in:
cctry
2026-06-17 17:11:05 -07:00
committed by GitHub
parent 343aeeef39
commit cfa4aa988f
2 changed files with 92 additions and 20 deletions
+21 -20
View File
@@ -595,27 +595,28 @@ class ModelConfig:
def _derive_model_shapes(self):
# Unify the config keys for hf_text_config
self.head_dim = getattr(
self.hf_text_config,
"head_dim",
self.hf_text_config.hidden_size // self.hf_text_config.num_attention_heads,
)
self.v_head_dim = getattr(
self.hf_text_config,
"v_head_dim",
self.head_dim,
)
self.head_dim = getattr(self.hf_text_config, "head_dim", None)
if self.head_dim is None:
self.head_dim = (
self.hf_text_config.hidden_size
// self.hf_text_config.num_attention_heads
)
setattr(self.hf_text_config, "head_dim", self.head_dim)
self.swa_head_dim = getattr(
self.hf_text_config,
"swa_head_dim",
self.head_dim,
)
self.swa_v_head_dim = getattr(
self.hf_text_config,
"swa_v_head_dim",
self.swa_head_dim,
)
self.v_head_dim = getattr(self.hf_text_config, "v_head_dim", None)
if self.v_head_dim is None:
self.v_head_dim = self.head_dim
setattr(self.hf_text_config, "v_head_dim", self.v_head_dim)
self.swa_head_dim = getattr(self.hf_text_config, "swa_head_dim", None)
if self.swa_head_dim is None:
self.swa_head_dim = self.head_dim
setattr(self.hf_text_config, "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
if (
"DeepseekV2ForCausalLM" in self.hf_config.architectures