llama : support running Mamba-Codestral-7B-v0.1

This commit is contained in:
Francis Couture-Harpin 2024-08-18 22:43:39 -04:00
parent dceff23fae
commit 2bfe9de6d3
2 changed files with 5 additions and 1 deletions

View file

@ -2843,6 +2843,10 @@ class Mamba2Model(Model):
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
del bid # unused
if name.startswith("model.backbone") or name.startswith("model.lm_head"):
# map Mamba-Codestral-7B-v0.1 tensor names to the names used by Mamba-2
name = name.removeprefix("model.")
if name.endswith(".dt_bias"):
name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"