Merge branch 'main_beta' into 'main'

Adam betas and eps

See merge request ADLR/megatron-lm!156

megatron/arguments.py

@@ -199,6 +199,15 @@ def _add_regularization_args(parser):
                        help='Weight decay coefficient for L2 regularization.')
     group.add_argument('--clip-grad', type=float, default=1.0,
                        help='Gradient clipping based on global L2 norm.')
+    group.add_argument('--adam-beta1', type=float, default=0.9,
+                       help='First coefficient for computing running averages of'
+                       'gradient and its square')
+    group.add_argument('--adam-beta2', type=float, default=0.999,
+                       help='Second coefficient for computing running averages of'
+                       'gradient and its square')
+    group.add_argument('--adam-eps', type=float, default=1e-08,
+                       help='Term added to the denominator to improve'
+                       'numerical stability')
 
     return parser
 

megatron/training.py

@@ -164,7 +164,8 @@ def get_optimizer(model):
                 param.model_parallel = False
 
     # Use Adam.
-    optimizer = Adam(param_groups, lr=args.lr, weight_decay=args.weight_decay)
+    optimizer = Adam(param_groups, lr=args.lr, weight_decay=args.weight_decay,
+        betas=(args.adam_beta1, args.adam_beta2), eps=args.adam_eps)
 
     # Wrap into fp16 optimizer.
     if args.fp16: