Unverified Commit e5ef4dfc authored by toncao's avatar toncao Committed by GitHub
Browse files

[Kimi-Linear] Correct prefixes and add compatibility to AWQ quants (#27834)


Signed-off-by: default avatartoncao <cpatonn@gmail.com>
Co-authored-by: default avatartoncao <cpatonn@gmail.com>
parent 36960501
...@@ -155,6 +155,7 @@ class KimiMoE(nn.Module): ...@@ -155,6 +155,7 @@ class KimiMoE(nn.Module):
hidden_act=config.hidden_act, hidden_act=config.hidden_act,
quant_config=quant_config, quant_config=quant_config,
reduce_results=False, reduce_results=False,
prefix=f"{prefix}.shared_experts",
) )
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
...@@ -340,7 +341,7 @@ class KimiDecoderLayer(nn.Module): ...@@ -340,7 +341,7 @@ class KimiDecoderLayer(nn.Module):
self.block_sparse_moe = KimiMoE( self.block_sparse_moe = KimiMoE(
config=config, config=config,
quant_config=quant_config, quant_config=quant_config,
prefix=f"{prefix}.mlp", prefix=f"{prefix}.block_sparse_moe",
) )
self.mlp = self.block_sparse_moe self.mlp = self.block_sparse_moe
else: else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment