aml datastore support using a data_path_prefix (deepspeedai#16)

Author: jeffra

bing_bert/bert_large_lamb_seq128.json

@@ -23,8 +23,8 @@
             "pretrain_type": "wiki_bc"
         },
         "datasets": {
-            "wiki_pretrain_dataset": "/data/bert/bnorick_format/128/wiki_pretrain",
-            "bc_pretrain_dataset": "/data/bert/bnorick_format/128/bookcorpus_pretrain"
+            "wiki_pretrain_dataset": "bnorick_format/128/wiki_pretrain",
+            "bc_pretrain_dataset": "bnorick_format/128/bookcorpus_pretrain"
         },
         "tp1pp_evalsets":
         {
@@ -45,6 +45,6 @@
         "total_training_steps": 187000
     },
     "validation": {
-        "path": "/data/bert/validation_set/"
+        "path": "validation_set/"
     }
 }

bing_bert/deepspeed_train.py

@@ -88,7 +88,7 @@ def pretrain_validation(args, index, model):
     logger = args.logger
 
     model.eval()
-    dataset = PreTrainingDataset(args.tokenizer, config['validation']['path'], args.logger,
+    dataset = PreTrainingDataset(args.tokenizer, os.path.join(args.data_path_prefix, config['validation']['path']), args.logger,
                                  args.max_seq_length, index, PretrainDataType.VALIDATION, args.max_predictions_per_seq)
     data_batches = get_dataloader(args, dataset, eval_set=True)
     eval_loss = 0
@@ -130,7 +130,7 @@ def get_train_dataset(args, index, finetune=False, shuffle=True):
             # Load Wiki Dataset
             wiki_pretrain_dataset = PreTrainingDataset(
                 args.tokenizer,
-                dataset_paths['wiki_pretrain_dataset'],
+                os.path.join(args.data_path_prefix, dataset_paths['wiki_pretrain_dataset']),
                 args.logger,
                 args.max_seq_length,
                 index,
@@ -145,7 +145,7 @@ def get_train_dataset(args, index, finetune=False, shuffle=True):
 
             bc_pretrain_dataset = PreTrainingDataset(
                 args.tokenizer,
-                dataset_paths['bc_pretrain_dataset'],
+                os.path.join(args.data_path_prefix, dataset_paths['bc_pretrain_dataset']),
                 args.logger,
                 args.max_seq_length,
                 index,

bing_bert/utils.py

@@ -80,6 +80,11 @@ def get_argument_parser():
                         default=100,
                         help='Interval to print training details.')
 
+    parser.add_argument('--data_path_prefix',
+                        type=str,
+                        default="",
+                        help="Path to prefix data loading, helpful for AML and other environments")
+
     return parser
 
 def is_time_to_exit(args, epoch_steps=0, global_steps=0):