internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: hpcaitech/ColossalAI

Issue:
[BUG]: Activation Checkpointing Failed Test with PyTorch 1.9
### 🐛 Describe the bug

When running unit testing with PyTorch 1.9, the following exception occurs for `tests/test_utils/test_activation_checkpointing.py`. 

## Screenshot

<img width="1003" alt="Lark20220411-162602" src="https://user-images.githubusercontent.com/31818963/162695922-dc88cf7b-fdb3-4063-af3f-2eb40917fa20.png">

## Reasons

This bug occurs because of the invalid use of `ctx.save_for_backward`. In torch autograd function, we need to provide either inputs or intermediate values to `ctx.save_for_backward`. However, when `activation_offload=True`, the current implementation will copy the tensors to cpu, creating new tensors which are neither inputs or intermediate activations (i.e. does not have gradient accumulator objects associated with them). 

```python
for i, arg in enumerate(args):
      if torch.is_tensor(arg):
          if activation_offload:
              tensor_inputs.append(copy_to_device(arg, 'cpu'))
          else:
              tensor_inputs.append(arg)
          ctx.tensor_indices.append(i)
          ctx.inputs.append(None)
      else:
          ctx.inputs.append(arg)

  ctx.save_for_backward(*tensor_inputs)
```

This will trigger the [assertion statement](https://github.com/pytorch/pytorch/blob/ed7ece389daa980a03a2f09848a6439a90af9782/torch/csrc/autograd/saved_variable.cpp#L196) as shown below.

<img width="660" alt="Screenshot 2022-04-11 at 4 27 50 PM" src="https://user-images.githubusercontent.com/31818963/162696232-9fee460e-3275-4914-bccd-4a579b20980f.png">


## Solution

Use `cxt.tensor_inputs = tensor_inputs` instead of `ctx.save_for_backward(*tensor_inputs)` if `activation_offload=True`.


### Environment

CUDA: 11.1
PyTorch: 1.9.1


Discussion:

Resolving patch:
diff --git a/colossalai/utils/activation_checkpoint.py b/colossalai/utils/activation_checkpoint.py
--- a/colossalai/utils/activation_checkpoint.py
+++ b/colossalai/utils/activation_checkpoint.py
@@ -68,7 +68,10 @@ def forward(ctx, run_function, activation_offload=False, *args):
             else:
                 ctx.inputs.append(arg)
 
-        ctx.save_for_backward(*tensor_inputs)
+        if activation_offload:
+            ctx.tensor_inputs = tensor_inputs
+        else:
+            ctx.save_for_backward(*tensor_inputs)
         return outputs
 
     @staticmethod
@@ -79,7 +82,11 @@ def backward(ctx, *args):
         # Copy the list to avoid modifying original list.
         inputs = list(ctx.inputs)
         tensor_indices = ctx.tensor_indices
-        tensors = ctx.saved_tensors
+
+        if ctx.activation_offload:
+            tensors = ctx.tensor_inputs
+        else:
+            tensors = ctx.saved_tensors
 
         # store the current states
         bwd_cpu_rng_state = torch.get_rng_state()


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[BUG]: Checkpointing Test Failed with PyTorch 1.9
### 🐛 Describe the bug

When running unit testing for model checkpointing, the following exception occurs.

<img width="1370" alt="Screenshot 2022-04-11 at 5 10 48 PM" src="https://user-images.githubusercontent.com/31818963/162704185-ad267034-2f49-4d18-82a1-70a08c59bcd6.png">

This is because that the `_rank_not_in_group` API is not exposed at the `torch.distributed` level with PyTorch 1.9.

### Environment

CUDA: 11.1
PyTorch 1.9.1


Discussion:

Resolving patch:
diff --git a/colossalai/communication/collective.py b/colossalai/communication/collective.py
--- a/colossalai/communication/collective.py
+++ b/colossalai/communication/collective.py
@@ -211,7 +211,7 @@ def reduce(tensor: Tensor,
 def scatter_object_list(scatter_object_output_list, scatter_object_input_list, src=0, group=None):
     r"""Modified from `torch.distributed.scatter_object_list <https://pytorch.org/docs/stable/_modules/torch/distributed/distributed_c10d.html#scatter_object_list>` to fix issues
     """
-    if dist._rank_not_in_group(group):
+    if dist.distributed_c10d._rank_not_in_group(group):
         return
 
     if (not isinstance(scatter_object_output_list, list) or len(scatter_object_output_list) < 1):
@@ -220,7 +220,7 @@ def scatter_object_list(scatter_object_output_list, scatter_object_input_list, s
     # set tensor device to cuda if backend is nccl
     device = torch.cuda.current_device() if dist.get_backend(group) == 'nccl' else torch.device("cpu")
 
-    my_rank = dist.get_rank()  # use global rank
+    my_rank = dist.get_rank()    # use global rank
     if my_rank == src:
         tensor_list, tensor_sizes = zip(
             *[dist.distributed_c10d._object_to_tensor(obj) for obj in scatter_object_input_list])


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[BUG]: Torch AMP Failed with Torch 1.8
### 🐛 Describe the bug

The torch AMP module failed with torch 1.8 in Colossal-AI. The error is shown below.

<img width="1029" alt="Screenshot 2022-04-12 at 3 14 39 PM" src="https://user-images.githubusercontent.com/31818963/162902471-16185dcd-b1e8-4a6e-b12e-2bbd26fdd444.png">
 
## Reason

The AMP module in torch 1.8 uses a different API in grad scaler.

<img width="624" alt="Screenshot 2022-04-12 at 3 15 53 PM" src="https://user-images.githubusercontent.com/31818963/162902730-9872bdc4-5ac7-48ea-bbfc-d263ad98de1f.png">



### Environment

CUDA: 11.1.1
PyTorch: 1.8.1


Discussion:

Resolving patch:
diff --git a/colossalai/amp/torch_amp/_grad_scaler.py b/colossalai/amp/torch_amp/_grad_scaler.py
--- a/colossalai/amp/torch_amp/_grad_scaler.py
+++ b/colossalai/amp/torch_amp/_grad_scaler.py
@@ -12,6 +12,7 @@
 import torch.distributed as dist
 from colossalai.core import global_context as gpc
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+from packaging import version
 
 
 class _MultiDeviceReplicator(object):
@@ -122,6 +123,14 @@ def __init__(self, init_scale=2.**16, growth_factor=2.0, backoff_factor=0.5, gro
         else:
             self._enabled = enabled
 
+        # check version
+        torch_version = version.parse(torch.__version__)
+        assert torch_version.major == 1
+        if torch_version.minor > 8:
+            self._higher_than_torch18 = True
+        else:
+            self._higher_than_torch18 = False
+
         if self._enabled:
             assert growth_factor > 1.0, "The growth factor must be > 1.0."
             assert backoff_factor < 1.0, "The backoff factor must be < 1.0."
@@ -404,8 +413,12 @@ def update(self, new_scale=None):
                 for i in range(1, len(found_infs)):
                     found_inf_combined += found_infs[i]
 
-            torch._amp_update_scale_(_scale, _growth_tracker, found_inf_combined, self._growth_factor,
-                                     self._backoff_factor, self._growth_interval)
+            if self._higher_than_torch18:
+                torch._amp_update_scale_(_scale, _growth_tracker, found_inf_combined, self._growth_factor,
+                                         self._backoff_factor, self._growth_interval)
+            else:
+                self._scale = torch._amp_update_scale(_growth_tracker, _scale, found_inf_combined, self._growth_factor,
+                                                      self._backoff_factor, self._growth_interval)
 
         # To prepare for next iteration, clear the data collected from optimizers this iteration.
         self._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[BUG]: Global Process Group is not Compatible with PyTorch 1.7
### 🐛 Describe the bug

In PyTorch 1.7, we cannot pass `group=None` to distributed communication operations as shown below.

<img width="718" alt="Screenshot 2022-04-14 at 2 41 23 PM" src="https://user-images.githubusercontent.com/31818963/163328621-740dd5df-1368-4fe8-8944-af9d1cd21e19.png">

When we do `dist.all_reduce(tensor, group=gpc.get_process_group(ParallelMode.GLOBAL))`, an exception will be raise as `gpc.get_process_group` returns None.

### Environment

_No response_


Discussion:

Resolving patch:
diff --git a/colossalai/context/parallel_context.py b/colossalai/context/parallel_context.py
--- a/colossalai/context/parallel_context.py
+++ b/colossalai/context/parallel_context.py
@@ -374,7 +374,7 @@ def init_global_dist(self, rank: int, world_size: int, backend: str, host: str,
         # None will give the default global process group for pytorch dist operations
         ranks = list(range(world_size))
         cpu_group = dist.new_group(ranks, backend='gloo') if dist.get_backend() != 'gloo' else None
-        self._register_dist(rank, world_size, None, cpu_group, ranks, ParallelMode.GLOBAL)
+        self._register_dist(rank, world_size, dist.GroupMember.WORLD, cpu_group, ranks, ParallelMode.GLOBAL)
         self.add_global_rank(ParallelMode.GLOBAL, rank)
 
     def _register_dist(self, local_rank, world_size, process_group, cpu_group, ranks_in_group, mode):


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[BUG]: --global-option="--no_cuda_ext" does not work
### 🐛 Describe the bug

When we install colossalai without cuda extension, we use the command below:

```shell
pip install colossalai --global-option="--no_cuda_ext"
```

However, it may lead to error as the option`--no_cuda_ext` will be applied to the installation of the dependency libraries, leading to "unknown option" error.

Instead, we should use environment variable to specify no cuda extension needed.

### Environment

_No response_


Discussion:

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -9,8 +9,7 @@
 build_cuda_ext = True
 ext_modules = []
 
-if '--no_cuda_ext' in sys.argv:
-    sys.argv.remove('--no_cuda_ext')
+if int(os.environ.get('NO_CUDA_EXT', '0')) == 1:
     build_cuda_ext = False
 
 


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[BUG]: Process significantly slowed down and huge memory taken when dealling with images with different sizes
### 🐛 Describe the bug

Describe the bug:
The process was slowed down and huge memory was taken during the training process of DeTr. We found that it might be because of the COCO dataset's variant shape of input images. By adding `cudnn_benchmark = False` into the config file, the situation was solved. 

### Environment

_No response_


Discussion:

Resolving patch:
diff --git a/colossalai/initialize.py b/colossalai/initialize.py
--- a/colossalai/initialize.py
+++ b/colossalai/initialize.py
@@ -270,7 +270,7 @@ def initialize(model: nn.Module,
             ranks=[0])
 
     # cudnn
-    cudnn_benchmark = config.get('cudnn_benchmark', True)
+    cudnn_benchmark = config.get('cudnn_benchmark', False)
     cudnn_deterministic = config.get('cudnn_deterministic', False)
     torch.backends.cudnn.benchmark = cudnn_benchmark
     torch.backends.cudnn.deterministic = cudnn_deterministic


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[BUG]:  Cannot synchronize grads of shared parameters cross pipeline stages when using ZERO-3
### 🐛 Describe the bug

 @FrankLeeeee @ver217 
Hi, in line 36 of _pipeline_parallel_gradient_handler.py:
https://github.com/hpcaitech/ColossalAI/blob/1aad903c1537eafb73fac1729b6df30b7006312f/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py#L36

the condition "param.grad is not None" will not work properly with ZERO-3, because after ZERO-3 synchronized grads, all parameters's grads were set to "colo_attr"，grads are None and buckets is empty here! 

This line also has the problem:
https://github.com/hpcaitech/ColossalAI/blob/1aad903c1537eafb73fac1729b6df30b7006312f/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py#L43

### Environment

colossalai latest version


Discussion:
Thanks for your response. I will fix this soon. And @YuliangLiu0306 is designing auto pipeline parallism, thus this API will be removed in the future.

Resolving patch:
diff --git a/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py b/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py
--- a/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py
+++ b/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py
@@ -33,14 +33,19 @@ def handle_gradient(self):
             # Pack the buckets.
             for param in self._model.parameters():
                 group = getattr(param, 'pipeline_shared_module_pg', None)
-                if param.requires_grad and param.grad is not None and group is not None:
+                if param.requires_grad and group is not None and (
+                    (hasattr(param, 'colo_attr') and not param.colo_attr.saved_grad.is_null())
+                        or param.grad is not None):
                     tp = param.data.type()
                     buckets[group][tp].append(param)
 
             # For each bucket, all-reduce and copy all-reduced grads.
             for group, group_buckets in buckets.items():
                 for tp, bucket in group_buckets.items():
-                    grads = [param.grad.data for param in bucket]
+                    grads = [
+                        param.colo_attr.grad_payload if hasattr(param, 'colo_attr') else param.grad.data
+                        for param in bucket
+                    ]
                     coalesced = _flatten_dense_tensors(grads).to(torch.cuda.current_device())
                     dist.all_reduce(coalesced, op=dist.ReduceOp.SUM, group=group)
                     for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)):


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[BUG]: FreqCacheEmbeddingBag._weight.ProcessGroup is initialized before _weight initialized
### 🐛 Describe the bug

When I init a DLRM with ParallelFreqAwareEmbeddingBag, the bug is reported as the following:
![image](https://user-images.githubusercontent.com/34452939/184055340-c99579bf-c9df-43ed-8552-7c9ae7aadffb.png)


I believe that is because the [ParallelFreqAwareEmbeddingBag](https://github.com/hpcaitech/ColossalAI/blob/039b7ed3bc33173e36c5c4decd41f8d7b1ec0f45/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py#L60) init its `_weight.ProcessGroup` before the `_weight` is initialized.

After I swap above line with its next line, the traceback shows another error:
![image](https://user-images.githubusercontent.com/34452939/184055612-de402e67-ac19-4860-ac12-9622f88a6aca.png)

It looks like some api update issue.

### Environment

_No response_


Discussion:
Thanks for the feedback. I will have a look.

Resolving patch:
diff --git a/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py b/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py
--- a/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py
+++ b/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py
@@ -7,7 +7,7 @@
 from torch.nn.parameter import Parameter
 from .._utils import dual_all_to_all
 
-from colossalai.tensor import ColoParameter, ShardSpec, ComputeSpec, ComputePattern, ProcessGroup
+from colossalai.tensor import ColoParameter, ShardSpec, ComputeSpec, ComputePattern, ProcessGroup, ColoTensorSpec
 
 
 def get_partition(embedding_dim, rank, world_size) -> Tuple[int, int, bool]:
@@ -57,13 +57,15 @@ def __init__(self,
         self.embedding_dim_per_partition = self.partition_end_index - self.partition_start_index
 
         if _weight is None:
-            self._weight.process_group = ProcessGroup(tp_degree=self.world_size)
+            colo_tensor_spec = ColoTensorSpec(pg=ProcessGroup(tp_degree=self.world_size),
+                                              dist_attr=ShardSpec(dims=[-1], num_partitions=[self.world_size]),
+                                              compute_attr=ComputePattern.TP1D)
             self._weight = ColoParameter.from_torch_tensor(torch.empty(self.num_embeddings,
                                                                        self.embedding_dim_per_partition,
                                                                        device='cpu',
                                                                        dtype=dtype),
                                                            requires_grad=True,
-                                                           spec=ShardSpec(dims=[-1], num_partitions=[self.world_size]))
+                                                           spec=colo_tensor_spec)
             self.init_parameters()
         else:
             assert isinstance(_weight, ColoParameter), "initialized weight must in type of ColoParameter"


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[BUG]: ModuleNotFoundError: No module named 'colossalai.nn.optimizer.zero_optimizer'
### 🐛 Describe the bug

I install colossalAI with the command `pip install colossalai==0.1.11rc3+torch1.12cu11.3 -f https://release.colossalai.org`
But I get an error when follow https://github.com/hpcaitech/ColossalAI/tree/main/examples/tutorial#-run-opt-finetuning-and-inference, I just run `bash ./run_clm_synthetic.sh` and get an error as follows:

```shell
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /home/he.yan/ColossalAI/examples/tutorial/opt/opt/run_clm.py:46 in <module>                      │
│                                                                                                  │
│    43 from colossalai.core import global_context as gpc                                          │
│    44 from colossalai.logging import disable_existing_loggers, get_dist_logger                   │
│    45 from colossalai.nn.optimizer import HybridAdam                                             │
│ ❱  46 from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer                           │
│    47 from colossalai.nn.parallel import ZeroDDP                                                 │
│    48 from colossalai.tensor import ProcessGroup                                                 │
│    49 from colossalai.utils import get_current_device, get_dataloader                            │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
ModuleNotFoundError: No module named 'colossalai.nn.optimizer.zero_optimizer'
```

### Environment

Python 3.8.15
torch1.12cu11.3


Discussion:
Hi, @heya5 , that's kind of strange, let me try to reproduce this error.
I think your latest examples `import colossalai.nn.optimizer.zero_optimizer`, but your code doesn't have it...
By the way, may I ask if the opt model in your library implements model parallelism (tensor parallelim or pipeline parallelism)?
Yes, the release seems to have some problems. I am initiating a new release, you should expect to download the correct version by the end of today. Meanwhile, the opt model does not have tensor parallelism and pipeline parallelism as it is implemented by huggingface.

Resolving patch:
diff --git a/colossalai/__init__.py b/colossalai/__init__.py
--- a/colossalai/__init__.py
+++ b/colossalai/__init__.py
@@ -7,4 +7,4 @@
     launch_from_torch,
 )
 
-__version__ = '0.1.11rc2'
+__version__ = '0.1.11rc4'


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[BUG]: Cannot build wheel with torch
### 🐛 Describe the bug

When I execute the pypi release workflow, an error occured due to the lack of dependency on torch.


<img width="803" alt="Screenshot 2023-01-05 at 13 47 13" src="https://user-images.githubusercontent.com/31818963/210710479-f1a615a7-26e1-43a4-80c8-be58de3e13eb.png">


### Environment

CUDA: 11.3
Python: 3.8


Discussion:

Resolving patch:
diff --git a/op_builder/builder.py b/op_builder/builder.py
--- a/op_builder/builder.py
+++ b/op_builder/builder.py
@@ -3,14 +3,18 @@
 from pathlib import Path
 from typing import List
 
-import torch
-
 
 def get_cuda_cc_flag() -> List:
     """get_cuda_cc_flag
 
     cc flag for your GPU arch
     """
+
+    # only import torch when needed
+    # this is to avoid importing torch when building on a machine without torch pre-installed
+    # one case is to build wheel for pypi release
+    import torch
+    
     cc_flag = []
     for arch in torch.cuda.get_arch_list():
         res = re.search(r'sm_(\d+)', arch)
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -15,8 +15,9 @@
     if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 10):
         raise RuntimeError("Colossal-AI requires Pytorch 1.10 or newer.\n"
                            "The latest stable release can be obtained from https://pytorch.org/")
+    TORCH_AVAILABLE = True
 except ImportError:
-    raise ModuleNotFoundError('torch is not found. You need to install PyTorch before installing Colossal-AI.')
+    TORCH_AVAILABLE = False
 
 
 # ninja build does not work unless include_dirs are abs path
@@ -24,7 +25,7 @@
 build_cuda_ext = True
 ext_modules = []
 
-if int(os.environ.get('NO_CUDA_EXT', '0')) == 1:
+if int(os.environ.get('NO_CUDA_EXT', '0')) == 1 or not TORCH_AVAILABLE:
     build_cuda_ext = False
 
 


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[FEATURE]: Make CUDA extensions default to optional
### Describe the feature

When installing `Colossal-AI` with `pip install .` from source or `pip install colossalai` from pypi, we always build CUDA extensions by default. However, this brings many problems for users who do not have a well-built environment with CUDA and torch. As a result, they failed at the first step, which is installation. I would suggest not to build extensions by default. The users can build extensions in two ways:
1. export `CUDA_EXT=1` when running `pip install`
2. build CUDA kernels only when the kernel is used


Discussion:

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -18,15 +18,22 @@
     TORCH_AVAILABLE = True
 except ImportError:
     TORCH_AVAILABLE = False
+    CUDA_HOME = None
 
 
 # ninja build does not work unless include_dirs are abs path
 this_dir = os.path.dirname(os.path.abspath(__file__))
-build_cuda_ext = True
+build_cuda_ext = False
 ext_modules = []
 
-if int(os.environ.get('NO_CUDA_EXT', '0')) == 1 or not TORCH_AVAILABLE:
-    build_cuda_ext = False
+if int(os.environ.get('CUDA_EXT', '0')) == 1:
+    if not TORCH_AVAILABLE:
+        raise ModuleNotFoundError("PyTorch is not found while CUDA_EXT=1. You need to install PyTorch first in order to build CUDA extensions")
+
+    if not CUDA_HOME:
+        raise RuntimeError("CUDA_HOME is not found while CUDA_EXT=1. You need to export CUDA_HOME environment vairable or install CUDA Toolkit first in order to build CUDA extensions")
+
+    build_cuda_ext = True
 
 
 def check_cuda_torch_binary_vs_bare_metal(cuda_dir):


Test patch: