convert-hf : simplify MoE weights stacking

convert-hf-to-gguf.py

@@ -23,7 +23,7 @@ if 'NO_LOCAL_GGUF' not in os.environ:
     sys.path.insert(1, str(Path(__file__).parent / 'gguf-py'))
 import gguf
 
-from convert import LlamaHfVocab, permute
+from convert import LlamaHfVocab
 
 
 ###### MODEL DEFINITIONS ######
@@ -165,10 +165,10 @@ class Model(Protocol):
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         return [(self.map_tensor_name(name), data_torch)]
 
-    def extra_f32_tensors(self, name: str, new_name: str, bid: int | None) -> bool:
+    def extra_f32_tensors(self, name: str, new_name: str, bid: int | None, n_dims: int) -> bool:
         return False
 
-    def extra_f16_tensors(self, name: str, new_name: str, bid: int | None) -> bool:
+    def extra_f16_tensors(self, name: str, new_name: str, bid: int | None, n_dims: int) -> bool:
         return False
 
     def write_tensors(self):
@@ -199,15 +199,16 @@ class Model(Protocol):
                     data = data.astype(np.float32)
 
                 # when both are true, the tensor keeps its original type
-                extra_f32 = self.extra_f32_tensors(name, new_name, bid)
+                extra_f32 = self.extra_f32_tensors(name, new_name, bid, n_dims)
-                extra_f16 = self.extra_f16_tensors(name, new_name, bid)
+                extra_f16 = self.extra_f16_tensors(name, new_name, bid, n_dims)
 
                 # 1d tensors need to be converted to float32
+                # Most of the codebase that takes in 1D tensors only handles F32 tensors
                 if self.ftype == 1 and data_dtype == np.float16 and (n_dims == 1 or extra_f32) and not extra_f16:
                     data = data.astype(np.float32)
 
                 # if f16 desired, convert any float32 2-dim weight tensors to float16
-                if self.ftype == 1 and data_dtype == np.float32 and (name.endswith(".weight") and n_dims == 2 or extra_f16) and not extra_f32:
+                if self.ftype == 1 and data_dtype == np.float32 and (name.endswith(".weight") and n_dims >= 2 or extra_f16) and not extra_f32:
                     data = data.astype(np.float16)
 
                 print(f"{new_name}, n_dims = {n_dims}, {old_dtype} --> {data.dtype}")
@@ -1038,8 +1039,8 @@ class PersimmonModel(Model):
         # self.gguf_writer.add_bos_token_id(71013)
         # self.gguf_writer.add_eos_token_id(71013)
 
-    def extra_f32_tensors(self, name: str, new_name: str, bid: int | None) -> bool:
+    def extra_f32_tensors(self, name: str, new_name: str, bid: int | None, n_dims: int) -> bool:
-        del name, new_name, bid  # unused
+        del name, new_name, bid, n_dims  # unused
 
         # TODO: FP16 conversion produces garbage outputs. (Q8_0 does not, so..?)
         return True
@@ -1072,90 +1073,73 @@ class StableLMModel(Model):
         self.gguf_writer.add_parallel_residual(hparams["use_parallel_residual"] if "use_parallel_residual" in hparams else True)
         self.gguf_writer.add_layer_norm_eps(self.find_hparam(["layer_norm_eps", "norm_eps"]))
 
+    _q_norms: list[dict[str, Tensor]] | None = None
+    _k_norms: list[dict[str, Tensor]] | None = None
+
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # FIXME
+        n_head = self.hparams["num_attention_heads"]
-        return super().modify_tensors(data_torch, name, bid)
+        n_kv_head = self.hparams["num_key_value_heads"]
 
-    def write_tensors(self):
-        block_count = self.hparams.get("n_layers", self.hparams.get("num_hidden_layers", self.hparams.get("n_layer")))
-        tensor_map = gguf.get_tensor_name_map(self.model_arch, block_count)
-        n_head = self.hparams.get("num_attention_heads")
-        n_kv_head = self.hparams.get("num_key_value_heads")
-        q_norms = dict()
-        k_norms = dict()
-        for name, data_torch in self.get_tensors():
-            # we don't need these
-            if name.endswith((".attention.masked_bias", ".attention.bias", ".attention.rotary_emb.inv_freq")):
-                continue
 
-            old_dtype = data_torch.dtype
-
-            # convert any unsupported data types to float32
-            if data_torch.dtype not in (torch.float16, torch.float32):
-                data_torch = data_torch.to(torch.float32)
-
-            data = data_torch.squeeze().numpy()
-            n_dims = len(data.shape)
         if name.find("q_layernorm.norms") != -1:
-                q_norms[name] = data
+            assert bid is not None
-                if len(q_norms) >= (block_count * n_head):
+
-                    self._stack_qk_norm(block_count, name, tensor_map, n_head, q_norms, n_dims, layer_name="q_layernorm")
+            if self._q_norms is None:
-                continue
+                self._q_norms = [{} for _ in range(self.block_count)]
+
+            self._q_norms[bid][name] = data_torch
+
+            if len(self._q_norms[bid]) >= n_head:
+                return self._stack_qk_norm(bid, n_head, self._q_norms[bid], "q_layernorm")
+            else:
+                return []
+
         if name.find("k_layernorm.norms") != -1:
-                k_norms[name] = data
+            assert bid is not None
-                if len(k_norms) >= (block_count * n_kv_head):
-                    self._stack_qk_norm(block_count, name, tensor_map, n_kv_head, k_norms, n_dims, layer_name="k_layernorm")
-                continue
 
-            # map tensor names
+            if self._k_norms is None:
-            new_name = tensor_map.get_name(name, try_suffixes=(".weight", ".bias"))
+                self._k_norms = [{} for _ in range(self.block_count)]
-            if new_name is None:
-                print(f"Can not map tensor {name!r}")
-                sys.exit()
 
-            n_dims = len(data.shape)
+            self._k_norms[bid][name] = data_torch
-            data_dtype = data.dtype
 
-            # if f32 desired, convert any float16 to float32
+            if len(self._k_norms[bid]) >= n_kv_head:
-            if self.ftype == 0 and data_dtype == np.float16:
+                return self._stack_qk_norm(bid, n_kv_head, self._k_norms[bid], "k_layernorm")
-                data = data.astype(np.float32)
+            else:
+                return []
 
-            # TODO: Why cant we use these float16 as-is? There should be not reason to store float16 as float32
+        return [(self.map_tensor_name(name), data_torch)]
-            if self.ftype == 1 and data_dtype == np.float16 and (n_dims == 1 or new_name.endswith("_norm.weight")):
-                data = data.astype(np.float32)
 
-            # if f16 desired, convert any float32 2-dim weight tensors to float16
+    def extra_f32_tensors(self, name: str, new_name: str, bid: int | None, n_dims: int) -> bool:
-            if self.ftype == 1 and data_dtype == np.float32 and name.endswith(".weight") and not new_name.endswith("_norm.weight") and n_dims == 2:
+        del name, bid, n_dims  # unused
-                data = data.astype(np.float16)
 
-            print(f"{new_name}, n_dims = {n_dims}, {old_dtype} --> {data.dtype}")
+        return new_name.endswith("_norm.weight")
 
-            self.gguf_writer.add_tensor(new_name, data)
+    def _stack_qk_norm(self, bid: int, n_head: int, norms: dict[str, Tensor], layer_name: str = "q_layernorm"):
-
+        datas: list[Tensor] = []
-    def _stack_qk_norm(self, block_count, name, tensor_map, n_head, norms, n_dims, layer_name="q_layernorm"):
+        # extract the norms in order
-        for bid in range(block_count):
-            datas = []
         for xid in range(n_head):
             ename = f"model.layers.{bid}.self_attn.{layer_name}.norms.{xid}.weight"
             datas.append(norms[ename])
             del norms[ename]
-            data = np.stack(datas, axis=0)
+        data_torch = torch.cat(datas, dim=0)
-            data_dtype = data.dtype
+
         merged_name = f"model.layers.{bid}.self_attn.{layer_name}.weight"
-            new_name = tensor_map.get_name(merged_name, try_suffixes=(".weight", ".bias"))
+        new_name = self.map_tensor_name(merged_name)
-            if new_name is None:
-                print(f"Can not map tensor {name!r}")
-                sys.exit()
-            if self.ftype == 1 and data_dtype == np.float16 and (n_dims == 1 or new_name.endswith("_norm.weight")):
-                data = data.astype(np.float32)
 
-            # if f16 desired, convert any float32 2-dim weight tensors to float16
+        return [(new_name, data_torch)]
-            if self.ftype == 1 and data_dtype == np.float32 and name.endswith(".weight") and not new_name.endswith("_norm.weight") and n_dims == 2:
-                data = data.astype(np.float16)
 
-            print(f"{new_name}, n_dims = {len(data.shape)}, shape = {data.shape} --> {data.dtype}")
+    def write_tensors(self):
+        super().write_tensors()
 
-            self.gguf_writer.add_tensor(new_name, data)
+        if self._q_norms is not None or self._k_norms is not None:
+            # flatten two `list[dict[str, Tensor]]` into a single `list[str]`
+            norms = (
+                [k for d in self._q_norms for k in d.keys()] if self._q_norms is not None else []
+            ) + (
+                [k for d in self._k_norms for k in d.keys()] if self._k_norms is not None else []
+            )
+            if len(norms) > 0:
+                raise ValueError(f"Unprocessed norms: {norms}")
 
 
 @Model.register("LlamaForCausalLM", "MistralForCausalLM", "MixtralForCausalLM")
@@ -1195,108 +1179,69 @@ class LlamaModel(Model):
                 self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR)
                 self.gguf_writer.add_rope_scaling_factor(self.hparams["rope_scaling"]["factor"])
 
+    @staticmethod
+    def permute(weights: Tensor, n_head: int, n_head_kv: int | None):
+        if n_head_kv is not None and n_head != n_head_kv:
+            n_head = n_head_kv
+        return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:])
+                .swapaxes(1, 2)
+                .reshape(weights.shape))
+
+    _experts: list[dict[str, Tensor]] | None = None
+
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # FIXME
+        n_head = self.hparams["num_attention_heads"]
-        return super().modify_tensors(data_torch, name, bid)
-
-    # Same as super class, but permuting q_proj, k_proj
-    def write_tensors(self):
-        block_count = self.hparams.get("n_layers", self.hparams.get("num_hidden_layers", self.hparams.get("n_layer")))
-        tensor_map = gguf.get_tensor_name_map(self.model_arch, block_count)
-        n_head = self.hparams.get("num_attention_heads")
         n_kv_head = self.hparams.get("num_key_value_heads")
-        n_experts = self.hparams.get("num_local_experts")
-        experts = dict()
-        for name, data_torch in self.get_tensors():
-            # we don't need these
-            if name.endswith((".attention.masked_bias", ".attention.bias", ".attention.rotary_emb.inv_freq")):
-                continue
-
-            old_dtype = data_torch.dtype
-
-            # convert any unsupported data types to float32
-            if data_torch.dtype not in (torch.float16, torch.float32):
-                data_torch = data_torch.to(torch.float32)
-
-            data = data_torch.numpy()
 
         if name.endswith("q_proj.weight"):
-                data = permute(data, n_head, n_head)
+            data_torch = LlamaModel.permute(data_torch, n_head, n_head)
         if name.endswith("k_proj.weight"):
-                data = permute(data, n_head, n_kv_head)
+            data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head)
-
-            data = data.squeeze()
 
         # process the experts separately
         if name.find("block_sparse_moe.experts") != -1:
-                experts[name] = data
+            n_experts = self.hparams["num_local_experts"]
-                if len(experts) >= n_experts:
+
+            assert bid is not None
+
+            if self._experts is None:
+                self._experts = [{} for _ in range(n_experts)]
+
+            self._experts[bid][name] = data_torch
+
+            if len(self._experts[bid]) >= n_experts * 3:
+                tensors: list[tuple[str, Tensor]] = []
+
                 # merge the experts into a single 3d tensor
-                    for bid in range(block_count):
+                for wid in ["w1", "w2", "w3"]:
-                        for wid in range(1, 4):
+                    datas: list[Tensor] = []
-                            full = True
+
                     for xid in range(n_experts):
-                                ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.w{wid}.weight"
+                        ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight"
-                                if ename not in experts:
+                        datas.append(self._experts[bid][ename])
-                                    full = False
+                        del self._experts[bid][ename]
-                                    break
-                            if not full:
-                                continue
 
-                            datas = []
+                    data_torch = torch.cat(datas, dim=0)
-                            for xid in range(n_experts):
-                                ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.w{wid}.weight"
-                                datas.append(experts[ename])
-                                del experts[ename]
 
-                            data = np.stack(datas, axis=0)
+                    merged_name = f"layers.{bid}.feed_forward.experts.{wid}.weight"
-                            data_dtype = data.dtype
 
-                            if self.ftype == 0 and data_dtype == np.float16:
+                    new_name = self.map_tensor_name(merged_name)
-                                data = data.astype(np.float32)
 
-                            if self.ftype == 1 and data_dtype == np.float32:
+                    tensors.append((new_name, data_torch))
-                                data = data.astype(np.float16)
+                return tensors
+            else:
+                return []
 
-                            merged_name = f"layers.{bid}.feed_forward.experts.w{wid}.weight"
+        return [(self.map_tensor_name(name), data_torch)]
 
-                            new_name = tensor_map.get_name(merged_name, try_suffixes=(".weight", ".bias"))
+    def write_tensors(self):
-                            if new_name is None:
+        super().write_tensors()
-                                print(f"Can not map tensor {name!r}")
-                                sys.exit()
-
-                            print(f"{new_name}, n_dims = {len(data.shape)}, shape = {data.shape} --> {data.dtype}")
-
-                            self.gguf_writer.add_tensor(new_name, data)
-                continue
-
-            # map tensor names
-            new_name = tensor_map.get_name(name, try_suffixes=(".weight", ".bias"))
-            if new_name is None:
-                print(f"Can not map tensor {name!r}")
-                sys.exit()
-
-            n_dims = len(data.shape)
-            data_dtype = data.dtype
-
-            # if f32 desired, convert any float16 to float32
-            if self.ftype == 0 and data_dtype == np.float16:
-                data = data.astype(np.float32)
-
-            # 1d tensors need to be converted to float32
-            if self.ftype == 1 and data_dtype == np.float16 and n_dims == 1:
-                data = data.astype(np.float32)
-
-            # if f16 desired, convert any float32 2-dim weight tensors to float16
-            if self.ftype == 1 and data_dtype == np.float32 and name.endswith(".weight") and n_dims == 2:
-                data = data.astype(np.float16)
-
-            print(f"{new_name}, n_dims = {n_dims}, {old_dtype} --> {data.dtype}")
-
-            self.gguf_writer.add_tensor(new_name, data)
 
+        if self._experts is not None:
+            # flatten `list[dict[str, Tensor]]` into `list[str]`
+            experts = [k for d in self._experts for k in d.keys()]
             if len(experts) > 0:
-            raise ValueError(f"Unprocessed experts: {experts.keys()}")
+                raise ValueError(f"Unprocessed experts: {experts}")
 
 
 @Model.register("GrokForCausalLM")
@@ -1313,95 +1258,44 @@ class GrokModel(Model):
         super().set_gguf_parameters()
         self.gguf_writer.add_name("Grok")
 
+    _experts: list[dict[str, Tensor]] | None = None
+
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # FIXME
-        return super().modify_tensors(data_torch, name, bid)
-
-    def write_tensors(self):
-        block_count = self.hparams.get("n_layers", self.hparams.get("num_hidden_layers", self.hparams.get("n_layer")))
-        tensor_map = gguf.get_tensor_name_map(self.model_arch, block_count)
-        n_experts = self.hparams.get("num_local_experts")
-        experts = dict()
-        for name, data_torch in self.get_tensors():
-            # we don't need these
-            if name.endswith((".attention.masked_bias", ".attention.bias", ".attention.rotary_emb.inv_freq")):
-                continue
-
-            old_dtype = data_torch.dtype
-
-            # convert any unsupported data types to float32
-            if data_torch.dtype not in (torch.float16, torch.float32):
-                data_torch = data_torch.to(torch.float32)
-
-            data = data_torch.squeeze().numpy()
-
         # process the experts separately
         if name.find(".moe.") != -1:
-                experts[name] = data
+            n_experts = self.hparams["num_local_experts"]
-                if len(experts) >= n_experts:
+
+            assert bid is not None
+
+            if self._experts is None:
+                self._experts = [{} for _ in range(n_experts)]
+
+            self._experts[bid][name] = data_torch
+
+            if len(self._experts[bid]) >= n_experts * 3:
+                tensors: list[tuple[str, Tensor]] = []
+
                 # merge the experts into a single 3d tensor
-                    for bid in range(block_count):
                 for wid in ["linear", "linear_1", "linear_v"]:
-                            full = True
+                    datas: list[Tensor] = []
+
                     for xid in range(n_experts):
                         ename = f"transformer.decoder_layer.{bid}.moe.{xid}.{wid}.weight"
-                                if ename not in experts:
+                        datas.append(self._experts[bid][ename])
-                                    full = False
+                        del self._experts[bid][ename]
-                                    break
-                            if not full:
-                                continue
 
-                            datas = []
+                    data_torch = torch.cat(datas, dim=0)
-                            for xid in range(n_experts):
-                                ename = f"transformer.decoder_layer.{bid}.moe.{xid}.{wid}.weight"
-                                datas.append(experts[ename])
-                                del experts[ename]
-
-                            data = np.stack(datas, axis=0)
-                            data_dtype = data.dtype
-
-                            if self.ftype == 0 and data_dtype == np.float16:
-                                data = data.astype(np.float32)
-
-                            if self.ftype == 1 and data_dtype == np.float32:
-                                data = data.astype(np.float16)
 
                     merged_name = f"transformer.decoder_layer.{bid}.moe.{wid}.weight"
 
-                            new_name = tensor_map.get_name(merged_name, try_suffixes=(".weight", ".bias"))
+                    new_name = self.map_tensor_name(merged_name)
-                            if new_name is None:
-                                print(f"Can not map tensor {name!r}")
-                                sys.exit()
 
-                            print(f"{new_name}, n_dims = {len(data.shape)}, shape = {data.shape} --> {data.dtype}")
+                    tensors.append((new_name, data_torch))
+                return tensors
+            else:
+                return []
 
-                            self.gguf_writer.add_tensor(new_name, data)
+        return [(self.map_tensor_name(name), data_torch)]
-                continue
-
-            # map tensor names
-            new_name = tensor_map.get_name(name, try_suffixes=(".weight", ".bias"))
-            if new_name is None:
-                print(f"Can not map tensor {name!r}")
-                sys.exit()
-
-            n_dims = len(data.shape)
-            data_dtype = data.dtype
-
-            # if f32 desired, convert any float16 to float32
-            if self.ftype == 0 and data_dtype == np.float16:
-                data = data.astype(np.float32)
-
-            # TODO: Why cant we use these float16 as-is? There should be not reason to store float16 as float32
-            if self.ftype == 1 and data_dtype == np.float16 and n_dims == 1:
-                data = data.astype(np.float32)
-
-            # if f16 desired, convert any float32 2-dim weight tensors to float16
-            if self.ftype == 1 and data_dtype == np.float32 and name.endswith(".weight") and n_dims == 2:
-                data = data.astype(np.float16)
-
-            print(f"{new_name}, n_dims = {n_dims}, {old_dtype} --> {data.dtype}")
-
-            self.gguf_writer.add_tensor(new_name, data)
 
 
 @Model.register("DbrxForCausalLM")
@@ -1435,13 +1329,8 @@ class DbrxModel(Model):
         print(f"gguf: file type = {self.ftype}")
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # FIXME
+        del bid  # unused
-        return super().modify_tensors(data_torch, name, bid)
 
-    def write_tensors(self):
-        block_count = self.hparams.get("n_layers")
-        tensor_map = gguf.get_tensor_name_map(self.model_arch, block_count)
-        for name, data_torch in self.get_tensors():
         n_expert = self.hparams["ffn_config"]["moe_num_experts"]
         n_ff = self.hparams["ffn_config"]["ffn_hidden_size"]
         n_embd = self.hparams["d_model"]
@@ -1455,6 +1344,7 @@ class DbrxModel(Model):
                             "ffn.experts.mlp.w2": (0, 2, 1),  # LLM_TENSOR_FFN_DOWN_EXPS ggml_tensor->ne{n_ff,   n_embd, n_expert}
                             "ffn.experts.mlp.v1": None}       # LLM_TENSOR_FFN_UP_EXPS   ggml_tensor->ne{n_embd, n_ff,   n_expert}
         experts = False
+
         for exp_tensor_name in exp_tensor_names.keys():
             if name.find(exp_tensor_name) != -1 and name.find(".weight") == -1:
                 experts = True
@@ -1463,45 +1353,20 @@ class DbrxModel(Model):
                     data_torch = data_torch.permute(*permute_tensor)
                 break
 
-            old_dtype = data_torch.dtype
-
-            # convert any unsupported data types to float32
-            if data_torch.dtype not in (torch.float16, torch.float32):
-                data_torch = data_torch.to(torch.float32)
-
-            data = data_torch.squeeze().numpy()
-
         # map tensor names
         # In MoE models the ffn tensors are typically most of the model weights,
         # and need to be quantizable. Quantize expects tensor names to be suffixed by .weight.
         # Every other model has the weight names ending in .weight,
         # let's assume that is the convention which is not the case for dbrx:
         # https://huggingface.co/databricks/dbrx-instruct/blob/main/model.safetensors.index.json#L15
-            new_name = tensor_map.get_name(name if not experts else name + ".weight", try_suffixes=(".weight",))
+        new_name = self.map_tensor_name(name if not experts else name + ".weight", try_suffixes=(".weight",))
-            if new_name is None:
-                print(f"Can not map tensor {name!r}")
-                sys.exit()
 
-            n_dims = len(data.shape)
+        return [(new_name, data_torch)]
-            data_dtype = data.dtype
 
-            # Most of the codebase that takes in 1D tensors only handles F32 tensors
+    def extra_f16_tensors(self, name: str, new_name: str, bid: int | None, n_dims: int) -> bool:
-            # and most of the outputs tensors are F32.
+        del name, new_name, bid  # unused
-            if data_dtype != np.float32 and n_dims == 1:
-                print(f"Can not map tensor {name!r}: all 1D tensors must be F32")
-                sys.exit()
 
-            # if f32 desired, convert any float16 to float32
+        return n_dims > 1;
-            if self.ftype == 0 and data_dtype == np.float16:
-                data = data.astype(np.float32)
-
-            # if f16 desired, convert any float32 2-dim weight tensors to float16
-            if self.ftype == 1 and data_dtype == np.float32 and n_dims > 1:
-                data = data.astype(np.float16)
-
-            print(f"{new_name}, n_dims = {n_dims}, shape = {data.shape}, {old_dtype} --> {data.dtype}")
-
-            self.gguf_writer.add_tensor(new_name, data)
 
 
 @Model.register("MiniCPMForCausalLM")
@@ -1611,98 +1476,57 @@ class Qwen2MoeModel(Model):
         if (n_experts := self.hparams.get("num_experts")) is not None:
             self.gguf_writer.add_expert_count(n_experts)
 
+    _experts: list[dict[str, Tensor]] | None = None
+
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # FIXME
-        return super().modify_tensors(data_torch, name, bid)
-
-    def write_tensors(self):
-        block_count = self.hparams.get("n_layers", self.hparams.get("num_hidden_layers", self.hparams.get("n_layer")))
-        tensor_map = gguf.get_tensor_name_map(self.model_arch, block_count)
-        n_experts = self.hparams.get("num_experts")
-        experts = dict()
-        for name, data_torch in self.get_tensors():
-            # we don't need these
-            if name.endswith((".attention.masked_bias", ".attention.bias", ".attention.rotary_emb.inv_freq")):
-                continue
-
-            old_dtype = data_torch.dtype
-
-            # convert any unsupported data types to float32
-            if data_torch.dtype not in (torch.float16, torch.float32):
-                data_torch = data_torch.to(torch.float32)
-
-            data = data_torch.squeeze().numpy()
-
         # process the experts separately
         if name.find("experts") != -1:
-                experts[name] = data
+            n_experts = self.hparams["num_experts"]
-                if len(experts) >= n_experts * 3:
+            assert bid is not None
+
+            if self._experts is None:
+                self._experts = [{} for _ in range(n_experts)]
+
+            self._experts[bid][name] = data_torch
+
+            if len(self._experts[bid]) >= n_experts * 3:
+                tensors: list[tuple[str, Tensor]] = []
+
                 # merge the experts into a single 3d tensor
-                    for bid in range(block_count):
                 for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                            full = True
+                    datas: list[Tensor] = []
+
                     for xid in range(n_experts):
                         ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                                if ename not in experts:
+                        datas.append(self._experts[bid][ename])
-                                    full = False
+                        del self._experts[bid][ename]
-                                    break
-                            if not full:
-                                continue
 
-                            datas = []
+                    data_torch = torch.cat(datas, dim=0)
-                            for xid in range(n_experts):
-                                ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                                datas.append(experts[ename])
-                                del experts[ename]
-
-                            data = np.stack(datas, axis=0)
-                            data_dtype = data.dtype
-
-                            if self.ftype == 0 and data_dtype == np.float16:
-                                data = data.astype(np.float32)
-
-                            if self.ftype == 1 and data_dtype == np.float32:
-                                data = data.astype(np.float16)
 
                     merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
 
-                            new_name = tensor_map.get_name(merged_name, try_suffixes=(".weight", ".bias"))
+                    new_name = self.map_tensor_name(merged_name)
-                            if new_name is None:
-                                print(f"Can not map tensor {name!r}")
-                                sys.exit()
 
-                            print(f"{new_name}, n_dims = {len(data.shape)}, shape = {data.shape} --> {data.dtype}")
+                    tensors.append((new_name, data_torch))
+                return tensors
+            else:
+                return []
 
-                            self.gguf_writer.add_tensor(new_name, data)
+        return [(self.map_tensor_name(name), data_torch)]
-                continue
 
-            # map tensor names
+    def extra_f32_tensors(self, name: str, new_name: str, bid: int | None, n_dims: int) -> bool:
-            new_name = tensor_map.get_name(name, try_suffixes=(".weight", ".bias"))
+        del name, bid, n_dims  # unused
-            if new_name is None:
-                print(f"Can not map tensor {name!r}")
-                sys.exit()
 
-            n_dims = len(data.shape)
+        return new_name.endswith("_norm.weight")
-            data_dtype = data.dtype
 
-            # if f32 desired, convert any float16 to float32
+    def write_tensors(self):
-            if self.ftype == 0 and data_dtype == np.float16:
+        super().write_tensors()
-                data = data.astype(np.float32)
-
-            # TODO: Why cant we use these float16 as-is? There should be not reason to store float16 as float32
-            if self.ftype == 1 and data_dtype == np.float16 and (n_dims == 1 or new_name.endswith("_norm.weight")):
-                data = data.astype(np.float32)
-
-            # if f16 desired, convert any float32 2-dim weight tensors to float16
-            if self.ftype == 1 and data_dtype == np.float32 and name.endswith(".weight") and n_dims == 2:
-                data = data.astype(np.float16)
-
-            print(f"{new_name}, n_dims = {n_dims}, shape = {data.shape}, {old_dtype} --> {data.dtype}")
-
-            self.gguf_writer.add_tensor(new_name, data)
 
+        if self._experts is not None:
+            # flatten `list[dict[str, Tensor]]` into `list[str]`
+            experts = [k for d in self._experts for k in d.keys()]
             if len(experts) > 0:
-            raise ValueError(f"Unprocessed experts: {experts.keys()}")
+                raise ValueError(f"Unprocessed experts: {experts}")
 
 
 @Model.register("GPT2LMHeadModel")
@@ -2152,8 +1976,8 @@ class BertModel(Model):
 
         return [(self.map_tensor_name(name), data_torch)]
 
-    def extra_f32_tensors(self, name: str, new_name: str, bid: int | None) -> bool:
+    def extra_f32_tensors(self, name: str, new_name: str, bid: int | None, n_dims: int) -> bool:
-        del new_name, bid  # unused
+        del new_name, bid, n_dims  # unused
 
         # not used with get_rows, must be F32
         return name == "embeddings.token_type_embeddings.weight"
@@ -2345,7 +2169,9 @@ class MambaModel(Model):
 
         return [(new_name, data_torch)]
 
-    def extra_f32_tensors(self, name: str, new_name: str, bid: int | None) -> bool:
+    def extra_f32_tensors(self, name: str, new_name: str, bid: int | None, n_dims: int) -> bool:
+        del n_dims  # unused
+
         return new_name in (self.format_tensor_name(n, bid, ".weight" if name.endswith(".weight") else "") for n in [
             gguf.MODEL_TENSOR.SSM_CONV1D,
             gguf.MODEL_TENSOR.SSM_X,
@@ -2386,54 +2212,17 @@ class OlmoModel(Model):
     # Same as super class, but permuting q_proj, k_proj
     # Copied from: LlamaModel
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # FIXME
+        del bid  # unused
-        return super().modify_tensors(data_torch, name, bid)
 
-    def write_tensors(self):
+        n_head = self.hparams["num_attention_heads"]
-        block_count = self.hparams.get("n_layers", self.hparams.get("num_hidden_layers", self.hparams.get("n_layer")))
-        tensor_map = gguf.get_tensor_name_map(self.model_arch, block_count)
-        n_head = self.hparams.get("num_attention_heads")
         n_kv_head = self.hparams.get("num_key_value_heads")
-        for name, data_torch in self.get_tensors():
-            old_dtype = data_torch.dtype
-
-            # convert any unsupported data types to float32
-            if data_torch.dtype not in (torch.float16, torch.float32):
-                data_torch = data_torch.to(torch.float32)
-
-            data = data_torch.numpy()
 
         if name.endswith("q_proj.weight"):
-                data = permute(data, n_head, n_head)
+            data_torch = LlamaModel.permute(data_torch, n_head, n_head)
         if name.endswith("k_proj.weight"):
-                data = permute(data, n_head, n_kv_head)
+            data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head)
 
-            data = data.squeeze()
+        return [(self.map_tensor_name(name), data_torch)]
-
-            # map tensor names
-            new_name = tensor_map.get_name(name, try_suffixes=(".weight", ".bias"))
-            if new_name is None:
-                print(f"Can not map tensor {name!r}")
-                sys.exit()
-
-            n_dims = len(data.shape)
-            data_dtype = data.dtype
-
-            # if f32 desired, convert any float16 to float32
-            if self.ftype == 0 and data_dtype == np.float16:
-                data = data.astype(np.float32)
-
-            # 1d tensors need to be converted to float32
-            if self.ftype == 1 and data_dtype == np.float16 and n_dims == 1:
-                data = data.astype(np.float32)
-
-            # if f16 desired, convert any float32 2-dim weight tensors to float16
-            if self.ftype == 1 and data_dtype == np.float32 and n_dims == 2:
-                data = data.astype(np.float16)
-
-            print(f"{new_name}, n_dims = {n_dims}, {old_dtype} --> {data.dtype}")
-
-            self.gguf_writer.add_tensor(new_name, data)
 
 
 ###### CONVERSION LOGIC ######