Repository: open-mmlab/mmengine

Issue:
[Bug] fsdp `_move_module_to_device()` will not change the special variable `device` of **Module** `DetDataPreprocessor` to GPU!!!
### 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).

### Environment

```python
OrderedDict([('sys.platform', 'linux'),
             ('Python', '3.9.16 (main, May 15 2023, 23:46:34) [GCC 11.2.0]'),
             ('CUDA available', True),
             ('numpy_random_seed', 2147483648),
             ('GPU 0,1,2,3,4,5,6,7', 'NVIDIA GeForce RTX 2080 Ti'),
             ('CUDA_HOME', '/usr/local/cuda'),
             ('NVCC', 'Cuda compilation tools, release 11.8, V11.8.89'),
             ('GCC', 'gcc (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0'),
             ('PyTorch', '2.0.0'),
             ('PyTorch compiling details',
              'PyTorch built with:\n'
              '  - GCC 9.3\n'
              '  - C++ Version: 201703\n'
              '  - Intel(R) oneAPI Math Kernel Library Version 2023.1-Product '
              'Build 20230303 for Intel(R) 64 architecture applications\n'
              '  - Intel(R) MKL-DNN v2.7.3 (Git Hash '
              '6dbeffbae1f23cbbeae17adb7b5b13f1f37c080e)\n'
              '  - OpenMP 201511 (a.k.a. OpenMP 4.5)\n'
              '  - LAPACK is enabled (usually provided by MKL)\n'
              '  - NNPACK is enabled\n'
              '  - CPU capability usage: AVX2\n'
              '  - CUDA Runtime 11.8\n'
              '  - NVCC architecture flags: '
              '-gencode;arch=compute_37,code=sm_37;-gencode;arch=compute_50,code=sm_50;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_61,code=sm_61;-gencode;arch=compute_70,code=sm_70;-gencode;arch=compute_75,code=sm_75;-gencode;arch=compute_80,code=sm_80;-gencode;arch=compute_86,code=sm_86;-gencode;arch=compute_90,code=sm_90;-gencode;arch=compute_37,code=compute_37\n'
              '  - CuDNN 8.7\n'
              '  - Magma 2.6.1\n'
              '  - Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, '
              'CUDA_VERSION=11.8, CUDNN_VERSION=8.7.0, '
              'CXX_COMPILER=/opt/rh/devtoolset-9/root/usr/bin/c++, CXX_FLAGS= '
              '-D_GLIBCXX_USE_CXX11_ABI=0 -fabi-version=11 -Wno-deprecated '
              '-fvisibility-inlines-hidden -DUSE_PTHREADPOOL -DNDEBUG '
              '-DUSE_KINETO -DLIBKINETO_NOROCTRACER -DUSE_FBGEMM -DUSE_QNNPACK '
              '-DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK '
              '-DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -Wall -Wextra '
              '-Werror=return-type -Werror=non-virtual-dtor '
              '-Werror=bool-operation -Wnarrowing '
              '-Wno-missing-field-initializers -Wno-type-limits '
              '-Wno-array-bounds -Wno-unknown-pragmas -Wunused-local-typedefs '
              '-Wno-unused-parameter -Wno-unused-function -Wno-unused-result '
              '-Wno-strict-overflow -Wno-strict-aliasing '
              '-Wno-error=deprecated-declarations -Wno-stringop-overflow '
              '-Wno-psabi -Wno-error=pedantic -Wno-error=redundant-decls '
              '-Wno-error=old-style-cast -fdiagnostics-color=always '
              '-faligned-new -Wno-unused-but-set-variable '
              '-Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math '
              '-Werror=format -Werror=cast-function-type '
              '-Wno-stringop-overflow, LAPACK_INFO=mkl, PERF_WITH_AVX=1, '
              'PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, '
              'TORCH_DISABLE_GPU_ASSERTS=ON, TORCH_VERSION=2.0.0, USE_CUDA=ON, '
              'USE_CUDNN=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, '
              'USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, '
              'USE_NCCL=ON, USE_NNPACK=ON, USE_OPENMP=ON, USE_ROCM=OFF, \n'),
             ('TorchVision', '0.15.0'),
             ('OpenCV', '4.7.0'),
             ('MMEngine', '0.8.1')])
```

### Reproduces the problem - code sample

```python
_base_ = 'detr_r50_8xb2-150e_coco.py'
# __file__ = 'detr_r50_fsdp.py

runner_type = 'FlexibleRunner'

# 指定 FSDPStrategy 并配置参数
strategy = dict(
    type='FSDPStrategy',
    model_wrapper=dict(
        type='MMFullyShardedDataParallel',
        cpu_offload=True,
        auto_wrap_policy='single_grad_auto_wrap_policy',  # manual
        mixed_precision=dict(
            param_dtype='float16',
            buffer_dtype='float32',
            reduce_dtype='float32'),
        use_orig_params=False),  # default
    state_dict_cfg=dict(
        state_dict_type='FULL_STATE_DICT',
        state_dict_config=dict(
            type='FullStateDictConfig', 
            offload_to_cpu=True),
        optim_state_dict_config=dict(
            type='FullOptimStateDictConfig', 
            offload_to_cpu=True)))

optim_wrapper = dict(
    type='AmpOptimWrapper',
    dtype='float16',
    use_fsdp=True)
```

```python
@FUNCTIONS.register_module()
def single_grad_auto_wrap_policy(module: nn.Module, recurse: bool, nonwrapped_numel: int) -> bool:
    if recurse:
        return True  # always recurse
    grad_types = {p.requires_grad for p in module.parameters()}
    return len(grad_types) <= 1
```

### Reproduces the problem - command or script

```shell
bash tools/dist_train.sh detr_r50_fsdp.py 2
```

### Reproduces the problem - error message

Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
![image](https://github.com/open-mmlab/mmengine/assets/49788935/ef30bc65-9379-4fde-967b-ffcf00ee502b)


### Additional information

1. In `torch.distributed.fsdp._init_utils._move_module_to_device()`, because module `DetDataPreprocessor` has no _parameter_, fsdp will **not move it to the GPU**. So it's spacial variable `device` will keep the value `device(type='cpu')`, and it's other two variable `mean` & `std` will still be on the **CPU**.
![image](https://github.com/open-mmlab/mmengine/assets/49788935/1e7eb5b0-e9ed-4678-b643-9b12fdbd1823)

2. But fsdp has `torch.distributed.fsdp._runtime_utils._lazy_init()` func, which will excute `_cast_buffers_to_dtype_and_device()` func. Then the buffer `mean` & `std` will be moved to GPU.
**However, the spacial variable `device` of `DetDataPreprocessor` will never be changed.**
_before cast_
![image](https://github.com/open-mmlab/mmengine/assets/49788935/4a7fdafc-2a30-4fc6-bcb7-dcdddfc4064a)
_after cast_
![image](https://github.com/open-mmlab/mmengine/assets/49788935/e822423b-c1f2-4336-80e0-bf17f9f4557b)

3. In `torch.distributed.fsdp.fully_sharded_data_parallel.FullyShardedDataParallel.forward()`, after excute `_root_pre_forward()`, the input data will be moved to GPU.
before
![image](https://github.com/open-mmlab/mmengine/assets/49788935/e5ce3ba2-1b68-44ea-b9a7-1a49837b2390)
after
![image](https://github.com/open-mmlab/mmengine/assets/49788935/049cb511-dca5-4ca1-814b-aa11226ea10c)

4. Then, in `BaseDataPreprocessor.cast_data()` func, because the special variable `device` of `DetDataPreprocessor` is `device(type='cpu')`, **the input data will be moved back to CPU**.
![image](https://github.com/open-mmlab/mmengine/assets/49788935/15f8821f-5d55-4c12-9830-65027a010da2)

5. suggestion. 
In `ImgDataPreprocessor.__init__()`
```python
# self.register_buffer('mean', torch.tensor(mean).view(-1, 1, 1), False)
self.mean = nn.Parameter(torch.tensor(mean).view(-1, 1, 1), False)  # Add
# self.register_buffer('std', torch.tensor(std).view(-1, 1, 1), False)
self.std = nn.Parameter(torch.tensor(std).view(-1, 1, 1), False)  # Add
```



Discussion:
Hi! Thank you for your feedback! This is actually a bug that occurs in a special case, and we are committed to fixing it as soon as possible.

The error you encountered can be raised when using a specific **wrap_policy** that shards all parameters of the root FSDP (Fully Sharded Data Parallelism) module during [`_recursive_wrap`](https://github.com/pytorch/pytorch/blob/92eb9d363a7110a43746cc3e704a393837d5047c/torch/distributed/fsdp/wrap.py#L320). In this scenario, when `_move_module_to_device` is called, it will get an `None` `params` because `_get_orig_params` ignores the sharded parameters and the parameters specified by `ignored_parameters`.

If you have defined a more fine-grained wrap policy, such as wrapping only certain transformer layers into an FSDP module, then only a portion of the parameters will be sharded. In this case, the `_get_orig_params` function in the root FSDP module will not return `None` for `params`, and the error you mentioned will not occur.

Regardless, we acknowledge that this is a bug in the current FSDPStrategy implementation, and we are committed to resolving it as soon as possible.

By the way, it's important not to define `mean` and `std` as `Parameter` for the following reasons:

1. The error you mentioned could still occur if `mean` and `std` are sharded using a special **wrap_policy**.
2. Semantically, `mean` and `std` should not be considered `Parameter` objects since they are not updated by the optimizer.
3. `mean` and `std` should be defined by the user and should not be overwritten when loading a saved checkpoint. If they are defined as `Parameter` objects, this rule could be violated.
A workaround is to insert the following code:

```python
        for module in model.modules():
            if isinstance(module, BaseDataPreprocessor):
                module.to(get_device())
```

at this location:

https://github.com/open-mmlab/mmengine/blob/b2295a258c4af3cdeb1a55fa5e49cddbd1c6c0bf/mmengine/_strategy/fsdp.py#L124

Will this solution address your problem?

Resolving patch:
diff --git a/mmengine/_strategy/fsdp.py b/mmengine/_strategy/fsdp.py
--- a/mmengine/_strategy/fsdp.py
+++ b/mmengine/_strategy/fsdp.py
@@ -19,8 +19,9 @@
 
 import mmengine
 from mmengine.config import Config, ConfigDict
+from mmengine.device import get_device
 from mmengine.dist import get_rank, is_main_process
-from mmengine.model import is_model_wrapper
+from mmengine.model import BaseDataPreprocessor, is_model_wrapper
 from mmengine.optim import (AmpOptimWrapper, BaseOptimWrapper, OptimWrapper,
                             OptimWrapperDict, _ParamScheduler,
                             build_optim_wrapper)
@@ -118,6 +119,10 @@ def _wrap_model(self, model: nn.Module) -> None:
             FullyShardedDataParallel: ``MMFullyShardedDataParallel``
             or subclass of ``FullyShardedDataParallel``.
         """
+        for module in model.modules():
+            if isinstance(module, BaseDataPreprocessor):
+                module.to(get_device())
+
         if is_model_wrapper(model):
             return
 


Test patch: