diff --git a/scripts/build_deepspeed_config.py b/scripts/build_deepspeed_config.py
index 728101a3..eff67c96 100644
--- 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")