conversion: only allow selected models

This commit is contained in:
ngxson 2024-07-10 00:23:07 +02:00
parent 03d24cae19
commit ee2b35c65f
2 changed files with 37 additions and 19 deletions

View file

@ -373,6 +373,9 @@ class Model:
except KeyError:
raise NotImplementedError(f'Architecture {arch!r} not supported!') from None
def support_lora(self) -> bool:
return False
# used for GPT-2 BPE and WordPiece vocabs
def get_vocab_base(self) -> tuple[list[str], list[int], str]:
tokens: list[str] = []
@ -1416,9 +1419,9 @@ class LlamaModel(Model):
n_head = self.hparams["num_attention_heads"]
n_kv_head = self.hparams.get("num_key_value_heads")
if name.endswith(("q_proj.weight", "q_proj.bias")):
if name.endswith(("q_proj.weight", "q_proj.bias", "q_proj.lora_B.weight")):
data_torch = LlamaModel.permute(data_torch, n_head, n_head)
if name.endswith(("k_proj.weight", "k_proj.bias")):
if name.endswith(("k_proj.weight", "k_proj.bias", "k_proj.lora_B.weight")):
data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head)
# process the experts separately
@ -1466,6 +1469,10 @@ class LlamaModel(Model):
if len(experts) > 0:
raise ValueError(f"Unprocessed experts: {experts}")
def support_lora(self) -> bool:
# TODO: support lora conversion for MOE
return "num_local_experts" not in self.hparams
@Model.register("BitnetForCausalLM")
class BitnetModel(Model):