Skip to content

modules_to_save doesn't work with PyTorch 2.6 #1882

@benjamin-marie

Description

@benjamin-marie
    modules_to_save = ['embed_tokens',"lm_head"],

in get_peft_model triggers the following error with Pytorch 2.6:

[/usr/local/lib/python3.11/dist-packages/unsloth/models/llama.py](https://localhost:8080/#) in get_peft_model(model, r, target_modules, lora_alpha, lora_dropout, bias, layers_to_transform, layers_pattern, use_gradient_checkpointing, random_state, max_seq_length, use_rslora, modules_to_save, init_lora_weights, loftq_config, temporary_location, **kwargs)
   2375             if train_embed_tokens:
   2376                 print("Unsloth: Offloading input_embeddings to disk to save VRAM")
-> 2377                 offload_input_embeddings(model, temporary_location)
   2378             pass
   2379 

[/usr/local/lib/python3.11/dist-packages/unsloth/models/_utils.py](https://localhost:8080/#) in offload_input_embeddings(model, temporary_location)
    765 
    766 def offload_input_embeddings(model, temporary_location : str = "_unsloth_temporary_saved_buffers"):
--> 767     offloaded_W = offload_to_disk(model.get_input_embeddings(), model, "input_embeddings", temporary_location)
    768     new_input_embeddings = torch.nn.Embedding.from_pretrained(offloaded_W)
    769     new_input_embeddings._offloaded_file_location = offloaded_W._offloaded_file_location

[/usr/local/lib/python3.11/dist-packages/unsloth/models/_utils.py](https://localhost:8080/#) in offload_to_disk(W, model, name, temporary_location)
    758     W = W.weight if hasattr(W, "weight") else W
    759     torch.save(W, filename, pickle_module = pickle, pickle_protocol = pickle.HIGHEST_PROTOCOL,)
--> 760     offloaded_W = torch.load(filename, map_location = "cpu", mmap = True)
    761     offloaded_W._offloaded_file_location = filename
    762     return offloaded_W

[/usr/local/lib/python3.11/dist-packages/torch/serialization.py](https://localhost:8080/#) in load(f, map_location, pickle_module, weights_only, mmap, **pickle_load_args)
   1468                         )
   1469                     except pickle.UnpicklingError as e:
-> 1470                         raise pickle.UnpicklingError(_get_wo_message(str(e))) from None
   1471                 return _load(
   1472                     opened_zipfile,

UnpicklingError: Weights only load failed. In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
Please file an issue with the following so that we can make `weights_only=True` compatible with your use case: WeightsUnpickler error: Unsupported operand 149

Check the documentation of torch.load to learn more about types accepted by default with weights_only https://pytorch.org/docs/stable/generated/torch.load.html.


Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions