Fix legacy PyTorch checkpoint scanning above 2 GiB #23

Problem

The published stable-diffusion-native v0.2.0-4 rejects a valid legacy PyTorch state-dict checkpoint above 2 GiB. The pinned pickle scanner narrows the remaining buffer span to int, causing valid GLOBAL/name scanning to fail. This is a confirmed native reader defect, separate from the large ZIP-checkpoint timeout.

Found during llamadart#1218. No runtime pin or upstream source was changed to obtain these results. This is not a newly established regression between releases.

Exact baseline

• Owner artifact: v0.2.0-4, macOS arm64, CPU/four threads.
• Bundled upstream: 3f8527a46c54ecf4cb4ed6003da8e8982283c73c.
• Runtime archive SHA-256: 840df6e27e1cac6d378dd5529f7c3b8d3b176417f950ef523c3535c09c92f528.
• Runtime dylib SHA-256: 79ed1ab7c0451596bcb18070d459826a8b04409875873e91f39ffd84b2ffe1ce.
• Published Dart consumer: llamadart 0.15.3, explicit ImageModelRole.checkpoint.

Root cause and model-free reproducer

skip_pickle_object uses find_char(p, (int)(end - p), '\n'), including for GLOBAL operands. read_torch_legacy_file supplies the remaining span of the whole checkpoint, which includes tensor storage. That span can exceed INT_MAX even when the pickle metadata itself is small.

Compiled against the unmodified pinned pickle_io.cpp and its ggml dependencies, this test reserves virtual address space and touches only the short prefix; it needs no model or 2 GiB data file:

#include "model_io/pickle_io.h"
#include <sys/mman.h>
#include <cstring>
#include <iostream>

int main() {
    const char prefix[] = "\x80\x02" "ccollections\nOrderedDict\n)R.";
    const size_t large = (size_t(1) << 31) + 4096;
    auto* data = static_cast<uint8_t*>(mmap(
        nullptr, large, PROT_READ | PROT_WRITE,
        MAP_PRIVATE | MAP_ANON, -1, 0));
    if (data == MAP_FAILED) return 3;
    memcpy(data, prefix, sizeof(prefix) - 1);
    size_t consumed = 0;
    const bool bounded = skip_pickle_object(data, sizeof(prefix) - 1, &consumed);
    const bool containing = skip_pickle_object(data, large, &consumed);
    std::cout << bounded << " " << containing << "\n";
    munmap(data, large);
    return bounded && !containing ? 0 : 2;
}

Observed: 1 0 for the identical valid pickle prefix, with containing span 2147487744. Expected after a fix: both calls succeed and report the same consumed prefix size. MAP_ANONYMOUS can replace MAP_ANON where needed.

Real-checkpoint control

A trusted SD-Turbo Safetensors file (1323 tensors: 709 F16, 614 F32) was repackaged with PyTorch 2.11.0, torch.save(state, ..., pickle_protocol=2, _use_new_zipfile_serialization=False). No checkpoint pickle was executed or loaded in Python.

• Source Safetensors: 2,612,796,864 bytes, SHA-256 4505ca3da482f3e01679f6301bf23cde136e0cb3ffeb66b7e809038f619a19b4.
• Legacy torch-save.pth: 2,612,865,076 bytes, SHA-256 2010fdf410e311b9d457be86e0a530735d9230c4380ba308ff273b4bf6e9b702.
• Public Dart load: typed LlamaModelException (native model-loader initialization failed).
• Direct new_sd_ctx in the same verified owner dylib: null / load failure, without Dart.
• Extracted legacy state-dict metadata: 214,579 bytes; the unmodified pinned metadata parser successfully reads all 1323 tensors when given the bounded metadata span.
• Equivalent full GGUF and Safetensors generate identical 128x128 CPU pixels.
• A 9,800,395-byte TAESD legacy checkpoint generated by the same torch.save procedure loads and generates successfully through both Dart and native controls.

Acceptance

• Fix span handling in the owning upstream reader and consume the fix through the owner workflow; do not patch upstream source in llamadart.
• Add a model-free regression for the same valid pickle prefix with bounded and greater-than-INT_MAX containing spans; keep malformed/truncated input checks.
• Requalify a trusted large legacy checkpoint and the small TAESD control through direct native and public Dart paths.
• Audit the other narrowing find_char call sites without assuming this explains ZIP timeouts.
• Publish/repin only as a separately authorized artifact-consumer change. The current docs/regression PR keeps runtime pins unchanged.

Downstream docs/regression tracking: llamadart draft PR #1237, with runtime pins unchanged.