fix(convert): Sanity check on merged FFN tensor sizes

Branch: GraniteMoE Co-Authored-By: [email protected] Signed-off-by: Gabe Goodhart <[email protected]>
ggerganov · Sep 23, 2024 · 3907795 · 3907795
1 parent b5b9e4b
commit 3907795
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py
@@ -4127,7 +4127,9 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter
         """
 
         if name.endswith("block_sparse_moe.input_linear.weight"):
-            gate, up = data_torch.chunk(2, dim=-2)
+            ffn_dim = self.hparams["intermediate_size"]
+            assert data_torch.shape[-2] == 2 * ffn_dim, "Merged FFN tensor size must be 2 * intermediate_size"
+            gate, up = data_torch[..., :ffn_dim, :], data_torch[..., ffn_dim:, :]
             return [
                 (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_EXP, bid), gate),
                 (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_EXP, bid), up),