Repository: deepspeedai/DeepSpeed
Pull request: fix cpu_adam memory leak on deepspeed re-use in the same process
As reported in https://github.com/microsoft/DeepSpeed/issues/879 currently deepspeed leaks some 8GB general RAM when re-running `deepspeed.initialize` in the same process (notebook or pytest worker). There is no problem with one  `deepspeed.initialize`  per process, since the memory is reclaimed on process exit.

With the help from @tjruwase we narrowed it down to cpu_adam creating objects, but not freeing them. It was happening only with zero2 or 3 with `cpu_offload` enabled.

This PR:
* extends the `cuda_adam.cpp` extension to add `destroy_adam_optimizer`
* adds `DeepSpeedCPUAdam.__del__` that calls the new cpp method

I haven't checked if other deepspeed optimizers have the same issue.

There is another small leak  of ~50MB general RAM remaining somewhere still, when re-running `deepspeed.initialize` in the same process, but it's far less of a problem than the current 8GB.

Fixes: https://github.com/microsoft/DeepSpeed/issues/879

@RezaYazdaniAminabadi 

Problem statement:
memory leak in an environment like notebook
In HF/DS tests I use a notebook-like environment for some of the tests, which means I don't fork a new process for each deepspeed run, but emulate the distributed env in the process and repeatedly run deepspeed. This makes it much easier/faster to test the values of the weights.

There must be either some global variable that holds onto memory, or the circular reference happening, since memory doesn't get released after deepspeed has done its work and I even explicitly deleted the engine/scheduler/optimizer variables.

I tested that if I remove deepspeed from the equation there is no leak.

Now I have a few dozens of those tests and I get some 10GB extra RAM used per deepspeed invocation. So it quickly grows to 100GBs.
 
This is with a test model of 2x 1 value weights y = wx+b - i.e. the model's memory foot print is ~0. 

Any idea where the objects might be held and not destroyed? If I repeat the same re-creation of the trainer in a loop the old one should get its memory freed up.

So this is an example of a cell in a jupyter notebook:
```
if trainer.deepspeed:
    print("reloading")
    trainer.deepspeed = None
    trainer.optimizer = None
    trainer.lr_scheduler = None
trainer = get_regression_trainer(output_dir=output_dir, deepspeed=ds_config_dict, skip_memory_metrics=True)
trainer.train()
```
of course the explicit `None`s should be needed as they should get overwritten by the new trainer, I was just doing a sanity check.

So there must be a circular reference where 2 or more internal variables refer to each other and thus the memory doesn't get freed.


I emulate a dist env with just:
```
dist_env_1_gpu = dict(
            MASTER_ADDR="localhost", MASTER_PORT="10999", RANK="0", LOCAL_RANK="0", WORLD_SIZE="1"
        )
for k,v in dist_env_1_gpu.items():
    os.environ[k]=v
```

I'd be happy to try to investigate this on my own if you could help me with pointers at potential suspects.

Thank you!

@jeffra, @samyam

Hints from issue discussion:

Commit message: receivet the optimizer Id from client side when running step (Default=0)

Patch to review:
diff --git a/csrc/adam/cpu_adam.cpp b/csrc/adam/cpu_adam.cpp
index d425dc3169ef..bb6bd13c0c77 100644
--- a/csrc/adam/cpu_adam.cpp
+++ b/csrc/adam/cpu_adam.cpp
@@ -672,6 +672,15 @@ int ds_adam_step_plus_copy(int optimizer_id,
     return 0;
 }
 
+int destroy_adam_optimizer(int optimizer_id)
+{
+    // std::cout << "Adam Optimizer #" << optimizer_id
+    //          << " is destroyed." << std::endl;
+    s_optimizers.erase(optimizer_id);
+
+    return 0;
+}
+
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
 {
     m.def("adam_update", &ds_adam_step, "DeepSpeed CPU Adam update (C++)");
@@ -679,4 +688,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
           &ds_adam_step_plus_copy,
           "DeepSpeed CPU Adam update and param copy (C++)");
     m.def("create_adam", &create_adam_optimizer, "DeepSpeed CPU Adam (C++)");
+    m.def("destroy_adam", &destroy_adam_optimizer, "DeepSpeed CPU Adam destroy (C++)");
 }
diff --git a/deepspeed/ops/adam/cpu_adam.py b/deepspeed/ops/adam/cpu_adam.py
index 7977d232b1fa..dd12e6238aab 100755
--- a/deepspeed/ops/adam/cpu_adam.py
+++ b/deepspeed/ops/adam/cpu_adam.py
@@ -85,13 +85,18 @@ def __init__(self,
                                      weight_decay,
                                      adamw_mode)
 
+    def __del__(self):
+        # need to destroy the C++ object explicitly to avoid a memory leak when deepspeed.initialize
+        # is used multiple times in the same process (notebook or pytest worker)
+        self.ds_opt_adam.destroy_adam(self.opt_id)
+
     def __setstate__(self, state):
         super(DeepSpeedCPUAdam, self).__setstate__(state)
         for group in self.param_groups:
             group.setdefault('amsgrad', False)
 
     @torch.no_grad()
-    def step(self, closure=None, fp16_param_groups=None):
+    def step(self, optId=0, closure=None, fp16_param_groups=None):
         """Update the model parameters.
 
         .. note::
@@ -142,7 +147,7 @@ def step(self, closure=None, fp16_param_groups=None):
 
                 if fp16_param_groups is not None:
                     self.ds_opt_adam.adam_update_copy(
-                        self.opt_id,
+                        optId,
                         state['step'],
                         group['lr'],
                         beta1,
@@ -156,7 +161,7 @@ def step(self, closure=None, fp16_param_groups=None):
                         state['exp_avg_sq'],
                         fp16_param_groups[group_id][param_id].data)
                 else:
-                    self.ds_opt_adam.adam_update(self.opt_id,
+                    self.ds_opt_adam.adam_update(optId,
                                                  state['step'],
                                                  group['lr'],
                                                  beta1,


Code review comments:
diff_hunk: @@ -96,7 +96,7 @@ def __setstate__(self, state):
             group.setdefault('amsgrad', False)
 
     @torch.no_grad()
-    def step(self, closure=None, fp16_param_groups=None):
+    def step(self, optId=0, closure=None, fp16_param_groups=None):
path: deepspeed/ops/adam/cpu_adam.py
text: @author:
this is python land, so camel case is typically reserved for class names only. so `opt_id`?

@author:
so, again, if this case is reduced to only ever having `opt_id=0`, this is now becoming a bit of a mess. Why do we then create `self.opt_id`, but not use it? i.e. you removed its usage in most places. 

This is very error-prone, since now a user may try to access an optimizer that doesn't exist if they call `step(opt_id=1)` and there is no checking that it exists.

What happens if someone calls `step(opt_id=0)` but it was deleted already?

Where does `step()` gets the id now from? It looks like you're relying on the default `0`.

How would an object created with ID 42, know to operate on `step(opt_id=5)`?

Clearly you guys need to have a bit of a design session around this.

If I may suggest, let's revert the change you added, so that this PR is just adding a proper cleanup. And then if you decided to kill the counter, or keeping it, but having the default of just one `opt_id=0`, this is a different change I'd say.

@user1:
I agree. I revert back my changes