github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[hotfix] fix chunk size can not be divided (#2867) 

* [hotfix] fix chunk size can not be divided

* [hotfix] use numpy for python3.8
pull/2876/head^2
HELSON 2023-02-22 15:04:46 +08:00 committed by GitHub
parent a4fc125c34
commit 6e4ac08172
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/gemini/chunk/manager.py
View File

@ -72,6 +72,9 @@ class ChunkManager:
if tensor.numel() > chunk_size:
chunk_size = tensor.numel()
dp_size = tensor.process_group.dp_world_size()
chunk_size = chunk_size + (-chunk_size % dp_size)
chunk = Chunk(
chunk_size=chunk_size,
process_group=tensor.process_group,

3
colossalai/gemini/chunk/search_utils.py
View File

@ -119,6 +119,7 @@ def search_chunk_configuration(
assert search_range_byte >= 0
params_dict = classify_params_by_dp_degree(param_order, strict_ddp_flag)
size_lcm = np.lcm.reduce(list(params_dict.keys()))
config_dict: Dict[int, Dict] = dict()
total_param_size = 0
@ -154,6 +155,8 @@ def search_chunk_configuration(
min_chunk_waste = temp_waste
best_chunk_size = chunk_size
# the chunk size needs to be divided by each groups sizes
best_chunk_size = best_chunk_size + (-best_chunk_size % size_lcm)
for dp_degree in params_dict:
if dp_degree in config_dict:
continue