github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-07 20:10:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse Source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot]
2024-05-16 07:26:19 +00:00
parent 82b25524ff
commit 6bbe956316
4 changed files with 32 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/zero/gemini/gemini_mgr.py
View File

@@ -1,6 +1,6 @@
import functools
from time import time
from typing import Dict, List, Optional, Tuple, Iterable
from typing import Dict, Iterable, List, Optional, Tuple
import torch
import torch.distributed as dist
@@ -101,7 +101,7 @@ class GeminiManager:
start = time()
self._record_warmup_chunks_order(chunks, record_anyway=record_anyway)
cuda_demand, can_evict_chunks = self._get_layout_info(self._compute_idx, self._warmup, chunks)
# don't evict chunks that are asynchronously fetched
# don't evict chunks that are asynchronously fetched
can_evict_chunks = [chunk for chunk in can_evict_chunks if chunk not in self._async_works]
self._layout_time += time() - start