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

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.<p><a href="https:&#x2F;&#x2F;huggingface.co&#x2F;blog&#x2F;zero-deepspeed-fairscale" rel="nofollow">https:&#x2F;&#x2F;huggingface.co&#x2F;blog&#x2F;zero-deepspeed-fairscale</a>

Support for this was also added to [Fairscale](<a href="https:&#x2F;&#x2F;fairscale.readthedocs.io&#x2F;en&#x2F;latest&#x2F;" rel="nofollow">https:&#x2F;&#x2F;fairscale.readthedocs.io&#x2F;en&#x2F;latest&#x2F;</a>) and [Fairseq](<a href="https:&#x2F;&#x2F;github.com&#x2F;pytorch&#x2F;fairseq" rel="nofollow">https:&#x2F;&#x2F;github.com&#x2F;pytorch&#x2F;fairseq</a>) last week. In particular, the Fairscale implementation can be used in any pyotrch project without requiring the use of the Deepspeed trainer.

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?

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.<p><a href="https:&#x2F;&#x2F;github.com&#x2F;EleutherAI&#x2F;gpt-neox" rel="nofollow">https:&#x2F;&#x2F;github.com&#x2F;EleutherAI&#x2F;gpt-neox</a>

ELI5? All this techno babble just sounds like &quot;it&#x27;s faster because we optimized it&quot;. What are the nontrivial, new fundamental tricks?

This is also being added to pytorch<p><a href="https:&#x2F;&#x2F;github.com&#x2F;pytorch&#x2F;pytorch&#x2F;pull&#x2F;46750" rel="nofollow">https:&#x2F;&#x2F;github.com&#x2F;pytorch&#x2F;pytorch&#x2F;pull&#x2F;46750</a>

Simple 10 min overview&#x2F;tutorial (official) if someone is interested - <a href="https:&#x2F;&#x2F;www.youtube.com&#x2F;watch?v=ovQC7FqXHXk" rel="nofollow">https:&#x2F;&#x2F;www.youtube.com&#x2F;watch?v=ovQC7FqXHXk</a>

See also zeroth order backpropagation which allows 300X faster training while not reducing throughput that much <a href="https:&#x2F;&#x2F;arxiv.org&#x2F;abs&#x2F;2011.08895" rel="nofollow">https:&#x2F;&#x2F;arxiv.org&#x2F;abs&#x2F;2011.08895</a> How much zero-3 affect accuracy?<p>See also <a href="https:&#x2F;&#x2F;github.com&#x2F;microsoft&#x2F;fastformers" rel="nofollow">https:&#x2F;&#x2F;github.com&#x2F;microsoft&#x2F;fastformers</a>

Alternatively, one could get rid of the memory used by optimizers entirely by switching to vanilla SGD.<p>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.

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

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

please hook this up to Jax!

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!