prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN 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
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py`
  Content:
  ```
  1 #!/usr/bin/env python
  2 
  3 from collections import defaultdict
  4 
  5 import torch
  6 import torch.distributed as dist
  7 from colossalai.core import global_context as gpc
  8 from colossalai.registry import GRADIENT_HANDLER
  9 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
  10 
  11 from ._base_gradient_handler import BaseGradientHandler
  12 
  13 
  14 @GRADIENT_HANDLER.register_module
  15 class PipelineSharedModuleGradientHandler(BaseGradientHandler):
  16     """A helper class to handle all-reduce operations in sub parallel groups.
  17     A all-reduce collective communication will be operated in 
  18     :func:`handle_gradient` among all sub pipeline parallel groups.
  19     For better performance, it bucketizes the gradients of all parameters that are 
  20     the same type to improve the efficiency of communication.
  21 
  22     Args:
  23         model (Module): Model where the gradients accumulate.
  24         optimizer (Optimizer): Optimizer for updating the parameters.
  25     """
  26 
  27     def handle_gradient(self):
  28         """A method running a all-reduce operation in sub pipeline parallel groups.
  29         """
  30         if gpc.pipeline_parallel_size > 1:
  31             # bucketize and all-reduce
  32             buckets = defaultdict(lambda: defaultdict(list))
  33             # Pack the buckets.
  34             for param in self._model.parameters():
  35                 group = getattr(param, 'pipeline_shared_module_pg', None)
  36                 if param.requires_grad and param.grad is not None and group is not None:
  37                     tp = param.data.type()
  38                     buckets[group][tp].append(param)
  39 
  40             # For each bucket, all-reduce and copy all-reduced grads.
  41             for group, group_buckets in buckets.items():
  42                 for tp, bucket in group_buckets.items():
  43                     grads = [param.grad.data for param in bucket]
  44                     coalesced = _flatten_dense_tensors(grads).to(torch.cuda.current_device())
  45                     dist.all_reduce(coalesced, op=dist.ReduceOp.SUM, group=group)
  46                     for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)):
  47                         buf.copy_(synced)
  48
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: 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 @@
               # 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)):
  
verification_info: {"golden_diff": "diff --git a/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py b/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py\n--- a/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py\n+++ b/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py\n@@ -33,14 +33,19 @@\n             # Pack the buckets.\n             for param in self._model.parameters():\n                 group = getattr(param, 'pipeline_shared_module_pg', None)\n-                if param.requires_grad and param.grad is not None and group is not None:\n+                if param.requires_grad and group is not None and (\n+                    (hasattr(param, 'colo_attr') and not param.colo_attr.saved_grad.is_null())\n+                        or param.grad is not None):\n                     tp = param.data.type()\n                     buckets[group][tp].append(param)\n \n             # For each bucket, all-reduce and copy all-reduced grads.\n             for group, group_buckets in buckets.items():\n                 for tp, bucket in group_buckets.items():\n-                    grads = [param.grad.data for param in bucket]\n+                    grads = [\n+                        param.colo_attr.grad_payload if hasattr(param, 'colo_attr') else param.grad.data\n+                        for param in bucket\n+                    ]\n                     coalesced = _flatten_dense_tensors(grads).to(torch.cuda.current_device())\n                     dist.all_reduce(coalesced, op=dist.ReduceOp.SUM, group=group)\n                     for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)):\n", "issue": "[BUG]:  Cannot synchronize grads of shared parameters cross pipeline stages when using ZERO-3\n### \ud83d\udc1b Describe the bug\r\n\r\n @FrankLeeeee @ver217 \r\nHi, in line 36 of _pipeline_parallel_gradient_handler.py:\r\nhttps://github.com/hpcaitech/ColossalAI/blob/1aad903c1537eafb73fac1729b6df30b7006312f/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py#L36\r\n\r\nthe 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\"\uff0cgrads are None and buckets is empty here! \r\n\r\nThis line also has the problem:\r\nhttps://github.com/hpcaitech/ColossalAI/blob/1aad903c1537eafb73fac1729b6df30b7006312f/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py#L43\r\n\r\n### Environment\r\n\r\ncolossalai latest version\n", "before_files": [{"content": "#!/usr/bin/env python\n\nfrom collections import defaultdict\n\nimport torch\nimport torch.distributed as dist\nfrom colossalai.core import global_context as gpc\nfrom colossalai.registry import GRADIENT_HANDLER\nfrom torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors\n\nfrom ._base_gradient_handler import BaseGradientHandler\n\n\n@GRADIENT_HANDLER.register_module\nclass PipelineSharedModuleGradientHandler(BaseGradientHandler):\n    \"\"\"A helper class to handle all-reduce operations in sub parallel groups.\n    A all-reduce collective communication will be operated in \n    :func:`handle_gradient` among all sub pipeline parallel groups.\n    For better performance, it bucketizes the gradients of all parameters that are \n    the same type to improve the efficiency of communication.\n\n    Args:\n        model (Module): Model where the gradients accumulate.\n        optimizer (Optimizer): Optimizer for updating the parameters.\n    \"\"\"\n\n    def handle_gradient(self):\n        \"\"\"A method running a all-reduce operation in sub pipeline parallel groups.\n        \"\"\"\n        if gpc.pipeline_parallel_size > 1:\n            # bucketize and all-reduce\n            buckets = defaultdict(lambda: defaultdict(list))\n            # Pack the buckets.\n            for param in self._model.parameters():\n                group = getattr(param, 'pipeline_shared_module_pg', None)\n                if param.requires_grad and param.grad is not None and group is not None:\n                    tp = param.data.type()\n                    buckets[group][tp].append(param)\n\n            # For each bucket, all-reduce and copy all-reduced grads.\n            for group, group_buckets in buckets.items():\n                for tp, bucket in group_buckets.items():\n                    grads = [param.grad.data for param in bucket]\n                    coalesced = _flatten_dense_tensors(grads).to(torch.cuda.current_device())\n                    dist.all_reduce(coalesced, op=dist.ReduceOp.SUM, group=group)\n                    for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)):\n                        buf.copy_(synced)\n", "path": "colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py"}], "after_files": [{"content": "#!/usr/bin/env python\n\nfrom collections import defaultdict\n\nimport torch\nimport torch.distributed as dist\nfrom colossalai.core import global_context as gpc\nfrom colossalai.registry import GRADIENT_HANDLER\nfrom torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors\n\nfrom ._base_gradient_handler import BaseGradientHandler\n\n\n@GRADIENT_HANDLER.register_module\nclass PipelineSharedModuleGradientHandler(BaseGradientHandler):\n    \"\"\"A helper class to handle all-reduce operations in sub parallel groups.\n    A all-reduce collective communication will be operated in \n    :func:`handle_gradient` among all sub pipeline parallel groups.\n    For better performance, it bucketizes the gradients of all parameters that are \n    the same type to improve the efficiency of communication.\n\n    Args:\n        model (Module): Model where the gradients accumulate.\n        optimizer (Optimizer): Optimizer for updating the parameters.\n    \"\"\"\n\n    def handle_gradient(self):\n        \"\"\"A method running a all-reduce operation in sub pipeline parallel groups.\n        \"\"\"\n        if gpc.pipeline_parallel_size > 1:\n            # bucketize and all-reduce\n            buckets = defaultdict(lambda: defaultdict(list))\n            # Pack the buckets.\n            for param in self._model.parameters():\n                group = getattr(param, 'pipeline_shared_module_pg', None)\n                if param.requires_grad and group is not None and (\n                    (hasattr(param, 'colo_attr') and not param.colo_attr.saved_grad.is_null())\n                        or param.grad is not None):\n                    tp = param.data.type()\n                    buckets[group][tp].append(param)\n\n            # For each bucket, all-reduce and copy all-reduced grads.\n            for group, group_buckets in buckets.items():\n                for tp, bucket in group_buckets.items():\n                    grads = [\n                        param.colo_attr.grad_payload if hasattr(param, 'colo_attr') else param.grad.data\n                        for param in bucket\n                    ]\n                    coalesced = _flatten_dense_tensors(grads).to(torch.cuda.current_device())\n                    dist.all_reduce(coalesced, op=dist.ReduceOp.SUM, group=group)\n                    for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)):\n                        buf.copy_(synced)\n", "path": "colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py"}]}