Merge pull request AI-Hypercomputer#3807 from AI-Hypercomputer:aireen/shard_by_row2

PiperOrigin-RevId: 938837636

Author: Google-ML-Automation

src/maxtext/input_pipeline/grain_data_processing.py

@@ -14,6 +14,7 @@
 
 """Input pipeline using Grain."""
 
+import math
 import glob
 from pathlib import Path
 import functools
@@ -179,35 +180,37 @@ def create_dataset_from_pattern(pattern):
           elastic=elastic,
       )
       return dataset
-  elif data_file_type == "tfrecord":
+  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(
+        len(data_files), dataloading_host_index, dataloading_host_count
+    )
+    if len(data_files) < dataloading_host_count:
+      max_logging.warning(
+          f"Data shard count ({len(data_files)}) < host count ({dataloading_host_count}). "
+          f"Each shard will be read by at most {math.ceil(dataloading_host_count / len(data_files))} hosts. "
+          f"Concurrent reading by multiple hosts may cause slow down."
+      )
+    min_files_per_host = len(data_files) // dataloading_host_count
+    if grain_worker_count > min_files_per_host:
+      raise ValueError(
+          f"grain_worker_count ({grain_worker_count}) exceeds the minimum number of {data_file_type} files "
+          f"per host ({min_files_per_host} = {len(data_files)} files / {dataloading_host_count} hosts). "
+          f"Lower grain_worker_count to at most {min_files_per_host}."
+      )
     dataset = grain.MapDataset.source(data_files)
     if shuffle:
       dataset = dataset.shuffle(seed=shuffle_seed)
     dataset = dataset.repeat(num_epoch)
-    dataset = dataset[dataloading_host_index::dataloading_host_count]  # sharding
-    dataset = dataset.map(input_pipeline_utils.make_tfrecord_iter_dataset)
-    files_per_host = max(len(data_files) // dataloading_host_count, 1)
+    dataset = dataset[file_slice]
+    if data_file_type == "tfrecord":
+      dataset = dataset.map(input_pipeline_utils.make_tfrecord_iter_dataset)
+    else:
+      dataset = dataset.map(grain.experimental.ParquetIterDataset)
     cycle_length = min(files_per_host, grain_num_threads)
     dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=cycle_length)
-    if shuffle:
-      dataset = grain.experimental.WindowShuffleIterDataset(dataset, window_size=shuffle_buffer_size, seed=shuffle_seed)
-    return dataset
-  elif data_file_type == "parquet":
-    data_files = find_data_files(data_file_pattern)
-    dataset = grain.MapDataset.source(data_files)
-    if shuffle:
-      dataset = dataset.shuffle(seed=shuffle_seed)
-    dataset = dataset.repeat(num_epoch)
-    dataset = dataset[dataloading_host_index::dataloading_host_count]  # sharding
-    assert grain_worker_count <= len(dataset), (
-        f"grain worker count is currently {grain_worker_count}, exceeding the max allowable value {len(dataset)} "
-        f"(file shard count of a data loading host) for your dataset. "
-        f"Please lower grain_worker_count or increase file shard count."
-    )
-    dataset = dataset.map(grain.experimental.ParquetIterDataset)
-    cycle_length = min(len(dataset) // num_epoch, grain_num_threads)
-    dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=cycle_length)
+    if row_shard is not None:
+      dataset = input_pipeline_utils.IndexShardIterDataset(dataset, host_index=row_shard[0], host_count=row_shard[1])
     if shuffle:
       dataset = grain.experimental.WindowShuffleIterDataset(dataset, window_size=shuffle_buffer_size, seed=shuffle_seed)
     return dataset
@@ -489,7 +492,7 @@ def make_grain_train_iterator(
         for x in train_dataloader_list
     ]
 
-  # Default non-colocated, non-expansion path
+  # Default non-colocated, expansion_factor_real_data >=1 path
   shard_index = process_indices.index(jax.process_index())
   shard_count = len(process_indices)
   train_ds = get_ds_fn(

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -496,6 +496,79 @@ def make_tfrecord_iter_dataset(path: str):
   return TFRecordIterDataset(path)
 
 
+def compute_file_sharding(file_count, host_index, host_count):
+  """Compute per-host file slicing and optional row-shard parameters.
+
+  When file_count >= host_count, each host reads a disjoint subset of files via the
+  standard `[host_index::host_count]` slice. When file_count < host_count, every file
+  is replicated across `ceil(host_count/file_count)` hosts (or `floor` for files past
+  the remainder), and those hosts shard records by index within the file. This bounds
+  concurrent readers per file at `ceil(host_count/file_count)`
+
+  Returns:
+    file_slice (slice): file slice this host should keep — also the (start, step)
+      pair `(host_index, host_count)` to feed `tf.distribute.InputContext` when applicable.
+    files_per_host (int): files this host's slice contains per epoch.
+    row_shard (tuple|None): `(row_shard_index, row_shard_count)` when host_count > file_count
+      and the file's group has >1 reader; otherwise None.
+  """
+  if file_count >= host_count:
+    return slice(host_index, None, host_count), max(file_count // host_count, 1), None
+  file_idx = host_index % file_count
+  row_shard_idx = host_index // file_count
+  row_shard_count = (host_count // file_count) + (1 if file_idx < (host_count % file_count) else 0)
+  row_shard = (row_shard_idx, row_shard_count) if row_shard_count > 1 else None
+  return slice(file_idx, None, file_count), 1, row_shard
+
+
+class _IndexShardDatasetIterator(grain.DatasetIterator):
+  """Iterator that yields every nth element of its parent (round-robin by index)."""
+
+  def __init__(self, parent: grain.DatasetIterator, host_index: int, host_count: int):
+    super().__init__(parent)
+    self._host_index = host_index
+    self._host_count = host_count
+    self._next_index = 0
+
+  def __next__(self):
+    while True:
+      value = next(self._parent)
+      current = self._next_index
+      self._next_index += 1
+      if current % self._host_count == self._host_index:
+        return value
+
+  def get_state(self):
+    return {
+        "next_index": self._next_index,
+        "parent": self._parent.get_state(),
+    }
+
+  def set_state(self, state):
+    self._next_index = state["next_index"]
+    self._parent.set_state(state["parent"])
+
+
+class IndexShardIterDataset(grain.IterDataset):
+  """Shards an IterDataset across hosts by element index (host i keeps records where idx % N == i).
+
+  Use when the upstream `IterDataset` order is deterministic and identical on every host;
+  this guarantees disjoint, balanced slices without per-file sharding.
+  """
+
+  def __init__(self, parent: grain.IterDataset, host_index: int, host_count: int):
+    super().__init__(parent)
+    self._host_index = host_index
+    self._host_count = host_count
+
+  def __iter__(self) -> _IndexShardDatasetIterator:
+    return _IndexShardDatasetIterator(
+        self._parent.__iter__(),
+        host_index=self._host_index,
+        host_count=self._host_count,
+    )
+
+
 @dataclasses.dataclass
 class ParseFeatures(grain.MapTransform):
   """Parse serialized tf.train.Example protos for arrayrecord/tfrecord datasets.

src/maxtext/input_pipeline/tfds_data_processing.py

@@ -14,8 +14,8 @@
 
 """Input pipeline for a LM1B dataset."""
 
-import warnings
 import functools
+import math
 
 import ml_collections
 
@@ -27,6 +27,7 @@
 from maxtext.input_pipeline import multihost_dataloading
 from maxtext.input_pipeline.packing import sequence_packing
 from maxtext.input_pipeline import input_pipeline_utils
+from maxtext.utils import max_logging
 
 AUTOTUNE = tf.data.experimental.AUTOTUNE
 
@@ -54,21 +55,31 @@ def get_datasets(
   else:
     read_config = tfds.ReadConfig()
 
-  if ds_builder.info.splits[data_split].num_shards >= dataloading_host_count:
-    read_config.input_context = tf.distribute.InputContext(
-        input_pipeline_id=dataloading_host_index,
-        num_input_pipelines=dataloading_host_count,
-    )
-    ds = ds_builder.as_dataset(split=data_split, read_config=read_config, shuffle_files=shuffle_files)
+  num_shards = ds_builder.info.splits[data_split].num_shards
+  _, _, row_shard = input_pipeline_utils.compute_file_sharding(num_shards, dataloading_host_index, dataloading_host_count)
+  if num_shards >= dataloading_host_count:
+    input_pipeline_id = dataloading_host_index
+    num_input_pipelines = dataloading_host_count
   else:
-    warnings.warn(
-        f"WARNING: Inefficient dataloading. Your {dataset_name} contains {ds_builder.info.splits[data_split].num_shards}"
-        f"shards, smaller than {dataloading_host_count=}. This is known to lead to inefficient dataloading."
-        "see https://github.com/google/maxtext/blob/main/getting_started/Data_Input_Pipeline.md"
-        "#multihost-dataloading-best-practice"
+    # When `num_shards < dataloading_host_count`, falls back to hybrid sharding: each TFDS
+    # shard is read by a `ceil(host_count/num_shards)` group of hosts, and within that group
+    # hosts split records via `Dataset.shard`. This bounds concurrent readers per shard,
+    # avoiding the thundering-herd I/O pattern of a "every host reads all shards" fallback.
+    input_pipeline_id = dataloading_host_index % num_shards
+    num_input_pipelines = num_shards
+    max_logging.warning(
+        f"{dataset_name}: shard count ({num_shards}) < host count ({dataloading_host_count}). "
+        f"Each shard will be read by at most {math.ceil(dataloading_host_count / num_shards)} hosts. "
+        f"Concurrent reading by multiple hosts may cause slow down."
     )
-    ds = ds_builder.as_dataset(split=data_split, read_config=read_config, shuffle_files=shuffle_files)
-    ds = ds.shard(num_shards=dataloading_host_count, index=dataloading_host_index)
+
+  read_config.input_context = tf.distribute.InputContext(
+      input_pipeline_id=input_pipeline_id,
+      num_input_pipelines=num_input_pipelines,
+  )
+  ds = ds_builder.as_dataset(split=data_split, read_config=read_config, shuffle_files=shuffle_files)
+  if row_shard is not None:
+    ds = ds.shard(num_shards=row_shard[1], index=row_shard[0])
 
   return ds
 

tests/unit/input_pipeline_utils_test.py

@@ -0,0 +1,113 @@
+# Copyright 2023–2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for input_pipeline_utils."""
+
+import pytest
+import unittest
+
+from maxtext.input_pipeline.input_pipeline_utils import compute_file_sharding
+
+
+@pytest.mark.cpu_only
+class ComputeFileShardingNormalCaseTest(unittest.TestCase):
+  """file_count >= host_count: disjoint file subsets, no row sharding."""
+
+  def test_even_split(self):
+    # 8 files, 4 hosts → interleaved assignment, 2 files each
+    file_slice, files_per_host, _ = compute_file_sharding(8, host_index=0, host_count=4)
+    self.assertEqual(list(range(8)[file_slice]), [0, 4])
+    self.assertEqual(files_per_host, 2)
+
+    file_slice, files_per_host, _ = compute_file_sharding(8, host_index=1, host_count=4)
+    self.assertEqual(list(range(8)[file_slice]), [1, 5])
+    self.assertEqual(files_per_host, 2)
+
+    file_slice, files_per_host, _ = compute_file_sharding(8, host_index=2, host_count=4)
+    self.assertEqual(list(range(8)[file_slice]), [2, 6])
+    self.assertEqual(files_per_host, 2)
+
+    file_slice, files_per_host, _ = compute_file_sharding(8, host_index=3, host_count=4)
+    self.assertEqual(list(range(8)[file_slice]), [3, 7])
+    self.assertEqual(files_per_host, 2)
+
+  def test_uneven_split(self):
+    # 5 files, 4 hosts → host 0 gets an extra file
+    file_slice, _, _ = compute_file_sharding(5, host_index=0, host_count=4)
+    self.assertEqual(list(range(5)[file_slice]), [0, 4])
+
+    file_slice, _, _ = compute_file_sharding(5, host_index=1, host_count=4)
+    self.assertEqual(list(range(5)[file_slice]), [1])
+
+    file_slice, _, _ = compute_file_sharding(5, host_index=2, host_count=4)
+    self.assertEqual(list(range(5)[file_slice]), [2])
+
+    file_slice, _, _ = compute_file_sharding(5, host_index=3, host_count=4)
+    self.assertEqual(list(range(5)[file_slice]), [3])
+
+  def test_single_host_gets_all_files(self):
+    file_slice, files_per_host, _ = compute_file_sharding(8, host_index=0, host_count=1)
+    self.assertEqual(list(range(8)[file_slice]), [0, 1, 2, 3, 4, 5, 6, 7])
+    self.assertEqual(files_per_host, 8)
+
+  def test_no_row_shard_in_normal_case(self):
+    for host_index in range(4):
+      _, _, row_shard = compute_file_sharding(8, host_index, host_count=4)
+      self.assertIsNone(row_shard)
+
+
+class ComputeFileShardingUndersizedCaseTest(unittest.TestCase):
+  """file_count < host_count: multiple hosts share a file, split by row."""
+
+  def test_single_file_four_hosts(self):
+    # All 4 hosts read the same file, each gets a quarter of the rows
+    _, _, row_shard = compute_file_sharding(1, host_index=0, host_count=4)
+    self.assertEqual(row_shard, (0, 4))  # row index 0 of 4
+
+    _, _, row_shard = compute_file_sharding(1, host_index=1, host_count=4)
+    self.assertEqual(row_shard, (1, 4))  # row index 1 of 4
+
+    _, _, row_shard = compute_file_sharding(1, host_index=2, host_count=4)
+    self.assertEqual(row_shard, (2, 4))  # row index 2 of 4
+
+    _, _, row_shard = compute_file_sharding(1, host_index=3, host_count=4)
+    self.assertEqual(row_shard, (3, 4))  # row index 3 of 4
+
+  def test_three_files_eight_hosts(self):
+    # 8 hosts round-robin across 3 files:
+    # hosts 0,3,6 → file 0 (3 readers); hosts 1,4,7 → file 1 (3 readers); hosts 2,5 → file 2 (2 readers)
+    expected = {
+        # host_index: (file_indices, row_shard)
+        0: ([0], (0, 3)),
+        1: ([1], (0, 3)),
+        2: ([2], (0, 2)),
+        3: ([0], (1, 3)),
+        4: ([1], (1, 3)),
+        5: ([2], (1, 2)),
+        6: ([0], (2, 3)),
+        7: ([1], (2, 3)),
+    }
+    for host_index, (exp_files, exp_row_shard) in expected.items():
+      file_slice, _, row_shard = compute_file_sharding(3, host_index, host_count=8)
+      self.assertEqual(list(range(3)[file_slice]), exp_files, f"host {host_index} file assignment")
+      self.assertEqual(row_shard, exp_row_shard, f"host {host_index} row shard")
+
+  def test_no_row_shard_when_only_one_reader(self):
+    # 2 files, 3 hosts: file 1 has only one reader (host 1) → no row split needed
+    _, _, row_shard = compute_file_sharding(2, host_index=1, host_count=3)
+    self.assertIsNone(row_shard)
+
+
+if __name__ == "__main__":
+  unittest.main()