Repository: pytorch
Issue: Torch.save with zip serialization causes memory bloat
Labels: module: serialization,triaged

## 🐛 Bug

When using torch.save with zip serialization turned on, I've noticed that the entire pickle file is loaded into memory before anything is written to disk. Only after the entire object has been pickled will torch.save actually start writing to disk. This very inefficient memory usage can cause problems when saving large objects/models on memory constrained systems.

When not using the zip serializer, I do not see such behavior with the reproducing script below.

## To Reproduce

1. Copy script (below) locally
2. Create 4GiB random file: `dd if=/dev/urandom of=random.dat bs=1M count=4`
3. Run script with call to main(use_zip=True)
4. Observe Memory usage of process using top
5. Observe filesize growth of CHECKPOINT_FILE(model.pt) using ls -la
6. Repeat 3-5 calling main(use_zip=False)

## Reproducing script
```
import pickle
import io

import torch

# create random.dat:
# below size is 4GiB but it can be smaller, it is easier
# to see the memory usage using top -u <username> this way

# dd if=/dev/urandom of=random.dat bs=1M count=4

INPUT_FILE="random.dat"
CHECKPOINT_FILE="model.pt"
DEFAULT_CHUNK_SIZE = 1024 * 1024

class FileChunkingIterator():
    def __init__(self, filename, chunk_size_bytes=DEFAULT_CHUNK_SIZE):
        self.filename = filename
        self.file = None
        self.chunk_size_bytes = chunk_size_bytes
        self.num_chunks_read = 0

    def __iter__(self):
        self.file = io.open(self.filename, "rb", buffering=0)
        self.num_chunks_read = 0
        return self

    def __next__(self):
        next_chunk = self.file.read(self.chunk_size_bytes)

        if len(next_chunk) == 0:
            raise StopIteration
        self.num_chunks_read += 1

        return LifetimeWrapper((self.num_chunks_read, next_chunk))

class DummyTensor():
    def __init__(self, filename=INPUT_FILE):
        self.filename = filename

    def __reduce_ex__(self, protocol):
        return type(self), (self.filename,), None, iter(FileChunkingIterator(self.filename))

"""
By default Python Pickler uses memoization to support
self-referencing objects. This causes an extra reference count
for every object to be pickled. This can be a problem if you're
trying to create temporary objects that should be cleaned up
ASAP to keep memory usage low.


To disable this, we have to create a pickler and set the
fast attribute to true. torch.save doesn't permit passing
in a Pickler object, only a pickle_module, so I had to
duck-type the necessary functions in this class.
"""

class DisableMemoizationPicklerModule():
    @classmethod
    def Pickler(cls, data_buf, protocol):
        p = pickle.Pickler(data_buf, protocol)
        p.fast = True
        return p

    @classmethod
    def dump(cls, obj, f, protocol):
        pickle.dump(obj, f, protocol)

"""
Just used to visualize temporary objects being created and
deleted
"""
class LifetimeWrapper():
    object_id = 0
    def __init__(self, i):
        self.id = LifetimeWrapper.object_id
        LifetimeWrapper.object_id += 1
        self.i = i

        print(f"Create obj: {self.id}")

    def __del__(self):
        print(f"Delete obj: {self.id}")

def main(use_zip):
    dummy_tensor = DummyTensor()
    print("start saving")
    torch.save({
        'model_state_dict': dummy_tensor,
        },
        CHECKPOINT_FILE,
        DisableMemoizationPicklerModule,
        _use_new_zipfile_serialization=use_zip)
    print("end saving")

if __name__ == "__main__":
    #main(use_zip=False) # Memory usage stays low and flat,
                        # model.pt filesize increases linearly

    main(use_zip=True)  # Memory usage keeps increasing,
                        # model.pt filesize remains 0
                        # until entire random.dat is loaded
```

## Expected behavior

I would expect main(use_zip=True) behavior to be very similar to main(use_zip=False). Meaning that Memory usage throughout the whole test remains relatively constant and the CHECKPOINT_FILE (model.pt) increases with size linearly throughout the whole test.

This is not the case however, when use_zip=True, memory usage continually increases and the CHECKPOINT_FILE is not written to until the end of the execution. 

## Environment

 - PyTorch Version (e.g., 1.0): 1.9.0+cu102
 - OS (e.g., Linux): CentOS Stream 8 (x86_64)
 - How you installed PyTorch (`conda`, `pip`, source): Conda
 - Python version: 3.9.6 (default, Jul 30 2021, 16:35:19)  [GCC 7.5.0] (64-bit runtime)
 - CUDA/cuDNN version: CUDA Version: 11.0
 - GPU models and configuration:
```
GPU 0: Tesla V100-SXM2-16GB
GPU 1: Tesla V100-SXM2-16GB
GPU 2: Tesla V100-SXM2-16GB
GPU 3: Tesla V100-SXM2-16GB
GPU 4: Tesla V100-SXM2-16GB
GPU 5: Tesla V100-SXM2-16GB
GPU 6: Tesla V100-SXM2-16GB
GPU 7: Tesla V100-SXM2-16GB
```


cc @mruberry