char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Sigbjørn Skjæretconvert : fix lora base model arch retrieval (#24621)dd4623a74

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
22.7 KiB546 linesraw
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3
4from __future__ import annotations
5
6from dataclasses import dataclass
7import logging
8import argparse
9import os
10import sys
11import json
12from math import prod
13from pathlib import Path
14from typing import TYPE_CHECKING, Any, Callable, Iterable, Iterator, Sequence, SupportsIndex, cast
15from transformers import AutoConfig, AutoTokenizer
16
17import torch
18
19if TYPE_CHECKING:
20    from torch import Tensor
21
22if 'NO_LOCAL_GGUF' not in os.environ:
23    sys.path.insert(1, str(Path(__file__).parent / 'gguf-py'))
24import gguf
25from gguf.constants import GGUFValueType
26
27# reuse model definitions from the conversion/ package
28from conversion import LazyTorchTensor, ModelBase, get_model_class, ModelType, get_model_architecture
29
30logger = logging.getLogger("lora-to-gguf")
31
32
33@dataclass
34class PartialLoraTensor:
35    A: Tensor | None = None
36    B: Tensor | None = None
37
38
39# magic to support tensor shape modifications and splitting
40class LoraTorchTensor:
41    _lora_A: Tensor  # (n_rank, row_size)
42    _lora_B: Tensor  # (col_size, n_rank)
43    _rank: int
44
45    def __init__(self, A: Tensor, B: Tensor):
46        assert len(A.shape) == len(B.shape)
47        assert A.shape[-2] == B.shape[-1]
48        if A.dtype != B.dtype:
49            A = A.to(torch.float32)
50            B = B.to(torch.float32)
51        self._lora_A = A
52        self._lora_B = B
53        self._rank = B.shape[-1]
54
55    def get_lora_A_B(self) -> tuple[Tensor, Tensor]:
56        return (self._lora_A, self._lora_B)
57
58    def __getitem__(
59        self,
60        indices: (
61            SupportsIndex
62            | slice
63            | tuple[SupportsIndex | slice | Tensor, ...]  # TODO: add ellipsis in the type signature
64        ),
65    ) -> LoraTorchTensor:
66        shape = self.shape
67        if isinstance(indices, SupportsIndex):
68            if len(shape) > 2:
69                return LoraTorchTensor(self._lora_A[indices], self._lora_B[indices])
70            else:
71                raise NotImplementedError  # can't return a vector
72        elif isinstance(indices, slice):
73            if len(shape) > 2:
74                return LoraTorchTensor(self._lora_A[indices], self._lora_B[indices])
75            else:
76                return LoraTorchTensor(self._lora_A, self._lora_B[indices])
77        elif isinstance(indices, tuple):
78            assert len(indices) > 0
79            if indices[-1] is Ellipsis:
80                return self[indices[:-1]]
81            # expand ellipsis
82            indices = tuple(
83                u
84                for v in (
85                    (
86                        (slice(None, None) for _ in range(len(indices) - 1))
87                        if i is Ellipsis
88                        else (i,)
89                    )
90                    for i in indices
91                )
92                for u in v
93            )
94
95            if len(indices) < len(shape):
96                indices = (*indices, *(slice(None, None) for _ in range(len(indices), len(shape))))
97
98            # TODO: make sure this is correct
99            indices_A = (
100                *(
101                    (
102                        j.__index__() % self._lora_A.shape[i]
103                        if isinstance(j, SupportsIndex)
104                        else slice(None, None)
105                    )
106                    for i, j in enumerate(indices[:-2])
107                ),
108                slice(None, None),
109                indices[-1],
110            )
111            indices_B = indices[:-1]
112            return LoraTorchTensor(self._lora_A[indices_A], self._lora_B[indices_B])
113        else:
114            raise NotImplementedError  # unknown indice type
115
116    @property
117    def dtype(self) -> torch.dtype:
118        assert self._lora_A.dtype == self._lora_B.dtype
119        return self._lora_A.dtype
120
121    @property
122    def shape(self) -> tuple[int, ...]:
123        assert len(self._lora_A.shape) == len(self._lora_B.shape)
124        return (*self._lora_B.shape[:-1], self._lora_A.shape[-1])
125
126    def size(self, dim=None):
127        assert dim is None
128        return self.shape
129
130    def contiguous(self) -> LoraTorchTensor:
131        return LoraTorchTensor(
132            self._lora_A.contiguous(),
133            self._lora_B.contiguous(),
134        )
135
136    def reshape(self, *shape: int | tuple[int, ...]) -> LoraTorchTensor:
137        if isinstance(shape[0], tuple):
138            new_shape: tuple[int, ...] = shape[0]
139        else:
140            new_shape = cast(tuple[int, ...], shape)
141        orig_shape = self.shape
142        if len(new_shape) < 2:
143            raise NotImplementedError  # can't become a vector
144
145        # expand -1 in the shape
146        if any(dim == -1 for dim in new_shape):
147            n_elems = prod(orig_shape)
148            n_new_elems = prod(dim if dim != -1 else 1 for dim in new_shape)
149            assert n_elems % n_new_elems == 0
150            new_shape = (*(dim if dim != -1 else n_elems // n_new_elems for dim in new_shape),)
151
152        if new_shape[-1] != orig_shape[-1]:
153            raise NotImplementedError  # can't reshape the row size trivially
154
155        shape_A = (*(1 for _ in new_shape[:-2]), self._rank, orig_shape[-1])
156        shape_B = (*new_shape[:-1], self._rank)
157        return LoraTorchTensor(
158            self._lora_A.reshape(shape_A),
159            self._lora_B.reshape(shape_B),
160        )
161
162    def reshape_as(self, other: Tensor) -> LoraTorchTensor:
163        return self.reshape(*other.shape)
164
165    def view(self, *size: int) -> LoraTorchTensor:
166        return self.reshape(*size)
167
168    def permute(self, *dims: int) -> LoraTorchTensor:
169        shape = self.shape
170        dims = tuple(dim - len(shape) if dim >= 0 else dim for dim in dims)
171        if dims[-1] == -1:
172            # TODO: support higher dimensional A shapes bigger than 1
173            assert all(dim == 1 for dim in self._lora_A.shape[:-2])
174            return LoraTorchTensor(self._lora_A, self._lora_B.permute(*dims))
175        if len(shape) == 2 and dims[-1] == -2 and dims[-2] == -1:
176            return LoraTorchTensor(self._lora_B.permute(*dims), self._lora_A.permute(*dims))
177        else:
178            # TODO: compose the above two
179            raise NotImplementedError
180
181    def transpose(self, dim0: int, dim1: int) -> LoraTorchTensor:
182        shape = self.shape
183        dims = [i for i in range(len(shape))]
184        dims[dim0], dims[dim1] = dims[dim1], dims[dim0]
185        return self.permute(*dims)
186
187    def swapaxes(self, axis0: int, axis1: int) -> LoraTorchTensor:
188        return self.transpose(axis0, axis1)
189
190    def split(self, split_size: int | Sequence[int], dim: int = 0) -> tuple[LoraTorchTensor, ...]:
191        shape = self.shape
192        ndim = len(shape)
193        if dim < 0:
194            dim += ndim
195        if dim == ndim - 1:
196            A_chunks = self._lora_A.split(split_size, dim=-1)
197            return tuple(LoraTorchTensor(a, self._lora_B) for a in A_chunks)
198        elif dim == ndim - 2:
199            B_chunks = self._lora_B.split(split_size, dim=-2)
200            return tuple(LoraTorchTensor(self._lora_A, b) for b in B_chunks)
201        else:
202            B_chunks = self._lora_B.split(split_size, dim=dim)
203            if self._lora_A.shape[dim] == 1:
204                return tuple(LoraTorchTensor(self._lora_A, b) for b in B_chunks)
205            A_chunks = self._lora_A.split(split_size, dim=dim)
206            return tuple(LoraTorchTensor(a, b) for a, b in zip(A_chunks, B_chunks))
207
208    def to(self, *args, **kwargs):
209        return LoraTorchTensor(self._lora_A.to(*args, **kwargs), self._lora_B.to(*args, **kwargs))
210
211    def __mul__(self, other) -> LoraTorchTensor:
212        # Only output-side multiplication for now
213        # W = B @ A, so M_out * W == (M_out * B) @ A
214        if not isinstance(other, (int, float)) and other.shape and other.shape[-1] != 1:
215            raise NotImplementedError
216        return LoraTorchTensor(self._lora_A, self._lora_B * other)
217
218    def __rmul__(self, other) -> LoraTorchTensor:
219        return self * other
220
221    @classmethod
222    def __torch_function__(cls, func: Callable, types, args=(), kwargs=None):
223        del types  # unused
224
225        if kwargs is None:
226            kwargs = {}
227
228        if func is torch.permute:
229            assert len(args)
230            return type(args[0]).permute(*args, **kwargs)
231        elif func is torch.reshape:
232            assert len(args)
233            return type(args[0]).reshape(*args, **kwargs)
234        elif func is torch.stack:
235            assert len(args)
236            assert isinstance(args[0], Sequence)
237            dim = kwargs.get("dim", 0)
238            assert dim == 0
239            return LoraTorchTensor(
240                torch.stack([a._lora_A for a in args[0]], dim),
241                torch.stack([b._lora_B for b in args[0]], dim),
242            )
243        elif func is torch.cat:
244            assert len(args)
245            assert isinstance(args[0], Sequence)
246            dim = kwargs.get("dim", 0)
247            assert dim == 0
248            if len(args[0][0].shape) > 2:
249                return LoraTorchTensor(
250                    torch.cat([a._lora_A for a in args[0]], dim),
251                    torch.cat([b._lora_B for b in args[0]], dim),
252                )
253            elif all(torch.equal(args[0][0]._lora_A, t._lora_A) for t in args[0][1:]):
254                return LoraTorchTensor(
255                    args[0][0]._lora_A,
256                    torch.cat([b._lora_B for b in args[0]], dim),
257                )
258            else:
259                raise NotImplementedError
260        elif func is torch.split:
261            assert len(args) and len(args) >= 2
262            tensor, split_size = args[0], args[1]
263            dim = args[2] if len(args) > 2 else kwargs.get("dim", 0)
264            return tensor.split(split_size, dim=dim)
265        else:
266            raise NotImplementedError
267
268
269def get_base_tensor_name(lora_tensor_name: str) -> str:
270    base_name = lora_tensor_name.replace("base_model.model.", "")
271    base_name = base_name.replace(".lora_A.weight", ".weight")
272    base_name = base_name.replace(".lora_B.weight", ".weight")
273    # models produced by mergekit-extract-lora have token embeddings in the adapter
274    base_name = base_name.replace(".lora_embedding_A", ".weight")
275    base_name = base_name.replace(".lora_embedding_B", ".weight")
276    return base_name
277
278
279def parse_args() -> argparse.Namespace:
280    parser = argparse.ArgumentParser(
281        description="Convert a Hugging Face PEFT LoRA adapter to a GGUF file")
282    parser.add_argument(
283        "--outfile", type=Path,
284        help="path to write to; default: based on input. {ftype} will be replaced by the outtype.",
285    )
286    parser.add_argument(
287        "--outtype", type=str, choices=["f32", "f16", "bf16", "q8_0", "auto"], default="f32",
288        help="output format - use f32 for float32, f16 for float16, bf16 for bfloat16, q8_0 for Q8_0, auto for the highest-fidelity 16-bit float type depending on the first loaded tensor type",
289    )
290    parser.add_argument(
291        "--bigendian", action="store_true",
292        help="model is executed on big endian machine",
293    )
294    parser.add_argument(
295        "--no-lazy", action="store_true",
296        help="use more RAM by computing all outputs before writing (use in case lazy evaluation is broken)",
297    )
298    parser.add_argument(
299        "--verbose", action="store_true",
300        help="increase output verbosity",
301    )
302    parser.add_argument(
303        "--dry-run", action="store_true",
304        help="only print out what will be done, without writing any new files",
305    )
306    parser.add_argument(
307        "--base", type=Path,
308        help="directory containing Hugging Face model config files (config.json, tokenizer.json) for the base model that the adapter is based on - only config is needed, actual model weights are not required. If base model is unspecified, it will be loaded from Hugging Face hub based on the adapter config",
309    )
310    parser.add_argument(
311        "--base-model-id", type=str,
312        help="the model ID of the base model, if it is not available locally or in the adapter config. If specified, it will ignore --base and load the base model config from the Hugging Face hub (Example: 'meta-llama/Llama-3.2-1B-Instruct')",
313    )
314    parser.add_argument(
315        "--trust-remote-code", default=False, action="store_true",
316        help="trust remote code in the model",
317    )
318    parser.add_argument(
319        "lora_path", type=Path,
320        help="directory containing Hugging Face PEFT LoRA config (adapter_model.json) and weights (adapter_model.safetensors or adapter_model.bin)",
321    )
322
323    return parser.parse_args()
324
325
326def load_hparams_from_hf(hf_model_id: str, trust_remote_code: bool) -> tuple[dict[str, Any], Path | None]:
327    from huggingface_hub import try_to_load_from_cache
328
329    # normally, adapter does not come with base model config, we need to load it from AutoConfig
330    config = AutoConfig.from_pretrained(hf_model_id, trust_remote_code=trust_remote_code)
331    cache_dir = try_to_load_from_cache(hf_model_id, "config.json")
332    cache_dir = Path(cache_dir).parent if isinstance(cache_dir, str) else None
333
334    return config.to_dict(), cache_dir
335
336
337if __name__ == '__main__':
338    args = parse_args()
339    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
340
341    ftype_map: dict[str, gguf.LlamaFileType] = {
342        "f32": gguf.LlamaFileType.ALL_F32,
343        "f16": gguf.LlamaFileType.MOSTLY_F16,
344        "bf16": gguf.LlamaFileType.MOSTLY_BF16,
345        "q8_0": gguf.LlamaFileType.MOSTLY_Q8_0,
346        "auto": gguf.LlamaFileType.GUESSED,
347    }
348
349    ftype = ftype_map[args.outtype]
350
351    dir_base_model: Path | None = args.base
352    dir_lora: Path = args.lora_path
353    base_model_id: str | None = args.base_model_id
354    lora_config = dir_lora / "adapter_config.json"
355    input_model = dir_lora / "adapter_model.safetensors"
356
357    if args.outfile is not None:
358        fname_out = args.outfile
359    else:
360        # output in the same directory as the model by default
361        fname_out = dir_lora
362
363    if os.path.exists(input_model):
364        # lazy import load_file only if lora is in safetensors format.
365        from safetensors.torch import load_file
366
367        lora_model = load_file(input_model, device="cpu")
368    else:
369        input_model = os.path.join(dir_lora, "adapter_model.bin")
370        lora_model = torch.load(input_model, map_location="cpu", weights_only=True)
371
372    # load LoRA config
373    with open(lora_config, "r") as f:
374        lparams: dict[str, Any] = json.load(f)
375
376    # load base model
377    if base_model_id is not None:
378        logger.info(f"Loading base model from Hugging Face: {base_model_id}")
379        hparams, dir_base_model = load_hparams_from_hf(base_model_id, args.trust_remote_code)
380    elif dir_base_model is None:
381        if "base_model_name_or_path" in lparams:
382            model_id = lparams["base_model_name_or_path"]
383            logger.info(f"Loading base model from Hugging Face: {model_id}")
384            try:
385                hparams, dir_base_model = load_hparams_from_hf(model_id, args.trust_remote_code)
386            except OSError as e:
387                logger.error(f"Failed to load base model config: {e}")
388                logger.error("Please try downloading the base model and add its path to --base")
389                sys.exit(1)
390        else:
391            logger.error("'base_model_name_or_path' is not found in adapter_config.json")
392            logger.error("Base model config is required. Please download the base model and add its path to --base")
393            sys.exit(1)
394    else:
395        logger.info(f"Loading base model: {dir_base_model.name}")
396        hparams = ModelBase.load_hparams(dir_base_model, False)
397
398    with torch.inference_mode():
399        model_arch = get_model_architecture(hparams, ModelType.TEXT)
400        try:
401            model_class = get_model_class(model_arch)
402            logger.info("Using model architecture: %s", model_arch)
403        except NotImplementedError:
404            logger.error(f"Model {model_arch} is not supported")
405            sys.exit(1)
406
407        class LoraModel(model_class):  # ty: ignore[unsupported-base]
408            model_arch = model_class.model_arch
409
410            lora_alpha: float
411
412            def __init__(self, *args, dir_lora_model: Path, lora_alpha: float, **kwargs):
413
414                super().__init__(*args, **kwargs)
415
416                self.dir_model_card = dir_lora_model
417                self.lora_alpha = float(lora_alpha)
418
419            def set_vocab(self):
420                pass
421
422            def set_type(self):
423                self.gguf_writer.add_type(gguf.GGUFType.ADAPTER)
424                self.gguf_writer.add_string(gguf.Keys.Adapter.TYPE, "lora")
425
426            def set_gguf_parameters(self):
427                logger.debug("GGUF KV: %s = %d", gguf.Keys.Adapter.LORA_ALPHA, self.lora_alpha)
428                self.gguf_writer.add_float32(gguf.Keys.Adapter.LORA_ALPHA, self.lora_alpha)
429                alora_invocation_tokens = lparams.get("alora_invocation_tokens")
430                invocation_string = lparams.get("invocation_string")
431                if invocation_string and not alora_invocation_tokens:
432                    logger.debug("Tokenizing invocation_string -> alora_invocation_tokens")
433                    base_model_path_or_id = hparams.get("_name_or_path")
434                    try:
435                        tokenizer = AutoTokenizer.from_pretrained(base_model_path_or_id)
436                    except ValueError:
437                        logger.error("Unable to load tokenizer from %s", base_model_path_or_id)
438                        raise
439                    # NOTE: There's an off-by-one with the older aLoRAs where
440                    # the invocation string includes the "<|start_of_turn|>"
441                    # token, but the adapters themselves were trained to
442                    # activate _after_ that first token, so we drop it here.
443                    alora_invocation_tokens = tokenizer(invocation_string)["input_ids"][1:]  # ty: ignore[call-non-callable]
444                if alora_invocation_tokens:
445                    logger.debug("GGUF KV: %s = %s", gguf.Keys.Adapter.ALORA_INVOCATION_TOKENS, alora_invocation_tokens)
446                    self.gguf_writer.add_key_value(
447                        gguf.Keys.Adapter.ALORA_INVOCATION_TOKENS,
448                        alora_invocation_tokens,
449                        GGUFValueType.ARRAY,
450                        GGUFValueType.UINT32,
451                    )
452
453            def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
454                # Never add extra tensors (e.g. rope_freqs) for LoRA adapters
455                return ()
456
457            def get_tensors(self) -> Iterator[tuple[str, Tensor]]:
458                tensor_map: dict[str, PartialLoraTensor] = {}
459
460                for name, tensor in lora_model.items():
461                    if self.lazy:
462                        tensor = LazyTorchTensor.from_eager(tensor)
463                    base_name = get_base_tensor_name(name)
464                    # filter base name, ignore tensor transformations for now
465                    data_gen = lambda g=tensor: g  # noqa: E731
466                    if (titem := self.filter_tensors((base_name, data_gen))) is None:
467                        continue
468                    base_name, _ = titem
469                    # note: mergekit-extract-lora also adds token embeddings to the adapter
470                    is_lora_a = ".lora_A.weight" in name or ".lora_embedding_A" in name
471                    is_lora_b = ".lora_B.weight" in name or ".lora_embedding_B" in name
472                    if not is_lora_a and not is_lora_b:
473                        if ".base_layer.weight" in name:
474                            continue
475                        # mergekit-extract-lora add these layernorm to the adapter, we need to keep them
476                        if "_layernorm" in name or ".norm" in name:
477                            yield (base_name, tensor)
478                            continue
479                        logger.error(f"Unexpected name '{name}': Not a lora_A or lora_B tensor")
480                        if ".embed_tokens.weight" in name or ".lm_head.weight" in name:
481                            logger.error("Embeddings is present in the adapter. This can be due to new tokens added during fine tuning")
482                            logger.error("Please refer to https://github.com/ggml-org/llama.cpp/pull/9948")
483                        sys.exit(1)
484
485                    if base_name in tensor_map:
486                        if is_lora_a:
487                            tensor_map[base_name].A = tensor
488                        else:
489                            tensor_map[base_name].B = tensor
490                    else:
491                        if is_lora_a:
492                            tensor_map[base_name] = PartialLoraTensor(A=tensor)
493                        else:
494                            tensor_map[base_name] = PartialLoraTensor(B=tensor)
495
496                for name, tensor in tensor_map.items():
497                    assert tensor.A is not None
498                    assert tensor.B is not None
499                    yield (name, cast(torch.Tensor, LoraTorchTensor(tensor.A, tensor.B)))
500
501            def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
502                dest = list(super().modify_tensors(data_torch, name, bid))
503                # some archs may have the same tensor for lm_head and output (tie word embeddings)
504                # in this case, adapters targeting lm_head will fail when using llama-export-lora
505                # therefore, we ignore them for now
506                # see: https://github.com/ggml-org/llama.cpp/issues/9065
507                if name == "lm_head.weight" and len(dest) == 0:
508                    raise ValueError("lm_head is present in adapter, but is ignored in base model")
509                for dest_name, dest_data in dest:
510                    # mergekit-extract-lora add these layernorm to the adapter
511                    if "_norm" in dest_name:
512                        assert dest_data.dim() == 1
513                        yield (dest_name, dest_data)
514                        continue
515
516                    # otherwise, we must get the lora_A and lora_B tensors
517                    assert isinstance(dest_data, LoraTorchTensor)
518                    lora_a, lora_b = dest_data.get_lora_A_B()
519
520                    # note: mergekit-extract-lora flip and transpose A and B
521                    # here we only need to transpose token_embd.lora_a, see llm_build_inp_embd()
522                    if "token_embd.weight" in dest_name:
523                        lora_a = lora_a.T
524
525                    yield (dest_name + ".lora_a", lora_a)
526                    yield (dest_name + ".lora_b", lora_b)
527
528        alpha: float = lparams["lora_alpha"]
529
530        model_instance = LoraModel(
531            dir_base_model,
532            ftype,
533            fname_out,
534            is_big_endian=args.bigendian,
535            use_temp_file=False,
536            eager=args.no_lazy,
537            dry_run=args.dry_run,
538            dir_lora_model=dir_lora,
539            lora_alpha=alpha,
540            hparams=hparams,
541            remote_hf_model_id=base_model_id,
542        )
543
544        logger.info("Exporting model...")
545        model_instance.write()
546        logger.info(f"Model successfully exported to {model_instance.fname_out}")