Repository: hpcaitech/ColossalAI

Issue:
[BUG]: colossalai check failed in NGC pytorch 21.07
### 🐛 Describe the bug

I built the docker images sucessfully following the [instructions](https://github.com/hpcaitech/ColossalAI#use-docker). And then I run `colossalai check -i` and check failed with the following trace:
```
Traceback (most recent call last):
  File "/opt/conda/bin/colossalai", line 5, in <module>
    from colossalai.cli import cli
  File "/opt/conda/lib/python3.8/site-packages/colossalai/__init__.py", line 1, in <module>
    from .initialize import (initialize, launch, launch_from_openmpi,
  File "/opt/conda/lib/python3.8/site-packages/colossalai/initialize.py", line 23, in <module>
    from colossalai.engine.schedule import NonPipelineSchedule, PipelineSchedule, InterleavedPipelineSchedule, get_tensor_shape
  File "/opt/conda/lib/python3.8/site-packages/colossalai/engine/__init__.py", line 1, in <module>
    from ._base_engine import Engine
  File "/opt/conda/lib/python3.8/site-packages/colossalai/engine/_base_engine.py", line 11, in <module>
    from colossalai.engine.schedule import BaseSchedule, NonPipelineSchedule, PipelineSchedule, InterleavedPipelineSchedule
  File "/opt/conda/lib/python3.8/site-packages/colossalai/engine/schedule/__init__.py", line 1, in <module>
    from ._base_schedule import BaseSchedule
  File "/opt/conda/lib/python3.8/site-packages/colossalai/engine/schedule/_base_schedule.py", line 10, in <module>
    from colossalai.utils import get_current_device
  File "/opt/conda/lib/python3.8/site-packages/colossalai/utils/__init__.py", line 15, in <module>
    from .model.colo_init_context import ColoInitContext
  File "/opt/conda/lib/python3.8/site-packages/colossalai/utils/model/colo_init_context.py", line 3, in <module>
    from colossalai.tensor import ColoTensor, ColoParameter, register_colo_module, init_colo_module, \
  File "/opt/conda/lib/python3.8/site-packages/colossalai/tensor/__init__.py", line 4, in <module>
    from .colo_tensor import ColoTensor
  File "/opt/conda/lib/python3.8/site-packages/colossalai/tensor/colo_tensor.py", line 7, in <module>
    from colossalai.tensor.dist_spec_mgr import DistSpecManager
  File "/opt/conda/lib/python3.8/site-packages/colossalai/tensor/dist_spec_mgr.py", line 2, in <module>
    from colossalai.nn.layer.utils import divide
  File "/opt/conda/lib/python3.8/site-packages/colossalai/nn/__init__.py", line 1, in <module>
    from .layer import *
  File "/opt/conda/lib/python3.8/site-packages/colossalai/nn/layer/__init__.py", line 7, in <module>
    from .moe import *
  File "/opt/conda/lib/python3.8/site-packages/colossalai/nn/layer/moe/__init__.py", line 1, in <module>
    from .experts import Experts, FFNExperts, TPExperts
  File "/opt/conda/lib/python3.8/site-packages/colossalai/nn/layer/moe/experts.py", line 8, in <module>
    from colossalai.zero.init_ctx import no_shard_zero_decrator
  File "/opt/conda/lib/python3.8/site-packages/colossalai/zero/__init__.py", line 6, in <module>
    from colossalai.zero.sharded_model.sharded_model_v2 import ShardedModelV2
  File "/opt/conda/lib/python3.8/site-packages/colossalai/zero/sharded_model/__init__.py", line 1, in <module>
    from .sharded_model_v2 import ShardedModelV2
  File "/opt/conda/lib/python3.8/site-packages/colossalai/zero/sharded_model/sharded_model_v2.py", line 5, in <module>
    from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX
ImportError: cannot import name '_EXTRA_STATE_KEY_SUFFIX' from 'torch.nn.modules.module' (/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py)

```

### Environment

NGC pytorch 21.07


Discussion:
I tried Dockerfile in the latest branch and colossalai check sucessed.
However, the same error occured when building colossalai from source in a conda env with colossalai-0.1.5+torch1.8cu11.6
> However, the same error occured when building colossalai from source in a conda env with colossalai-0.1.5+torch1.8cu11.6

This occurs perhaps because of API incompatibility of PyTorch, you can try a higher version PyTorch like 1.10 and 1.11.


Resolving patch:
diff --git a/colossalai/zero/sharded_model/sharded_model_v2.py b/colossalai/zero/sharded_model/sharded_model_v2.py
--- a/colossalai/zero/sharded_model/sharded_model_v2.py
+++ b/colossalai/zero/sharded_model/sharded_model_v2.py
@@ -2,7 +2,6 @@
 from collections import OrderedDict
 from typing import Any, Optional, Iterator, Tuple
 from copy import deepcopy
-from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX
 import itertools
 import torch
 import torch.distributed as dist
@@ -28,6 +27,11 @@
 from ._utils import (cast_float_arguments, cast_tensor_to_fp16, cast_tensor_to_fp32, chunk_and_pad, free_storage,
                      get_gradient_predivide_factor)
 
+try:
+    from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX
+except ImportError:
+    _EXTRA_STATE_KEY_SUFFIX = '_extra_state'
+
 
 class ShardedModelV2(nn.Module):
     """


Test patch: