Zero-3 Offload: Scale DL models to trillion parameters without code changes

[FL33TW00D]

Huggingface has been working on implementing this into their library, and it has some pretty amazing effects on the size of models you can train on a simple Colab.

https://huggingface.co/blog/zero-deepspeed-fairscale

[stephenroller]

Support for this was also added to [Fairscale](https://fairscale.readthedocs.io/en/latest/) and [Fairseq](https://github.com/pytorch/fairseq) last week. In particular, the Fairscale implementation can be used in any pyotrch project without requiring the use of the Deepspeed trainer.

[diptanu]

What are the relevant commits in Fairseq for this? I couldn't figure out the changes by looking at the commits from last week.

[ansk]

Question for someone knowledgable about this: if I have a model which is large -- but small enough that I can fit a single training example on GPU -- does this approach offer speedups compared to simple gradient accumulation? Or is this only useful for models which are so large that the model parameters themselves are overwhelming GPU memory?

[joshlk]

GPT-NeoX is an example project that is using deepspeed and Zero-3 offloading. The wider project intend to train a GPT-3 sized model and release it freely to the world.

https://github.com/EleutherAI/gpt-neox

[ma2rten]

It seems like Zero-3 doesn't work for them:

https://github.com/EleutherAI/gpt-neox/issues/171

[dataangel]

ELI5? All this techno babble just sounds like "it's faster because we optimized it". What are the nontrivial, new fundamental tricks?

[jonbaer]

I think there is some explanation (on the previous model?) here, https://www.youtube.com/watch?v=tC01FRB0M7w

[jiofih]

Third paragraph or so in the overview:

> ZeRO removes the memory redundancies across data-parallel processes by partitioning the three model states (optimizer states, gradients, and parameters) across data-parallel processes instead of replicating them. By doing this, it boosts memory efficiency compared to classic data-parallelism while retaining its computational granularity and communication efficiency

[bevenky]

This is also being added to pytorch

https://github.com/pytorch/pytorch/pull/46750

[minimaxir]

I don't think that's the Stage 3 announced in this blog post, but it's def a framework for it.

[alphagrep12345]

Simple 10 min overview/tutorial (official) if someone is interested - https://www.youtube.com/watch?v=ovQC7FqXHXk

[The_rationalist]

See also zeroth order backpropagation which allows 300X faster training while not reducing throughput that much https://arxiv.org/abs/2011.08895 How much zero-3 affect accuracy?

See also https://github.com/microsoft/fastformers

[vladf]

Alternatively, one could get rid of the memory used by optimizers entirely by switching to vanilla SGD.

I haven’t tried this on transformers and maybe that’s what breaks down here but in “classic” supervised settings I’ve found SGD with schedule tuning just as fast as Adam.

[gwern]

SGD doesn't work on large Transformers, no. You need something like AdamW.

[andrewprock]

How much data do you need to mitigate the risk of over fitting a trillion parameter model?

[gwern]

You ideally need ~500GB of text, or so. EleutherAI's The Pile was designed to be just big enough to fit a 1t GPT efficiently, and you can get the various scaling curves out of the OA-related scaling papers. (You want the amount of data that fits into a single epoch, because if you reuse data, you get less bang for the FLOPs buck, and FLOPS constraints are right now much more binding than data or model size.)

[singhrac]

For those searching, DeepSpeed is implemented as a set of C++/CUDA extensions on top of PyTorch (compiled using their JIT).

[bionhoward]

please hook this up to Jax!

[mchusma]

This is super impressive. I could not figure out for a while who exactly was running this project, but it looks like its Microsoft. Great work!