[Tensor] init ColoParameter (#914)

colossalai/tensor/colo_parameter.py

@@ -0,0 +1,28 @@
+from .colo_tensor import ColoTensor
+from .const import TensorType
+import torch
+
+
+class ColoParameter(ColoTensor):
+    r"""A kind of ColoTensor to be considered as a module parameter.
+
+    """
+
+    def __init__(self, *args, **kargs):
+        super().__init__(*args, **kargs)
+        self._type = TensorType.MODEL
+
+    def __new__(cls, *args, **kwargs):
+        t = super(ColoParameter, cls).__new__(cls)
+        t._type = TensorType.MODEL
+        return t
+
+    @staticmethod
+    def init_from_torch_tensor(tensor: torch.Tensor, save_payload=True) -> 'ColoParameter':
+        colo_p = ColoParameter(*tensor.size(),
+                               dtype=tensor.dtype,
+                               requires_grad=tensor.requires_grad,
+                               pin_memory=tensor.is_pinned(),
+                               device=tensor.device,
+                               torch_tensor=tensor if save_payload else torch.empty(0))
+        return colo_p