[feat] refactored extension module (#5298)

* [feat] refactored extension module * polish * polish * polish * polish * polish * polish * polish * polish * polish * polish
2025-09-04 10:34:41 +00:00 · 2024-01-25 17:01:48 +08:00
parent d7f8db8e21
commit 7cfed5f076
157 changed files with 1353 additions and 8966 deletions
--- a/colossalai/zero/gemini/chunk/chunk.py
+++ b/colossalai/zero/gemini/chunk/chunk.py
@@ -190,8 +190,10 @@ class Chunk:
    def device_type(self) -> str:
        if self.chunk_temp is not None:
            return self.chunk_temp.device.type
-        else:
+        elif self.is_gathered or self.cuda_shard is not None:
            return get_accelerator().name
+        else:
+            return "cpu"

    @property
    def payload(self) -> torch.Tensor: