Repository: hpcaitech/ColossalAI

Issue:
[BUG]: The problem of running 2D, 2.5D, and 3D on vit model (related to https://github.com/hpcaitech/ColossalAI-Examples/tree/main/image/vision_transformer/hybrid_parallel)
### 🐛 Describe the bug

I have the issue to run high dimensional tensor parallelism, such as 2D, 2.5D, 3D. 

Declaration: norm_forward(Tensor x, Tensor mean, Tensor sqr_mean, Tensor weight, Tensor bias, Tensor eps) -> ((Tensor, Tensor, Tensor))
Cast error details: Unable to cast Python instance to C++ type (compile in debug mode for details)
    return self._forward_func(*args)
  File "/usr/local/lib/python3.7/dist-packages/colossalai-0.1.12+torch1.10cu11.3-py3.7-linux-x86_64.egg/colossalai/nn/layer/parallel_3d/layers.py", line 160, in forward
    self.input_x_weight_parallel_mode,
  File "/usr/local/lib/python3.7/dist-packages/colossalai-0.1.12+torch1.10cu11.3-py3.7-linux-x86_64.egg/colossalai/nn/layer/parallel_3d/_operation.py", line 397, in layernorm_3d
    input_x_weight_parallel_mode,
  File "/usr/local/lib/python3.7/dist-packages/torch/cuda/amp/autocast_mode.py", line 94, in decorate_fwd
    return fwd(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/colossalai-0.1.12+torch1.10cu11.3-py3.7-linux-x86_64.egg/colossalai/nn/layer/parallel_3d/_operation.py", line 330, in forward
    output, mu, sigma = norm_forward(input_, mean, sqr_mean, weight, bias, eps)
RuntimeError: norm_forward() Expected a value of type 'Tensor (inferred)' for argument 'eps' but instead found type 'float'.
Inferred 'eps' to be of type 'Tensor' because it was not annotated with an explicit type.
Position: 5
Value: 1e-06
Declaration: norm_forward(Tensor x, Tensor mean, Tensor sqr_mean, Tensor weight, Tensor bias, Tensor eps) -> ((Tensor, Tensor, Tensor))
Cast error details: Unable to cast Python instance to C++ type (compile in debug mode for details)

### Environment

I am running using 4 GPUs and 8 GPUs, where each is Tesla-V100-PCIE-32GB. 
Colossal-AI version: 0.1.12
----------------------------
PyTorch Version: 1.10.0
PyTorch Version required by Colossal-AI: 1.10
PyTorch version match: ✓
----------------------------
System CUDA Version: 11.3
CUDA Version required by PyTorch: 11.3
CUDA Version required by Colossal-AI: 11.3
CUDA Version Match: ✓
----------------------------
CUDA Extension: ✓

[BUG]: The problem of running 2D, 2.5D, and 3D on vit model (related to https://github.com/hpcaitech/ColossalAI-Examples/tree/main/image/vision_transformer/hybrid_parallel)
### 🐛 Describe the bug

I have the issue to run high dimensional tensor parallelism, such as 2D, 2.5D, 3D. 

Declaration: norm_forward(Tensor x, Tensor mean, Tensor sqr_mean, Tensor weight, Tensor bias, Tensor eps) -> ((Tensor, Tensor, Tensor))
Cast error details: Unable to cast Python instance to C++ type (compile in debug mode for details)
    return self._forward_func(*args)
  File "/usr/local/lib/python3.7/dist-packages/colossalai-0.1.12+torch1.10cu11.3-py3.7-linux-x86_64.egg/colossalai/nn/layer/parallel_3d/layers.py", line 160, in forward
    self.input_x_weight_parallel_mode,
  File "/usr/local/lib/python3.7/dist-packages/colossalai-0.1.12+torch1.10cu11.3-py3.7-linux-x86_64.egg/colossalai/nn/layer/parallel_3d/_operation.py", line 397, in layernorm_3d
    input_x_weight_parallel_mode,
  File "/usr/local/lib/python3.7/dist-packages/torch/cuda/amp/autocast_mode.py", line 94, in decorate_fwd
    return fwd(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/colossalai-0.1.12+torch1.10cu11.3-py3.7-linux-x86_64.egg/colossalai/nn/layer/parallel_3d/_operation.py", line 330, in forward
    output, mu, sigma = norm_forward(input_, mean, sqr_mean, weight, bias, eps)
RuntimeError: norm_forward() Expected a value of type 'Tensor (inferred)' for argument 'eps' but instead found type 'float'.
Inferred 'eps' to be of type 'Tensor' because it was not annotated with an explicit type.
Position: 5
Value: 1e-06
Declaration: norm_forward(Tensor x, Tensor mean, Tensor sqr_mean, Tensor weight, Tensor bias, Tensor eps) -> ((Tensor, Tensor, Tensor))
Cast error details: Unable to cast Python instance to C++ type (compile in debug mode for details)

### Environment

I am running using 4 GPUs and 8 GPUs, where each is Tesla-V100-PCIE-32GB. 
Colossal-AI version: 0.1.12
----------------------------
PyTorch Version: 1.10.0
PyTorch Version required by Colossal-AI: 1.10
PyTorch version match: ✓
----------------------------
System CUDA Version: 11.3
CUDA Version required by PyTorch: 11.3
CUDA Version required by Colossal-AI: 11.3
CUDA Version Match: ✓
----------------------------
CUDA Extension: ✓



Discussion:
@tongping Thank you for reporting the bug. Type hints are missing here. We will fix it asap.
@tongping Thank you for reporting the bug. Type hints are missing here. We will fix it asap.

Resolving patch:
diff --git a/colossalai/nn/layer/parallel_3d/_operation.py b/colossalai/nn/layer/parallel_3d/_operation.py
old mode 100644
new mode 100755
--- a/colossalai/nn/layer/parallel_3d/_operation.py
+++ b/colossalai/nn/layer/parallel_3d/_operation.py
@@ -281,7 +281,7 @@ def vocab_parallel_classifier_3d(
 
 
 @torch.jit.script
-def norm_forward(x, mean, sqr_mean, weight, bias, eps):
+def norm_forward(x: Tensor, mean: Tensor, sqr_mean: Tensor, weight: Tensor, bias: Tensor, eps: float):
     mu = x - mean
     var = sqr_mean - mean**2
     sigma = torch.sqrt(var + eps)
@@ -292,7 +292,7 @@ def norm_forward(x, mean, sqr_mean, weight, bias, eps):
 
 
 @torch.jit.script
-def norm_backward(grad, mu, sigma, weight):
+def norm_backward(grad: Tensor, mu: Tensor, sigma: Tensor, weight: Tensor):
     # dbias, dweight = grad, grad * mu / sigma
     dz = grad * weight
     dmu = dz / sigma


Test patch: