[zero] cuda margin space for OS (#418)

colossalai/utils/commons/memory.py

@@ -0,0 +1,10 @@
+import imp
+import torch
+from colossalai.utils import get_current_device
+
+
+def col_cuda_memory_capacity():
+    """
+    Get cuda memory capacity of the current cuda.
+    """
+    return torch.cuda.get_device_properties(get_current_device()).total_memory