Repository: hpcaitech/ColossalAI

Issue:
How to release a Chunk when the keep_gathered is True[BUG]: 
### 🐛 Describe the bug

Hello Guys, I am looking through the CAI code and have a doubt about the chunk management. I saw the keep_gathered is True when `group_acc_size < min_chunk_size_byte`. Here is the code refernce. [Ref](https://github.com/hpcaitech/ColossalAI/blob/f477a14f4aeb49f7a30ee0f46775040391a96e1c/colossalai/gemini/chunk/search_utils.py#L133)
```python
# let small parameters keep gathered in CUDA all the time
if group_acc_size < min_chunk_size_byte:
    config_dict[dp_degree] = dict(chunk_size=group_acc_size, keep_gathered=True)
else:
    size_dict[dp_degree] = size_list
```
However, the chunk is not able to release to do a validation after training and get an `AssertionError`. Here are the errors.
```
File "/root/miniconda3/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/root/miniconda3/envs/test/lib/python3.7/site-packages/colossalai/nn/parallel/data_parallel.py", line 287, in forward
    self._post_forward()
  File "/root/miniconda3/envs/test/lib/python3.7/site-packages/colossalai/nn/parallel/data_parallel.py", line 264, in _post_forward
    assert chunk.can_release
AssertionError

```
I wonder if this is a bug and how to solve such problem. Many thanks.

### Environment

```
#### Installation Report ####

------------ Environment ------------
Colossal-AI version: 0.2.0
PyTorch version: 1.12.1
CUDA version: 11.3
CUDA version required by PyTorch: 11.3

Note:
1. The table above checks the versions of the libraries/tools in the current environment
2. If the CUDA version is N/A, you can set the CUDA_HOME environment variable to locate it

------------ CUDA Extensions AOT Compilation ------------
Found AOT CUDA Extension: ✓
PyTorch version used for AOT compilation: 1.12
CUDA version used for AOT compilation: 11.3

Note:
1. AOT (ahead-of-time) compilation of the CUDA kernels occurs during installation when the environment varialbe CUDA_EXT=1 is set
2. If AOT compilation is not enabled, stay calm as the CUDA kernels can still be built during runtime

------------ Compatibility ------------
PyTorch version match: ✓
System and PyTorch CUDA version match: ✓
System and Colossal-AI CUDA version match: ✓

```
[tensor] fix some unittests

[tensor] fix some unittests



Discussion:
The root reason happens at this function: 

```
    @property
    def can_release(self) -> bool:
        if self.keep_gathered:
            return False
        else:
            return self.tensor_state_cnter[TensorState.HOLD] + \
                   self.tensor_state_cnter[TensorState.HOLD_AFTER_BWD] == self.num_tensors
```

I believe it is a bug on this [line of code](https://github.com/hpcaitech/ColossalAI/blob/f566b0ce6b96c0045322c2a7623013ec0c23b6dc/colossalai/nn/parallel/data_parallel.py#L254) which does not distinguish keep_gathered. We will fix it soon.
@JThh Thanks for the reply. I think so. The `keep_gathered` is set by the chunk init, but no modification anymore after that,



Resolving patch:
diff --git a/colossalai/gemini/chunk/manager.py b/colossalai/gemini/chunk/manager.py
--- a/colossalai/gemini/chunk/manager.py
+++ b/colossalai/gemini/chunk/manager.py
@@ -140,6 +140,14 @@ def reduce_chunk(self, chunk: Chunk) -> bool:
         self.__add_memory_usage(chunk.memory_usage)
         return True
 
+    def fake_release_chunk(self, chunk: Chunk) -> None:
+        """Release gathered chunk in a fake mode.
+        This function is used for keep-gathered chunk in the inference mode.
+        """
+        assert chunk.keep_gathered
+        assert chunk.tensor_state_cnter[TensorState.HOLD] == chunk.num_tensors
+        self.__sub_accessed_chunk(chunk)
+
     def copy_tensor_to_chunk_slice(self, tensor: torch.Tensor, data: torch.Tensor) -> None:
         """
         Copy data to the chunk.
diff --git a/colossalai/nn/parallel/data_parallel.py b/colossalai/nn/parallel/data_parallel.py
--- a/colossalai/nn/parallel/data_parallel.py
+++ b/colossalai/nn/parallel/data_parallel.py
@@ -257,8 +257,11 @@ def _post_forward(self):
         access_list = list(self.chunk_manager.accessed_chunks)
         # we need to scatter all accessed chunks and move them to their original places
         for chunk in access_list:
-            assert chunk.can_release
-            self.chunk_manager.release_chunk(chunk)
+            if chunk.keep_gathered:
+                self.chunk_manager.fake_release_chunk(chunk)
+            else:
+                assert chunk.can_release
+                self.chunk_manager.release_chunk(chunk)
             first_param = next(iter(chunk.tensors_info))
             self.chunk_manager.move_chunk(chunk, self.grads_device[first_param])
         assert self.chunk_manager.accessed_mem == 0


Test patch:
diff --git a/tests/test_gemini/update/test_inference.py b/tests/test_gemini/update/test_inference.py
--- a/tests/test_gemini/update/test_inference.py
+++ b/tests/test_gemini/update/test_inference.py
@@ -1,4 +1,5 @@
 from functools import partial
+from typing import Callable
 
 import pytest
 import torch
@@ -13,7 +14,7 @@
 from colossalai.gemini.gemini_mgr import GeminiManager
 from colossalai.nn.optimizer import HybridAdam
 from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer
-from colossalai.nn.parallel import ZeroDDP
+from colossalai.nn.parallel import ZeroDDP, zero_model_wrapper
 from colossalai.testing import parameterize, rerun_if_address_is_in_use
 from colossalai.utils import free_port
 from colossalai.utils.cuda import get_current_device
@@ -36,9 +37,35 @@ def check_param(model: ZeroDDP, torch_model: torch.nn.Module):
         assert_close(value, temp_zero_value, rtol=1e-3, atol=4e-3)
 
 
+def multi_chunk_init(model: torch.nn.Module, placement_policy: str):
+    world_size = dist.get_world_size()
+    config_dict, *_ = search_chunk_configuration(model, search_range_mb=1, search_interval_byte=100)
+    config_dict[world_size]['chunk_size'] = 5000
+    config_dict[world_size]['keep_gathered'] = False
+    if placement_policy != 'cuda':
+        init_device = torch.device('cpu')
+    else:
+        init_device = None
+    chunk_manager = ChunkManager(config_dict, init_device=init_device)
+    gemini_manager = GeminiManager(placement_policy, chunk_manager)
+    model = ZeroDDP(model, gemini_manager, pin_memory=True)
+    return model
+
+
+def single_chunk_init(model: torch.nn.Module, placement_policy: str):
+    gemini_config = dict(
+        device=get_current_device(),
+        placement_policy=placement_policy,
+        pin_memory=True,
+    )
+    model = zero_model_wrapper(model=model, zero_stage=3, gemini_config=gemini_config)
+    return model
+
+
 @parameterize('placement_policy', ['cuda', 'cpu', 'auto', 'const'])
 @parameterize('model_name', ['gpt2'])
-def exam_inference(placement_policy, model_name: str):
+@parameterize('model_init_func', [single_chunk_init, multi_chunk_init])
+def exam_inference(placement_policy: str, model_name: str, model_init_func: Callable):
     set_seed(19360226)
     get_components_func = non_distributed_component_funcs.get_callable(model_name)
     model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
@@ -56,18 +83,7 @@ def exam_inference(placement_policy, model_name: str):
     for torch_p, p in zip(torch_model.parameters(), model.parameters()):
         p.data.copy_(torch_p.data)
 
-    world_size = torch.distributed.get_world_size()
-    config_dict, *_ = search_chunk_configuration(model, search_range_mb=1, search_interval_byte=100)
-    config_dict[world_size]['chunk_size'] = 5000
-    config_dict[world_size]['keep_gathered'] = False
-    if placement_policy != 'cuda':
-        init_device = torch.device('cpu')
-    else:
-        init_device = None
-    chunk_manager = ChunkManager(config_dict, init_device=init_device)
-    gemini_manager = GeminiManager(placement_policy, chunk_manager)
-    model = ZeroDDP(model, gemini_manager, pin_memory=True)
-
+    model = model_init_func(model, placement_policy)
     optimizer = HybridAdam(model.parameters(), lr=1e-3)
     zero_optim = ZeroOptimizer(optimizer, model, initial_scale=128)