[BUG] Memory overhead issue with Zero2

Describe the bug

In Zero Stage2, deepspeed will move all the parameters to cpu first, flatten the data, and then move back to GPU. In some large model cases, this operation will use 2x memory and leads to OOMs.

Some Explanation

According to PyTorch memory management, when deepspeed move the parameters to cpu,

https://github.com/microsoft/DeepSpeed/blob/cda7c718954e22f22df222872400e29eaf5708c1/deepspeed/runtime/zero/stage2.py#L271

Pytorch reserves theses memory instead of release it. Then deepspeed move the parameters back to GPU here:

https://github.com/microsoft/DeepSpeed/blob/cda7c718954e22f22df222872400e29eaf5708c1/deepspeed/runtime/zero/stage2.py#L295

In some large model cases, pytorch will malloc new GPU memory instead of using the reserved memory in cache pool. This will lead to 2x gpu memory usage, which is improper，and face the risk of OOMs.

To Reproduce
This is a simple demo of pytorch memory management:

import torch
from deepspeed.ops.op_builder import UtilsBuilder

util_ops = UtilsBuilder().load()
flatten_func = util_ops.flatten
print("Finish utils load")

data_1 = torch.zeros(10000, 10000).cuda()
data_2 = torch.zeros(10000, 10000).cuda()
print(f"MA {round(torch.cuda.memory_allocated() / (1024 * 1024 * 1024),2 )} GB \
    Max_MA {round(torch.cuda.max_memory_allocated() / (1024 * 1024 * 1024),2)} GB \
    CA {round(torch.cuda.memory_reserved() / (1024 * 1024 * 1024),2)} GB \
    Max_CA {round(torch.cuda.max_memory_reserved() / (1024 * 1024 * 1024))} GB ")

data_1.data = data_1.data.cpu()
data_2.data = data_2.data.cpu()
print(f"MA {round(torch.cuda.memory_allocated() / (1024 * 1024 * 1024),2 )} GB \
    Max_MA {round(torch.cuda.max_memory_allocated() / (1024 * 1024 * 1024),2)} GB \
    CA {round(torch.cuda.memory_reserved() / (1024 * 1024 * 1024),2)} GB \
    Max_CA {round(torch.cuda.max_memory_reserved() / (1024 * 1024 * 1024))} GB ")
#data = torch.add(data,data_2)
result = flatten_func([data_1, data_2])
result = result.cuda()
print(f"MA {round(torch.cuda.memory_allocated() / (1024 * 1024 * 1024),2 )} GB \
    Max_MA {round(torch.cuda.max_memory_allocated() / (1024 * 1024 * 1024),2)} GB \
    CA {round(torch.cuda.memory_reserved() / (1024 * 1024 * 1024),2)} GB \
    Max_CA {round(torch.cuda.max_memory_reserved() / (1024 * 1024 * 1024))} GB ")

The output is:

MA 0.75 GB     Max_MA 0.75 GB     CA 0.75 GB     Max_CA 1 GB
MA 0.0 GB     Max_MA 0.75 GB     CA 0.75 GB     Max_CA 1 GB
MA 0.75 GB     Max_MA 0.75 GB     CA 1.49 GB     Max_CA 1 GB

After move data_1 and data_2 to cpu, the allocated memory is 0.0GB and these memory is reserved in memory cache. While move the flattened data back to gpu, pytorch malloc new memory and double the memory cache, leads to unnecessary memory allocation.

If the model is 15.9 G, here needs 15.9 *2 = 31.8G memory usage, pytorch won't trigger the OOMs. However, the following NCCL or cublas call will easily trigger the OOMs. However, if the model is 17G, the memory malloc of another 17G will fail. The pytorch will flash the cache by itself, the cached memory will down to 17G, and the OOMs won't happen.

Suggestion
Add torch.cuda.empty_cache() after move data to cpu, flashing the cache manually.

Additional context
According to the code, Zero3 seems meet the same problem.

Start with deepspeed/runtime/zero/stage2.py at the parameter CPU transfer around lines 271 and 295, then review the reproduction script's CUDA memory measurements. Check whether the same behavior applies to Zero3 as reported. Done means the transfer and flattening path no longer causes unnecessary peak GPU memory or OOM risk, with memory behavior verified on the provided example.