We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests

[BUG]: Wrong import in `zero/sharded_optim/_utils.py`
### 🐛 Describe the bug

In issue #2774 , thanks @malfet for pointing out that we should not use `torch._six` to import `inf` and use `torch` to import `inf` instead, however, there is a small mistake in PR #2775 use an invalid `torch.six` module to import `inf`. We should fix this typo.

### Environment

_No response_

--- 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/zero/sharded_optim/_utils.py`
Content:
```
1 import math
2 from typing import Optional
3 
4 import torch
5 import torch.distributed as dist
6 from torch.six import inf
7 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
8 
9 from colossalai.tensor import ColoParameter
10 from colossalai.utils import is_model_parallel_parameter
11 
12 
13 def flatten(input_):
14     return _flatten_dense_tensors(input_)
15 
16 
17 def unflatten(flat, tensors):
18     return _unflatten_dense_tensors(flat, tensors)
19 
20 
21 def count_numel(tensor_list):
22     res = 0
23     for tensor in tensor_list:
24         res += tensor.numel()
25     return res
26 
27 
28 def calculate_padding(numel, unit_size):
29     remainder = numel % unit_size
30     return unit_size - remainder if remainder else remainder
31 
32 
33 def shuffle_by_round_robin(tensor_list, num_partitions):
34     partitions = dict()
35 
36     for tensor_idx, tensor in enumerate(tensor_list):
37         partition_to_go = tensor_idx % num_partitions
38         if partition_to_go not in partitions:
39             partitions[partition_to_go] = []
40         partitions[partition_to_go].append(dict(tensor=tensor, index=tensor_idx))
41 
42     partitions_count = len(partitions)
43     new_tensor_list = []
44     tensor_index_mapping = dict()
45 
46     for partition_id in range(partitions_count):
47         partition_tensors = partitions[partition_id]
48         for item in partition_tensors:
49             tensor_index_mapping[item['index']] = len(new_tensor_list)
50             new_tensor_list.append(item['tensor'])
51 
52     return new_tensor_list, tensor_index_mapping
53 
54 
55 # create a flat tensor aligned at the alignment boundary
56 def flatten_dense_tensors_with_padding(tensor_list, unit_size):
57     num_elements = count_numel(tensor_list)
58     padding = calculate_padding(num_elements, unit_size=unit_size)
59 
60     if padding > 0:
61         pad_tensor = torch.zeros(padding, device=tensor_list[0].device, dtype=tensor_list[0].dtype)
62         padded_tensor_list = tensor_list + [pad_tensor]
63     else:
64         padded_tensor_list = tensor_list
65 
66     return flatten(padded_tensor_list)
67 
68 
69 def is_nccl_aligned(tensor):
70     return tensor.data_ptr() % 4 == 0
71 
72 
73 def get_grad_accumulate_object(tensor):
74     """
75     Return the AccumulateGrad of the input tensor
76     """
77 
78     # grad_fn reference:
79     # https://discuss.pytorch.org/t/in-the-grad-fn-i-find-a-next-functions-but-i-dont-understand-the-meaning-of-the-attribute/24463
80     # expand_as reference: https://pytorch.org/docs/stable/generated/torch.Tensor.expand.html#torch.Tensor.expand
81     #
82     # `next_functions` will return the backward graph where
83     # the first element is the AccumulateGrad of the leaf nodes.
84     # we want to get the AccumulateGrad of the input tensor instead of the leaf
85     # node in the whole computation graph.
86     # Therefore, we call expand_as to create a dummy graph
87     # where tensor_tmp and tensor indeed point to the same object.
88     # You can check this by print(tensor.data_ptr() == tensor_tmp.data_ptr())
89     tensor_tmp = tensor.expand_as(tensor)
90     grad_acc_obj = tensor_tmp.grad_fn.next_functions[0][0]
91     return grad_acc_obj
92 
93 
94 def split_half_float_double(tensor_list):
95     dtypes = ["torch.cuda.HalfTensor", "torch.cuda.FloatTensor", "torch.cuda.DoubleTensor", "torch.cuda.BFloat16Tensor"]
96     buckets = []
97     for i, dtype in enumerate(dtypes):
98         bucket = [t for t in tensor_list if t.type() == dtype]
99         if bucket:
100             buckets.append(bucket)
101     return buckets
102 
103 
104 def reduce_tensor_dp_group(tensor: torch.Tensor,
105                            dtype: Optional[torch.dtype] = None,
106                            dst_local_rank: Optional[int] = None,
107                            dst_global_rank: Optional[int] = None,
108                            group: Optional[dist.ProcessGroup] = None):
109     """
110     Reduce the tensor in the data parallel process group
111 
112     :param tensor: A tensor object to reduce/all-reduce
113     :param dtype: The data type used in communication
114     :param dst_rank: The source rank for reduce. If dst_rank is None,
115     :param parallel_mode: Communication parallel mode
116     all-reduce will be used instead of reduce. Default is None.
117 
118     :type tensor: torch.Tensor
119     :type dtype: torch.dtype, optional
120     :type dst_rank: int, optional
121     :type pg: ProcessGroup, optional
122     """
123     # use the original dtype
124     if dtype is None:
125         dtype = tensor.dtype
126 
127     # cast the data to specified dtype for reduce/all-reduce
128     if tensor.dtype != dtype:
129         tensor_to_reduce = tensor.to(dtype)
130     else:
131         tensor_to_reduce = tensor
132 
133     world_size = dist.get_world_size(group=group)
134     tensor_to_reduce.div_(world_size)
135 
136     # if rank is None, all reduce will be used
137     # else, reduce is used
138     use_all_reduce = dst_local_rank is None
139 
140     if use_all_reduce:
141         dist.all_reduce(tensor_to_reduce, group=group)
142     else:
143         dist.reduce(tensor=tensor_to_reduce, dst=dst_global_rank, group=group)
144 
145     # recover the original dtype
146     if tensor.dtype != dtype and tensor is not tensor_to_reduce:
147         local_rank = dist.get_rank(group=group)
148         if use_all_reduce or dst_local_rank == local_rank:
149             tensor.copy_(tensor_to_reduce)
150 
151     return tensor
152 
153 
154 def has_inf_or_nan(tensor):
155     try:
156         # if tensor is half, the .float() incurs an additional deep copy, but it's necessary if
157         # Pytorch's .sum() creates a one-element tensor of the same type as tensor
158         # (which is true for some recent version of pytorch).
159         tensor_sum = float(tensor.float().sum())
160         # More efficient version that can be used if .sum() returns a Python scalar
161         # tensor_sum = float(tensor.sum())
162     except RuntimeError as instance:
163         # We want to check if inst is actually an overflow exception.
164         # RuntimeError could come from a different error.
165         # If so, we still want the exception to propagate.
166         if "value cannot be converted" not in instance.args[0]:
167             raise
168         return True
169     else:
170         if tensor_sum == float('inf') or tensor_sum == -float('inf') or tensor_sum != tensor_sum:
171             return True
172         return False
173 
174 
175 def release_param_grad(tensor_list):
176     for tensor in tensor_list:
177         tensor.grad = None
178 
179 
180 def calculate_global_norm_from_list(norm_list):
181     """ Compute total from a list of norms
182     """
183     total_norm = 0.0
184     for norm in norm_list:
185         total_norm += norm**2.0
186     return math.sqrt(total_norm)
187 
188 
189 def compute_norm(gradients, params, dp_group, mp_group, norm_type=2):
190     """Clips gradient norm of an iterable of parameters.
191     This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and
192     added functionality to handle model parallel parameters. Note that
193     the gradients are modified in place.
194     Arguments:
195         parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
196             single Tensor that will have gradients normalized
197         max_norm (float or int): max norm of the gradients
198         norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
199             infinity norm.
200     Returns:
201         Total norm of the parameters (viewed as a single vector).
202     """
203 
204     if mp_group is None:
205         mp_rank = 0
206     else:
207         mp_rank = dist.get_rank(mp_group)
208 
209     norm_type = float(norm_type)
210     if norm_type == inf:
211         total_norm = max(g.data.abs().max() for g in gradients)
212         total_norm_cuda = torch.cuda.FloatTensor([float(total_norm)])
213         dist.all_reduce(total_norm_cuda, op=torch.distributed.ReduceOp.MAX, group=dp_group)
214 
215         # Take max across all GPUs.
216         if mp_group is not None:
217             dist.all_reduce(tensor=total_norm_cuda, op=torch.distributed.ReduceOp.MAX)
218         total_norm = total_norm_cuda[0].item()
219     else:
220         total_norm = 0.0
221         # if dist.get_rank() == 0:
222         #    logger.info(f"Total Norm beginning {total_norm}")
223 
224         for g, p in zip(gradients, params):
225             # Pipeline parallelism may replicate parameters. Avoid multi-counting.
226             tp_param_flag = False
227             if is_model_parallel_parameter(p) or (isinstance(p, ColoParameter) and not p.is_replicate()):
228                 tp_param_flag = True
229             if tp_param_flag or mp_rank == 0:
230                 param_norm = g.data.double().norm(2)
231                 total_norm += param_norm.item()**2
232 
233         # Sum across all model parallel GPUs.
234         total_norm_cuda = torch.cuda.FloatTensor([float(total_norm)])
235         torch.distributed.all_reduce(total_norm_cuda, op=torch.distributed.ReduceOp.SUM, group=dp_group)
236 
237         if mp_group is not None:
238             dist.all_reduce(tensor=total_norm_cuda, op=torch.distributed.ReduceOp.SUM, group=mp_group)
239 
240         total_norm = total_norm_cuda[0].item()**(1. / norm_type)
241 
242     if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm:
243         total_norm = -1
244 
245     return total_norm
246 
247 
248 def sync_param(flat_tensor, tensor_list):
249     """
250     Synchronize the flattened tensor and unflattened tensor list. When
251     a list of tensor are flattened with `torch._utils._unflatten_dense_tensors`,
252     a new tensor is created. Thus, the flat tensor and original tensor list do not
253     share the same memory space. This function will update the tensor list so that
254     they point to the same value.
255 
256     :param flat_tensor: A flat tensor obtained by calling `torch._utils._unflatten_dense_tensors` on a tensor lsit
257     :param tensor_list: A list of tensors corresponding to the flattened tensor
258     :type flat_tensor: torch.Tensor
259     :type tensor_list: List[torch.Tensor]
260     """
261     updated_params = unflatten(flat_tensor, tensor_list)
262 
263     # update the tensor data
264     for p, q in zip(tensor_list, updated_params):
265         p.data = q.data
266
```
--- 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()

```

diff --git a/colossalai/zero/sharded_optim/_utils.py b/colossalai/zero/sharded_optim/_utils.py
--- a/colossalai/zero/sharded_optim/_utils.py
+++ b/colossalai/zero/sharded_optim/_utils.py
@@ -3,7 +3,7 @@
 
 import torch
 import torch.distributed as dist
-from torch.six import inf
+from torch import inf
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 
 from colossalai.tensor import ColoParameter

{"golden_diff": "diff --git a/colossalai/zero/sharded_optim/_utils.py b/colossalai/zero/sharded_optim/_utils.py\n--- a/colossalai/zero/sharded_optim/_utils.py\n+++ b/colossalai/zero/sharded_optim/_utils.py\n@@ -3,7 +3,7 @@\n \n import torch\n import torch.distributed as dist\n-from torch.six import inf\n+from torch import inf\n from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors\n \n from colossalai.tensor import ColoParameter\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[BUG]: Wrong import in `zero/sharded_optim/_utils.py`\n### \ud83d\udc1b Describe the bug\n\nIn issue #2774 , thanks @malfet for pointing out that we should not use `torch._six` to import `inf` and use `torch` to import `inf` instead, however, there is a small mistake in PR #2775 use an invalid `torch.six` module to import `inf`. We should fix this typo.\n\n### Environment\n\n_No response_\n", "before_files": [{"content": "import math\nfrom typing import Optional\n\nimport torch\nimport torch.distributed as dist\nfrom torch.six import inf\nfrom torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors\n\nfrom colossalai.tensor import ColoParameter\nfrom colossalai.utils import is_model_parallel_parameter\n\n\ndef flatten(input_):\n    return _flatten_dense_tensors(input_)\n\n\ndef unflatten(flat, tensors):\n    return _unflatten_dense_tensors(flat, tensors)\n\n\ndef count_numel(tensor_list):\n    res = 0\n    for tensor in tensor_list:\n        res += tensor.numel()\n    return res\n\n\ndef calculate_padding(numel, unit_size):\n    remainder = numel % unit_size\n    return unit_size - remainder if remainder else remainder\n\n\ndef shuffle_by_round_robin(tensor_list, num_partitions):\n    partitions = dict()\n\n    for tensor_idx, tensor in enumerate(tensor_list):\n        partition_to_go = tensor_idx % num_partitions\n        if partition_to_go not in partitions:\n            partitions[partition_to_go] = []\n        partitions[partition_to_go].append(dict(tensor=tensor, index=tensor_idx))\n\n    partitions_count = len(partitions)\n    new_tensor_list = []\n    tensor_index_mapping = dict()\n\n    for partition_id in range(partitions_count):\n        partition_tensors = partitions[partition_id]\n        for item in partition_tensors:\n            tensor_index_mapping[item['index']] = len(new_tensor_list)\n            new_tensor_list.append(item['tensor'])\n\n    return new_tensor_list, tensor_index_mapping\n\n\n# create a flat tensor aligned at the alignment boundary\ndef flatten_dense_tensors_with_padding(tensor_list, unit_size):\n    num_elements = count_numel(tensor_list)\n    padding = calculate_padding(num_elements, unit_size=unit_size)\n\n    if padding > 0:\n        pad_tensor = torch.zeros(padding, device=tensor_list[0].device, dtype=tensor_list[0].dtype)\n        padded_tensor_list = tensor_list + [pad_tensor]\n    else:\n        padded_tensor_list = tensor_list\n\n    return flatten(padded_tensor_list)\n\n\ndef is_nccl_aligned(tensor):\n    return tensor.data_ptr() % 4 == 0\n\n\ndef get_grad_accumulate_object(tensor):\n    \"\"\"\n    Return the AccumulateGrad of the input tensor\n    \"\"\"\n\n    # grad_fn reference:\n    # https://discuss.pytorch.org/t/in-the-grad-fn-i-find-a-next-functions-but-i-dont-understand-the-meaning-of-the-attribute/24463\n    # expand_as reference: https://pytorch.org/docs/stable/generated/torch.Tensor.expand.html#torch.Tensor.expand\n    #\n    # `next_functions` will return the backward graph where\n    # the first element is the AccumulateGrad of the leaf nodes.\n    # we want to get the AccumulateGrad of the input tensor instead of the leaf\n    # node in the whole computation graph.\n    # Therefore, we call expand_as to create a dummy graph\n    # where tensor_tmp and tensor indeed point to the same object.\n    # You can check this by print(tensor.data_ptr() == tensor_tmp.data_ptr())\n    tensor_tmp = tensor.expand_as(tensor)\n    grad_acc_obj = tensor_tmp.grad_fn.next_functions[0][0]\n    return grad_acc_obj\n\n\ndef split_half_float_double(tensor_list):\n    dtypes = [\"torch.cuda.HalfTensor\", \"torch.cuda.FloatTensor\", \"torch.cuda.DoubleTensor\", \"torch.cuda.BFloat16Tensor\"]\n    buckets = []\n    for i, dtype in enumerate(dtypes):\n        bucket = [t for t in tensor_list if t.type() == dtype]\n        if bucket:\n            buckets.append(bucket)\n    return buckets\n\n\ndef reduce_tensor_dp_group(tensor: torch.Tensor,\n                           dtype: Optional[torch.dtype] = None,\n                           dst_local_rank: Optional[int] = None,\n                           dst_global_rank: Optional[int] = None,\n                           group: Optional[dist.ProcessGroup] = None):\n    \"\"\"\n    Reduce the tensor in the data parallel process group\n\n    :param tensor: A tensor object to reduce/all-reduce\n    :param dtype: The data type used in communication\n    :param dst_rank: The source rank for reduce. If dst_rank is None,\n    :param parallel_mode: Communication parallel mode\n    all-reduce will be used instead of reduce. Default is None.\n\n    :type tensor: torch.Tensor\n    :type dtype: torch.dtype, optional\n    :type dst_rank: int, optional\n    :type pg: ProcessGroup, optional\n    \"\"\"\n    # use the original dtype\n    if dtype is None:\n        dtype = tensor.dtype\n\n    # cast the data to specified dtype for reduce/all-reduce\n    if tensor.dtype != dtype:\n        tensor_to_reduce = tensor.to(dtype)\n    else:\n        tensor_to_reduce = tensor\n\n    world_size = dist.get_world_size(group=group)\n    tensor_to_reduce.div_(world_size)\n\n    # if rank is None, all reduce will be used\n    # else, reduce is used\n    use_all_reduce = dst_local_rank is None\n\n    if use_all_reduce:\n        dist.all_reduce(tensor_to_reduce, group=group)\n    else:\n        dist.reduce(tensor=tensor_to_reduce, dst=dst_global_rank, group=group)\n\n    # recover the original dtype\n    if tensor.dtype != dtype and tensor is not tensor_to_reduce:\n        local_rank = dist.get_rank(group=group)\n        if use_all_reduce or dst_local_rank == local_rank:\n            tensor.copy_(tensor_to_reduce)\n\n    return tensor\n\n\ndef has_inf_or_nan(tensor):\n    try:\n        # if tensor is half, the .float() incurs an additional deep copy, but it's necessary if\n        # Pytorch's .sum() creates a one-element tensor of the same type as tensor\n        # (which is true for some recent version of pytorch).\n        tensor_sum = float(tensor.float().sum())\n        # More efficient version that can be used if .sum() returns a Python scalar\n        # tensor_sum = float(tensor.sum())\n    except RuntimeError as instance:\n        # We want to check if inst is actually an overflow exception.\n        # RuntimeError could come from a different error.\n        # If so, we still want the exception to propagate.\n        if \"value cannot be converted\" not in instance.args[0]:\n            raise\n        return True\n    else:\n        if tensor_sum == float('inf') or tensor_sum == -float('inf') or tensor_sum != tensor_sum:\n            return True\n        return False\n\n\ndef release_param_grad(tensor_list):\n    for tensor in tensor_list:\n        tensor.grad = None\n\n\ndef calculate_global_norm_from_list(norm_list):\n    \"\"\" Compute total from a list of norms\n    \"\"\"\n    total_norm = 0.0\n    for norm in norm_list:\n        total_norm += norm**2.0\n    return math.sqrt(total_norm)\n\n\ndef compute_norm(gradients, params, dp_group, mp_group, norm_type=2):\n    \"\"\"Clips gradient norm of an iterable of parameters.\n    This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and\n    added functionality to handle model parallel parameters. Note that\n    the gradients are modified in place.\n    Arguments:\n        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a\n            single Tensor that will have gradients normalized\n        max_norm (float or int): max norm of the gradients\n        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for\n            infinity norm.\n    Returns:\n        Total norm of the parameters (viewed as a single vector).\n    \"\"\"\n\n    if mp_group is None:\n        mp_rank = 0\n    else:\n        mp_rank = dist.get_rank(mp_group)\n\n    norm_type = float(norm_type)\n    if norm_type == inf:\n        total_norm = max(g.data.abs().max() for g in gradients)\n        total_norm_cuda = torch.cuda.FloatTensor([float(total_norm)])\n        dist.all_reduce(total_norm_cuda, op=torch.distributed.ReduceOp.MAX, group=dp_group)\n\n        # Take max across all GPUs.\n        if mp_group is not None:\n            dist.all_reduce(tensor=total_norm_cuda, op=torch.distributed.ReduceOp.MAX)\n        total_norm = total_norm_cuda[0].item()\n    else:\n        total_norm = 0.0\n        # if dist.get_rank() == 0:\n        #    logger.info(f\"Total Norm beginning {total_norm}\")\n\n        for g, p in zip(gradients, params):\n            # Pipeline parallelism may replicate parameters. Avoid multi-counting.\n            tp_param_flag = False\n            if is_model_parallel_parameter(p) or (isinstance(p, ColoParameter) and not p.is_replicate()):\n                tp_param_flag = True\n            if tp_param_flag or mp_rank == 0:\n                param_norm = g.data.double().norm(2)\n                total_norm += param_norm.item()**2\n\n        # Sum across all model parallel GPUs.\n        total_norm_cuda = torch.cuda.FloatTensor([float(total_norm)])\n        torch.distributed.all_reduce(total_norm_cuda, op=torch.distributed.ReduceOp.SUM, group=dp_group)\n\n        if mp_group is not None:\n            dist.all_reduce(tensor=total_norm_cuda, op=torch.distributed.ReduceOp.SUM, group=mp_group)\n\n        total_norm = total_norm_cuda[0].item()**(1. / norm_type)\n\n    if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm:\n        total_norm = -1\n\n    return total_norm\n\n\ndef sync_param(flat_tensor, tensor_list):\n    \"\"\"\n    Synchronize the flattened tensor and unflattened tensor list. When\n    a list of tensor are flattened with `torch._utils._unflatten_dense_tensors`,\n    a new tensor is created. Thus, the flat tensor and original tensor list do not\n    share the same memory space. This function will update the tensor list so that\n    they point to the same value.\n\n    :param flat_tensor: A flat tensor obtained by calling `torch._utils._unflatten_dense_tensors` on a tensor lsit\n    :param tensor_list: A list of tensors corresponding to the flattened tensor\n    :type flat_tensor: torch.Tensor\n    :type tensor_list: List[torch.Tensor]\n    \"\"\"\n    updated_params = unflatten(flat_tensor, tensor_list)\n\n    # update the tensor data\n    for p, q in zip(tensor_list, updated_params):\n        p.data = q.data\n", "path": "colossalai/zero/sharded_optim/_utils.py"}], "after_files": [{"content": "import math\nfrom typing import Optional\n\nimport torch\nimport torch.distributed as dist\nfrom torch import inf\nfrom torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors\n\nfrom colossalai.tensor import ColoParameter\nfrom colossalai.utils import is_model_parallel_parameter\n\n\ndef flatten(input_):\n    return _flatten_dense_tensors(input_)\n\n\ndef unflatten(flat, tensors):\n    return _unflatten_dense_tensors(flat, tensors)\n\n\ndef count_numel(tensor_list):\n    res = 0\n    for tensor in tensor_list:\n        res += tensor.numel()\n    return res\n\n\ndef calculate_padding(numel, unit_size):\n    remainder = numel % unit_size\n    return unit_size - remainder if remainder else remainder\n\n\ndef shuffle_by_round_robin(tensor_list, num_partitions):\n    partitions = dict()\n\n    for tensor_idx, tensor in enumerate(tensor_list):\n        partition_to_go = tensor_idx % num_partitions\n        if partition_to_go not in partitions:\n            partitions[partition_to_go] = []\n        partitions[partition_to_go].append(dict(tensor=tensor, index=tensor_idx))\n\n    partitions_count = len(partitions)\n    new_tensor_list = []\n    tensor_index_mapping = dict()\n\n    for partition_id in range(partitions_count):\n        partition_tensors = partitions[partition_id]\n        for item in partition_tensors:\n            tensor_index_mapping[item['index']] = len(new_tensor_list)\n            new_tensor_list.append(item['tensor'])\n\n    return new_tensor_list, tensor_index_mapping\n\n\n# create a flat tensor aligned at the alignment boundary\ndef flatten_dense_tensors_with_padding(tensor_list, unit_size):\n    num_elements = count_numel(tensor_list)\n    padding = calculate_padding(num_elements, unit_size=unit_size)\n\n    if padding > 0:\n        pad_tensor = torch.zeros(padding, device=tensor_list[0].device, dtype=tensor_list[0].dtype)\n        padded_tensor_list = tensor_list + [pad_tensor]\n    else:\n        padded_tensor_list = tensor_list\n\n    return flatten(padded_tensor_list)\n\n\ndef is_nccl_aligned(tensor):\n    return tensor.data_ptr() % 4 == 0\n\n\ndef get_grad_accumulate_object(tensor):\n    \"\"\"\n    Return the AccumulateGrad of the input tensor\n    \"\"\"\n\n    # grad_fn reference:\n    # https://discuss.pytorch.org/t/in-the-grad-fn-i-find-a-next-functions-but-i-dont-understand-the-meaning-of-the-attribute/24463\n    # expand_as reference: https://pytorch.org/docs/stable/generated/torch.Tensor.expand.html#torch.Tensor.expand\n    #\n    # `next_functions` will return the backward graph where\n    # the first element is the AccumulateGrad of the leaf nodes.\n    # we want to get the AccumulateGrad of the input tensor instead of the leaf\n    # node in the whole computation graph.\n    # Therefore, we call expand_as to create a dummy graph\n    # where tensor_tmp and tensor indeed point to the same object.\n    # You can check this by print(tensor.data_ptr() == tensor_tmp.data_ptr())\n    tensor_tmp = tensor.expand_as(tensor)\n    grad_acc_obj = tensor_tmp.grad_fn.next_functions[0][0]\n    return grad_acc_obj\n\n\ndef split_half_float_double(tensor_list):\n    dtypes = [\"torch.cuda.HalfTensor\", \"torch.cuda.FloatTensor\", \"torch.cuda.DoubleTensor\", \"torch.cuda.BFloat16Tensor\"]\n    buckets = []\n    for i, dtype in enumerate(dtypes):\n        bucket = [t for t in tensor_list if t.type() == dtype]\n        if bucket:\n            buckets.append(bucket)\n    return buckets\n\n\ndef reduce_tensor_dp_group(tensor: torch.Tensor,\n                           dtype: Optional[torch.dtype] = None,\n                           dst_local_rank: Optional[int] = None,\n                           dst_global_rank: Optional[int] = None,\n                           group: Optional[dist.ProcessGroup] = None):\n    \"\"\"\n    Reduce the tensor in the data parallel process group\n\n    :param tensor: A tensor object to reduce/all-reduce\n    :param dtype: The data type used in communication\n    :param dst_rank: The source rank for reduce. If dst_rank is None,\n    :param parallel_mode: Communication parallel mode\n    all-reduce will be used instead of reduce. Default is None.\n\n    :type tensor: torch.Tensor\n    :type dtype: torch.dtype, optional\n    :type dst_rank: int, optional\n    :type pg: ProcessGroup, optional\n    \"\"\"\n    # use the original dtype\n    if dtype is None:\n        dtype = tensor.dtype\n\n    # cast the data to specified dtype for reduce/all-reduce\n    if tensor.dtype != dtype:\n        tensor_to_reduce = tensor.to(dtype)\n    else:\n        tensor_to_reduce = tensor\n\n    world_size = dist.get_world_size(group=group)\n    tensor_to_reduce.div_(world_size)\n\n    # if rank is None, all reduce will be used\n    # else, reduce is used\n    use_all_reduce = dst_local_rank is None\n\n    if use_all_reduce:\n        dist.all_reduce(tensor_to_reduce, group=group)\n    else:\n        dist.reduce(tensor=tensor_to_reduce, dst=dst_global_rank, group=group)\n\n    # recover the original dtype\n    if tensor.dtype != dtype and tensor is not tensor_to_reduce:\n        local_rank = dist.get_rank(group=group)\n        if use_all_reduce or dst_local_rank == local_rank:\n            tensor.copy_(tensor_to_reduce)\n\n    return tensor\n\n\ndef has_inf_or_nan(tensor):\n    try:\n        # if tensor is half, the .float() incurs an additional deep copy, but it's necessary if\n        # Pytorch's .sum() creates a one-element tensor of the same type as tensor\n        # (which is true for some recent version of pytorch).\n        tensor_sum = float(tensor.float().sum())\n        # More efficient version that can be used if .sum() returns a Python scalar\n        # tensor_sum = float(tensor.sum())\n    except RuntimeError as instance:\n        # We want to check if inst is actually an overflow exception.\n        # RuntimeError could come from a different error.\n        # If so, we still want the exception to propagate.\n        if \"value cannot be converted\" not in instance.args[0]:\n            raise\n        return True\n    else:\n        if tensor_sum == float('inf') or tensor_sum == -float('inf') or tensor_sum != tensor_sum:\n            return True\n        return False\n\n\ndef release_param_grad(tensor_list):\n    for tensor in tensor_list:\n        tensor.grad = None\n\n\ndef calculate_global_norm_from_list(norm_list):\n    \"\"\" Compute total from a list of norms\n    \"\"\"\n    total_norm = 0.0\n    for norm in norm_list:\n        total_norm += norm**2.0\n    return math.sqrt(total_norm)\n\n\ndef compute_norm(gradients, params, dp_group, mp_group, norm_type=2):\n    \"\"\"Clips gradient norm of an iterable of parameters.\n    This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and\n    added functionality to handle model parallel parameters. Note that\n    the gradients are modified in place.\n    Arguments:\n        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a\n            single Tensor that will have gradients normalized\n        max_norm (float or int): max norm of the gradients\n        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for\n            infinity norm.\n    Returns:\n        Total norm of the parameters (viewed as a single vector).\n    \"\"\"\n\n    if mp_group is None:\n        mp_rank = 0\n    else:\n        mp_rank = dist.get_rank(mp_group)\n\n    norm_type = float(norm_type)\n    if norm_type == inf:\n        total_norm = max(g.data.abs().max() for g in gradients)\n        total_norm_cuda = torch.cuda.FloatTensor([float(total_norm)])\n        dist.all_reduce(total_norm_cuda, op=torch.distributed.ReduceOp.MAX, group=dp_group)\n\n        # Take max across all GPUs.\n        if mp_group is not None:\n            dist.all_reduce(tensor=total_norm_cuda, op=torch.distributed.ReduceOp.MAX)\n        total_norm = total_norm_cuda[0].item()\n    else:\n        total_norm = 0.0\n        # if dist.get_rank() == 0:\n        #    logger.info(f\"Total Norm beginning {total_norm}\")\n\n        for g, p in zip(gradients, params):\n            # Pipeline parallelism may replicate parameters. Avoid multi-counting.\n            tp_param_flag = False\n            if is_model_parallel_parameter(p) or (isinstance(p, ColoParameter) and not p.is_replicate()):\n                tp_param_flag = True\n            if tp_param_flag or mp_rank == 0:\n                param_norm = g.data.double().norm(2)\n                total_norm += param_norm.item()**2\n\n        # Sum across all model parallel GPUs.\n        total_norm_cuda = torch.cuda.FloatTensor([float(total_norm)])\n        torch.distributed.all_reduce(total_norm_cuda, op=torch.distributed.ReduceOp.SUM, group=dp_group)\n\n        if mp_group is not None:\n            dist.all_reduce(tensor=total_norm_cuda, op=torch.distributed.ReduceOp.SUM, group=mp_group)\n\n        total_norm = total_norm_cuda[0].item()**(1. / norm_type)\n\n    if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm:\n        total_norm = -1\n\n    return total_norm\n\n\ndef sync_param(flat_tensor, tensor_list):\n    \"\"\"\n    Synchronize the flattened tensor and unflattened tensor list. When\n    a list of tensor are flattened with `torch._utils._unflatten_dense_tensors`,\n    a new tensor is created. Thus, the flat tensor and original tensor list do not\n    share the same memory space. This function will update the tensor list so that\n    they point to the same value.\n\n    :param flat_tensor: A flat tensor obtained by calling `torch._utils._unflatten_dense_tensors` on a tensor lsit\n    :param tensor_list: A list of tensors corresponding to the flattened tensor\n    :type flat_tensor: torch.Tensor\n    :type tensor_list: List[torch.Tensor]\n    \"\"\"\n    updated_params = unflatten(flat_tensor, tensor_list)\n\n    # update the tensor data\n    for p, q in zip(tensor_list, updated_params):\n        p.data = q.data\n", "path": "colossalai/zero/sharded_optim/_utils.py"}]}