RFC: split checkpoint load/save for huge models
Author: stas00
Pull request: no
Labels: name: WIP

# 🚀 Feature request

While [discussing with pytorch devs adding the ability to load/save state_dict on the finer granularity level and not needing to manifest the whole state_dict in memory](https://github.com/pytorch/pytorch/issues/64327), we have an additional issue of the  model file just being too large. I'd like to propose for `transformers` to support multi-part checkpoints.

Reasons for the need: 
- the hub limitation: Cloudfront does not support >20GB files so downloads via s3 can't be fast with those large files
- the current pytorch issue loading the whole state_dict into memory and requiring 2x model size in memory - checkpoint conversion is quite demanding on memory as well for the same reason.
- in general it's a potential issue for users with imperfect up/down internet connection. uploading/downloading 25GB files is still not easy for all.

Possible solutions:

1. as mentioned [here](https://github.com/pytorch/pytorch/issues/64327), [SplitCheckpoint](https://github.com/finetuneanon/transformers/blob/ca5d90ac1965982db122a649c2c9c902bde74a03/src/transformers/modeling_utils.py#L417-L443) already implements a possible solution which saves each state_dict's key separately
2. as solution 1 but we may save groups of these - e.g. save each layer's keys together in one pickled state_dict per layer. I looked at some large models and they will have a huge amount of keys, e.g. even t5-small is ~150 keys. But this approach would be more complicated since we now need to define the container block and it'll be different from model to model. May be by sub-module? So perhaps the first solution is much more simple.

The only addition I'd propose to actually name the files with the full key name rather than obscure files like `m18.pt` as implemented by  [SplitCheckpoint](https://github.com/finetuneanon/transformers/blob/ca5d90ac1965982db122a649c2c9c902bde74a03/src/transformers/modeling_utils.py#L417-L443) , and which require an extra file to do look ups.

So my proposal is:
```
config.json
merges.txt
README.md
tokenizer.json
vocab.json
pytorch_model/map.pt
pytorch_model/shared.weight.pt
pytorch_model/encoder.embed_tokens.weight.pt
[...]
pytorch_model/encoder.block.3.layer.0.SelfAttention.v.weight.pt
[...]
pytorch_model/decoder.block.5.layer.1.EncDecAttention.q.weight.pt
[...]
pytorch_model/lm_head.weight
```
and these are all raw files not belonging to any archive. and map just has the list of keys in their order for when `OrderedDict` is important.

the cost of the 1st solution is somewhat slower save/load. I haven't benchmarked, but the IO will be the bottleneck here, and the ZIP structure currently gets unravelled one tensor at a time anyway, so the difference is likely to be negligible. 

other solutions are welcome.

Other examples of split checkpoints:
- Deepspeed's pipeline (PP) saves each layer as a separate checkpoint, which allows to quickly change the PP degree at run time. 

Threshold:
- need to define the threshold at which we automatically switch to this multi-part format unless the user overrides the default. Probably can use the size of the model as the measurement. I think it should be 3B or even less.  if model size == 3B the resulting file size are:

1. 6GB in fp16/bf16
2. 12GB in fp32.

@patrickvonplaten, @patil-suraj, @LysandreJik, @sgugger, @julien-c

Comments:
Thanks for the write-up! 

I'm happy with a multi-part checkpoint solution I think if it's clearly the cleanest solution to reduce the 2x-model-size-memory-requirement-at-loading problem. If there is a rather simple solution to solve the 2x model size memory problem while keeping a single checkpoint file I'd be in favor of this because:
- people are used to it
- it's more in line with TF and PT
- it's less complex as for large models 500+ checkpoints can really clutter the model repo

IMO, the points: 
- the hub limitation: Cloudfront does not support >20GB files so downloads via s3 can't be fast with those large files
- in general it's a potential issue for users with imperfect up/down internet connection. uploading/downloading 25GB files is still not easy for all.

are not very important since when dealing with such large models, download times will take some time anyways and even a speed-up in download from 10min to 3min is no that important since the files are cached afterwards anyways IMO. Also I'm not sure if it's an advantage to being able to load layers individually since I can see people getting lost in for which layers the upload was successful / unsuccessful. 

Thanks for your feedback, @patrickvonplaten 

>     * it's less complex as for large models 500+ checkpoints can really clutter the model repo

Have you meant checkpoints or files? The proposal is to have a separate dir so from the top level nothing has changed.

I'm with you that when there are many hundreds of params it may become somewhat hard to manage. As you flagged watching that all files got correctly uploaded and more chances for a timeout on a single file during the download. I was sitting on the fence about posting this proposal because of these possible issues and after encountering multiple problems with uploading a 26GB model I decided to go ahead with it.

So I'm hoping that perhaps some alternative solution will come from this discussion.



How about a compromise that regroups weights in a state dict until we get to the biggest file possible still < 20GB. This way we don't clutter the repo and only split the checkpoint when necessary, while also avoiding the annoying timeouts. We can then have a load in several part to recover the whole weights. It would require a bit of refactor in `from_pretrained` but not too much, the bulk of the work would be in `save_pretrained` to split the checkpoint in several parts when needed.

Well, the idea is to make models a bit lighter in size, and 20GB is still a huge size! Though your suggestion is sensible.

But to put the issue in perspective consider our BigScience's plan to train a 200B model in a few months. That means the model will be 400GB in fp16, 800GB in fp32.

One of the possible approaches there would be to split it by layer, each of which would still be huge. so there perhaps separate param per file would be just right.



BTW, there was another related discussion of compressing the models: https://github.com/huggingface/transformers/issues/10008

And here is a non-completed PR from @anthon-l, where the checkpoint is 19GB https://github.com/huggingface/transformers/pull/10301

I guess t5-11b is the only other one that's very large on the hub - 45GB in size, or do we have more examples?


Think the biggest models we have on the hub are the 45GB T5 models

if we have more and more large models, for instance the ones from Bigscience, i think we need to split checkpoints indeed. 

> download times will take some time anyways and even a speed-up in download from 10min to 3min is no that important
(quote from @patrickvonplaten)

Depending on where you (or your server) are in the world and your kind of connection, the delta can be way bigger – 10x to 100x download speed is likely in a lot of cases in my opinion.

If we end up having 400GB artefacts in the hub, I think strictly no-one will be able to load them if we stick to single file (which, as @stas00 mentioned, _implies_ download from `us-east` S3)

More precisely, the speed difference between Cloudfront and S3 increases with how "far" (in networking topology) you are from the AWS `us-east` region.

So the BigScience 176B model is training and the bf16 checkpoint is 329GB.

So please let's resume this discussion and bring it to completion soon as we need to start working on converting the checkpoints since the group desires to already start doing inference on an even incompletely trained model.

While at it let's also incorporate the discussion of `torch.load` due to the `pickle` security issues.

Thank you!

@patrickvonplaten, @patil-suraj, @LysandreJik, @sgugger, @julien-c, @thomwolf 

also @Narsil

And @julien-c's comment to remind:

> shards should be in the 5-30GB range if possible (in all case, always smaller than 30GB)

(that was changed recently from 20GB upper limit)

Here is the breakdown of the 176B model's layers in BF16 weights:

```
$ NHIDDEN=14336; NLAYERS=70; SEQ_LEN=2048; VOCAB_SIZE=250680; python -c "h=$NHIDDEN; l=$NLAYERS; \
s=$SEQ_LEN; v=$VOCAB_SIZE; t=2*(12*h**2 + 13*h); r=2*(v*h + s*h + 2*h); \
print(f'BF16 Transformer block size: {t/2**30:.02f}GB, the rest is: {r/2**30:.02f}GB, Total {(l*t+r)/2**30:.02f}GB')"
BF16 Transformer block size: 4.59GB, the rest is: 6.75GB, Total 328.34GB
```

So basically, save each layer separately to the tune of 4.59GB per checkpoint - 70 of those, one additional checkpoint with the rest of the weights of 6.75GB.

A total of 71 shards.

This is very close to how Megatron-Deepspeed's pipeline shards are already saved, Except it saves each layer in a separate file, perhaps we should follow suite, whereas I dumped all the non-transformer layers into one in the breakdown above. It makes it very easy to load one layer at a time following the model's layers one by one.

While working out the design it's good to look at the concrete model, but also to consider smaller and bigger models, as those surely are coming. That is to validate that one file per layer is a good solution, if you scroll back up to the beginning of this discussion this was my original proposal.

I'm okay with a system that will group weights until we get to the max < 5GB (or any number) then creates a new shard automatically for the next weights. Saving each layer separately has several drawbacks:
- you can't decide to do it automatically, as for smaller models it doesn't necessarily make sense
- it will result in multiple calls to fetch every file (even when they are downloaded, we will still call the HF API to see if there is an update for each of the files since it's in the core of the system in Transformers)

so I'd avoid that solution.

The `save_pretrained` method should then probably do this by default (so smaller models would be saved as before and large models would be automatically split) even if it's a breaking change, as it interfaces better with the Hub. We would of course have an argument to save the model as a whole for users that want to reuse their large model directly, without using `from_pretrained`.

The main problem when reloading is then going to be how to handle the calls to the Hub to know when `from_pretrained` should ask for one file `pytorch_model.bin` or several (probably named `pytorch_model_xxx.bin`. Since we can't call the API that returns the list of model files as it's not optimized (feel free to correct me @julien-c), we should have a config attribute that tells us when there are multiple checkpoints so we  know the list of files we have to download.

I think the discussion around `torch.load` is orthogonal to this discussion (at least the way everything is written right now in Transformers) so I would leave it for a different issue.


The problem of "has-it-changed" can be solved by saving an additional file which is the index to all other files, which could also provide a map of how to load the rest of the files. i.e. it'll tell which specific weights are located in which file. Since your proposal of `pytorch_model_xxx.bin` is somewhat difficult on its own since there is no indication of what's inside `pytorch_model_xxx.bin`

If there is just one index file to check, the hub functionality can check if there is `pytorch_model.bin` (single part model) or ` pytorch_model_index.bin` (multi-part model) - so only 2 files to check.

The problem of how to slice can be solved by a threshold setting at which the model starts sharding - probably 5GB is a good choice, since downloading more than 5GB can be an issue for those w/o bad connectivity.

And of course the math would be different for a model with bf16/fp16 weights vs fp32 weights as the latter will use 2x bytes. So the grouping should then happen by `params*dtype_bytes < gb_threshold`

So as long as the model is smaller than 5GB, it will remain a single file always.

But I'm mostly repeating Sylvain here, just added the component of dtype size  (2 vs 4).

The idea of an index that provides a map is very good, I like it!

Should the threshold be configurable by the user with the default set to what we prefer for the hub?

I'm asking since there will be users who will not use the hub and will perhaps want a 100GB model file.

While we are at it, does the hub uploader check that the file is not bigger than 20GB before attempting the upload?

Adding my two cents:
- I agree that threshold should be configurable. We implemented a similar threshold in the `datasets`' `push_to_hub` [method](https://github.com/huggingface/datasets/pull/3098/files#diff-247ce37d2f96c4fc8e77ce55dc9922675a5f7f561f8df0cbc563d13048c59abdR3319) and it should provide the most flexibility. Having 5GB as a default also sounds good to me. 
- I don't think the hub checks that the file is bigger than 20GB. Ideally we'd send out a very visible warning when a checkpoint is saved with larger shards than 20GB.
- In `datasets`, for a file split over three shards, we use the following format:
    - `<filename>-00000-of-00003.<extension>`
    - `<filename>-00001-of-00003.<extension>`
    - `<filename>-00002-of-00003.<extension>`
    
        This naming can be adopted if each individual shard are valid files in the `<extension>` format. If we split semantically (on weights) rather than on bytes, this naming makes sense, otherwise, the extension would be moved inside: `<filename>.<extension>-00000-of-00003` .
    
        If there are no other constraints on the naming, I'd vote to have a similar format, for both consistency's sake and to be able to quickly identify which files should be replaced in case of a `push_to_hub` that would need to overwrite a previously uploaded checkpoint, with a possibly differing number of shards.

And my 2 cents.

Can the config actually hold the location of its various weight files (relative to the repo) in something like `config.weight_files` which could default to `["pytorch_model.bin"]` and `["tf_model.h5"]` depending on the Arch.. ?

I think the config has to be downloaded for any model download so it could become the index itself no ? And users can choose whichever splitting they seem is more appropriate ?

Having an index is very important IMO to allow resharding existing repos (it does make them inaccessible to the old versions of the library that don't have the sharding logic though)

(nit) `pytorch_model.bin.index` is a better filename than `pytorch_model_index.bin` in my opinion (this scheme could also apply to other files and will be very easy to check on the hub side)

> * If we split semantically (on weights) rather than on bytes, this naming makes sense, otherwise, the extension would be moved inside: `<filename>.<extension>-00000-of-00003` .

absolutely by layer, and not bytes, which makes it possible to load parts of the model directly on different gpus. e.g. think Pipeline which will load for example file 0+1 on gpu 0, file 2+3 on gpu 1, etc.

> Can the config actually hold the location of its various weight files

that's exactly what we have been discussing - there will be an index file (not config) which will map weights to files, which of course then allows you to point to whatever file you want, including a different format.

> which will map weights to files,

Why do you need that ? If the files are in a nice format, looking up their weights names should be extremely fast (and never load the tensors). The main issue I see with an external index containing weight_names is the potential for it being outdated/inconsistent with what's actually within the other files. Recovering from the inconsistent state is going to add more code logic.

If we keep the index in the config (withtout weight names) then, at most a file is missing which is sort of the same issue as `pytorch_model.bin` missing.

Then we can just recreate the weight_name index when looking into the files and so there's no option for it being inconsistent because we don't commit it to disk.


Even in pure zip/pickle from PT I think it's possible to recreate that index on the fly with very little overhead but I will have to check.



The index files is needed for the following reasons:

1. only one file to look up for is-modified check - to see if the model needs to be re-downloaded - I don't see any reason why the code that generates the index file and the shards should lead to out of sync situation. It can also include md5 checksums to able to always verify that the files are in sync, which we probably need to do anyway.
2. in the case of pipeline parallelism it allows one to load just the weights each pipeline stage needs and this can't be derived from a file name and will require loading all files for each stage, rather than one file. It'd be different if each layer had its own file, then it'd be self-naming - but it looks like the proposal so far is that we are going to pack multiple layers together for smaller models. The Deepspeed pipeline saves each layer in its own file named after that layer id and thus avoids the need for an index file.

BTW, git LFS already provides us with checksums: e.g., https://huggingface.co/gpt2/blob/main/pytorch_model.bin
```
SHA256: 7c5d3f4b8b76583b422fcb9189ad6c89d5d97a094541ce8932dce3ecabde1421
```
so it should be instant to validate that they match w/o actually running a slow hashing program.

So download the index file - compare the hash of LFS files it points to - verification done.

> only one file to look up for is-modified check - to see if the model needs to be re-downloaded - I don't see any reason why the code that generates the index file and the shards should lead to out of sync situation. It can also include md5 checksums to able to always verify that the files are in sync, which we probably need to do anyway.

You cannot protect from users doing things without you expecting it, like uploading single files instead of using the API we designed and therefore breaking consistency. My experience with any system, is that if it CAN be out of sync, it will end up out of sync. And recovering is usually much harder than preventing it altogether.

> So download the index file - compare the hash of LFS files it points to - verification done.

This means that the files look coherent, it doesn't mean that weight `X` is actually within file `Y`.

> In the case of pipeline parallelism it allows one to load just the weights each pipeline stage needs

Checking the weight names is very fast:
```
Found 148 tensors, took 0:00:00.088987 on file of len 474.9MiB
Found 436 tensors, took 0:00:00.107209 on file of len 2.9GiB
```
(Corresponding code at the end)

Downloading the files selectively still requires careful sharding on your end (if the sharding is poorly done, and some layers are spread across multiple files, then you still need to download everything).

This is very reasonable, and enabling downloading only what's necessary seems like a great feature for such large models.

Having no weight_index and good internal sharding could enable you as a user to do that (`shard1.h5`, `shard2.h5`, `common_weights.h5` maybe ?).

What happens if the sharding changes based on the hardware your are running ? (Like moving from 8 A100 to 64 G4 maybe).
Are layers even the good amount of sharding in case of Deepspeed parallelism you can have tensor being sharded themselves no ?
Just asking questions to understand the requirements better for enabling downloading only what's necessary, it seems more is at play than pure file sharding.
Btw, using something like `h5` could enable streaming only necessary parts of even a large file probably (but does require logic and a server, I would need to check but maybe HTTP ranges is doable, not sure how that plays with Cloudfront cache though)

The other idea I have is if that index could be somehow inferred by the hub itself, it would save quite a bit of headaches of inconsistent files
@julien-c Do you think deriving that index per weight file automatically on the hub is doable ?




```python
from huggingface_hub import hf_hub_download
import datetime
import h5py
import os


def _load(f, tensors, prefix=""):
    for k in f.keys():
        if isinstance(f[k], h5py._hl.dataset.Dataset):
            tensors[f"{prefix}_{k}"] = f[k].shape
        else:
            tensors.update(_load(f[k], tensors, prefix=f"{prefix}_{k}"))
    return tensors


def filelen(filename, suffix="B"):
    num = os.path.getsize(filename)
    for unit in ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"]:
        if abs(num) < 1024.0:
            return f"{num:3.1f}{unit}{suffix}"
        num /= 1024.0
    return f"{num:.1f}Yi{suffix}"


def check(filename):
    start = datetime.datetime.now()
    shapes_gpt2 = {}
    with h5py.File(filename, "r") as f:
        _load(f, shapes_gpt2)
    print(
        f"Found {len(shapes_gpt2)} tensors, took {datetime.datetime.now() - start} on file of len {filelen(filename)}"
    )


def main():
    filename = hf_hub_download("gpt2", filename="tf_model.h5")
    filename_large = hf_hub_download("gpt2-large", filename="tf_model.h5")
    check(filename)
    check(filename_large)


if __name__ == "__main__":
    main()
```


Jumping up because I might be helping to implement this while adding the 176B parameters model

Note that there are 2 notions of sharding here:
- sharding on the hub: goal is to keep all the files at less than 5-20GB to get cloudfront/fast-downloads
- sharding the model locally: goal is to be able to load the model when there are less CPU memory than the size of the model for instance (which happens currently for the 176B model on Jean Zay) - this could be related to ZeRO3 at some point but probably we should not depend on DeepSpeed save file format and be more flexible to accomodate TF/JAX/etc

A threshold of 5-10GB per file should allow both sharding limit. On my side I would stay quite open in how the weights are spread among files (don't force per layer) and use a mapping file as mentioned above.

So we could go toward this organization I guess inspired by @stas00 first proposal:

```
README.md
config.json
[tokenizer files....]
pytorch_model.bin.index => pickle dictionnaire mapping stat_dict keys to file name in pytorch_model folder
pytorch_model/pytorch_model-00000-of-00003.pt => pickle of dict of weights
pytorch_model/pytorch_model-00001-of-00003.pt => pickle of dict of weights
pytorch_model/pytorch_model-00002-of-00003.pt => pickle of dict of weights
```

The index mapping file `pytorch_model.bin.index` could also be a human readable file (JSON for instance) by the way. Not sure we gain much by having it stored in binary.

> The index mapping file pytorch_model.bin.index could also be a human readable file (JSON for instance) by the way. Not sure we gain much by having it stored in binary.

Plus, we're trying to stay away from pickle those days :-)

Just replying on this for now:

> HTTP ranges is doable, not sure how that plays with Cloudfront cache though

yes Cloudfront supports range requests 👍 

Are we considering my proposal to completely part ways from torch checkpoint and using a tied dbm file instead? https://github.com/pytorch/pytorch/issues/64327

This makes the checkpoint manipulation completely transparent and removes all the custom framework specific formatting and allows for a single checkpoint that could be loaded by any framework.

Now you can load or save any layer or even a single weight w/o caring about how it's implemented underneath.

So in this case it'll just use multiple dbm files to handle shards.

The proposal includes a fully working code prototype and would just require logic to handle sharding.

#16343 has been merged so I believe this may be closed?

Just posting a small script I used to shard any model: https://gist.github.com/younesbelkada/382016361580b939a87edcddc94c6593 people may want to use it in the future to push sharded models ! 

> Just posting a small script I used to shard any model: https://gist.github.com/younesbelkada/382016361580b939a87edcddc94c6593 people may want to use it in the future to push sharded models !

@younesbelkada would this work for an optimizer? Currently my optimizer .pt is the issue, it needs to be less than 50GB per reqs specified here https://huggingface.co/docs/hub/repositories-recommendations#explanations

```
checkpoint-6000$ ls -lah
total 105G
drwxr-sr-x 3 brando9 root 4.0K Oct  9 13:34 .
drwxr-sr-x 7 brando9 root 4.0K Oct  9 12:38 ..
-rw-r--r-- 1 brando9 root  619 Oct  9 13:04 config.json
-rw-r--r-- 1 brando9 root  116 Oct  9 13:04 generation_config.json
drwxr-sr-x 9 brando9 root 4.0K Oct  9 13:53 .git
-rw-r--r-- 1 brando9 root 1.5K Oct  9 12:38 .gitattributes
-rw-r--r-- 1 brando9 root  54G Oct  9 13:04 optimizer.pt
-rw------- 1 brando9 root  38G Oct  9 13:58 optimizer.pt.gz
-rw-r--r-- 1 brando9 root 9.3G Oct  9 13:05 pytorch_model-00001-of-00002.bin
-rw-r--r-- 1 brando9 root 4.3G Oct  9 13:05 pytorch_model-00002-of-00002.bin
-rw-r--r-- 1 brando9 root  24K Oct  9 13:05 pytorch_model.bin.index.json
-rw-r--r-- 1 brando9 root  15K Oct  9 13:05 rng_state.pth
-rw-r--r-- 1 brando9 root  627 Oct  9 13:05 scheduler.pt
-rw-r--r-- 1 brando9 root  15K Oct  9 13:05 trainer_state.json
-rw-r--r-- 1 brando9 root 4.0K Oct  9 13:05 training_args.bin
```

hi @brando90 
Hmm not sure if this supports optimizer sharding out of the box, you might need to create your own custom logic for it