Merge pull request #290 from jonathanking/add_missing_warmupdecaylr_s…

…etting Add missing warmup_num_steps parameters to DeepSpeed config generation script.
aqlaboratory · Mar 14, 2023 · 7c32b79 · 7c32b79
2 parents 886494f + 9b2a4f8
commit 7c32b79
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/scripts/build_deepspeed_config.py b/scripts/build_deepspeed_config.py
@@ -248,9 +248,10 @@
         params["warmup_max_lr"] = args.warmup_max_lr
         params["warmup_num_steps"] = args.warmup_num_steps
     elif(args.scheduler == "WarmupDecayLR"):
-        params["total_num_steps"] = args.warmup_decay_total_num_steps
         params["warmup_min_lr"] = args.warmup_decay_min_lr
         params["warmup_max_lr"] = args.warmup_decay_max_lr
+        params["warmup_num_steps"] = args.warmup_decay_num_steps
+        params["total_num_steps"] = args.warmup_decay_total_num_steps
     else:
         raise ValueError("Invalid scheduler")