github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-08 12:30:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor the code structure to solve the circular import

Browse Source
This commit is contained in:
genghaozhe
2024-05-20 07:25:24 +00:00
parent a280517dd9
commit bfcb2d1ff8
3 changed files with 32 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/zero/gemini/gemini_hook.py
View File

@@ -51,7 +51,12 @@ class GeminiZeROHook(ColoParamOpHook):
self._chunk_manager.access_chunk(chunk)
# get possible chunks to prefetch
chunks_fetch_async = self._gemini_manager.placement_policy.get_prefetch_chunks()
chunks_fetch_async = self._gemini_manager.placement_policy.get_prefetch_chunks(
is_warmup=self._gemini_manager.is_warmup(),
compute_list=self._gemini_manager.compute_list,
compute_idx=self._gemini_manager.compute_idx,
async_works=self._gemini_manager.async_works,
)
# prefetch
for chunk in chunks_fetch_async: