mirror of https://github.com/razor-ai/soup.git
124 lines
4.5 KiB
Python
124 lines
4.5 KiB
Python
"""Multipack helper — task allowlist + sampler builder for trainer wrappers.
|
|
|
|
This is the high-level dispatcher that trainer wrappers call to get a
|
|
configured :class:`MultipackBatchSampler` from a tokenized dataset's per-
|
|
sample lengths. Keeping the construction logic out of the trainer modules
|
|
means the v0.37.0 wiring touches each trainer in ~3 lines.
|
|
|
|
Tasks wired in v0.37.0: ``sft``, ``pretrain``. Preference / RLHF / reward-
|
|
model trainers operate on paired data where multipack's index-bag
|
|
abstraction does not cleanly apply — those land in a future release.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from typing import TYPE_CHECKING
|
|
|
|
from soup_cli.utils.multipack_sampler import MultipackBatchSampler
|
|
|
|
if TYPE_CHECKING:
|
|
from soup_cli.config.schema import TrainingConfig
|
|
|
|
|
|
# v0.37.0 wiring — only flat-sequence tasks. Preference tasks deferred.
|
|
_MULTIPACK_SUPPORTED_TASKS: frozenset[str] = frozenset({"sft", "pretrain"})
|
|
|
|
|
|
def supports_multipack(task: str) -> bool:
|
|
"""Return True if ``task`` has multipack sampler wiring.
|
|
|
|
Mirrors :func:`soup_cli.utils.v028_features.supports_v028_features`
|
|
naming and shape so ``SoupConfig`` validators stay symmetric.
|
|
"""
|
|
if not isinstance(task, str):
|
|
return False
|
|
return task in _MULTIPACK_SUPPORTED_TASKS
|
|
|
|
|
|
def build_multipack_sampler_for_lengths(
|
|
*,
|
|
lengths: Sequence[int],
|
|
tcfg: TrainingConfig,
|
|
max_seq_length: int,
|
|
real_batches: bool = True,
|
|
seed: int = 0,
|
|
) -> MultipackBatchSampler:
|
|
"""Build a :class:`MultipackBatchSampler` from per-sample token lengths.
|
|
|
|
Args:
|
|
lengths: token counts for each sample in the (already tokenized)
|
|
training dataset.
|
|
tcfg: the run's :class:`TrainingConfig`. Must have ``multipack=True``.
|
|
max_seq_length: per-sample maximum length (matches
|
|
``DataConfig.max_length``).
|
|
real_batches: when True, yields list-of-bins per call (collator
|
|
stacks bins along the batch dim). When False, yields one flat
|
|
bin at a time and ``batch_max_len = batch_size * max_seq_length``
|
|
— Axolotl's "micro-batch as flat sequence" trick that maximises
|
|
packing density.
|
|
seed: deterministic shuffle seed (mirrored across DDP ranks).
|
|
|
|
Returns:
|
|
A configured :class:`MultipackBatchSampler`.
|
|
|
|
Raises:
|
|
ValueError: when ``multipack`` is not enabled on ``tcfg`` or when
|
|
``max_seq_length`` is non-positive.
|
|
"""
|
|
if not getattr(tcfg, "multipack", False):
|
|
raise ValueError(
|
|
"build_multipack_sampler_for_lengths called but "
|
|
"tcfg.multipack=False — guard upstream"
|
|
)
|
|
if isinstance(max_seq_length, bool):
|
|
raise TypeError(
|
|
f"max_seq_length must not be bool, got {max_seq_length!r}"
|
|
)
|
|
if not isinstance(max_seq_length, int):
|
|
raise TypeError(
|
|
f"max_seq_length must be int, got "
|
|
f"{type(max_seq_length).__name__}"
|
|
)
|
|
if max_seq_length <= 0:
|
|
raise ValueError(
|
|
f"max_seq_length must be positive, got {max_seq_length}"
|
|
)
|
|
|
|
raw_batch_size = getattr(tcfg, "batch_size", 1) or 1
|
|
# batch_size=='auto' must be resolved upstream before the sampler is
|
|
# built — multipack needs an explicit integer to size each bin.
|
|
if isinstance(raw_batch_size, str):
|
|
raise ValueError(
|
|
f"multipack requires an explicit integer batch_size, got "
|
|
f"{raw_batch_size!r}. Resolve auto-batch-size before calling "
|
|
"build_multipack_sampler_for_lengths (e.g. via "
|
|
"soup_cli.utils.batch_probe.pick_batch_size)."
|
|
)
|
|
# bool is a subclass of int — reject explicitly (matches v0.30.0 policy).
|
|
if isinstance(raw_batch_size, bool):
|
|
raise TypeError(
|
|
f"tcfg.batch_size must not be bool, got {raw_batch_size!r}"
|
|
)
|
|
if not isinstance(raw_batch_size, int):
|
|
raise TypeError(
|
|
f"tcfg.batch_size must be int, got {type(raw_batch_size).__name__}"
|
|
)
|
|
batch_size = raw_batch_size
|
|
|
|
if real_batches:
|
|
batch_max_len = max_seq_length
|
|
else:
|
|
# Axolotl trick: flat mode treats one bin as "batch_size copies of
|
|
# max_seq_length stitched together" so FFD has more room to pack.
|
|
batch_max_len = batch_size * max_seq_length
|
|
|
|
return MultipackBatchSampler(
|
|
lengths=lengths,
|
|
batch_max_len=batch_max_len,
|
|
batch_size=batch_size if real_batches else 1,
|
|
real_batches=real_batches,
|
|
seed=seed,
|
|
drop_last=False,
|
|
)
|