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):
# Unify the config keys for hf_text_config
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.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.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)
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,
)
# FIXME: temporary special judge for MLA architecture
if (
"DeepseekV2ForCausalLM" in self.hf_config.architectures