feat: group training rows by length to improve DP and batching efficiency

tests/test_length_grouping.py

@@ -0,0 +1,129 @@
+"""Tests for length grouping."""
+
+import importlib
+import sys
+from dataclasses import dataclass
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from torchspec.data.utils import length_grouped_order, sample_length_hint
+
+
+def samples(*lengths):
+    return [{"data_id": f"s{i}", "seq_len": n} for i, n in enumerate(lengths)]
+
+
+def lengths_of(entries):
+    return [sample_length_hint(e) for e in entries]
+
+
+@pytest.mark.parametrize(
+    "sample, expected",
+    [
+        ({"seq_len": 512}, 512),  # offline replay row
+        ({"input_ids": torch.zeros(128), "formatted_prompt": "ignored"}, 128),  # tokenized
+        ({"formatted_prompt": "hello"}, 5),  # defer_tokenization: chars stand in
+    ],
+)
+def test_length_hint_reads_whichever_field_is_present(sample, expected):
+    assert sample_length_hint(sample) == expected
+
+
+def test_sorts_longest_first_within_each_chunk():
+    # Chunks are positional, so the 100 in the second chunk stays there.
+    assert lengths_of(length_grouped_order(samples(1, 2, 100, 3), 2)) == [2, 1, 100, 3]
+
+
+def test_trailing_partial_chunk_is_sorted_too():
+    assert lengths_of(length_grouped_order(samples(1, 2, 3, 4, 9, 5), 4)) == [4, 3, 2, 1, 9, 5]
+
+
+def test_small_dataset_becomes_a_full_sort():
+    assert lengths_of(length_grouped_order(samples(1, 5, 3), 1024)) == [5, 3, 1]
+
+
+@pytest.mark.parametrize("group_size", [0, 1])
+def test_group_size_of_one_or_less_disables_grouping(group_size):
+    entries = samples(1, 5, 3)
+    assert length_grouped_order(entries, group_size) == entries
+
+
+def test_leaves_the_input_list_untouched():
+    entries = samples(1, 5, 3)
+    length_grouped_order(entries, 3)
+    assert lengths_of(entries) == [1, 5, 3]
+
+
+def test_groups_similar_lengths_into_neighbouring_dispatches():
+    # The property that matters: a dispatch-sized window should span a narrow
+    # range of lengths after grouping.
+    entries = samples(*((i * 37) % 1000 for i in range(600)))
+
+    def mean_spread(ordered, dispatch=4):
+        windows = [lengths_of(ordered[i : i + dispatch]) for i in range(0, 600, dispatch)]
+        return sum(max(w) - min(w) for w in windows) / len(windows)
+
+    assert mean_spread(length_grouped_order(entries, 128)) < mean_spread(entries) / 5
+
+
+# --- controller wiring ------------------------------------------------------
+
+
+@dataclass
+class MockArgs:
+    per_dp_rank_batch_size: int = 1
+    max_sample_pool_size: int = 0
+    seed: int = 0
+    shuffle_dataset: bool = True
+    length_group_size: int = 8
+
+
+def controller(dataset, **overrides):
+    module_name = "torchspec.controller.training_controller"
+    sys.modules.pop(module_name, None)
+    with patch("ray.remote", lambda cls: cls):
+        module = importlib.import_module(module_name)
+    actor = module.AsyncTrainingController(MockArgs(**overrides), dp_size=2)
+    actor._stored_dataset = dataset
+    return actor
+
+
+SHUFFLE_ME = samples(*((i * 13) % 100 for i in range(32)))
+
+
+def test_grouping_runs_after_the_shuffle():
+    data = controller(SHUFFLE_ME)._prepare_dataset()
+
+    assert len(data) == 32
+    for start in range(0, 32, 8):
+        chunk = lengths_of(data[start : start + 8])
+        assert chunk == sorted(chunk, reverse=True)
+
+
+def test_disabled_grouping_leaves_the_shuffled_order():
+    grouped = controller(SHUFFLE_ME)._prepare_dataset()
+    plain = controller(SHUFFLE_ME, length_group_size=0)._prepare_dataset()
+    assert lengths_of(grouped) != lengths_of(plain)
+
+
+def test_same_epoch_rebuilds_the_same_order():
+    # Mid-epoch resume slices this order, so it has to be reproducible.
+    first = controller(SHUFFLE_ME)._prepare_dataset()
+    second = controller(SHUFFLE_ME)._prepare_dataset()
+    assert [e["data_id"] for e in first] == [e["data_id"] for e in second]
+
+
+def test_skip_slices_the_grouped_order():
+    actor = controller(SHUFFLE_ME)
+    assert actor._prepare_dataset(skip=10) == actor._prepare_dataset()[10:]
+
+
+def test_each_epoch_groups_different_samples_together():
+    actor = controller(SHUFFLE_ME)
+    epoch0 = actor._prepare_dataset()
+    actor._dataset_epoch = 1
+    epoch1 = actor._prepare_dataset()
+
+    assert {e["data_id"] for e in epoch0[:8]} != {e["data_id"] for e in epoch1[:8]}

tests/test_offline_replay.py

@@ -227,7 +227,7 @@ def test_controller_sources_manifest_ids_without_retokenizing(tmp_path):
     assert controller.load_dataset(args) == 1
     assert controller.load_eval_dataset(args) == 1
     assert controller._stored_dataset == [
-        {"data_id": "train-1", "metadata": {"offline_replay": True}}
+        {"data_id": "train-1", "metadata": {"offline_replay": True}, "seq_len": 4}
     ]
 
     controller.submit_eval_chunk(0, 1)

torchspec/config/train_config.py

@@ -40,6 +40,7 @@ class DatasetConfig:
     eval_micro_batch_size: Optional[int] = None
     eval_prompt_key: Optional[str] = None
     last_turn_loss_only: Any = "auto"  # bool or "auto"
+    length_group_size: int = 1024
     min_loss_tokens: int = 0  # DFlash: skip sequences with < N supervised tokens (use 2*block_size)
     prompt_key: str = "conversations"
     shuffle_dataset: bool = True

torchspec/controller/training_controller.py

@@ -57,6 +57,7 @@
 import ray
 from ray.util.queue import Queue
 
+from torchspec.data.utils import length_grouped_order
 from torchspec.training.data_fetcher import TrainSample
 from torchspec.utils.logging import logger
 from torchspec.utils.memory import estimate_tensor_bytes
@@ -159,6 +160,7 @@ def __init__(self, args, dp_size: int):
         self._dataset_epoch: int = 0
         self._dataset_seed: int = getattr(args, "seed", 42)
         self._shuffle_dataset: bool = getattr(args, "shuffle_dataset", True)
+        self._length_group_size: int = getattr(args, "length_group_size", 1024)
 
         self._start_time = time.time()
         self._inference_monitor = SpeedMonitor(window_seconds=10.0)
@@ -214,7 +216,11 @@ def _load_dataset_split(self, args, split: str) -> list:
 
             dataset = OfflineDataset(args.offline_data_path)
             return [
-                {"data_id": str(row["data_id"]), "metadata": {"offline_replay": True}}
+                {
+                    "data_id": str(row["data_id"]),
+                    "metadata": {"offline_replay": True},
+                    "seq_len": row["seq_len"],
+                }
                 for row in dataset.rows(split)
             ]
 
@@ -250,7 +256,7 @@ def load_dataset(self, args) -> int:
         return len(self._stored_dataset)
 
     def _prepare_dataset(self, skip: int = 0) -> list:
-        """Return dataset for the current epoch, optionally shuffled.
+        """Return dataset for the current epoch, optionally shuffled and length-grouped.
 
         When shuffle is enabled the ordering is deterministic from
         (seed + epoch), so resume can reconstruct the same epoch ordering
@@ -265,6 +271,8 @@ def _prepare_dataset(self, skip: int = 0) -> list:
             rng = random.Random(self._dataset_seed + self._dataset_epoch)
             rng.shuffle(data)
 
+        data = length_grouped_order(data, self._length_group_size)
+
         if skip > 0:
             skip = min(skip, len(data))
             data = data[skip:]
@@ -275,8 +283,8 @@ def _prepare_dataset(self, skip: int = 0) -> list:
             else "shuffle disabled"
         )
         logger.info(
-            f"Prepared dataset ({len(data)} samples, {shuffle_tag}"
-            + (f", skipped {skip})" if skip > 0 else ")")
+            f"Prepared dataset ({len(data)} samples, {shuffle_tag}, "
+            f"length group {self._length_group_size}" + (f", skipped {skip})" if skip > 0 else ")")
         )
         return data
 

torchspec/data/utils.py

@@ -95,7 +95,7 @@ def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
             # Round up to nearest bucket to reduce unique shapes for torch.compile.
             # Without this, every batch gets a different padded length, causing
             # FlexAttention recompilation (~1s overhead per new shape).
-            _BUCKET = 256
+            _BUCKET = 128
             max_length = ((max_length + _BUCKET - 1) // _BUCKET) * _BUCKET
             attention_masks = [torch.ones_like(item["input_ids"]).long() for item in features]
 
@@ -381,6 +381,38 @@ def flatten_multimodal_content(messages, image_placeholder="<image>"):
     return messages
 
 
+def sample_length_hint(sample: Dict[str, Any]) -> int:
+    """Sequence length of a stored dataset entry, for length grouping.
+
+    Exactly one of these three is present: ``seq_len`` on offline replay rows,
+    the tokenized ``input_ids``, or — under ``defer_tokenization`` — only the
+    formatted text, whose character count is monotone in token count and free.
+    """
+    if "seq_len" in sample:
+        return sample["seq_len"]
+    if "input_ids" in sample:
+        return sample["input_ids"].numel()
+    return len(sample["formatted_prompt"])
+
+
+def length_grouped_order(samples: List[Any], group_size: int) -> List[Any]:
+    """Sort *samples* longest-first within fixed chunks of *group_size*.
+
+    Chunks are positional, so only the order inside each one changes. That keeps
+    a dispatch's sequences similar in length without turning length into a
+    curriculum the way a global sort would.
+    """
+    if group_size <= 1:
+        return list(samples)
+
+    ordered = []
+    for start in range(0, len(samples), group_size):
+        chunk = samples[start : start + group_size]
+        chunk.sort(key=sample_length_hint, reverse=True)
+        ordered.extend(chunk)
+    return ordered
+
+
 def estimate_row_count(data_path):
     if not os.path.isfile(data_path):
         return None

torchspec/offline/dataset.py

@@ -82,7 +82,12 @@ def __init__(
                     path = (self.root / row["file"]).resolve()
                     if self.root not in path.parents or not path.is_file():
                         raise FileNotFoundError(f"Offline sample not found: {path}")
-                    item = {"split": split, "data_id": data_id, "file": row["file"]}
+                    item = {
+                        "split": split,
+                        "data_id": data_id,
+                        "file": row["file"],
+                        "seq_len": row["seq_len"],
+                    }
                     self._rows[split].append(item)
                     self._by_id[data_id] = item
 
@@ -152,7 +157,12 @@ def append(
             if os.path.exists(temporary):
                 os.unlink(temporary)
 
-        row = {"split": split, "data_id": data_id, "file": relative.as_posix()}
+        row = {
+            "split": split,
+            "data_id": data_id,
+            "file": relative.as_posix(),
+            "seq_len": int(saved["input_ids"].numel()),
+        }
         with (self.root / "manifest.jsonl").open("a", encoding="utf-8") as stream:
             stream.write(json.dumps(row) + "\n")
             stream.flush()