distributed merge of per-rank Megatron data files

megatron/data/indexed_dataset.py

@@ -417,21 +417,21 @@ def __init__(self, path, skip_warmup=False):
                 offset = stream.tell()
 
             if not skip_warmup:
-                print_rank_0("    warming up index mmap file...")
+#                print_rank_0("    warming up index mmap file...")
                 _warmup_mmap_file(path)
 
             self._bin_buffer_mmap = np.memmap(path, mode='r', order='C')
             self._bin_buffer = memoryview(self._bin_buffer_mmap)
-            print_rank_0("    reading sizes...")
+#            print_rank_0("    reading sizes...")
             self._sizes = np.frombuffer(
                 self._bin_buffer,
                 dtype=np.int32,
                 count=self._len,
                 offset=offset)
-            print_rank_0("    reading pointers...")
+#            print_rank_0("    reading pointers...")
             self._pointers = np.frombuffer(self._bin_buffer, dtype=np.int64, count=self._len,
                                            offset=offset + self._sizes.nbytes)
-            print_rank_0("    reading document index...")
+#            print_rank_0("    reading document index...")
             self._doc_idx = np.frombuffer(self._bin_buffer, dtype=np.int64, count=self._doc_count,
                                           offset=offset + self._sizes.nbytes + self._pointers.nbytes)
 
@@ -576,8 +576,8 @@ def merge_file_(self, another_file):
         index = MMapIndexedDataset.Index(index_file_path(another_file))
         assert index.dtype == self._dtype
 
-        total_len = len(index.sizes)+len(self._sizes)
-        print(f"    concat {another_file} size={len(index.sizes)} for a total size of {total_len}")
+#        total_len = len(index.sizes)+len(self._sizes)
+#        print(f"    concat {another_file} size={len(index.sizes)} for a total size of {total_len}")
 
         offset = len(self._sizes)
         self._sizes.extend(index.sizes)
@@ -592,3 +592,213 @@ def finalize(self, index_file):
 
         with MMapIndexedDataset.Index.writer(index_file, self._dtype) as index:
             index.write(self._sizes, self._doc_idx)
+
+
+def mpi_get_sum(val, mpi, comm):
+    """Compute sum of val, and return total on all ranks."""
+    insize = np.array([val], dtype=np.int64)
+    outsize = np.zeros_like(insize)
+    comm.Allreduce(insize, outsize, op=mpi.SUM)
+    return outsize[0]
+
+
+def mpi_get_offset(val, mpi, comm):
+    """Compute preifx sum (exclusive scan) of val, and return offset of each rank."""
+    insize = np.array([val], dtype=np.int64)
+    outsize = np.zeros_like(insize)
+    comm.Scan(insize, outsize, op=mpi.SUM)
+    offset = outsize[0] - insize[0]
+    return offset
+
+
+def mpi_get_min(val, mpi, comm):
+    """Return minimum val to all ranks."""
+    insize = np.array([val], dtype=np.int64)
+    outsize = np.zeros_like(insize)
+    comm.Allreduce(insize, outsize, op=mpi.MIN)
+    return outsize[0]
+
+
+def mpi_create_file(filename, mpi, comm):
+    """Create, truncate, and open a file shared by all ranks."""
+    # Don't truncate file until all ranks reach this point
+    comm.barrier()
+
+    # Wait for rank 0 to open (and truncate) file,
+    # then have all ranks open file for writing.
+    rank = comm.Get_rank()
+    if rank == 0:
+        f = open(filename, 'wb')
+    comm.barrier()
+    if rank != 0:
+        f = open(filename, 'r+b')
+
+    # TODO: verify that all ranks successfully opened the file
+    comm.barrier()
+
+    return f
+
+
+# To create the binary files given a set of per-rank binary
+# files, one simply concatenates the data from the per-rank
+# binary files in rank order.  We stat each rank file to determine
+# its size, execute a scan to compute the byte offset where
+# the calling rank should write its data, seek to proper
+# spot, and copy the full file.
+def merge_files_mpi_bin(outfile, infile, mpi, comm):
+    """Concatenate per-rank binary files into a new file given by outfile"""
+    import stat
+    import shutil
+
+    comm.barrier()
+
+    # Create shared output file.
+    f = mpi_create_file(outfile, mpi, comm)
+
+    # get file size of binary file for this rank
+    filesize = os.stat(infile)[stat.ST_SIZE]
+
+    # compute offset this rank should start copying
+    # its data into the merged file
+    offset = mpi_get_offset(filesize, mpi, comm)
+
+    # seek to appropriate offset and copy data
+    f.seek(offset)
+    with open(infile, "rb") as fsrc:
+        shutil.copyfileobj(fsrc, f)
+
+    f.close()
+
+    # TODO: check that all ranks wrote successfully
+    comm.barrier()
+
+
+def merge_files_mpi_idx(outfile, infile, mpi, comm, dtype):
+    rank = comm.Get_rank()
+    numranks = comm.Get_size()
+
+    comm.barrier()
+
+    # Create shared output file
+    f = mpi_create_file(outfile, mpi, comm)
+
+    # Read the index file for the calling rank
+    index = MMapIndexedDataset.Index(infile)
+    sizes = index.sizes
+    if rank == 0:
+        docs = index.doc_idx
+    if rank != 0:
+        docs = index.doc_idx[1:]
+
+    # Compute total number of size and document index
+    # values across all ranks.  Also compute the offset
+    # of the calling rank for each value considering
+    # the values of sizes/docs for all ranks before the
+    # calling rank.
+    numsizes = len(sizes)
+    numdocs = len(docs)
+    size_count = mpi_get_sum(numsizes, mpi, comm)
+    docs_count = mpi_get_sum(numdocs, mpi, comm)
+    size_offset = mpi_get_offset(numsizes, mpi, comm)
+    docs_offset = mpi_get_offset(numdocs, mpi, comm)
+
+    # Have rank 0 write the file header
+    pos = 0
+    if rank == 0:
+        f.write(MMapIndexedDataset.Index._HDR_MAGIC)
+        f.write(struct.pack('<Q', 1))
+        f.write(struct.pack('<B', code(dtype)))
+        f.write(struct.pack('<Q', size_count))
+        f.write(struct.pack('<Q', docs_count))
+        pos = f.tell()
+
+    # Broadcast value of pos from rank 0,
+    # and advance file position past file header on all ranks.
+    pos = mpi_get_sum(pos, mpi, comm)
+
+    # The list of size values from each rank are
+    # concatenated and stored as int32.
+    f.seek(pos + size_offset * np.int32().itemsize)
+    sizes32 = np.array(sizes, dtype=np.int32)
+    f.write(sizes32.tobytes(order='C'))
+    del sizes32
+    pos += size_count * np.int32().itemsize
+
+    # The pointer values store the byte offset to each sentence.
+    # A sentence has a variable number of tokens, given by
+    # its corresponding entry in the size array.  Each token
+    # of a sentence is stored in units of type dtype, which consumes
+    # dtype().itemsize bytes (often a standard type that is just
+    # large enough to represent all elements of the vocabulary).
+
+    # First compute byte offsets for each sentence of our
+    # local set of sentences.
+    pointers = np.array(sizes, dtype=np.int64)
+    pointer_last = 0
+    if len(sizes) > 0:
+        np.cumsum(pointers, axis=0, out=pointers)
+        pointers *= dtype().itemsize
+        pointer_last = pointers[-1]
+
+    # Then account for bytes for all sentences on ranks
+    # before the calling rank.
+    pointer_offset = mpi_get_offset(pointer_last, mpi, comm)
+    pointers += pointer_offset
+
+    # Finally, zero-base the offset values by subtracting
+    # the number of bytes of the first sentence.  To do that
+    # we first need to find the rank having the first sentence,
+    # then bcast that size to all ranks.
+    if size_count > 0:
+        # There is at least one sentence across all ranks,
+        # figure out which rank has the first sentence which
+        # is not necessarily rank 0.
+        minrank = numranks
+        if len(sizes) > 0:
+            minrank = rank
+        minrank = mpi_get_min(minrank, mpi, comm)
+
+        # Broadcast size of the first sentence from minrank.
+        # We "bcast" using all allreduce.
+        pointers_shift = 0
+        if minrank == rank:
+            pointers_shift = pointers[0]
+        pointers_shift = mpi_get_sum(pointers_shift, mpi, comm)
+
+        # Zero-base pointers by subtracting size of first
+        # sentence from all values.
+        pointers -= pointers_shift
+
+    # Seek to proper offset for this rank and write
+    # pointer values into file, stored as int64.
+    f.seek(pos + size_offset * np.int64().itemsize)
+    f.write(pointers.tobytes(order='C'))
+    del pointers
+    pos += size_count * np.int64().itemsize
+
+    # The document index points to the position in the sizes
+    # array for the starting sentence of each document.
+    # A variable number of sentences can be in each document.
+    # Adjust document index for number of sentences that
+    # come before the calling rank.
+    doc_idx = np.array(docs, dtype=np.int64)
+    doc_idx += size_offset
+
+    # Seek to proper offset for this rank and write
+    # document index into file, stored as int64.
+    f.seek(pos + docs_offset * np.int64().itemsize)
+    f.write(doc_idx.tobytes(order='C'))
+    del doc_idx
+
+    f.close()
+
+    # TODO: check that all ranks wrote successfully
+    comm.barrier()
+
+
+def merge_files_mpi(filemain, filerank, mpi, comm, dtype=np.int64):
+    # Concatenate the data files
+    merge_files_mpi_bin(data_file_path(filemain), data_file_path(filerank), mpi, comm)
+
+    # Combine index files into a single index file
+    merge_files_mpi_idx(index_file_path(filemain), index_file_path(filerank), mpi, comm, dtype)