mergekit-extract-lora does not extract - the destination is empty

[raulod]

Tried to extracting LORA from gemma-2-it

mergekit-extract-lora models/gemma-2-2b-it/ models/gemma-2-2b models/lora/gemma-2-it-lora
The installed version of bitsandbytes was compiled without GPU support. 8-bit optimizers, 8-bit multiplication, and GPU quantization are unavailable.
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:00<00:00, 11.28it/s]
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:01<00:00,  1.66it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▌| 288/289 [08:55<00:01,  1.86s/it]
Traceback (most recent call last):
  File "/usr/local/bin/mergekit-extract-lora", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/usr/local/lib/python3.11/site-packages/click/core.py", line 1157, in __call__
    return self.main(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/click/core.py", line 1078, in main
    rv = self.invoke(ctx)
         ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/click/core.py", line 1434, in invoke
    return ctx.invoke(self.callback, **ctx.params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/click/core.py", line 783, in invoke
    return __callback(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/code/thirdparty/mergekit/mergekit/scripts/extract_lora.py", line 576, in main
    lora_weights, ranks = extract_lora(
                          ^^^^^^^^^^^^^
  File "/code/thirdparty/mergekit/mergekit/scripts/extract_lora.py", line 246, in extract_lora
    base_weight = base_loader.get_tensor(f"{module_name}.weight")
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/code/thirdparty/mergekit/mergekit/io/lazy_tensor_loader.py", line 127, in get_tensor
    raise KeyError(key)
KeyError: 'lm_head.weight'

It look long time to complete but the destination folder is empty.

[Touch-Night]

Gemma 2 does not have lm_head.weight

[Touch-Night]

Try this:

    for module_type, module_name in tqdm(module_details):
        try:
            base_weight = base_loader.get_tensor(f"{module_name}.weight")
            finetuned_weight = finetuned_loader.get_tensor(f"{module_name}.weight")
        except KeyError:
            logging.warning(f"Module '{module_name}' not found in the model!")
            continue