Repository: open-mmlab/mmengine

Issue:
[Bug] MMSeparateDistributedDataParallel skip init_weights
### Prerequisite

- [X] I have searched [Issues](https://github.com/open-mmlab/mmengine/issues) and [Discussions](https://github.com/open-mmlab/mmengine/discussions) but cannot get the expected help.
- [X] The bug has not been fixed in the latest version(https://github.com/open-mmlab/mmengine).


### Reproduces the problem - code sample

I encountered this problem when running my own custom project based on `mmengine`, which is too complicated to present.


### Additional information

The `MMSeparateDistributedDataParallel` is a **model wrapper of model wrapper**, whose `module` may contain the `MMDistributedDataParallel`.

https://github.com/open-mmlab/mmengine/blob/4d7a29342c02d78cb3ffabe036607e801d6713d9/mmengine/model/wrappers/seperate_distributed.py#L82-L87

When initializing the model_wrapper weights, the `runner` initialize the weight of the `runner.model.module`:
 https://github.com/open-mmlab/mmengine/blob/4d7a29342c02d78cb3ffabe036607e801d6713d9/mmengine/runner/runner.py#L893-L902
The above codes will cause the `init_weights` functions in model not work because the children module of the `runner.model.module` in this case is of type `MMDistributedDataParallel` which does NOT have `init_weights` function.


Discussion:
Hi, the module is defined in https://github.com/open-mmlab/mmengine/blob/4d7a29342c02d78cb3ffabe036607e801d6713d9/mmengine/model/wrappers/seperate_distributed.py#L71

Therefore, if you implement the `init_weights` method in the module, the runner can call it as expected.
@zhouzaida But the sub_modules in `self.module` are **modified** to be model_wrapper `MMDistributedDataParallel`:
https://github.com/open-mmlab/mmengine/blob/4d7a29342c02d78cb3ffabe036607e801d6713d9/mmengine/model/wrappers/seperate_distributed.py#L75-L87
Therefore, the `MMSeparateDistributedDataParallel.module` does have `init_weights` function but its sub-modules ( `MMDistributedDataParallel`) do not have `init_weights` function, which causes the `init_weights` functions of the module wrapped in the `MMDistributedDataParallel` not accessible during the initialzation.
Oh, the `init_weights` method is defined in the sub-module, so there might be an issue with that.

Resolving patch:
diff --git a/mmengine/model/base_module.py b/mmengine/model/base_module.py
--- a/mmengine/model/base_module.py
+++ b/mmengine/model/base_module.py
@@ -11,6 +11,7 @@
 from mmengine.dist import master_only
 from mmengine.logging import MMLogger, print_log
 from .weight_init import initialize, update_init_info
+from .wrappers.utils import is_model_wrapper
 
 
 class BaseModule(nn.Module, metaclass=ABCMeta):
@@ -123,6 +124,8 @@ def init_weights(self):
                 initialize(self, other_cfgs)
 
             for m in self.children():
+                if is_model_wrapper(m) and not hasattr(m, 'init_weights'):
+                    m = m.module
                 if hasattr(m, 'init_weights'):
                     m.init_weights()
                     # users may overload the `init_weights`


Test patch:
diff --git a/tests/test_model/test_base_module.py b/tests/test_model/test_base_module.py
--- a/tests/test_model/test_base_module.py
+++ b/tests/test_model/test_base_module.py
@@ -4,6 +4,7 @@
 import os.path as osp
 import tempfile
 from unittest import TestCase
+from unittest.mock import patch
 
 import torch
 from torch import nn
@@ -145,7 +146,6 @@ def test_init_weights(self):
             ├──conv1d (FooConv1d, weight=3, bias=4)
         ├──reg (nn.Linear, weight=1, bias=2)
         """
-
         self.model.init_weights()
 
         assert torch.equal(
@@ -222,6 +222,22 @@ def __init__(self,
         self.assertTrue((ori_layer_weight != model.linear.linear.weight).any())
         self.assertTrue((ori_layer_bias != model.linear.linear.bias).any())
 
+        class FakeDDP(nn.Module):
+
+            def __init__(self, module) -> None:
+                super().__init__()
+                self.module = module
+
+        # Test initialization of nested modules in DDPModule which define
+        # `init_weights`.
+        with patch('mmengine.model.base_module.is_model_wrapper',
+                   lambda x: isinstance(x, FakeDDP)):
+            model = FOOMODELS.build(model_cfg)
+            model.ddp = FakeDDP(CustomLinear())
+            model.init_weights()
+            self.assertTrue((model.ddp.module.linear.weight == 1).all())
+            self.assertTrue((model.ddp.module.linear.bias == 2).all())
+
     def test_dump_init_info(self):
         import os
         import shutil