Struct LegacyMultiStorageSource 

pub struct LegacyMultiStorageSource { /* private fields */ }

Legacy multi-storage source for old PyTorch format (pre-1.6)

Format Analysis

Based on research into PyTorch’s serialization.py and the legacy TAR format:

1. Storage Layout: PyTorch legacy format (0.1.10-1.5) stores data as:

   • Pickle metadata containing tensor definitions
   • A list of storage keys in order
   • Raw binary data with all storages concatenated

2. Boundary Detection Challenge: After extensive research, I found that:

   • PyTorch does NOT store explicit storage boundaries in the file
   • Storages are concatenated in the order specified by the storage keys list
   • Each tensor references its storage by key and specifies offset/size

3. Why True Lazy Loading is Difficult:

   • To determine storage boundaries, we would need to: a. Parse ALL tensor metadata to find which storage each uses b. Track the maximum extent of each storage based on tensor usage c. Infer boundaries from the gaps between storages
   • However, the TensorSnapshot abstraction hides storage keys in closures
   • This would require deep modifications to the pickle parsing logic

Current Implementation

This implementation provides a best-effort approach:

• Supports setting a storage map if boundaries can be determined externally
• Falls back to loading the entire blob if boundaries are unknown

Implementations

impl LegacyMultiStorageSource

pub fn new(path: PathBuf, data_offset: u64, data_size: u64) -> Self

Create a new legacy multi-storage source

pub fn set_storage_keys(&self, keys: Vec<String>)

Set the ordered storage keys from the pickle

pub fn track_storage_usage(&self, storage_key: &str, offset: usize, size: usize)

Track storage usage from tensor access This is called from within tensor loading closures

pub fn read(&self, key: &str) -> Result<Vec<u8>>

Read data for a specific storage key Only loads the specific storage portion, never the entire blob