Repository: hpcaitech/ColossalAI

Issue:
[BUG]: Memory consumption by fp16 is not normal when using Engine.
### 🐛 Describe the bug

when using colossalai.amp.convert_to_torch_amp to wrap the model, optimizer and criterion, 
```
if not use_colossai_engine:
    model, optimizer, criterion =  colossalai.amp.convert_to_torch_amp(model, optimizer, criterion)
```
and then train normally, which also only consumes 4700M of memory.

```
output, _ = model(img, label)
train_loss = criterion(output, label)
optimizer.backward(train_loss)
optimizer.step()
optimizer.zero_grad()
```

------

But if you use colossalai.initialize to initialize, it will consume 7700M of memory. But we did see that by reading the fp16 parameter in config, in the initialization code of ```colossalai.initialize```, the conversion of process  ```colossalai.amp.convert_to_torch_amp``` is performed, and then we use the Engine for training, but it needs to consume 7700M of memory at this time. This is where I get confused.
```
engine.zero_grad()
output, _ = engine(img, label)
train_loss = engine.criterion(output, label)
engine.backward(train_loss)
engine.step()   
```

### Environment

_No response_


Discussion:
Thanks! Give me some time to verify this. :)

Resolving patch:
diff --git a/colossalai/engine/_base_engine.py b/colossalai/engine/_base_engine.py
--- a/colossalai/engine/_base_engine.py
+++ b/colossalai/engine/_base_engine.py
@@ -92,7 +92,10 @@ def __init__(self,
             self._schedule = NonPipelineSchedule()
         if self.uses_pipeline:
             self._schedule.pre_processing(self)
-        register_ophooks_recursively(self._model, self._ophook_list)
+
+        #register hook if any
+        if len(self._ophook_list) > 0:
+            register_ophooks_recursively(self._model, self._ophook_list)
 
     @property
     def ophooks(self):
diff --git a/colossalai/engine/ophooks/utils.py b/colossalai/engine/ophooks/utils.py
--- a/colossalai/engine/ophooks/utils.py
+++ b/colossalai/engine/ophooks/utils.py
@@ -85,11 +85,15 @@ def backward(ctx, *args):
 
 
 def register_ophooks_recursively(module: torch.nn.Module,
-                                 ophook_list: List[BaseOpHook] = None,
+                                 ophook_list: List[BaseOpHook],
                                  name: str = "",
                                  filter_fn: Optional[Callable] = None):
     r"""Recursilvely register pre/post hooks for all submodules in the module in FWD and BWD."""
     assert isinstance(module, torch.nn.Module)
+    assert isinstance(ophook_list, (list, tuple))
+    assert len(ophook_list) > 0, 'expected at least 1 hook in the argument ophook_list but found 0'
+    for hook in ophook_list:
+            assert (isinstance(hook, BaseOpHook))
 
     # Add hooks for submodules
     for child_name, child in module.named_children():
@@ -103,10 +107,6 @@ def register_ophooks_recursively(module: torch.nn.Module,
     if filter_fn is not None and filter_fn(module):
         return
 
-    if ophook_list is not None:
-        for hook in ophook_list:
-            assert (isinstance(hook, BaseOpHook))
-
     def _pre_forward_module_hook(submodule, *args):
         for hook in ophook_list:
             assert isinstance(submodule, torch.nn.Module)


Test patch: