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


--- 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/tensor/param_op_hook.py`
Content:
```
1 from abc import ABC, abstractmethod
2 from contextlib import contextmanager
3 from typing import Any, List, Tuple
4 
5 import torch
6 from torch.utils._pytree import TreeSpec, tree_flatten, tree_unflatten
7 
8 
9 class ColoParamOpHook(ABC):
10     """
11     Hook which is triggered by each operation when operands contain ColoParameter.
12     To customize it, you must inherit this abstract class, and implement ``pre_forward``,
13     ``post_forward``, ``pre_backward`` and ``post_backward``.
14     These four methods apply a list of ColoParameter as input args.
15     """
16 
17     @abstractmethod
18     def pre_forward(self, params: List[torch.Tensor]) -> None:
19         pass
20 
21     @abstractmethod
22     def post_forward(self, params: List[torch.Tensor]) -> None:
23         pass
24 
25     @abstractmethod
26     def pre_backward(self, params: List[torch.Tensor]) -> None:
27         pass
28 
29     @abstractmethod
30     def post_backward(self, params: List[torch.Tensor]) -> None:
31         pass
32 
33 
34 class ColoParamOpHookManager:
35     """
36     Manage your param op hooks. It only has static methods.
37     The only static method you should call is ``use_hooks(*hooks)``.
38     """
39 
40     hooks: Tuple[ColoParamOpHook, ...] = tuple()
41 
42     @staticmethod
43     @contextmanager
44     def use_hooks(*hooks: ColoParamOpHook):
45         """Change the param op hooks you use. Nested calling is allowed.
46 
47         Example:
48             >>> with ColoParamOpHookManager.use_hooks(*hooks):
49             >>>     do_something()
50             >>>     with ColoParamOpHookManager.use_hooks():
51             >>>         // clear hooks
52             >>>         do_something()
53         """
54         try:
55             old_param_op_hooks = ColoParamOpHookManager.hooks
56             ColoParamOpHookManager.hooks = hooks
57             yield
58         finally:
59             ColoParamOpHookManager.hooks = old_param_op_hooks
60 
61     @staticmethod
62     def _trigger_pre_forward(params: List[torch.Tensor]) -> None:
63         for hook in ColoParamOpHookManager.hooks:
64             hook.pre_forward(params)
65 
66     @staticmethod
67     def _trigger_post_forward(params: List[torch.Tensor]) -> None:
68         for hook in ColoParamOpHookManager.hooks:
69             hook.post_forward(params)
70 
71     @staticmethod
72     def _trigger_pre_backward(params: List[torch.Tensor]) -> None:
73         for hook in ColoParamOpHookManager.hooks:
74             hook.pre_backward(params)
75 
76     @staticmethod
77     def _trigger_post_backward(params: List[torch.Tensor]) -> None:
78         for hook in ColoParamOpHookManager.hooks:
79             hook.post_backward(params)
80 
81     @staticmethod
82     def pre_op(params: List[torch.Tensor], *args: Any) -> list:
83         ColoParamOpHookManager._trigger_pre_forward(params)
84         # auto grad function can only recognize torch.Tensor, thus we have to flatten the input
85         # if one of the input requires grad, all the output will be treated as requires grad
86         # and will have grad fn even the corresponding input does not require grad
87         # we have to extract tensors requiring grad into flat list and then merge them back
88         grad_args, other_args, grad_flags, spec = _flatten_grad_args(args)
89         new_grad_args = PreFwdPostBwd.apply(params, *grad_args)
90         return _merge_args(new_grad_args, other_args, grad_flags, spec)
91 
92     @staticmethod
93     def post_op(params: List[torch.Tensor], arg: Any) -> Any:
94         ColoParamOpHookManager._trigger_post_forward(params)
95         return PostFwdPreBwd.apply(params, arg)
96 
97     @staticmethod
98     def has_hook() -> bool:
99         return len(ColoParamOpHookManager.hooks) > 0
100 
101 
102 class PreFwdPostBwd(torch.autograd.Function):
103     @staticmethod
104     def forward(ctx, params, *args):
105         ctx.params = params
106         return args
107 
108     @staticmethod
109     def backward(ctx, *grads):
110         ColoParamOpHookManager._trigger_post_backward(ctx.params)
111         return (None,) + grads
112 
113 
114 class PostFwdPreBwd(torch.autograd.Function):
115     @staticmethod
116     def forward(ctx, params, args):
117         ctx.params = params
118         return args
119 
120     @staticmethod
121     def backward(ctx, *grads):
122         ColoParamOpHookManager._trigger_pre_backward(ctx.params)
123         return (None,) + grads
124 
125 
126 def _is_grad_tensor(obj) -> bool:
127     if torch.is_tensor(obj):
128         if obj.grad_fn is not None or obj.requires_grad:
129             return True
130     return False
131 
132 
133 def _flatten_grad_args(args) -> Tuple[list, list, List[bool], TreeSpec]:
134     flat_args, spec = tree_flatten(args)
135     grad_args = []
136     other_args = []
137     grad_flags = []
138     for arg in flat_args:
139         flag = _is_grad_tensor(arg)
140         grad_flags.append(flag)
141         if flag:
142             grad_args.append(arg)
143         else:
144             other_args.append(arg)
145     assert len(grad_args) > 0
146     return grad_args, other_args, grad_flags, spec
147 
148 
149 def _merge_args(grad_args, other_args, grad_flags, spec):
150     grad_iter = iter(grad_args)
151     other_iter = iter(other_args)
152     flat_args = [next(grad_iter) if flag else next(other_iter) for flag in grad_flags]
153     return tree_unflatten(flat_args, spec)
154
```

Path: `colossalai/tensor/colo_parameter.py`
Content:
```
1 from typing import Optional
2 
3 import torch
4 
5 from colossalai.tensor.colo_tensor import ColoTensor
6 from colossalai.tensor.param_op_hook import ColoParamOpHookManager
7 
8 from .colo_tensor import _convert_output
9 
10 WHITE_LIST_FUNCS = {torch.Tensor.__getitem__, torch.Tensor.is_floating_point}
11 
12 
13 def is_no_hook_op(func) -> bool:
14     return func.__name__.startswith("__") and func not in WHITE_LIST_FUNCS
15 
16 
17 def filter_colo_parameters(*args, **kwargs):
18     param_list = []
19 
20     def get_colo_parameters(element) -> None:
21         if isinstance(element, list) or isinstance(element, tuple):
22             for e in element:
23                 get_colo_parameters(e)
24         elif isinstance(element, dict):
25             raise RuntimeError("Found Dict: ColoParameter can't deal with complicated arguments.")
26         elif isinstance(element, ColoParameter):
27             param_list.append(element)
28         return
29 
30     for a in args:
31         get_colo_parameters(a)
32     for v in kwargs.values():
33         get_colo_parameters(v)
34 
35     return param_list
36 
37 
38 def replace_args(args, kwargs, new_args):
39     args = new_args[: len(args)]
40     for k, v in zip(kwargs.keys(), new_args[len(args) :]):
41         kwargs[k] = v
42     return tuple(args), kwargs
43 
44 
45 class ColoParameter(ColoTensor, torch.nn.Parameter):
46     r"""A kind of ColoTensor to be considered as a module parameter."""
47 
48     def __new__(cls, data: Optional[torch.Tensor] = None, requires_grad: bool = True) -> "ColoParameter":
49         if data is None:
50             data = torch.empty(0)
51         return torch.Tensor._make_subclass(cls, data, requires_grad)
52 
53     @classmethod
54     def __torch_function__(cls, func, types, args=..., kwargs=None):
55         if kwargs is None:
56             kwargs = {}
57         if ColoParamOpHookManager.has_hook() and not is_no_hook_op(func):
58             params = filter_colo_parameters(*args, **kwargs)
59             if len(params) > 0:
60                 with torch._C.DisableTorchFunction():
61                     new_args = ColoParamOpHookManager.pre_op(params, *args, *kwargs.values())
62                 args, kwargs = replace_args(args, kwargs, new_args)
63                 ret = super().__torch_function__(func, types, args, kwargs)
64                 with torch._C.DisableTorchFunction():
65                     ret = ColoParamOpHookManager.post_op(params, ret)
66                 return _convert_output(ret, func)
67         return super().__torch_function__(func, types, args, kwargs)
68 
69     def __deepcopy__(self, memo):
70         if id(self) in memo:
71             return memo[id(self)]
72         else:
73             with torch._C.DisableTorchFunction():
74                 data = self.data.clone()
75             tensor = ColoParameter(data, self.requires_grad)
76             memo[id(self)] = tensor
77             return tensor
78 
79     def __reduce_ex__(self, proto):
80         # Adapted from torch._utils._rebuild_parameter
81         # def _rebuild_colo_parameter(data, requires_grad, backward_hooks):
82         #     colo_param = ColoParameter(data, requires_grad)
83         #     colo_param._backward_hooks = backward_hooks
84         #     return colo_param
85 
86         # return (
87         #     _rebuild_colo_parameter,
88         #     (self.data, self.requires_grad, OrderedDict())
89         # )
90 
91         # TODO(jzy) we don't support object reflection now.
92         # distspec cannot be pickled or rebuilt because it's tightly connected to runtime attribute `process_group`.
93         raise NotImplementedError
94
```
--- 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/tensor/colo_parameter.py b/colossalai/tensor/colo_parameter.py
--- a/colossalai/tensor/colo_parameter.py
+++ b/colossalai/tensor/colo_parameter.py
@@ -7,11 +7,12 @@
 
 from .colo_tensor import _convert_output
 
-WHITE_LIST_FUNCS = {torch.Tensor.__getitem__, torch.Tensor.is_floating_point}
+WHITE_LIST_FUNCS = {torch.Tensor.__getitem__}
+NO_HOOK_FUNCS = {torch.Tensor.is_floating_point}
 
 
 def is_no_hook_op(func) -> bool:
-    return func.__name__.startswith("__") and func not in WHITE_LIST_FUNCS
+    return (func.__name__.startswith("__") and func not in WHITE_LIST_FUNCS) or func in NO_HOOK_FUNCS
 
 
 def filter_colo_parameters(*args, **kwargs):
diff --git a/colossalai/tensor/param_op_hook.py b/colossalai/tensor/param_op_hook.py
--- a/colossalai/tensor/param_op_hook.py
+++ b/colossalai/tensor/param_op_hook.py
@@ -92,7 +92,10 @@
     @staticmethod
     def post_op(params: List[torch.Tensor], arg: Any) -> Any:
         ColoParamOpHookManager._trigger_post_forward(params)
-        return PostFwdPreBwd.apply(params, arg)
+        # incase the output is a tuple, we have to flatten it
+        grad_args, other_args, grad_flags, spec = _flatten_grad_args(arg)
+        new_grad_args = PostFwdPreBwd.apply(params, *grad_args)
+        return _merge_args(new_grad_args, other_args, grad_flags, spec)
 
     @staticmethod
     def has_hook() -> bool:
@@ -113,7 +116,7 @@
 
 class PostFwdPreBwd(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, params, args):
+    def forward(ctx, params, *args):
         ctx.params = params
         return args
 
@@ -142,7 +145,6 @@
             grad_args.append(arg)
         else:
             other_args.append(arg)
-    assert len(grad_args) > 0
     return grad_args, other_args, grad_flags, spec

{"golden_diff": "diff --git a/colossalai/tensor/colo_parameter.py b/colossalai/tensor/colo_parameter.py\n--- a/colossalai/tensor/colo_parameter.py\n+++ b/colossalai/tensor/colo_parameter.py\n@@ -7,11 +7,12 @@\n \n from .colo_tensor import _convert_output\n \n-WHITE_LIST_FUNCS = {torch.Tensor.__getitem__, torch.Tensor.is_floating_point}\n+WHITE_LIST_FUNCS = {torch.Tensor.__getitem__}\n+NO_HOOK_FUNCS = {torch.Tensor.is_floating_point}\n \n \n def is_no_hook_op(func) -> bool:\n-    return func.__name__.startswith(\"__\") and func not in WHITE_LIST_FUNCS\n+    return (func.__name__.startswith(\"__\") and func not in WHITE_LIST_FUNCS) or func in NO_HOOK_FUNCS\n \n \n def filter_colo_parameters(*args, **kwargs):\ndiff --git a/colossalai/tensor/param_op_hook.py b/colossalai/tensor/param_op_hook.py\n--- a/colossalai/tensor/param_op_hook.py\n+++ b/colossalai/tensor/param_op_hook.py\n@@ -92,7 +92,10 @@\n     @staticmethod\n     def post_op(params: List[torch.Tensor], arg: Any) -> Any:\n         ColoParamOpHookManager._trigger_post_forward(params)\n-        return PostFwdPreBwd.apply(params, arg)\n+        # incase the output is a tuple, we have to flatten it\n+        grad_args, other_args, grad_flags, spec = _flatten_grad_args(arg)\n+        new_grad_args = PostFwdPreBwd.apply(params, *grad_args)\n+        return _merge_args(new_grad_args, other_args, grad_flags, spec)\n \n     @staticmethod\n     def has_hook() -> bool:\n@@ -113,7 +116,7 @@\n \n class PostFwdPreBwd(torch.autograd.Function):\n     @staticmethod\n-    def forward(ctx, params, args):\n+    def forward(ctx, params, *args):\n         ctx.params = params\n         return args\n \n@@ -142,7 +145,6 @@\n             grad_args.append(arg)\n         else:\n             other_args.append(arg)\n-    assert len(grad_args) > 0\n     return grad_args, other_args, grad_flags, spec\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "from abc import ABC, abstractmethod\nfrom contextlib import contextmanager\nfrom typing import Any, List, Tuple\n\nimport torch\nfrom torch.utils._pytree import TreeSpec, tree_flatten, tree_unflatten\n\n\nclass ColoParamOpHook(ABC):\n    \"\"\"\n    Hook which is triggered by each operation when operands contain ColoParameter.\n    To customize it, you must inherit this abstract class, and implement ``pre_forward``,\n    ``post_forward``, ``pre_backward`` and ``post_backward``.\n    These four methods apply a list of ColoParameter as input args.\n    \"\"\"\n\n    @abstractmethod\n    def pre_forward(self, params: List[torch.Tensor]) -> None:\n        pass\n\n    @abstractmethod\n    def post_forward(self, params: List[torch.Tensor]) -> None:\n        pass\n\n    @abstractmethod\n    def pre_backward(self, params: List[torch.Tensor]) -> None:\n        pass\n\n    @abstractmethod\n    def post_backward(self, params: List[torch.Tensor]) -> None:\n        pass\n\n\nclass ColoParamOpHookManager:\n    \"\"\"\n    Manage your param op hooks. It only has static methods.\n    The only static method you should call is ``use_hooks(*hooks)``.\n    \"\"\"\n\n    hooks: Tuple[ColoParamOpHook, ...] = tuple()\n\n    @staticmethod\n    @contextmanager\n    def use_hooks(*hooks: ColoParamOpHook):\n        \"\"\"Change the param op hooks you use. Nested calling is allowed.\n\n        Example:\n            >>> with ColoParamOpHookManager.use_hooks(*hooks):\n            >>>     do_something()\n            >>>     with ColoParamOpHookManager.use_hooks():\n            >>>         // clear hooks\n            >>>         do_something()\n        \"\"\"\n        try:\n            old_param_op_hooks = ColoParamOpHookManager.hooks\n            ColoParamOpHookManager.hooks = hooks\n            yield\n        finally:\n            ColoParamOpHookManager.hooks = old_param_op_hooks\n\n    @staticmethod\n    def _trigger_pre_forward(params: List[torch.Tensor]) -> None:\n        for hook in ColoParamOpHookManager.hooks:\n            hook.pre_forward(params)\n\n    @staticmethod\n    def _trigger_post_forward(params: List[torch.Tensor]) -> None:\n        for hook in ColoParamOpHookManager.hooks:\n            hook.post_forward(params)\n\n    @staticmethod\n    def _trigger_pre_backward(params: List[torch.Tensor]) -> None:\n        for hook in ColoParamOpHookManager.hooks:\n            hook.pre_backward(params)\n\n    @staticmethod\n    def _trigger_post_backward(params: List[torch.Tensor]) -> None:\n        for hook in ColoParamOpHookManager.hooks:\n            hook.post_backward(params)\n\n    @staticmethod\n    def pre_op(params: List[torch.Tensor], *args: Any) -> list:\n        ColoParamOpHookManager._trigger_pre_forward(params)\n        # auto grad function can only recognize torch.Tensor, thus we have to flatten the input\n        # if one of the input requires grad, all the output will be treated as requires grad\n        # and will have grad fn even the corresponding input does not require grad\n        # we have to extract tensors requiring grad into flat list and then merge them back\n        grad_args, other_args, grad_flags, spec = _flatten_grad_args(args)\n        new_grad_args = PreFwdPostBwd.apply(params, *grad_args)\n        return _merge_args(new_grad_args, other_args, grad_flags, spec)\n\n    @staticmethod\n    def post_op(params: List[torch.Tensor], arg: Any) -> Any:\n        ColoParamOpHookManager._trigger_post_forward(params)\n        return PostFwdPreBwd.apply(params, arg)\n\n    @staticmethod\n    def has_hook() -> bool:\n        return len(ColoParamOpHookManager.hooks) > 0\n\n\nclass PreFwdPostBwd(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, params, *args):\n        ctx.params = params\n        return args\n\n    @staticmethod\n    def backward(ctx, *grads):\n        ColoParamOpHookManager._trigger_post_backward(ctx.params)\n        return (None,) + grads\n\n\nclass PostFwdPreBwd(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, params, args):\n        ctx.params = params\n        return args\n\n    @staticmethod\n    def backward(ctx, *grads):\n        ColoParamOpHookManager._trigger_pre_backward(ctx.params)\n        return (None,) + grads\n\n\ndef _is_grad_tensor(obj) -> bool:\n    if torch.is_tensor(obj):\n        if obj.grad_fn is not None or obj.requires_grad:\n            return True\n    return False\n\n\ndef _flatten_grad_args(args) -> Tuple[list, list, List[bool], TreeSpec]:\n    flat_args, spec = tree_flatten(args)\n    grad_args = []\n    other_args = []\n    grad_flags = []\n    for arg in flat_args:\n        flag = _is_grad_tensor(arg)\n        grad_flags.append(flag)\n        if flag:\n            grad_args.append(arg)\n        else:\n            other_args.append(arg)\n    assert len(grad_args) > 0\n    return grad_args, other_args, grad_flags, spec\n\n\ndef _merge_args(grad_args, other_args, grad_flags, spec):\n    grad_iter = iter(grad_args)\n    other_iter = iter(other_args)\n    flat_args = [next(grad_iter) if flag else next(other_iter) for flag in grad_flags]\n    return tree_unflatten(flat_args, spec)\n", "path": "colossalai/tensor/param_op_hook.py"}, {"content": "from typing import Optional\n\nimport torch\n\nfrom colossalai.tensor.colo_tensor import ColoTensor\nfrom colossalai.tensor.param_op_hook import ColoParamOpHookManager\n\nfrom .colo_tensor import _convert_output\n\nWHITE_LIST_FUNCS = {torch.Tensor.__getitem__, torch.Tensor.is_floating_point}\n\n\ndef is_no_hook_op(func) -> bool:\n    return func.__name__.startswith(\"__\") and func not in WHITE_LIST_FUNCS\n\n\ndef filter_colo_parameters(*args, **kwargs):\n    param_list = []\n\n    def get_colo_parameters(element) -> None:\n        if isinstance(element, list) or isinstance(element, tuple):\n            for e in element:\n                get_colo_parameters(e)\n        elif isinstance(element, dict):\n            raise RuntimeError(\"Found Dict: ColoParameter can't deal with complicated arguments.\")\n        elif isinstance(element, ColoParameter):\n            param_list.append(element)\n        return\n\n    for a in args:\n        get_colo_parameters(a)\n    for v in kwargs.values():\n        get_colo_parameters(v)\n\n    return param_list\n\n\ndef replace_args(args, kwargs, new_args):\n    args = new_args[: len(args)]\n    for k, v in zip(kwargs.keys(), new_args[len(args) :]):\n        kwargs[k] = v\n    return tuple(args), kwargs\n\n\nclass ColoParameter(ColoTensor, torch.nn.Parameter):\n    r\"\"\"A kind of ColoTensor to be considered as a module parameter.\"\"\"\n\n    def __new__(cls, data: Optional[torch.Tensor] = None, requires_grad: bool = True) -> \"ColoParameter\":\n        if data is None:\n            data = torch.empty(0)\n        return torch.Tensor._make_subclass(cls, data, requires_grad)\n\n    @classmethod\n    def __torch_function__(cls, func, types, args=..., kwargs=None):\n        if kwargs is None:\n            kwargs = {}\n        if ColoParamOpHookManager.has_hook() and not is_no_hook_op(func):\n            params = filter_colo_parameters(*args, **kwargs)\n            if len(params) > 0:\n                with torch._C.DisableTorchFunction():\n                    new_args = ColoParamOpHookManager.pre_op(params, *args, *kwargs.values())\n                args, kwargs = replace_args(args, kwargs, new_args)\n                ret = super().__torch_function__(func, types, args, kwargs)\n                with torch._C.DisableTorchFunction():\n                    ret = ColoParamOpHookManager.post_op(params, ret)\n                return _convert_output(ret, func)\n        return super().__torch_function__(func, types, args, kwargs)\n\n    def __deepcopy__(self, memo):\n        if id(self) in memo:\n            return memo[id(self)]\n        else:\n            with torch._C.DisableTorchFunction():\n                data = self.data.clone()\n            tensor = ColoParameter(data, self.requires_grad)\n            memo[id(self)] = tensor\n            return tensor\n\n    def __reduce_ex__(self, proto):\n        # Adapted from torch._utils._rebuild_parameter\n        # def _rebuild_colo_parameter(data, requires_grad, backward_hooks):\n        #     colo_param = ColoParameter(data, requires_grad)\n        #     colo_param._backward_hooks = backward_hooks\n        #     return colo_param\n\n        # return (\n        #     _rebuild_colo_parameter,\n        #     (self.data, self.requires_grad, OrderedDict())\n        # )\n\n        # TODO(jzy) we don't support object reflection now.\n        # distspec cannot be pickled or rebuilt because it's tightly connected to runtime attribute `process_group`.\n        raise NotImplementedError\n", "path": "colossalai/tensor/colo_parameter.py"}], "after_files": [{"content": "from abc import ABC, abstractmethod\nfrom contextlib import contextmanager\nfrom typing import Any, List, Tuple\n\nimport torch\nfrom torch.utils._pytree import TreeSpec, tree_flatten, tree_unflatten\n\n\nclass ColoParamOpHook(ABC):\n    \"\"\"\n    Hook which is triggered by each operation when operands contain ColoParameter.\n    To customize it, you must inherit this abstract class, and implement ``pre_forward``,\n    ``post_forward``, ``pre_backward`` and ``post_backward``.\n    These four methods apply a list of ColoParameter as input args.\n    \"\"\"\n\n    @abstractmethod\n    def pre_forward(self, params: List[torch.Tensor]) -> None:\n        pass\n\n    @abstractmethod\n    def post_forward(self, params: List[torch.Tensor]) -> None:\n        pass\n\n    @abstractmethod\n    def pre_backward(self, params: List[torch.Tensor]) -> None:\n        pass\n\n    @abstractmethod\n    def post_backward(self, params: List[torch.Tensor]) -> None:\n        pass\n\n\nclass ColoParamOpHookManager:\n    \"\"\"\n    Manage your param op hooks. It only has static methods.\n    The only static method you should call is ``use_hooks(*hooks)``.\n    \"\"\"\n\n    hooks: Tuple[ColoParamOpHook, ...] = tuple()\n\n    @staticmethod\n    @contextmanager\n    def use_hooks(*hooks: ColoParamOpHook):\n        \"\"\"Change the param op hooks you use. Nested calling is allowed.\n\n        Example:\n            >>> with ColoParamOpHookManager.use_hooks(*hooks):\n            >>>     do_something()\n            >>>     with ColoParamOpHookManager.use_hooks():\n            >>>         // clear hooks\n            >>>         do_something()\n        \"\"\"\n        try:\n            old_param_op_hooks = ColoParamOpHookManager.hooks\n            ColoParamOpHookManager.hooks = hooks\n            yield\n        finally:\n            ColoParamOpHookManager.hooks = old_param_op_hooks\n\n    @staticmethod\n    def _trigger_pre_forward(params: List[torch.Tensor]) -> None:\n        for hook in ColoParamOpHookManager.hooks:\n            hook.pre_forward(params)\n\n    @staticmethod\n    def _trigger_post_forward(params: List[torch.Tensor]) -> None:\n        for hook in ColoParamOpHookManager.hooks:\n            hook.post_forward(params)\n\n    @staticmethod\n    def _trigger_pre_backward(params: List[torch.Tensor]) -> None:\n        for hook in ColoParamOpHookManager.hooks:\n            hook.pre_backward(params)\n\n    @staticmethod\n    def _trigger_post_backward(params: List[torch.Tensor]) -> None:\n        for hook in ColoParamOpHookManager.hooks:\n            hook.post_backward(params)\n\n    @staticmethod\n    def pre_op(params: List[torch.Tensor], *args: Any) -> list:\n        ColoParamOpHookManager._trigger_pre_forward(params)\n        # auto grad function can only recognize torch.Tensor, thus we have to flatten the input\n        # if one of the input requires grad, all the output will be treated as requires grad\n        # and will have grad fn even the corresponding input does not require grad\n        # we have to extract tensors requiring grad into flat list and then merge them back\n        grad_args, other_args, grad_flags, spec = _flatten_grad_args(args)\n        new_grad_args = PreFwdPostBwd.apply(params, *grad_args)\n        return _merge_args(new_grad_args, other_args, grad_flags, spec)\n\n    @staticmethod\n    def post_op(params: List[torch.Tensor], arg: Any) -> Any:\n        ColoParamOpHookManager._trigger_post_forward(params)\n        # incase the output is a tuple, we have to flatten it\n        grad_args, other_args, grad_flags, spec = _flatten_grad_args(arg)\n        new_grad_args = PostFwdPreBwd.apply(params, *grad_args)\n        return _merge_args(new_grad_args, other_args, grad_flags, spec)\n\n    @staticmethod\n    def has_hook() -> bool:\n        return len(ColoParamOpHookManager.hooks) > 0\n\n\nclass PreFwdPostBwd(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, params, *args):\n        ctx.params = params\n        return args\n\n    @staticmethod\n    def backward(ctx, *grads):\n        ColoParamOpHookManager._trigger_post_backward(ctx.params)\n        return (None,) + grads\n\n\nclass PostFwdPreBwd(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, params, *args):\n        ctx.params = params\n        return args\n\n    @staticmethod\n    def backward(ctx, *grads):\n        ColoParamOpHookManager._trigger_pre_backward(ctx.params)\n        return (None,) + grads\n\n\ndef _is_grad_tensor(obj) -> bool:\n    if torch.is_tensor(obj):\n        if obj.grad_fn is not None or obj.requires_grad:\n            return True\n    return False\n\n\ndef _flatten_grad_args(args) -> Tuple[list, list, List[bool], TreeSpec]:\n    flat_args, spec = tree_flatten(args)\n    grad_args = []\n    other_args = []\n    grad_flags = []\n    for arg in flat_args:\n        flag = _is_grad_tensor(arg)\n        grad_flags.append(flag)\n        if flag:\n            grad_args.append(arg)\n        else:\n            other_args.append(arg)\n    return grad_args, other_args, grad_flags, spec\n\n\ndef _merge_args(grad_args, other_args, grad_flags, spec):\n    grad_iter = iter(grad_args)\n    other_iter = iter(other_args)\n    flat_args = [next(grad_iter) if flag else next(other_iter) for flag in grad_flags]\n    return tree_unflatten(flat_args, spec)\n", "path": "colossalai/tensor/param_op_hook.py"}, {"content": "from typing import Optional\n\nimport torch\n\nfrom colossalai.tensor.colo_tensor import ColoTensor\nfrom colossalai.tensor.param_op_hook import ColoParamOpHookManager\n\nfrom .colo_tensor import _convert_output\n\nWHITE_LIST_FUNCS = {torch.Tensor.__getitem__}\nNO_HOOK_FUNCS = {torch.Tensor.is_floating_point}\n\n\ndef is_no_hook_op(func) -> bool:\n    return (func.__name__.startswith(\"__\") and func not in WHITE_LIST_FUNCS) or func in NO_HOOK_FUNCS\n\n\ndef filter_colo_parameters(*args, **kwargs):\n    param_list = []\n\n    def get_colo_parameters(element) -> None:\n        if isinstance(element, list) or isinstance(element, tuple):\n            for e in element:\n                get_colo_parameters(e)\n        elif isinstance(element, dict):\n            raise RuntimeError(\"Found Dict: ColoParameter can't deal with complicated arguments.\")\n        elif isinstance(element, ColoParameter):\n            param_list.append(element)\n        return\n\n    for a in args:\n        get_colo_parameters(a)\n    for v in kwargs.values():\n        get_colo_parameters(v)\n\n    return param_list\n\n\ndef replace_args(args, kwargs, new_args):\n    args = new_args[: len(args)]\n    for k, v in zip(kwargs.keys(), new_args[len(args) :]):\n        kwargs[k] = v\n    return tuple(args), kwargs\n\n\nclass ColoParameter(ColoTensor, torch.nn.Parameter):\n    r\"\"\"A kind of ColoTensor to be considered as a module parameter.\"\"\"\n\n    def __new__(cls, data: Optional[torch.Tensor] = None, requires_grad: bool = True) -> \"ColoParameter\":\n        if data is None:\n            data = torch.empty(0)\n        return torch.Tensor._make_subclass(cls, data, requires_grad)\n\n    @classmethod\n    def __torch_function__(cls, func, types, args=..., kwargs=None):\n        if kwargs is None:\n            kwargs = {}\n        if ColoParamOpHookManager.has_hook() and not is_no_hook_op(func):\n            params = filter_colo_parameters(*args, **kwargs)\n            if len(params) > 0:\n                with torch._C.DisableTorchFunction():\n                    new_args = ColoParamOpHookManager.pre_op(params, *args, *kwargs.values())\n                args, kwargs = replace_args(args, kwargs, new_args)\n                ret = super().__torch_function__(func, types, args, kwargs)\n                with torch._C.DisableTorchFunction():\n                    ret = ColoParamOpHookManager.post_op(params, ret)\n                return _convert_output(ret, func)\n        return super().__torch_function__(func, types, args, kwargs)\n\n    def __deepcopy__(self, memo):\n        if id(self) in memo:\n            return memo[id(self)]\n        else:\n            with torch._C.DisableTorchFunction():\n                data = self.data.clone()\n            tensor = ColoParameter(data, self.requires_grad)\n            memo[id(self)] = tensor\n            return tensor\n\n    def __reduce_ex__(self, proto):\n        # Adapted from torch._utils._rebuild_parameter\n        # def _rebuild_colo_parameter(data, requires_grad, backward_hooks):\n        #     colo_param = ColoParameter(data, requires_grad)\n        #     colo_param._backward_hooks = backward_hooks\n        #     return colo_param\n\n        # return (\n        #     _rebuild_colo_parameter,\n        #     (self.data, self.requires_grad, OrderedDict())\n        # )\n\n        # TODO(jzy) we don't support object reflection now.\n        # distspec cannot be pickled or rebuilt because it's tightly connected to runtime attribute `process_group`.\n        raise NotImplementedError\n", "path": "colossalai/tensor/colo_parameter.py"}]}