Repository: pytorch
Issue: RFC: multi-part `torch.load`/`torch.save` to support huge models and/or low CPU memory
Labels: module: nn,module: serialization,triaged,onnx-triaged,onnx-needs-info

## 🚀 Feature

## The need

Models are getting bigger and there are times when loading all the params  from external storage into CPU memory at once is either not possible or calls for some extra craftiness to make it work.

Examples:
- users wanting to use large models that can't fit into 12GB free colab, yet which fit just fine into the provided GPU (this issue is definitely solvable with current tools, but it can be made easier.)
- I encounter this issue quite often when working with multi-node training, or when [converting multi-gpu partitioned param checkpoints into normal checkpoints](https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/utils/zero_to_fp32.py) like what Deepspeed does. Here I had to reconstitute params partitioned across many state dicts the main SD in CPU memory before I could save it, and the input can easily be 100GB or bigger, calling for quite a lot of RAM to complete. And needing 2x RAM. This is an example of needing a leaner `torch.save`.

   Another immediate use-case is converting from one framework to another. e.g. at BigScience we will shortly need to be able to convert a 200B param model from Megatron-Deepspeed to Megatron-LM to HF Transformers, so currently such conversion will require at least 800GB CPU RAM for fp16, and 1600GB CPU RAM for fp32.
   
   In the same context reshaping a huge checkpoint to a different TP/PP layout too may require a huge amount of CPU memory.

## Proposal

The current storage facility is pickled data, which is a zip file with multiple separate files internally, but it's currently not possible to load these parts selectively - it is all-in-one. Same for saving data, one has to first build up a full state dict and only then can save it (important when params are spread out across gpus, e.g. in model parallelism).

This RFC proposes switching to an optional different format, while supporting the existing one, where there is a higher granularity than all tensors in the saved state dict, which reduces CPU memory requirements.

The proposal is to add support for tied to filesystem dicts, so that the dict is never fully loaded into memory, but instead gets loaded on demand, one key at a time. Same with saving, one can "grow" the saved state dict one key at a time.

One possible solution is using [dbm](https://docs.python.org/3/library/dbm.html) because it's a python built-in and it provides the required functionality. So a working prototype follows.

## Possible Implementation

You can play with the code directly using this notebook: https://github.com/stas00/porting/blob/master/pytorch/torch-save-load-multi-part-dbm.ipynb

```

import dbm
import pickle
import torch
from torch import nn
from collections.abc import MutableMapping

class DBMStateDict(MutableMapping):
    def __init__(self, path):
        self.path = path
        self.db = dbm.open(path, 'c')
        
    def __len__(self):
        return len(self.db.keys())
    
    def __getitem__(self, key):
        return pickle.loads(self.db[key])
    
    def __setitem__(self, key, value):
        self.db[key] = pickle.dumps(value)
        # it looks like dbm syncs immediately
    
    def __delitem__(self, key):
        return self.db.pop(key)
    
    def keys(self):
        return [k.decode() for k in self.db.keys()]
    
    def __iter__(self):    
        return iter(self.db)
            
    def copy(self):
        return DBMStateDict(self.path)
    
    def __del__(self):
        self.db.close()
                            
def save_new(sd_dict, path):
    sd = DBMStateDict(path)
    for k,v in sd_dict.items():
        sd[k] = v
        
def load_new(path):
    # this doesn't load the whole sd into memory!
    return DBMStateDict(path)

class SubNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1, 1)
        self.fc2 = nn.Linear(1, 1)

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = SubNet()

# original
m = Net()
path = "model1.pt"
sd_dict = m.state_dict()
torch.save(sd_dict, path)
sd_dict = torch.load(path)
m.load_state_dict(sd_dict)

# same but loading / saving one key at a time   
m = Net()
path = "model1.dbm"
sd_dict = m.state_dict()
save_new(sd_dict, path)
sd_new = load_new(path)
m.load_state_dict(sd_new)
```

This is a rough prototype and doesn't pretend to be complete.

And of course, it'll have to be made back-compatible. But the point here is to demonstrate the feature and not to provide a complete new version.

## Problems

- I can see this is not `OrderedDict` which could be a problem in some circumstances (but good enough as a prototype).

## Other approaches

- multi-file checkpoint splitting https://github.com/finetuneanon/transformers/#checkpoint-splitting / [code](https://github.com/finetuneanon/transformers/blob/ca5d90ac1965982db122a649c2c9c902bde74a03/src/transformers/modeling_utils.py#L417-L443)
- Deepspeed's pipeline parallelism saves each layer in a separate file

## Credits

The main class has been inspired by [SplitCheckpoint](https://github.com/finetuneanon/transformers/blob/ca5d90ac1965982db122a649c2c9c902bde74a03/src/transformers/modeling_utils.py#L417-L443)

Thank you!

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki