add dataset preparation to main train scripts

configs/base_model.json

@@ -1,5 +1,8 @@
 {
- "data_path": "./data/enwik8.tar.gz",
+ "dataset": {
+ "name": "enwik8",
+ "path": "./data/enwik8.tar.gz"
+ },
  "num_epochs": 10,
  "vocab_size": 256,
  "batch_size": 4,

scripts/train_enwik8.sh

@@ -1,3 +1,2 @@
 mkdir logs
-python3 prepare_dataset.py --model enwik8
 NCCL_SHM_DISABLE=1 NCCL_DEBUG=info MASTER_ADDR=127.0.0.1 MASTER_PORT=2000 deepspeed train_enwik8.py --deepspeed --deepspeed_config configs/base_deepspeed.json

scripts/train_gpt3small.sh

@@ -1,3 +1,2 @@
 mkdir logs
-python3 prepare_dataset.py 
 NCCL_SHM_DISABLE=1 NCCL_DEBUG=info MASTER_ADDR=127.0.0.1 MASTER_PORT=2000 deepspeed train.py --deepspeed --deepspeed_config configs/base_deepspeed.json

train.py

@@ -35,6 +35,7 @@
 dset_params = params["dataset"]
 assert dset_params is not None
 
+deepspeed.init_distributed(dist_backend='nccl')
 torch.distributed.barrier() # barrier will force processes to stop until *all* processes have reached the barrier
 if is_main(train_args):
  prepare_data(dset_params["name"])
@@ -56,7 +57,11 @@
 val_loader = iter(val_loader)
 
 # optimizer
-optim = torch.optim.Adam(model.parameters(), lr=params["learning_rate"])
+if train_args.local_rank == -1: # non-deepspeed
+ optim = torch.optim.Adam(model.parameters(), lr=params["learning_rate"])
+else:
+ optim = None # deepspeed will prepare the optimizer for us
+
 
 # training
 ds_model_params = prepare_optimizer_parameters(model)
@@ -68,6 +73,7 @@
  model_parameters=ds_model_params,
  training_data=train_dataset)
 
+print("OPTIMIZER:", optim)
 pbar = trange(params.get("train_steps", 1), mininterval=10., desc='Training Model', dynamic_ncols=True)
 for _ in pbar:
  for i, data in enumerate(train_loader):

train_enwik8.py

@@ -9,7 +9,7 @@
 from tqdm.auto import trange
 
 from gpt_neox import (GPTNeoX, AutoregressiveWrapper, TextSamplerDataset, download_dataset,
- cycle, prepare_optimizer_parameters, decode_tokens, read_enwik8_data, is_main)
+ cycle, prepare_optimizer_parameters, decode_tokens, read_enwik8_data, is_main, prepare_data)
 
 
 def get_args():
@@ -44,9 +44,17 @@ def get_params(model):
 )
 
 model = AutoregressiveWrapper(model)
+dset_params = params["dataset"]
+deepspeed.init_distributed(dist_backend='nccl')
+torch.distributed.barrier() # barrier will force processes to stop until *all* processes have reached the barrier
+if is_main(train_args):
+ prepare_data(dset_params["name"])
+ torch.distributed.barrier() # barrier will force processes to stop until *all* processes have reached the barrier
+else:
+ torch.distributed.barrier()
 
 # prepare enwik8 data
-data_train, data_val = read_enwik8_data(params["data_path"])
+data_train, data_val = read_enwik8_data(dset_params["data_path"])
 train_dataset = TextSamplerDataset(data_train, params["seq_len"])
 val_dataset = TextSamplerDataset(data_val, params["seq_len"])
 val_loader = cycle(DataLoader(val_dataset, batch_size=params["batch_size"]))