Mlm adaptation

megatron/arguments.py

@@ -925,6 +925,9 @@ def __call__(self, parser, args, values, option_string=None):
                        'specific positions. This option tries to un-bias the loss by reweighting loss on specific '
                        'positions based on how frequently we train on that position.'
                        'This is mostly used for prefix_lm training')
+    group.add_argument("--noise_density", type=float, default=None, help="Span corruption noise density")
+    group.add_argument("--mean_noise_span_length", type=int, default=None, help="Span corruption mean noise span length")
+
 
     return parser
 

megatron/data/gpt_dataset.py

@@ -35,7 +35,7 @@ def build_train_valid_test_datasets(data_prefix, data_impl, splits_string,
 
     # Single dataset.
     if len(data_prefix) == 1:
-        all_train_datasets, all_valid_datasets, all_test_datasets =  _build_train_valid_test_datasets(data_prefix[0],
+        all_train_datasets, all_valid_datasets, all_test_datasets = _build_train_valid_test_datasets(data_prefix[0],
                                                 data_impl, splits_string,
                                                 train_valid_test_num_samples,
                                                 seq_length, seed, skip_warmup)