CUDA OOM when saving checkpoint (in consolidate_state_dict()) using OSS

I am experiencing CUDA out of memory crashes when consolidating my optimizer state dict before saving it. I am training on 32 40GB A100s, four nodes with eight GPUs each, using PyTorch Lightning's `'ddp_sharded'` strategy, which is OSS. I get the OOM crash *in the middle* of running consolidate_state_dict(). I have tried adding `del` statements, `gc.collect()` and `torch.cuda.empty_cache()` inside the loop to no avail. I am using a custom optimizer class, a modified AdamW that also saves an exponential moving average of the weights, and I need optimizer state sharding because the extra memory overhead for the EMA weights is so onerous. Here is the custom optimizer code: https://gist.github.com/crowsonkb/ea0ed1f6e88594046c72735f3cef1d05. I don't understand how I am running out of GPU memory partway through consolidate_state_dict() (I put in print statements and it got through 27 of 32 ranks) since it moves the tensors to CPU after each broadcast. I am using NCCL so it has to broadcast on GPU but it copies to CPU right afterwards.

Thank you,
Katherine Crowson

Begin at consolidate_state_dict() in the OSS/ddp_sharded path and reproduce the failure with the linked custom optimizer on the described multi-node setup. Track GPU memory through each rank's NCCL broadcast and CPU transfer; done means checkpoint consolidation completes without a CUDA OOM while retaining the optimizer state.

Written by the indexing model from the issue text.