[Feature] Single-file zip archives for memory-mapped TensorDicts

[vmoens]

Description

This PR upgrades the memory-mapped serialization story so that a saved TensorDict can travel as a single file instead of a tree of small files -- without introducing a new format. A memmap archive is a standard zip whose entries mirror the memmap directory layout byte-for-byte (meta.json per level, one *.memmap payload per leaf). It is not a third serialization scheme next to memmap directories and consolidation: it is the memmap directory format, in one file. unzip archive.tdz yields a directory that load_memmap() accepts today, and zip -0 -r over an existing memmap directory yields a loadable archive.

td.save("data.tdz")                                   # single file (or archive=True)
td2 = TensorDict.load_memmap("data.tdz")              # zero-copy, lazy
sub = TensorDict.load_memmap("data.tdz", subpath="module/0")  # partial load
rw = TensorDict.load_memmap("data.tdz", mode="r+")    # write-through leaves
buf = data.memmap_like("buffer.tdz")                  # preallocated writable archive
pack_memmap("saved_dir", "saved.tdz")                 # convert existing saves
unpack_memmap("saved.tdz", "saved_dir")               # back to a directory

Motivation

A memmap directory is ideal to work with locally (per-leaf files, in-place writable, partially loadable), but shipping it means copying thousands of small files -- slow on NFS/S3 and easy to corrupt with a partial copy. The consolidated format is a single file but is opaque and disconnected from the memmap tooling. This PR keeps the directory semantics and makes the copyable artifact one file.

The design follows the same playbook as torch.save (uncompressed zip64, aligned records, mmap=True loading) and HuggingFace safetensors (parse small metadata, mmap the file, expose tensors as offset views). Compared to safetensors, the archive natively supports nesting, lazy stacks, tensorclasses, non-tensor data and nested tensors, because the payload schema is the existing memmap one.

How it works

Write path (save/dumps/memmap with a .tdz suffix or archive=True):

def save_as_archive(td, path):
    if td is already memmapped in a directory:
        dir = td.saved_prefix                  # pack directly, no staging
    else:
        dir = td.memmap_like(temp dir next to path)
        # ^ metadata-only staging: meta.json files + empty SPARSE payloads,
        #   no tensor data is written (zero data I/O)
    with ZipFile(path, "w") as zf:
        for file in walk(dir):                 # meta.json first at each level
            if file is a *.memmap payload:
                pad the local header "extra" field so the payload
                starts on a 64-byte boundary   # same trick as torch.save
                stream the bytes STRAIGHT FROM THE SOURCE TENSORS
            else:
                copy the metadata file
    remove staging dir

Tensor data is written exactly once (single pass, plus a cheap CRC-32), so archive saving is on par with a directory save for large leaves (see benchmarks). The staging directory only ever holds metadata and sparse files. Tensordicts containing nested tensors fall back to full staging (memmap_like does not support them).

Alignment is achieved with a private zip "extra" block, so the file remains a perfectly standard zip. compression="deflate"|"bzip2"|"lzma" is available as an explicit opt-in for archival storage (uint8 images, masks and index tensors compress well; float weights barely do). Compressed archives load with load_memmap(path, num_threads=N) inflating the deflate entries in parallel (pread at known offsets + raw zlib, GIL released): ~4x end-to-end at 8 threads (512MB / 2000 leaves: 952ms -> 239ms).

Read path (load_memmap on a file):

def load_memmap(path, mode="r"):
    if path is a zip file:
        parse the central directory once      # entry -> (payload offset, size)
        storage = mmap the whole file as uint8   # MAP_PRIVATE, or MAP_SHARED if mode == "r+"
        path = ArchivePath(root)              # Path-like shim over the entries
    metadata = read meta.json                 # unchanged from here on
    for each leaf: tensor = storage[offset : offset+nbytes].view(dtype).view(shape)
    for each nested entry: recurse

The key implementation move is the _ArchivePath shim: it implements the small Path surface the existing loaders use (/, exists, iterdir, is_dir, open, with_suffix), so the recursive _load_memmap implementations of TensorDict, LazyStackedTensorDict and tensorclasses traverse an archive exactly as they traverse a directory -- no per-class loader duplication. Only the leaf materialization dispatches: real files still produce MemoryMappedTensor.from_filename, archive entries produce views into the shared mapping (the same mechanism from_consolidated uses).

Loading is lazy twice over: only the requested subpath subtree is parsed, and below that, mmap only faults in the pages of leaves that are actually accessed. device="meta" works unchanged. Misaligned or compressed entries (e.g. foreign zip -r archives) transparently fall back to a copying read instead of a view.

Partial loading. subpath accepts a NestedKey (arbitrary nesting normalized as usual) or a "/"-separated string path; it resolves the nested meta.json inside the archive and loads only that subtree. The same argument works for directory prefixes.

Write-through loading. By default an archive-loaded tensordict behaves like from_consolidated(): leaves are copy-on-write views and in-place writes stay in memory. mode="r+" maps the file shared instead, restoring directory semantics (t.add_(1) reaches the file). Because zip stores a CRC-32 per entry, in-place writes leave checksums stale: load_memmap never verifies them, and refresh_archive_checksums(path) re-stamps them (local headers, central directory, data descriptors) before the archive is handed to tools that do verify (unzip, unpack_memmap). mode="r+" refuses layouts where write-through is impossible -- compressed, misaligned (foreign zips) or nested-tensor entries -- rather than silently diverting some writes to copies.

Preallocated writable archives. memmap_like accepts a .tdz target (or archive=True) and creates a zero-filled archive, returned loaded with mode="r+": a single-file dataset buffer. The expanded-tensordict preallocation idiom works unchanged (datum.expand(1_000_000).memmap_like("data.tdz")) since the writer streams from the staging layout, not from materialized data.

Semantics and limitations

• Archives are snapshots layout-wise: tensor values can be updated in place with mode="r+" (or through a memmap_like buffer), but no appending, reshaping or dtype changes after the fact (inherent to any single-file design).
• Nested-tensor tensordicts save and load fine, but fall back to full staging at write time and cannot use mode="r+" / memmap_like archives.

Tests, docs, benchmarks

• test/tensordict/test_methods.py: round-trip across all TestTensorDicts fixtures, plus a TestMemmapArchive class covering payload alignment, zero-copy/single-storage, copy-on-write vs mode="r+" write-through, checksum refresh, memmap_like archives (incl. expanded/stride-0 sources), foreign unaligned zips, pack/unpack round trip, compression, subpath (dir and archive, string and nested-key forms), meta-device loading, nested tensors, load_memmap_, pickling and error paths.
• docs/source/saving.rst gains a "Single-file memmap archives" section with a benchmark figure; pack_memmap/unpack_memmap/is_memmap_archive/refresh_archive_checksums are referenced in docs/source/reference/td.rst.
• benchmarks/common/memmap_benchmarks_test.py gains archive save/load benchmarks; benchmarks/scripts/serialization_formats_bench.py reproduces the format-comparison sweep and the docs figure.

Found in passing (pre-existing, not addressed here): _SubTensorDict memmap round trips are lossy for directories and archives alike (wrong batch size, broken idx on reload), and load_memmap's cross-class dispatch drops the device argument.

Performance

Measured on an Apple-silicon laptop (APFS SSD, warm page cache, min over reps; flat = 8 large leaves, small = 2000 leaves in 100 nested groups). torch.save and safetensors are measured on their fastest paths -- flat tensor dicts (dotted keys), loaded with torch.load(mmap=True, weights_only=True) / mmap-backed load_file, nesting rebuilt on load. Reproduce with python benchmarks/scripts/serialization_formats_bench.py --plot fig.png.

op	layout/size	memmap dir	consolidated	tdz	torch.save	safetensors
save	flat 1GB	135 ms	118 ms	114 ms	637 ms	89 ms
save	small 1GB	231 ms	285 ms	454 ms	615 ms	107 ms
open (lazy)	flat (any size)	~360 us	~49 us	~620 us	~1.6 ms	~35 us
open (lazy)	small (any size)	~88 ms	~5 ms	~21 ms	~44 ms	~820 us
open+read all	flat 1GB	22 ms	19 ms	25 ms	25 ms	30 ms
copy artifact	small 100MB	269 ms	21 ms	22 ms	20 ms	20 ms
copy artifact	flat 1GB	186 ms	160 ms	181 ms	181 ms	174 ms

Takeaways: opening is lazy and size-independent for every format and scales with entry count -- the archive sits between the directory (one file open per leaf) and the flat single-blob formats (one metadata parse). Bulk read throughput is bandwidth-bound and near-identical everywhere. Archive saving matches a directory save for large leaves; many-small layouts pay a per-entry zip overhead. torch.save is the slowest writer at scale (~5x). safetensors is the fastest flat-tensor format, as expected from its minimalism (no structure, no checksums, no writability) -- what tdz buys for its overhead is native nesting/lazy stacks/tensorclass/non-tensor/NJT support, zip-standard tooling, directory interconvertibility, subpath loading and in-place writability. Copying is where all single-file formats beat the directory (13x on the many-small layout locally, more on network filesystems). Multithreaded saves (num_threads, measurable with --num-threads) speed up directory saves of large flat leaves (~1.6x at 8 threads here) but not archives: the zip payload is written sequentially and num_threads only parallelizes the metadata-only staging. A parallel archive writer (precomputed STORED offsets + pwrite) is possible if GPU-resident or network-filesystem saves ever need it.

🤖 Generated with Claude Code