fix issue with data

EleutherAI · Jun 21, 2023 · fba32d5 · fba32d5
1 parent c00ce70
commit fba32d5
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 6 deletions.
diff --git a/configs/1-3B.yml b/configs/1-3B.yml
@@ -3,7 +3,7 @@
  # parallelism settings ( you will want to change these based on your cluster setup, ideally scheduling pipeline stages
  # across the node boundaries )
  "pipe_parallel_size": 1,
- "model_parallel_size": 1,
+ "model_parallel_size": 2,
 
  # model settings
  "num_layers": 24,
@@ -74,7 +74,7 @@
  },
 
  # misc. training settings
- "train_iters": 320000,
+ "train_iters": 320,
  "lr_decay_iters": 320000,
  "distributed_backend": "nccl",
  "lr_decay_style": "cosine",

diff --git a/data/enwik8/enwik8.zip b/data/enwik8/enwik8.zip
diff --git a/megatron/data/data_utils.py b/megatron/data/data_utils.py
@@ -67,6 +67,7 @@ def build_the_dataset(
  indexed_dataset = make_indexed_dataset(data_prefix, data_impl, skip_warmup)
  if label_prefix is None:
  label_dataset = None
+
  else:
  label_dataset = make_indexed_dataset(label_prefix, data_impl, skip_warmup)
 
@@ -207,7 +208,7 @@ def build_weighted_datasets(
  for i, (train_path, label_path, valid_path, test_path) in enumerate(
  zip_longest(
  neox_args.train_data_paths,
- neox_args.label_data_paths if neox_args.label_data_paths else [],
+ neox_args.label_data_paths if neox_args.label_data_paths else None,
  neox_args.valid_data_paths,
  neox_args.test_data_paths,
  )

diff --git a/megatron/training.py b/megatron/training.py
@@ -273,7 +273,7 @@ def pretrain(neox_args):
 
 def _get_batch(neox_args, tokenizer, keys, data, datatype):
  """Support function for get_batch / get_batch pipe (to avoid code repetition)"""
- _keys = [k for k in keys if k in data] if data else keys
+ _keys = [k for k in keys if k in data.keys()] if data else keys
  data_b = mpu.broadcast_data(_keys, data, datatype)
 
  # Unpack.
@@ -304,7 +304,7 @@ def get_batch(neox_args, data_iterator):
  """Generate a batch"""
 
  # Items and their type.
- keys = ["text", "label"]
+ keys = ["text", "label"] if neox_args.label_data_paths else ["text"]
  datatype = torch.int64
 
  # Broadcast data.
@@ -324,7 +324,7 @@ def get_batch(neox_args, data_iterator):
 def get_batch_pipe(data, neox_args, curr_scheduler=None):
  """A modification of get_batch() to work with the latest batch instead of an iterator."""
  # Items and their type.
- keys = ["text", "label"]
+ keys = ["text", "label"] if neox_args.label_data_paths else ["text"]
  datatype = torch.int64
 
  tokens, labels, loss_mask, attention_mask, position_ids = _get_batch(