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

[gemini] fix argument naming during chunk configuration searching

Browse Source
This commit is contained in:
Baizhou Zhang
2023-06-25 13:34:15 +08:00
parent b463651f3e
commit 0bb0b481b4
17 changed files with 62 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_auto_parallel/test_offload/test_perf.py
View File

@@ -60,7 +60,7 @@ def exam_fwd_bwd(model_name: str, memory_budget: float, solver_name: str):
placement_policy='cpu',
pin_memory=True,
hidden_dim=8192,
search_range_mb=128)
search_range_m=128)
gemini_model = zero_model_wrapper(gemini_model, 3, gemini_config)
optim_config = dict(reduce_bucket_size=12 * 1024 * 1024, overlap_communication=True, verbose=True)
gemini_optim = zero_optim_wrapper(gemini_model, hybrid_optimizer, optim_config=optim_config)

2
tests/test_auto_parallel/test_tensor_shard/test_compatibility_with_gemini.py
View File

@@ -75,7 +75,7 @@ def check_auto_parallel_with_gemini(rank, world_size, port):
device=get_current_device(),
placement_policy='cpu',
pin_memory=True,
search_range_mb=128)
search_range_m=128)
post_process_colo_init_ctx(gm, device=get_current_device(), default_pg=dp_process_group)
gm = zero_model_wrapper(gm, zero_stage=3, gemini_config=gemini_config)