[QUESTION] checkpointing/loading memory overhead

It seems that in current implementation the torch_dist checkpointing and loading will introduce around 2GB GPU memory overhead for rank 0 (for a 400m model), which will cause OOM if the original GPU untilization is high.

Is there a way to free this GPU memory allocation after successfully checkpointing/loading? so that such overhead will not impact the later training process, otherwise it's quite annoying as it requires an additional test (and optimization if OOM) for checkpointing and loading.

(tried with adding torch.cuda.empty_cache() in the checkpointing function but not working. And the load_checkpoint function alr has a "torch.cuda.empty_cache()" in the end)

Thank you!