Improving packed sequences SFT for large datasets

src/megatron/bridge/data/datasets/packed_sequence.py

@@ -15,6 +15,10 @@
 import logging
 from dataclasses import dataclass
 from pathlib import Path
+from multiprocessing import Pool
+import multiprocessing as mp
+from tqdm import tqdm
+from typing import Literal, List
 
 import numpy as np
 from megatron.core.msc_utils import MultiStorageClientFeature
@@ -26,6 +30,20 @@
 
 logger = logging.getLogger(__name__)
 
+_shared_dataset = None
+def _tokenize_get_item(i):
+    return _shared_dataset[i]
+
+def _tokenize_init_worker(dataset):
+    global _shared_dataset
+    _shared_dataset = dataset
+
+def _retrieve_tokenized(dataset, num_workers):
+    if num_workers == 1:
+        return np.array([dataset[i] for i in tqdm(range(len(dataset)))])
+    num_workers = num_workers if num_workers > 0 else mp.cpu_count()
+    with Pool(num_workers, initializer=_tokenize_init_worker, initargs=(dataset,)) as pool:
+        return np.array(list(tqdm(pool.imap(_tokenize_get_item, range(len(dataset))), total=len(dataset))))
 
 def tokenize_dataset(
     path: Path,
@@ -34,6 +52,7 @@ def tokenize_dataset(
     seed: int,
     dataset_kwargs: dict | None = None,
     pad_seq_to_mult: int | None = 1,
+    num_tokenizer_workers: int = -1,
 ):
     """
     Tokenizes a dataset from the provided path using the specified tokenizer
@@ -88,7 +107,7 @@ def tokenize_dataset(
     pad_id = dataset.tokenizer.eod
     pad_seq_length_to_mult = dataset.pad_seq_length_to_mult
     max_seq_length = dataset.max_seq_length
-    dataset = np.array([dataset[i] for i in range(len(dataset))])
+    dataset = _retrieve_tokenized(dataset, num_tokenizer_workers)
 
     if pad_seq_to_mult > 1:
 
@@ -120,7 +139,6 @@ def pre_pad_dataset(data, max_seq_length, max_length_to_pad, pad_id):
 
     return dataset
 
-
 def prepare_packed_sequence_data(
     input_path: Path,
     output_path: Path,
@@ -132,6 +150,7 @@ def prepare_packed_sequence_data(
     packing_algorithm: str = "first_fit_shuffle",
     dataset_kwargs: dict | None = None,
     pad_seq_to_mult: int | None = 1,
+    num_tokenizer_workers: int = -1,
 ):
     """
     Prepares a packed sequence dataset from a given input file and saves it to an output file.
@@ -162,6 +181,7 @@ def prepare_packed_sequence_data(
         seed,
         dataset_kwargs,
         pad_seq_to_mult=pad_seq_to_mult,
+        num_tokenizer_workers=num_tokenizer_workers
     )
     sequences, histogram = create_hist(dataset, max_seq_length)
 
@@ -220,6 +240,12 @@ class PackedSequenceSpecs:
     This field is set by llm.finetune api.
     """
 
+    num_tokenizer_workers: int = -1
+    """
+    The number of worker processes to use for tokenization when preparing the packed sequence dataset.
+    If -1, the number of workers will be set to the number of CPU cores available
+    """
+
     packed_train_data_path: str = None
     """
     If specified, use this file for the packed training dataset instead of the default path.

src/megatron/bridge/data/datasets/packing_utils.py

@@ -25,7 +25,7 @@
 logger = logging.getLogger(__name__)
 
 
-def find_first_bin_that_fits(bins: List[List[int]], s: int, bin_size: int) -> int:
+def find_first_bin_that_fits(bin_sums: List[int], s: int, bin_size: int) -> int:
     """
     Finds the first bin in a list of bins that has enough space to fit a sequence of size 's'.
 
@@ -37,9 +37,9 @@ def find_first_bin_that_fits(bins: List[List[int]], s: int, bin_size: int) -> in
     Returns:
       The index of the first bin that can fit the sequence 's', or -1 if no such bin exists.
     """
-    for i, abin in enumerate(bins):
-        if sum(abin) + s <= bin_size:
-            return i
+    for i,cur_sum in enumerate(bin_sums):
+        if cur_sum + s <= bin_size:
+             return i
     return -1
 
 
@@ -56,12 +56,15 @@ def first_fit(seqlens: List[int], pack_size: int) -> List[List[int]]:
         of the sequences assigned to that bin.
     """
     res = []
-    for s in seqlens:
-        first_bin = find_first_bin_that_fits(res, s, pack_size)
+    res_sums = []
+    for s in tqdm(seqlens):
+        first_bin = find_first_bin_that_fits(res_sums, s, pack_size)
         if first_bin == -1:  # open a new bin
             res.append([s])
+            res_sums.append(s)
         else:
             res[first_bin].append(s)
+            res_sums[first_bin] += s
     return res