Merge pull request EleutherAI#771 from zhuzilin/optimize-data-preprocess

Optimize data preprocessing by using numpy

megatron/data/indexed_dataset.py

@@ -287,12 +287,13 @@ def __init__(self, out_file, dtype=np.int32):
  self.element_size = self.element_sizes[self.dtype]
  self.doc_idx = [0]
 
- def add_item(self, tensor):
- bytes = self.out_file.write(np.array(tensor.numpy(), dtype=self.dtype))
+ def add_item(self, np_array):
+ assert isinstance(np_array, np.ndarray) and np_array.dtype == self.dtype
+ bytes = self.out_file.write(np_array)
  self.data_offsets.append(self.data_offsets[-1] + bytes / self.element_size)
- for s in tensor.size():
+ for s in np_array.shape:
  self.sizes.append(s)
- self.dim_offsets.append(self.dim_offsets[-1] + len(tensor.size()))
+ self.dim_offsets.append(self.dim_offsets[-1] + len(np_array.shape))
 
  def end_document(self):
  self.doc_idx.append(len(self.sizes))
@@ -360,14 +361,11 @@ def __enter__(self):
 
  @staticmethod
  def _get_pointers(sizes):
- dtype_size = dtype().itemsize
- address = 0
- pointers = []
-
- for size in sizes:
- pointers.append(address)
- address += size * dtype_size
+ pointers = np.zeros(len(sizes), dtype=np.int64)
+ sizes = np.array(sizes, dtype=np.int64)
 
+ np.cumsum(sizes[:-1], out=pointers[1:])
+ pointers = pointers * dtype().itemsize
  return pointers
 
  def write(self, sizes, doc_idx):
@@ -568,8 +566,12 @@ def __init__(self, out_file, dtype=np.int64):
  self._sizes = []
  self._doc_idx = [0]
 
- def add_item(self, tensor):
- np_array = np.array(tensor.numpy(), dtype=self._dtype)
+ @property
+ def dtype(self):
+ return self._dtype
+
+ def add_item(self, np_array):
+ assert isinstance(np_array, np.ndarray) and np_array.dtype == self.dtype
  self._data_file.write(np_array.tobytes(order="C"))
  self._sizes.append(np_array.size)
 

tools/preprocess_data.py

@@ -23,6 +23,7 @@
 import sys
 
 import lm_dataformat as lmd
+import numpy as np
 
 sys.path.append(
  os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir))
@@ -217,7 +218,7 @@ def main():
  # add each tokenized document / sentence
  for key, sentences in doc.items():
  for sentence in sentences:
- builders[key].add_item(torch.IntTensor(sentence))
+ builders[key].add_item(np.array(sentence, dtype=builders[key].dtype))
  # separate with eos token
  builders[key].end_document()