Add concat_then_split packing in Grain pipeline

Author: aireenmei

src/MaxText/configs/base.yml

@@ -595,6 +595,7 @@ grain_train_files: ''
 grain_eval_files: ''
 grain_train_mixture_config_path: '' # Path to a JSON file specifying the mixture weights for Grain training data.
 grain_file_type: 'arrayrecord' # arrayrecord or parquet
+grain_packing_type: 'first_fit' # 'first_fit' or 'concat_then_split'
 grain_worker_count: 1 # Set to -1 to enable auto-tuning: automatically determines optimal worker count. See https://google-grain.readthedocs.io/en/latest/_autosummary/grain.experimental.pick_performance_config.html
 grain_per_worker_buffer_size: 1
 # num_threads and prefetch_buffer_size are per-worker per-dataset. Used in ReadOptions (https://google-grain.readthedocs.io/en/latest/tutorials/data_loader_tutorial.html#per-worker-readoptions)

src/MaxText/configs/types.py

@@ -803,6 +803,10 @@ class DatasetGeneral(BaseModel):
       True,
       description="Whether to pack multiple short examples into a single sequence.",
   )
+  grain_packing_type: Literal["first_fit", "concat_then_split"] = Field(
+      "first_fit",
+      description="Packing type when using Grain pipeline. 'first_fit' or 'concat_then_split'.",
+  )
   max_segments_per_seq: int = Field(
       32, description="Maximum number of segments that can be packed into a single sequence."
   )

src/MaxText/input_pipeline/_grain_data_processing.py

@@ -171,6 +171,8 @@ def pretrain_preprocessing_pipeline(
   if config.grain_file_type == "arrayrecord":
     dataset = dataset.map(_input_pipeline_utils.ParseFeatures(data_columns, tokenize))
     dataset = dataset.map(_input_pipeline_utils.NormalizeFeatures(data_columns, tokenize))
+  else:
+    dataset = dataset.map(_input_pipeline_utils.KeepFeatures(feature_names=data_columns))
 
   assert len(data_columns) == 1
   text_column = data_columns[0]
@@ -207,11 +209,26 @@ def pretrain_preprocessing_pipeline(
     # But when using Grain, we want to keep the batch_size consistent with that in the checkpoint.
     # We revert the batch_size expansion here, but load multiple batches per step in multihost_dataloading.py.
     batch_size = batch_size // config.expansion_factor_real_data
+
   if config.packing:
     length_struct = {col: config.max_target_length for col in data_columns}
-    dataset = grain.experimental.FirstFitPackIterDataset(
-        dataset, length_struct=length_struct, num_packing_bins=batch_size
-    )
+    if config.grain_packing_type == "first_fit":
+      dataset = grain.experimental.FirstFitPackIterDataset(
+          dataset, length_struct=length_struct, num_packing_bins=batch_size
+      )
+    elif config.grain_packing_type == "concat_then_split":
+      if config.add_bos and hasattr(tokenizer_model, "bos_id"):
+        dataset = grain.experimental.ConcatThenSplitIterDataset(
+            dataset,
+            length_struct=length_struct,
+            bos_handling=grain.experimental.BOSHandling.REPLACE_FIRST_TOKEN_WITH_BOS,
+            bos_token_id=tokenizer_model.bos_id,
+        )
+      else:
+        dataset = grain.experimental.ConcatThenSplitIterDataset(dataset, length_struct=length_struct)
+    else:
+      raise ValueError(f"Unkown packing type: {config.packing}")
+
     rekey_dict = {
         "targets_segmentation": "targets_segment_ids",
         "inputs_segmentation": "inputs_segment_ids",

src/MaxText/input_pipeline/_input_pipeline_utils.py

@@ -389,6 +389,16 @@ def map(self, element):
       return {col: element[col].numpy() for col in self.column_names}
 
 
+@dataclasses.dataclass
+class KeepFeatures(grain.MapTransform):
+
+  def __init__(self, feature_names):
+    self.feature_names = feature_names
+
+  def map(self, element):
+    return {k: v for k, v in element.items() if k in self.feature_names}
+
+
 @dataclasses.dataclass
 class Rekey(grain.MapTransform):
   """Rename keys according to a mapping dict"""