| 775 | |
| 776 | @Model.register("BitnetForCausalLM") |
| 777 | class BitnetModel(Model): |
| 778 | model_arch = gguf.MODEL_ARCH.BITNET |
| 779 | params: str = "" |
| 780 | |
| 781 | def set_params(self, params: str): |
| 782 | self.params = params |
| 783 | hp_config = model_config[self.params] |
| 784 | self.hparams["hidden_size"] = hp_config["hidden_size"] |
| 785 | self.hparams["intermediate_size"] = hp_config["intermediate_size"] |
| 786 | self.hparams["num_hidden_layers"] = hp_config["num_hidden_layers"] |
| 787 | self.hparams["num_attention_heads"] = hp_config["num_attention_heads"] |
| 788 | self.hparams["num_key_value_heads"] = hp_config["num_attention_heads"] |
| 789 | self.block_count = self.find_hparam(["n_layers", "num_hidden_layers", "n_layer"]) |
| 790 | self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) |
| 791 | |
| 792 | |
| 793 | def set_vocab(self): |
| 794 | self._set_vocab_sentencepiece() |
| 795 | |
| 796 | def set_gguf_parameters(self): |
| 797 | super().set_gguf_parameters() |
| 798 | |
| 799 | self.gguf_writer.add_vocab_size(self.hparams["vocab_size"]) |
| 800 | |
| 801 | self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) |
| 802 | self.gguf_writer.add_rope_scaling_factor(1.0) |
| 803 | |
| 804 | def weight_quant(self, weight): |
| 805 | dtype = weight.dtype |
| 806 | weight = weight.float() |
| 807 | s = 1 / weight.abs().mean().clamp(min=1e-5) |
| 808 | result = (weight * s).round().clamp(-1, 1) / s |
| 809 | return result.type(dtype) |
| 810 | |
| 811 | def transform_to_tl1(self, x: np.ndarray): |
| 812 | scale = np.max(np.abs(x)) |
| 813 | # res = np.round(x / scale + 2).astype(np.uint8) |
| 814 | res = preprocess_weights_tl1(x) |
| 815 | return res, scale |
| 816 | |
| 817 | def transform_to_tl2(self, x: np.ndarray): |
| 818 | scale = np.max(np.abs(x)) |
| 819 | # res = np.round(x / scale + 2).astype(np.uint8) |
| 820 | res = preprocess_weights_tl2(x) |
| 821 | return res, scale |
| 822 | |
| 823 | # generate dummy model |
| 824 | def generate_tensors(self) -> Iterator[tuple[str, np.ndarray]]: |
| 825 | hp_config = model_config[self.params] |
| 826 | hidden_size = hp_config["hidden_size"] |
| 827 | intermediate_size = hp_config["intermediate_size"] |
| 828 | num_hidden_layers = hp_config["num_hidden_layers"] |
| 829 | num_attention_heads = hp_config["num_attention_heads"] |
| 830 | |
| 831 | # generate dummy tensors |
| 832 | tensor = torch.randn((32002, hidden_size), dtype=torch.float32) |
| 833 | yield ("model.embed_tokens.weight", tensor) |
| 834 | for i in range(num_hidden_layers): |
nothing calls this directly
no outgoing calls
no test coverage detected