Sequential/Partial unpickling and loading of models

[ollelogdahl]

🚀 The feature, motivation and pitch

Given a 90GB model file that should not be run, but rather just inspected, you might want to load it partially.

Example usage:

descriptors = torch.load_tensors_desc(filename)
# at this state, the memory usage should be close to zero.
for desc in descriptors:
    tensor = torch.load_tensor_from_desc(filename, desc)
    # at this state, the mamory usage should be about the size of the tensor.

This is not a suggestion regarding syntax. At the moment, this is hardly possible:

• A normal torch.load will use about 90GB to load the model.
• Implementing a custom pickle.Unpickler still wants to construct the torch.storage's for all tensors, which conceptually should not be neccessary (also, doesn't work as some tensors may share storage)
• Reimplementing the pickle-format is tedious.

I'm not sure this is an easy patch, or if it even is possible. Please leave any suggestions in the comments.

Alternatives

No response

Additional context

No response

[vadimkantorov]

Related: #79967