minor fixes

megatron/learning_rates.py

@@ -128,8 +128,8 @@ class AnnealingLR(object):
         new_lr = self.get_lr()
         new_wd = self.get_wd()
         for group in self.optimizer.param_groups:
-            group['lr'] = new_lr * group['lr_mult']
-            group['weight_decay'] = new_wd * group['wd_mult']
+            group['lr'] = new_lr * group.get('lr_mult', 1.0)
+            group['weight_decay'] = new_wd * group.get('wd_mult', 1.0)
 
 
     def state_dict(self):

megatron/model/transformer.py

@@ -608,7 +608,7 @@ class ParallelTransformer(MegatronModule):
         self.num_layers = mpu.get_num_layers(
             args, args.model_type == ModelType.encoder_and_decoder)
 
-        self.dpr = [x.item() for x in torch.linspace(0, self.drop_path_rate, self.num_layers)]
+        self.dpr = [x.item() for x in torch.linspace(0, self.drop_path_rate, args.num_layers)]
 
         # Transformer layers.
         def build_layer(layer_number):