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

[Gemini] Use async stream to prefetch and h2d data moving (#5781)

Browse Source 
* use async stream to prefetch and h2d data moving

* Remove redundant code
This commit is contained in:
Haze188
2024-06-12 15:48:52 +08:00
committed by GitHub
parent 8554585a5f
commit d9dddf574f
4 changed files with 12 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/zero/gemini/chunk/utils.py
View File

@@ -21,6 +21,7 @@ def init_chunk_manager(
hidden_dim: Optional[int] = None,
reuse_fp16_chunk: bool = True,
verbose: bool = False,
max_prefetch: int = 0,
**kwargs,
) -> ChunkManager:
if hidden_dim:
@@ -51,9 +52,5 @@ def init_chunk_manager(
)
dist.barrier()
chunk_manager = ChunkManager(
config_dict,
init_device,
reuse_fp16_chunk=reuse_fp16_chunk,
)
chunk_manager = ChunkManager(config_dict, init_device, reuse_fp16_chunk=reuse_fp16_chunk, max_prefetch=max_prefetch)
return chunk_manager