[gemini] update ddp strict mode (#2518)

* [zero] add strict ddp mode for chunk init

* [gemini] update gpt example

colossalai/tensor/colo_tensor.py

@@ -1,3 +1,4 @@
+import math
 from copy import copy
 from functools import lru_cache
 from typing import Callable, Optional, Set
@@ -303,6 +304,11 @@ class ColoTensor(torch.Tensor):
         else:
             return size_list[args[0]]
 
+    def numel_global(self):
+        """Returns the number of elements in the tensor when it's replicated.
+        """
+        return math.prod(self.size_global())
+
     # Some API for dist spec check
 
     def is_replicate(self):