[BUG] memory overhead issue with optimizer leading to OOM

Could it be that unfused_optimizer is not careful at how it allocates memory when this condition occurs:

[2021-11-08 18:30:01,688] [INFO] [unfused_optimizer.py:275:_update_scale] Grad overflow on iteration: 2983
[2021-11-08 18:30:01,688] [INFO] [unfused_optimizer.py:276:_update_scale] Reducing dynamic loss scale from 65536.0 to 32768.0
[2021-11-08 18:30:01,689] [INFO] [unfused_optimizer.py:199:step] [deepspeed] fp16 dynamic loss scale overflow! Skipping step. Attempted loss scale: 65536.0, reducing to 32768.0

it tries to allocate more memory - a whooping 1GB!

    fp32_param.grad = fp16_param.grad.to(fp32_param.dtype)
RuntimeError: CUDA out of memory. Tried to allocate 1.00 GiB (GPU 1; 31.75 GiB total capacity; 24.28 GiB already allocated; 256.00 MiB free; 30.01 GiB reserved in total by PyTorch)RuntimeError
: CUDA out of memory. Tried to allocate 1.00 GiB (GPU 3; 31.75 GiB total capacity; 24.28 GiB already allocated; 276.00 MiB free; 30.01 GiB reserved in total by PyTorch)

I was running at 31 out of 32GB gpu memory used for many hours until the above occurred and it threw it into OOM.

Perhaps that logic needs to free memory it doesn't need any more first before allocating new memory?

e.g. this problem happened on a single SLURM job after running for 15h just fine.

This is with ZeRO-1 with Megatron-Deepspeed.

Thank you!

@tjruwase

Start in unfused_optimizer.py around _update_scale at lines 275-276 and step at line 199, then inspect the fp32_param.grad conversion shown in the report. Reproduce the overflow path with ZeRO-1 and Megatron-Deepspeed while monitoring GPU allocation. Done means the overflow handling no longer causes the reported unexpected 1 GiB allocation or OOM.