Skip to content

fix: pickle Quantized through its class, so torch.load can allow it - #101

Merged
cursor[bot] merged 4 commits into
mainfrom
akshey/torch-load-pickles-fefb
Sep 28, 2026
Merged

cursor[bot] merged 4 commits into
mainfrom
akshey/torch-load-pickles-fefb

Conversation

@aksheyd

@aksheyd aksheyd commented Sep 28, 2026

Copy link
Copy Markdown
Owner

torch.save({"layer": q}, "model.pt") worked, but torch.load("model.pt") refused it. Since PyTorch 2.6, torch.load defaults to weights_only=True, which only calls globals it trusts. Pickles called the static method from_bytes, which pickles as builtins.getattr(Quantized, "from_bytes"). So torch.serialization.add_safe_globals([Quantized]) didn't help, and allowing getattr, which PyTorch's message suggested, would defeat weights_only.

  • Quantized(data) loads the bytes that to_bytes saved, like from_bytes, and pickles now call it. Allowing the class is enough
  • pickles made through from_bytes still load, since from_bytes stays. A test loads one saved by main before this change
  • the class docstring and a README line say to call add_safe_globals([Quantized]) before torch.load
  • the test uses an unpickler that trusts only Quantized, like torch.load does after add_safe_globals, across every pickle protocol. So CI doesn't need torch

Checked with torch 2.14 (CPU): a checkpoint that holds every kind and scale type loads under the default weights_only=True once the class is allowed. Before that, PyTorch's error now suggests allowing quantize.Quantized instead of getattr.

Worth landing before 0.3.0, since every pickle saved by a release keeps whatever it calls. A follow-up PR for PyTorch inputs is stacked on this one.

just lint, just test, and just python pass.

Open in Web Open in Cursor 

torch.load defaults to weights_only=True, which refuses any global it doesn't trust. Pickles called the static method from_bytes, which pickles as builtins.getattr, so add_safe_globals([Quantized]) didn't help, and allowing getattr would defeat weights_only. Quantized(data) now loads the bytes that to_bytes saved, like from_bytes, and pickles call it, so allowing the class is enough. Pickles made through from_bytes still load, since from_bytes stays.
@aksheyd
aksheyd marked this pull request as ready for review September 28, 2026 04:47
@cursor
cursor Bot merged commit 1baf7d8 into main Sep 28, 2026
@cursor
cursor Bot deleted the akshey/torch-load-pickles-fefb branch September 28, 2026 04:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant