diff --git a/tests/test_length_grouping.py b/tests/test_length_grouping.py new file mode 100644 index 00000000..eee48e6c --- /dev/null +++ b/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]} diff --git a/tests/test_offline_replay.py b/tests/test_offline_replay.py index 3decd3ae..436e1fcf 100644 --- a/tests/test_offline_replay.py +++ b/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) diff --git a/torchspec/config/train_config.py b/torchspec/config/train_config.py index 8e160286..4dc0f219 100644 --- a/torchspec/config/train_config.py +++ b/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 diff --git a/torchspec/controller/training_controller.py b/torchspec/controller/training_controller.py index 30b74c18..920c3650 100644 --- a/torchspec/controller/training_controller.py +++ b/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 diff --git a/torchspec/data/utils.py b/torchspec/data/utils.py index 226123ef..ad4e38ea 100644 --- a/torchspec/data/utils.py +++ b/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=""): 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 diff --git a/torchspec/offline/dataset.py b/torchspec/offline/dataset.py index f025e24d..2f1cdfa7 100644 --- a/torchspec/offline/dataset.py +++ b/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()