From 4412987f3aa7dcee25b1e8e13764cd6aea2b26dc Mon Sep 17 00:00:00 2001 From: lc_pro Date: Fri, 16 Jun 2023 19:50:13 +0800 Subject: [PATCH] fix bucket size in example of gpt gemini --- examples/language/gpt/gemini/train_gpt_demo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/language/gpt/gemini/train_gpt_demo.py b/examples/language/gpt/gemini/train_gpt_demo.py index 4b78624f0110..a7b552c9e23d 100644 --- a/examples/language/gpt/gemini/train_gpt_demo.py +++ b/examples/language/gpt/gemini/train_gpt_demo.py @@ -250,7 +250,7 @@ def main(): plugin = None if args.distplan.startswith("CAI_ZeRO"): plugin = LowLevelZeroPlugin(stage=zero_stage, - reduce_bucket_size_in_m=12 * 1024 * 1024, + reduce_bucket_size_in_m=12, overlap_communication=True, verbose=True) elif args.distplan == "CAI_Gemini":