Fix incorrect default get_hidden_dim logic (#8987)

This commit is contained in:
Lifu Huang
2025-08-09 00:25:38 -07:00
committed by GitHub
parent a47baff12c
commit 6e2151183b
7 changed files with 36 additions and 143 deletions
-24
View File
@@ -330,30 +330,6 @@ class Qwen3ForCausalLM(nn.Module):
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.get_input_embeddings(input_ids)
def get_hidden_dim(self, module_name: str) -> Tuple[int]:
# return input_dim, output_dim
if module_name in ["q_proj", "qkv_proj"]:
return (
self.config.hidden_size,
self.config.head_dim * self.config.num_attention_heads,
)
elif module_name in ["o_proj"]:
return (
self.config.head_dim * self.config.num_attention_heads,
self.config.hidden_size,
)
elif module_name in ["kv_proj"]:
return (
self.config.hidden_size,
self.config.head_dim * self.config.num_key_value_heads,
)
elif module_name == "gate_up_proj":
return self.config.hidden_size, self.config.intermediate_size
elif module_name == "down_proj":
return self.config.intermediate_size, self.config.hidden_size
else:
raise NotImplementedError()
@torch.no_grad()
def forward(
self,