stage=0: params=2.0 GB, grads=2.0 GB, optimizer=12.0 GB, total=16.0 GB
stage=1: params=2.0 GB, grads=2.0 GB, optimizer=3.0 GB, total=7.0 GB
stage=2: params=2.0 GB, grads=0.5 GB, optimizer=3.0 GB, total=5.5 GB
stage=3: params=0.5 GB, grads=0.5 GB, optimizer=3.0 GB, total=4.0 GB
