🔧 ZeRO by hand with a 4-parameter model

Table of Contents


Motivation
Why study ZeRO
Setup
Why make a copy of the model weights for optimizer states
Why do we have to make copies of the momentum and variance and not just recompute them on the fly?
Stage 1
Stage 2
Stage 3







Motivation




This article seeks to implement by hand the Zero Redundancy Optimizer (ZeRO). The goal...