Add Bagz format support

src/dependencies/requirements/base_requirements/requirements.txt

@@ -12,6 +12,7 @@ google-cloud-aiplatform
 google-cloud-mldiagnostics
 google-cloud-monitoring
 grain[parquet]
+bagz
 huggingface_hub>=1.14.0
 jax
 jaxlib

src/dependencies/requirements/requirements.txt

@@ -12,6 +12,7 @@ google-cloud-aiplatform
 google-cloud-mldiagnostics>=0.5.10
 google-cloud-monitoring
 grain[parquet]
+bagz
 huggingface_hub
 jax!=0.7.1
 jaxlib!=0.7.1

src/maxtext/configs/base.yml

@@ -754,7 +754,7 @@ hf_access_token: ''
 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_file_type: 'arrayrecord' # arrayrecord, parquet, tfrecord, or bagz
 grain_packing_type: 'first_fit' # 'first_fit', 'best_fit' or 'concat_then_split'. See details of the corresponding module in https://google-grain.readthedocs.io/en/latest/grain.experimental.html
 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

src/maxtext/configs/types.py

@@ -1302,7 +1302,7 @@ class GrainDataset(BaseModel):
   )
   grain_file_type: str = Field(
       "arrayrecord",
-      description="File type for Grain data. Supported: arrayrecord, tfrecord, parquet.",
+      description="File type for Grain data. Supported: arrayrecord, tfrecord, parquet, bagz.",
   )
   grain_use_elastic_iterator: bool = Field(
       False,
@@ -3146,9 +3146,9 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           "Colocated python data input is only supported with Pathways (single"
           " controller) enabled (`enable_single_controller=True`)."
       )
-    if self.grain_use_elastic_iterator and self.grain_file_type != "arrayrecord":
+    if self.grain_use_elastic_iterator and self.grain_file_type not in ("arrayrecord", "bagz"):
       raise ValueError(
-          "`grain_use_elastic_iterator=True` only supports `grain_file_type=arrayrecord`. "
+          "`grain_use_elastic_iterator=True` only supports `grain_file_type=arrayrecord` or `bagz`. "
           "tfrecord and parquet pipelines use `InterleaveIterDataset` (a many-to-one "
           "IterDataset transform), which `ElasticIterator` forbids. "
           f"Got grain_file_type={self.grain_file_type}."

src/maxtext/input_pipeline/data_processing_utils.py

@@ -27,7 +27,7 @@
 
 def parse_and_keep_features(dataset, config, data_columns, tokenize):
   """Parse arrayrecord features or keep specified columns for other formats."""
-  if config.grain_file_type in ("arrayrecord", "tfrecord"):
+  if config.grain_file_type in ("arrayrecord", "tfrecord", "bagz"):
     dataset = dataset.map(input_pipeline_utils.ParseFeatures(data_columns, tokenize))
     dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(data_columns, tokenize))
   else:

src/maxtext/input_pipeline/grain_data_processing.py

@@ -26,6 +26,7 @@
 
 import grain.python as grain
 from grain.experimental import ElasticIterator
+import bagz
 
 from maxtext.input_pipeline import data_processing_utils
 from maxtext.input_pipeline import input_pipeline_utils
@@ -180,6 +181,76 @@ def create_dataset_from_pattern(pattern):
           elastic=elastic,
       )
       return dataset
+  elif data_file_type == "bagz":
+    def create_dataset_from_pattern(pattern):
+      files = find_data_files(pattern)
+      source = input_pipeline_utils.BagzDataSource(files)
+      return grain.MapDataset.source(source)
+
+    if mixture_config_path:
+      with open(mixture_config_path, "r", encoding="utf-8") as f:
+        mixture_config = json.load(f)
+      paths = [config["path"] for config in mixture_config.values()]
+      weights = [float(config["weight"]) for config in mixture_config.values()]
+
+      executor = futures.ThreadPoolExecutor(max_workers=grain_data_source_max_workers)
+      dataset_list = list(executor.map(create_dataset_from_pattern, paths))
+      executor.shutdown(wait=True)
+
+      datasets_dict = dict(zip(mixture_config.keys(), dataset_list))
+      for name, ds in datasets_dict.items():
+        datasets_dict[name] = _apply_mapdataset_transforms(
+            ds,
+            shuffle,
+            shuffle_seed,
+            num_epoch,
+            dataloading_host_index,
+            dataloading_host_count,
+            grain_num_threads,
+            grain_prefetch_buffer_size,
+        )
+
+      total_weight = sum(weights)
+      weights_dict = {name: weight / total_weight for name, weight in zip(mixture_config.keys(), weights)}
+      dataset = grain.IterDataset.mix(datasets_dict, weights_dict)
+      return dataset
+    elif ";" in data_file_pattern:
+      data_file_patterns, weights = zip(*[pattern.split(",") for pattern in data_file_pattern.split(";")])
+      assert len(data_file_patterns) == len(weights)
+      weights = [float(w) for w in weights]
+      weights = [round(w / sum(weights), 4) for w in weights]
+
+      executor = futures.ThreadPoolExecutor(max_workers=grain_data_source_max_workers)
+      dataset_list = list(executor.map(create_dataset_from_pattern, data_file_patterns))
+      executor.shutdown(wait=True)
+
+      for d, _ in enumerate(dataset_list):
+        dataset_list[d] = _apply_mapdataset_transforms(
+            dataset_list[d],
+            shuffle,
+            shuffle_seed,
+            num_epoch,
+            dataloading_host_index,
+            dataloading_host_count,
+            grain_num_threads,
+            grain_prefetch_buffer_size,
+        )
+      dataset = grain.IterDataset.mix(dataset_list, weights)
+      return dataset
+    else:
+      dataset = create_dataset_from_pattern(data_file_pattern)
+      dataset = _apply_mapdataset_transforms(
+          dataset,
+          shuffle,
+          shuffle_seed,
+          num_epoch,
+          dataloading_host_index,
+          dataloading_host_count,
+          grain_num_threads,
+          grain_prefetch_buffer_size,
+          elastic=elastic,
+      )
+      return dataset
   elif data_file_type in ("tfrecord", "parquet"):
     data_files = find_data_files(data_file_pattern)
     file_slice, files_per_host, row_shard = input_pipeline_utils.compute_file_sharding(
@@ -215,7 +286,7 @@ def create_dataset_from_pattern(pattern):
     return dataset
   else:
     raise ValueError(
-        f"grain pipeline supports (arrayrecord, tfrecord, parquet) as grain_file_type, but got {data_file_type}"
+        f"grain pipeline supports (arrayrecord, tfrecord, parquet, bagz) as grain_file_type, but got {data_file_type}"
     )
 
 

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -1070,3 +1070,43 @@ def map(self, element: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
     element[f"{self.data_column}_mrope_deltas"] = mrope_position_deltas
 
     return element
+
+
+import bagz
+
+class BagzDataSource:
+  """A fully picklable and fork-safe RandomAccessDataSource for Bagz files."""
+  def __init__(self, paths: list[str] | str):
+    if isinstance(paths, (list, tuple)):
+      self._path = ",".join(sorted(paths))
+    else:
+      self._path = paths
+    self._len = None
+    self._reader = None
+
+  @property
+  def reader(self):
+    if self._reader is None:
+      self._reader = bagz.Reader(self._path)
+    return self._reader
+
+  def __len__(self) -> int:
+    if self._len is None:
+      self._len = len(bagz.Reader(self._path))
+    return self._len
+
+  def __getitem__(self, record_key: int):
+    return self.reader[record_key]
+
+  def __repr__(self) -> str:
+    return f"BagzDataSource(path={self._path!r})"
+
+  def __getstate__(self):
+    state = self.__dict__.copy()
+    state["_reader"] = None
+    return state
+
+  def __setstate__(self, state):
+    self.__dict__.update(state)
+    self._reader = None
+

src/maxtext/trainers/tokenizer/train_tokenizer.py

@@ -41,6 +41,7 @@
 
 import jax
 import grain.python as grain
+import bagz
 
 from maxtext.input_pipeline import input_pipeline_utils
 from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT
@@ -93,6 +94,12 @@ def build_grain_iterator(data_file_pattern: str, data_file_type: str, data_keys:
     dataset = dataset.map(input_pipeline_utils.ParseFeatures(list(data_keys), tokenize=True))
     dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(list(data_keys), tokenize=True))
     return iter(dataset)
+  elif data_file_type == "bagz":
+    source = input_pipeline_utils.BagzDataSource(data_files)
+    dataset = grain.MapDataset.source(source)
+    dataset = dataset.map(input_pipeline_utils.ParseFeatures(list(data_keys), tokenize=True))
+    dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(list(data_keys), tokenize=True))
+    return iter(dataset)
   elif data_file_type == "tfrecord":
     dataset = grain.MapDataset.source(data_files)
     dataset = dataset.map(input_pipeline_utils.make_tfrecord_iter_dataset)
@@ -103,7 +110,7 @@ def build_grain_iterator(data_file_pattern: str, data_file_type: str, data_keys:
     dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(list(data_keys), tokenize=True))
     return iter(dataset)
   else:
-    raise ValueError(f"Unsupported grain_file_type: {data_file_type!r}. Use 'parquet', 'arrayrecord', or 'tfrecord'.")
+    raise ValueError(f"Unsupported grain_file_type: {data_file_type!r}. Use 'parquet', 'arrayrecord', 'tfrecord', or 'bagz'.")
 
 
 def _dump_chars_to_textfile(dataset_iter: Iterator, maxchars: int = int(1e7), data_keys=("text",)) -> tuple[str, int]: