Skip to content

fix: carry the pickle line-span fix for legacy checkpoints above 2 GiB - #26

Merged
leehack merged 2 commits into
mainfrom
fix/legacy-pickle-span-carry
Oct 22, 2026
Merged

leehack merged 2 commits into
mainfrom
fix/legacy-pickle-span-carry

Conversation

@leehack

@leehack leehack commented Oct 22, 2026

Copy link
Copy Markdown
Owner

Fixes the legacy PyTorch checkpoint reader for files above 2 GiB by carrying a one-file patch to upstream's pickle scanner until a pin includes a fix. Refs #23. Stacked on #25: its test commit (d9bacef) is the first commit here.

Decision needed

AGENTS.md said to update the pin instead of patching vendored sources. Upstream master (0518b34) still has the bug, and upstream's own agent rules forbid agent-opened PRs, so a pin bump is not available. This PR adds a narrow exception: the submodule is never edited; a fix the pin lacks may be carried as a patches/ entry with a test that fails without it. If you would rather not carry patches, merge #25 alone and close this.

Root cause

skip_pickle_object finds the end of GLOBAL, INST, PERSID and other line operands with find_char(p, (int)(end - p), '\n'). read_torch_legacy_file passes it the rest of the whole checkpoint, tensor storage included. Above INT_MAX the cast goes negative and nothing is found; above UINT32_MAX only the low 32 bits survive.

Change

File Change
patches/pickle_line_span.patch find_line_end caps the scan at INT_MAX (no pickle line is that long); all 13 narrowing casts in pickle_io.cpp use it. Only pickle_io.cpp calls find_char.
CMakeLists.txt Applies each patch to a build-tree copy of the sources it names and compiles the copy. Configuring fails if a patch no longer applies or changes nothing. Minimum CMake 3.18 (TARGET_DIRECTORY).
tests/test_pickle_span.py Compiles the patched scanner; the expected failure is gone. A new control fails once the unpatched pin passes, so the patch is dropped when upstream fixes it.
tools/build.py, tools/package_release.py build-info.json and manifest.json list carried patches under upstream.patches.
AGENTS.md, README.md, patches/README.md The carried-patch rule and inventory.

The patch also applies cleanly to upstream 0518b34 (the pin the upstream-982 lane moves to).

Validation (macOS arm64, Apple clang)

Check Result
python3 -m unittest discover -s tests 60 tests OK, 7 skipped
tools/build.py build --target macos-arm64 + validate_artifacts.py macos-arm64 builds the patched copy (Carrying patches/pickle_line_span.patch); ok
Mutants of the patch: no cap, cap at UINT32_MAX, scan ignoring end, and each of the 6 skip_pickle_object call sites reverted 9/9 killed (8 by the large-span test, the overread by the bounded and ASan tests)

Direct native new_sd_ctx + one-step 128x128 CPU generation (4 threads, seed 42). Legacy files were written from local weights with a pickle writer, without torch and without loading any pickle:

Model TAESD Pinned runtime Patched runtime
SD-Turbo F16 GGUF safetensors 5ab8864b20e3cd03 5ab8864b20e3cd03
SD-Turbo legacy .pth, 2,612,843,011 bytes safetensors load failed 5ab8864b20e3cd03
SD-Turbo F16 GGUF legacy .pth, 9,797,897 bytes 5ab8864b20e3cd03 5ab8864b20e3cd03

Not run here: Linux, Windows, Android and iOS builds (hosted CI covers them), the public Dart load path, and an upstream submission (needs the maintainer). The ZIP-checkpoint timeout in #24 is a separate cause and untouched. No release, tag or runtime repin.

The pinned stable-diffusion.cpp narrows the remaining span to int before
scanning a pickle operand for its newline, and read_torch_legacy_file passes
the rest of the whole checkpoint, so a valid legacy checkpoint above 2 GiB is
rejected (#23). Compile the pinned pickle_io.cpp and check valid and malformed
pickles in exact-size buffers and in read-only spans above INT_MAX and
UINT32_MAX. The valid large-span case is an expected failure until an upstream
fix is pinned; a bump that fixes it fails the run until the marker goes.
The pinned stable-diffusion.cpp narrows the rest of a legacy checkpoint to
int before scanning a pickle operand for its newline, so a valid legacy
.pth above INT_MAX is rejected (#23). Upstream master has the same code.

Carry patches/pickle_line_span.patch until a pin includes a fix: CMake
applies it to a build-tree copy of pickle_io.cpp and compiles the copy, so
the submodule stays pristine, and configuring fails once the patch no longer
applies or changes nothing. The scan is capped at INT_MAX; no pickle line is
that long.

The span tests now compile the patched scanner, so valid pickles in spans
above INT_MAX and UINT32_MAX pass, and a control fails once the pinned
upstream no longer needs the patch. Build info and the release manifest list
carried patches next to the upstream commit.
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