Revert "revert the head_dim assignment from PR 23862" (#28583)
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user