_pickle.UnpicklingError: Unsupported type torch._tensor._rebuild_from_type_v2

[rangan2510]

Getting an issue while merging BioMistral and Zephyr. Trying to replicate the code here.

Here's my yaml.

slices:
  - sources:
      - model: BioMistral/BioMistral-7B
        layer_range: [0, 32]
      - model: HuggingFaceH4/zephyr-7b-beta
        layer_range: [0, 32]
merge_method: slerp
base_model: BioMistral/BioMistral-7B
parameters:
  t:
    - filter: self_attn
      value: [0, 0.5, 0.3, 0.7, 1]
    - filter: mlp
      value: [1, 0.5, 0.7, 0.3, 0]
    - value: 0.5
dtype: bfloat16

This is the error.

File "D:\miniconda3\envs\mergekit\Scripts\mergekit-yaml.exe\__main__.py", line 7, in <module>
  File "D:\miniconda3\envs\mergekit\Lib\site-packages\click\core.py", line 1157, in __call__
    return self.main(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\miniconda3\envs\mergekit\Lib\site-packages\click\core.py", line 1078, in main
    rv = self.invoke(ctx)
         ^^^^^^^^^^^^^^^^
  File "D:\miniconda3\envs\mergekit\Lib\site-packages\click\core.py", line 1434, in invoke
    return ctx.invoke(self.callback, **ctx.params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\miniconda3\envs\mergekit\Lib\site-packages\click\core.py", line 783, in invoke
    return __callback(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\Github\mergekit-mistral\mergekit\mergekit\options.py", line 78, in wrapper
    f(*args, **kwargs)
  File "D:\Github\mergekit-mistral\mergekit\mergekit\scripts\run_yaml.py", line 47, in main
    run_merge(
  File "D:\Github\mergekit-mistral\mergekit\mergekit\merge.py", line 87, in run_merge
    for _task, value in exec.run():
  File "D:\Github\mergekit-mistral\mergekit\mergekit\graph.py", line 191, in run
    res = task.execute(**arguments)
          ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\Github\mergekit-mistral\mergekit\mergekit\io\tasks.py", line 81, in execute
    x = loader.get_tensor(name, device=self.device or "cpu")
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\Github\mergekit-mistral\mergekit\mergekit\io\lazy_tensor_loader.py", line 143, in get_tensor
    self.current_shard = TensorLoader.get(
                         ^^^^^^^^^^^^^^^^^
  File "D:\Github\mergekit-mistral\mergekit\mergekit\io\loader.py", line 49, in get
    return LazyPickleLoader(shard_path, device=device)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\Github\mergekit-mistral\mergekit\mergekit\io\loader.py", line 64, in __init__
    self.index = torch.load(path)
                 ^^^^^^^^^^^^^^^^
  File "D:\miniconda3\envs\mergekit\Lib\site-packages\torch\serialization.py", line 1026, in load
    return _load(opened_zipfile,
           ^^^^^^^^^^^^^^^^^^^^^
  File "D:\miniconda3\envs\mergekit\Lib\site-packages\torch\serialization.py", line 1438, in _load
    result = unpickler.load()
             ^^^^^^^^^^^^^^^^
  File "D:\miniconda3\envs\mergekit\Lib\site-packages\torch\serialization.py", line 1431, in find_class
    return super().find_class(mod_name, name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\Github\mergekit-mistral\mergekit\mergekit\io\lazy_unpickle.py", line 108, in find_class
    raise pickle.UnpicklingError(f"Unsupported type {module}.{name}")
_pickle.UnpicklingError: Unsupported type torch._tensor._rebuild_from_type_v2

What am I doing wrong? Right now, I am using Python v3.11.8. CPU only, no cuda. Here are the packages in the system.

accelerate==0.27.2
annotated-types==0.6.0
certifi==2024.2.2
charset-normalizer==3.3.2
click==8.1.7
colorama==0.4.6
filelock==3.13.3
fsspec==2024.3.1
huggingface-hub==0.22.2
idna==3.6
immutables==0.20
Jinja2==3.1.3
MarkupSafe==2.1.5
-e git+https://github.com/cg123/mergekit.git@ef7f5850f2658892b87f83d4955f817d398cd421#egg=mergekit
mpmath==1.3.0
networkx==3.2.1
numpy==1.26.4
packaging==24.0
peft==0.10.0
protobuf==5.26.1
psutil==5.9.8
pydantic==2.6.2
pydantic_core==2.16.3
PyYAML==6.0.1
regex==2023.12.25
requests==2.31.0
safetensors==0.4.2
sentencepiece==0.2.0
sympy==1.12
tokenizers==0.15.2
torch==2.2.2
tqdm==4.66.2
transformers==4.39.3
typing_extensions==4.10.0
urllib3==2.2.1

[cg123]

Looks like the lazy unpickler doesn't support one of those models. There was a bug making the lazy unpickler be always used that was just resolved by #247 - if you try again on main it should work.

I'll also look into if I can handle _rebuild_from_type_v2 safely in the lazy unpickler.

[rangan2510]

Hi, the issue persists.

[cg123]

Have you tried with the --lazy-unpickle argument?

[joreyolo]

hello, I had the same error while merging BioMistral/BioMistral-7B and Rakuten/RakutenAI-7B.

There was a bug making the lazy unpickler be always used that was just resolved by https://github.com/arcee-ai/mergekit/pull/247 - if you try again on main it should work.

I followed this and tried with --lazy-unpickle argument, but the error still occurred.

[AmineBechar07]

Any Solutions?