Merge branch 'main' into sync/npu

colossalai/legacy/zero/gemini/__init__.py

+from .colo_init_context import ColoInitContext, post_process_colo_init_ctx
 from .ophooks import BaseOpHook, register_ophooks_recursively
 from .stateful_tensor import StatefulTensor
 from .stateful_tensor_mgr import StatefulTensorMgr
@@ -11,4 +12,6 @@ __all__ = [
     "AutoTensorPlacementPolicy",
     "register_ophooks_recursively",
     "BaseOpHook",
+    "ColoInitContext",
+    "post_process_colo_init_ctx",
 ]