Repository: microsoft/DeepSpeed

Issue:
[BUG] zero3 memory leak on return from training loop
**Describe the bug**

Hmm, I thought I fixed the leak here https://github.com/microsoft/DeepSpeed/pull/2665 but there is one more of a similar nature, where the model gets gathered, but not ungathered when the weights are no longer needed.

When entering the training loop the weights are sharded. But on exit they are no longer sharded, so the memory isn't getting freed. Say, you want to do something else with the gpu - e.g. additional eval - it's going to OOM most likely.

**To Reproduce**


```
$ cat test.py
import deepspeed
from deepspeed.ops.op_builder import CPUAdamBuilder

from unit.common import DistributedTest, DistributedFixture
from unit.simple_model import *
from unit.util import required_minimum_torch_version

from unit.checkpoint.common import *

import pytest

class TestZeROCheckpointNew(DistributedTest):
    world_size = 1

    def test_1_load_module_only_new(self):
        tmpdir = 'zero_stage'
        zero_stage = 3
        config_dict = {
            "train_batch_size": 2,
            "optimizer": {
                "type": 'Adam'
            },
            "fp16": {
                "enabled": True,
                "initial_scale_power": 8
            },
            "zero_optimization": {
                "stage": zero_stage,
            }
        }
        hidden_dim = 10

        with deepspeed.zero.Init(config_dict_or_path=config_dict):
            models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]

        print("CREATE", models[0].linears[0].weight[:5])

        dtype = torch.half
        ds_model = create_deepspeed_model(config_dict=config_dict,
                                          model=models[0],
                                          base_optimizer=None)

        data_loader = random_dataloader(model=ds_model,
                                        total_samples=5,
                                        hidden_dim=hidden_dim,
                                        device=ds_model.device,
                                        dtype=dtype)

        print("ORIGIN 1", models[0].linears[0].weight)
        for _, batch in enumerate(data_loader):
            loss = ds_model(batch[0], batch[1])
            ds_model.backward(loss)
            ds_model.step()

        print("ORIGIN 2", models[0].linears[0].weight[0,:5])
```

```
PYTHONPATH=. pytest --torch_ver 1.13.1 --cuda_ver 11.7 test.py -sv
```
gives:
```
ORIGIN 1 Parameter containing:
tensor([], device='cuda:0', dtype=torch.float16, requires_grad=True)
ORIGIN 2 tensor([ 0.0699, -0.0974, -0.1665, -0.1251, -0.1195], device='cuda:0',
       dtype=torch.float16, grad_fn=<SliceBackward0>)
```
you can see that the model wasn't partitioned and memory leaked for the remainder of the program.

It should be `tensor([], device='cuda:0', ...)` after `step` maybe?

As far as what the solution should be, I think `step` should call `partition` internally, similar to the fix in https://github.com/microsoft/DeepSpeed/pull/2665 - it scatters the updated weights anyway.

@tjruwase 


Discussion:
@stas00, @satpalsr this issue is a bit more nuanced. The intended behavior for training is to cache parameters in GPU memory in anticipation of the next forward pass. This is why we don't explicitly release the parameters after step. This cache size and behavior is controlled by configuration knobs like `stage3_max_live_parameters`, `stage3_prefetch_bucket_size`, and `stage3_max_reuse_distance`. In contrast, `GatherParameters` designed for one-time computation with these parameters and so implicitly releases them afterwards. 

@stas00, do you have a specific need for this behavior. Perhaps the solution is to provide an API for client to explicit flush the parameter cache. 

@satpalsr, apologies for not responding before your PR. 
> @stas00, do you have a specific need for this behavior. Perhaps the solution is to provide an API for client to explicit flush the parameter cache.

Yes, I'm thinking if the model needs to do something on the gpu after the training has finished it's now stuck with not having free space.

I think such an API would work well. 

> But on exit they are no longer sharded, so the memory isn't getting freed.

Sorry for naive question, I was looking at this [post video](https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/), After optimiser step, I see fp16 weights to be sharded across GPUs. Can we not cache parameters while they are being sharded?
@satpalsr, by caching the parameters I mean not sharding them so that they available to use immediately without an allgather. Hope that helps. 

Resolving patch:
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -9,6 +9,7 @@
 import hashlib
 from collections import defaultdict, OrderedDict, deque
 from shutil import copyfile
+import gc
 
 from torch.nn.modules import Module
 from torch.nn.parameter import Parameter
@@ -3546,3 +3547,12 @@ def save_16bit_model(self, save_dir, save_filename="pytorch_model.bin"):
         self.checkpoint_engine.commit(tag)
 
         return True
+
+    def empty_partition_cache(self):
+        """
+        Release GPU memory consumed by offloaded model parameters.
+        """
+        if hasattr(self.optimizer, 'empty_partition_cache'):
+            self.optimizer.empty_partition_cache()
+            gc.collect()
+            get_accelerator().empty_cache()
diff --git a/deepspeed/runtime/zero/parameter_offload.py b/deepspeed/runtime/zero/parameter_offload.py
--- a/deepspeed/runtime/zero/parameter_offload.py
+++ b/deepspeed/runtime/zero/parameter_offload.py
@@ -259,6 +259,9 @@ def get_param_coordinator(self, training):
 
         return self.param_coordinators[training]
 
+    def empty_partition_cache(self):
+        self.partition_all_parameters()
+
     def _convert_to_zero_parameters(self, ds_config, module, mpu):
         non_zero_params = [p for p in module.parameters() if not is_zero_param(p)]
         if non_zero_params:
@@ -321,7 +324,7 @@ def mark_persistent_parameters(self, param_threshold, model_threshold):
             if param.ds_numel + total_persistent_parameters > model_threshold:
                 continue
 
-            if param.ds_numel < param_threshold:
+            if param.ds_numel <= param_threshold:
                 params_count += 1
                 param.ds_persist = True
                 persistent_params.append(param)
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -2467,6 +2467,9 @@ def checkpoint_event_epilogue(self):
         if len(self.persistent_parameters) > 0:
             self.persistent_parameters[0].all_gather(self.persistent_parameters)
 
+    def empty_partition_cache(self):
+        self.parameter_offload.empty_partition_cache()
+
 
 def _handle_overflow(cpu_sum, x, i):
     import math


Test patch:
diff --git a/tests/unit/runtime/zero/test_zero.py b/tests/unit/runtime/zero/test_zero.py
--- a/tests/unit/runtime/zero/test_zero.py
+++ b/tests/unit/runtime/zero/test_zero.py
@@ -1422,3 +1422,49 @@ def test(self, force_ds_optim):
             model, _, _, _ = deepspeed.initialize(model=model,
                                                   optimizer=optimizer,
                                                   config=config_dict)
+
+
+@pytest.mark.parametrize('training', [True, False])
+class TestZeroPartitionCache(DistributedTest):
+    world_size = 1
+
+    def test_training_partition_cache(self, training):
+        hidden_dim = 10
+        config_dict = {
+            "train_batch_size": 2,
+            "fp16": {
+                "enabled": True,
+                "initial_scale_power": 8
+            },
+            "zero_optimization": {
+                "stage": 3,
+                "stage3_param_persistence_threshold": hidden_dim
+            }
+        }
+        if training:
+            config_dict["optimizer"] = {"type": "Adam"}
+
+        with deepspeed.zero.Init(config_dict_or_path=config_dict):
+            model = SimpleModel(hidden_dim, empty_grad=False)
+
+        model, _, _, _ = deepspeed.initialize(model=model, config=config_dict)
+
+        dtype = torch.half
+        data_loader = random_dataloader(model=model,
+                                        total_samples=6,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=dtype)
+
+        for _, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            if training:
+                model.backward(loss)
+                model.step()
+
+        persist_param_size = sum([p.numel() for p in model.parameters() if p.ds_persist])
+
+        assert persist_param_size >= sum([p.numel() for p in model.parameters()])
+
+        model.empty_partition_cache()
+        assert sum([p.numel() for p in model.parameters()]) == 0