github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-04 02:57:20 +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-21 02:09:14 +00:00
parent 137a7c341b
commit b3c0e6d871
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -42,7 +42,9 @@ class GeminiManager:
self._mem_stats_collector = ( self._mem_stats_collector = (
ChunkMemStatsCollector(chunk_manager, self._memstats) if policy_cls.need_mem_stats else None ChunkMemStatsCollector(chunk_manager, self._memstats) if policy_cls.need_mem_stats else None
) )
self._placement_policy = policy_cls(chunk_manager=chunk_manager, mem_stats_collector=self._mem_stats_collector, **placement_kwargs) self._placement_policy = policy_cls(
chunk_manager=chunk_manager, mem_stats_collector=self._mem_stats_collector, **placement_kwargs
)
self._compute_list: List[Tuple[Chunk, ...]] = [] self._compute_list: List[Tuple[Chunk, ...]] = []
self._compute_idx: int = -1 self._compute_idx: int = -1
self._async_works: Dict[Chunk, dist.Work] = {} self._async_works: Dict[Chunk, dist.Work] = {}