Repository: microsoft/DeepSpeed

Issue:
[REQUEST] efficiently deal with frozen weights during training
**Is your feature request related to a problem? Please describe.**

This is a need for HF `accelerate` and `m4` projects.

We have more and more situations where a large part of the model that's being trained is frozen. 

At the moment the only way to emulate frozen weights in deepspeed is to create a param group with `lr=0`, which works but a huge amount of memory is wasted to allocate states that aren't being trained (grads too?) and the compute is wasted too.

**Describe the solution you'd like**

Ideally, if `requires_grad == False` those weights shouldn't be in the optimizer and not allocating any optim state/grad memory.

Thank you!

@tjruwase 

@pacman100 - would you like to add/expand anything to this request?


Discussion:
Current workaround is to create a param_group with all the frozen parameters and set the lr=0, an example is shown below:
```
if getattr(accelerator.state, "deepspeed_plugin", None) is not None:
        optimizer_grouped_parameters = [
            {
                "params": [p for p in model.parameters() if p.requires_grad],
                "lr": lr
            },
            {
                "params": [p for p in model.parameters() if not p.requires_grad],
                "lr": 0
            },
        ]
        
        for param in model.parameters():
            param.requires_grad = True
        optimizer = torch.optim.AdamW(optimizer_grouped_parameters)
```

Below are the screenshots of memory usage (in MB) when using DS Stage 3 (GPU: 65 GB, CPU: 327 GB :exploding_head:) with CPU offloading vs plain PyTorch (GPU: 56 GB, CPU: 3 GB). The model is MT0-XXL (13B params) and is using LoRA for parameter efficient fine-tuning.

`Plain Pytorch`
<img width="1512" alt="Screenshot 2022-12-09 at 4 51 13 PM" src="https://user-images.githubusercontent.com/13534540/208633163-2b775f0a-9c0c-494a-8349-3f2d0b96e185.png">

`DS ZeRO-3 with CPU offloading`
<img width="1510" alt="Screenshot 2022-12-09 at 4 42 15 PM" src="https://user-images.githubusercontent.com/13534540/208633188-08682f13-fbe4-4f18-b88e-9225f0a7ce3e.png">

**Request**:  Efficiently deal with frozen weights during training such that large models could be offloaded on CPUs/sharded across GPUs properly with storage of the optimizer state only for the trainable parameters, e.g., we can see that using plain PyTorch, mt0-xxl (13B params) model takes up 56GB on GPU, now, it would be really helpful if one could do CPU offloading such that training could work on a 16GB or 24GB GPU. 

@pacman100, this is a great demo of the issue. Could you please share the repro steps?

Resolving patch:
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
@@ -252,11 +252,15 @@ def __init__(self,
         self.sub_group_size = sub_group_size
 
         self.sub_group_to_group_id = {}
-        see_memory_usage("Before creating fp16 partitions", force=False)
-        self._create_fp16_partitions_with_defragmentation()
+
+        # Trainable parameters
+        self.trainable_param_groups = self._get_trainable_parameter_groups()
+
+        see_memory_usage("Before creating fp16 partitions", force=True)
+        self._create_fp16_partitions_with_defragmentation(self.trainable_param_groups)
         num_fp16_subgroups = len(self.fp16_partitioned_groups_flat)
         see_memory_usage(f"After creating fp16 partitions: {num_fp16_subgroups}",
-                         force=False)
+                         force=True)
 
         # Optimizer tensor swapping
         if self.swap_optimizer:
@@ -350,19 +354,28 @@ def __init__(self,
     def destroy(self):
         self.parameter_offload.destroy()
 
+    def _get_trainable_parameter_groups(self):
+        param_groups = []
+        for param_group in self.optimizer.param_groups:
+            trainable_params = {
+                "params": [p for p in param_group["params"] if p.requires_grad]
+            }
+            param_groups.append(trainable_params)
+        return param_groups
+
     def _setup_for_real_optimizer(self):
-        see_memory_usage("Before creating fp32 partitions", force=False)
+        see_memory_usage("Before creating fp32 partitions", force=True)
         self._create_fp32_partitions()
-        see_memory_usage("After creating fp32 partitions", force=False)
+        see_memory_usage("After creating fp32 partitions", force=True)
         dist.barrier()
 
         # To support pipelined optimizer swapping
         self._create_next_swappable_fp32_groups()
 
-        see_memory_usage("Before initializing optimizer states", force=False)
+        see_memory_usage("Before initializing optimizer states", force=True)
 
         self.initialize_optimizer_states()
-        see_memory_usage("After initializing optimizer states", force=False)
+        see_memory_usage("After initializing optimizer states", force=True)
         dist.barrier()
 
         if dist.get_rank() == 0:
@@ -523,7 +536,7 @@ def _create_param_groups_fp16_flat_cpu_memory(self):
 
         aggregate_params_count = 0
 
-        for j, param_group in enumerate(self.optimizer.param_groups):
+        for j, param_group in enumerate(self.trainable_param_groups):
             params_in_group = sum([p.partition_numel() for p in param_group['params']])
 
             flat_buffer_size = params_in_group
@@ -552,11 +565,12 @@ def _create_param_groups_fp16_flat_cpu_memory(self):
                     torch.empty(1,
                                 dtype=self.dtype))
 
-    def _create_fp16_partitions_with_defragmentation(self):
+    def _create_fp16_partitions_with_defragmentation(self, fp16_param_groups):
         dist.barrier()
+
         param_groups: List[List[Parameter]] = tuple(
             self._create_fp16_sub_groups(param_group["params"])
-            for param_group in self.optimizer.param_groups)
+            for param_group in fp16_param_groups)
 
         # bookkeeping related to param groups
         for param_group_idx, param_group in enumerate(param_groups):
@@ -884,7 +898,6 @@ def initialize_optimizer_states(self):
                                       dtype=gradient_dtype,
                                       device=self.device)
 
-        timers = self.timers
         timer_names = set()
 
         if self.swap_optimizer:
@@ -2122,6 +2135,7 @@ def _get_param_groups(self):
 
     def _set_param_groups(self, value):
         self.optimizer.param_groups = value
+        self.trainable_param_groups = self._get_trainable_parameter_groups()
 
     param_groups = property(_get_param_groups, _set_param_groups)
 


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
@@ -1313,3 +1313,63 @@ def test(self, zero_stage):
                         state = optimizer.optimizer.state[param]
                         step_counts.append(state['step'])
                 assert all(step == step_counts[0] for step in step_counts)
+
+
+class TestZeroFrozenWeights(DistributedTest):
+    world_size = 1
+
+    def test(self):
+        config_dict = {
+            "train_batch_size": 4,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-4
+                }
+            },
+            "fp16": {
+                "enabled": True
+            },
+            "zero_optimization": {
+                "stage": 3
+            }
+        }
+        hidden_dim = 10
+
+        class MyModel(torch.nn.Module):
+            def __init__(self, hidden_dim):
+                super(MyModel, self).__init__()
+                self.l1 = torch.nn.Linear(hidden_dim, hidden_dim)
+                self.l2 = torch.nn.Linear(hidden_dim, hidden_dim)
+                self.act = torch.nn.ReLU()
+                self.cel = torch.nn.CrossEntropyLoss()
+
+                # freeze one fc
+                self.l2.weight.requires_grad = False
+                self.l2.bias.requires_grad = False
+
+            def forward(self, x, y):
+                x = self.l1(x)
+                x = self.act(x)
+                x = self.l2(x)
+                loss = self.cel(x, y)
+                val = (x, loss)
+                return val
+
+        with deepspeed.zero.Init(config_dict_or_path=config_dict):
+            model = MyModel(hidden_dim)
+
+        model, _, _, _ = deepspeed.initialize(model=model,
+                                              model_parameters=model.parameters(),
+                                              config=config_dict)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=50,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device)
+        dist.barrier()
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            loss = loss[1]
+            model.backward(loss)
+            model.step()