Support ZeRO-Infinity

Is your feature request related to a problem? Please describe.

I´m frustrated because I can´t use my Geforce MX 250 to train a 13B GPT-NeoX.

Describe the solution you’d like

In DeepSpeed v0.3.15 they introduces ZeRO-Infinity, a ZeRO-3 that also checkpoints to disk.

If I understand correctly, after approving https://github.com/EleutherAI/gpt-neox/pull/199 , a merge of https://github.com/EleutherAI/DeeperSpeed with upstream and exposing the parameters:

"offload_optimizer": {
      "device": "cpu/gpu/nvme"
    },
    "offload_param": {
      "device": "cpu/gpu/nvme"
    }

Would be enough

[StellaAthena]

@bratao Basically that the marginal improvement for Z-3 and Z-infinity over Z-1 + pipeline parallelism was minimal for training at the 100B scale.

[bratao]

Sounds fair @StellaAthena , hopefully in my upcoming vacation I can work on it. But I’m curious, can you share the Nvidia feedback?