soup/soup_cli/utils/multipack_trainer.py

"""HF Trainer / TRL SFTTrainer subclass that swaps ``_get_train_sampler``.

The schema field ``training.multipack: bool`` shipped in v0.37.0 along
with the FFD bin-packing :class:`MultipackBatchSampler` and the
arch-allowlist validator. The live HF Trainer wiring was deliberately
deferred (mirrors the v0.27.0 MII / v0.38.0 quant-menu / v0.39.0 ReLoRA
stub-then-live pattern). v0.40.3 (#65) wires it up.

Usage from a trainer wrapper::

    from soup_cli.utils.multipack_trainer import (
        attach_multipack_state, lengths_from_dataset,
        make_multipack_trainer_class,
    )
    if tcfg.multipack:
        validate_multipack_architecture(arch)
        TrainerCls = make_multipack_trainer_class(SFTTrainer)
        trainer = TrainerCls(**trainer_kwargs)
        attach_multipack_state(
            trainer,
            lengths=lengths_from_dataset(train_ds),
            max_seq_len=cfg.data.max_length,
            batch_size=batch_size,
            seed=cfg.training.seed or 0,
        )
"""

from __future__ import annotations

import functools
import logging
from typing import Any, Optional, Sequence

logger = logging.getLogger(__name__)

# Hidden attrs on the trainer instance — picked up by _get_train_sampler.
_LENGTHS_ATTR = "_soup_multipack_lengths"
_MAX_SEQ_ATTR = "_soup_multipack_max_seq_len"
_BATCH_SIZE_ATTR = "_soup_multipack_batch_size"
_SEED_ATTR = "_soup_multipack_seed"


def lengths_from_dataset(
    dataset: Any,
    *,
    key: str = "input_ids",
) -> list[int]:
    """Extract per-sample token lengths from a tokenized dataset.

    Falls back to a ``"length"`` column when ``input_ids`` is absent. A
    sample with neither contributes ``0`` (which the sampler will reject
    upstream — surfaces the misconfiguration loudly).
    """
    if dataset is None:
        return []
    lengths: list[int] = []
    try:
        iterator = iter(dataset)
    except TypeError:
        return []
    for row in iterator:
        if not isinstance(row, dict):
            lengths.append(0)
            continue
        value = row.get(key)
        if value is not None:
            try:
                lengths.append(len(value))
                continue
            except TypeError:
                pass
        fallback = row.get("length")
        if isinstance(fallback, int) and not isinstance(fallback, bool):
            lengths.append(fallback)
        else:
            lengths.append(0)

    # Surface the case where every length is 0 — the multipack sampler
    # would otherwise silently produce phantom batches and the trainer
    # would NaN-loss with no obvious cause. Loud-fail policy mirrors the
    # v0.37.0 multipack-arch allowlist.
    if lengths and not any(lengths):
        logger.warning(
            "lengths_from_dataset(key=%r) returned all zeros — dataset "
            "rows have neither '%s' nor 'length' keys. Multipack sampling "
            "will not work; check tokenisation pipeline.",
            key, key,
        )
    return lengths


def detect_arch_name(model: Any) -> Optional[str]:
    """Return the model's architecture class name, or ``None``.

    Probes ``model.config.architectures[0]`` first (HF convention), then
    falls back to ``type(model).__name__``.
    """
    if model is None:
        return None
    config = getattr(model, "config", None)
    if config is not None:
        archs = getattr(config, "architectures", None)
        if archs:
            try:
                first = archs[0]
            except (IndexError, TypeError):
                first = None
            if isinstance(first, str) and first:
                return first
    cls = type(model).__name__
    return cls or None


@functools.lru_cache(maxsize=None)
def make_multipack_trainer_class(base_cls: type) -> type:
    """Return a subclass of ``base_cls`` overriding ``_get_train_sampler``.

    The override returns a :class:`MultipackBatchSampler` built from
    instance attrs set by :func:`attach_multipack_state`. If state is
    missing the override delegates to the base implementation — so the
    subclass is safe to instantiate even when multipack is later disabled.

    Cached via ``functools.lru_cache`` so two calls with the same
    ``base_cls`` return the SAME subclass — this keeps ``isinstance``
    checks consistent across sweep runs and avoids confusing pickle.

    .. note::
       v0.40.4 (#65) wires this factory into the SFT / Pretrain trainer
       wrappers via a ``get_train_dataloader`` override. HF Trainer's
       ``_get_train_sampler`` returns a ``Sampler[int]`` which the
       DataLoader then consumes as scalar indices, while
       :class:`MultipackBatchSampler` yields ``list[list[int]]``. The
       solution is to bypass ``_get_train_sampler`` for the live path and
       install the multipack sampler as the DataLoader's
       ``batch_sampler=``. The ``_get_train_sampler`` override stays as a
       defensive no-op fallback (delegates to super) so the subclass
       remains safe to instantiate even when state was never attached.
    """
    from soup_cli.utils.multipack_sampler import MultipackBatchSampler

    class MultipackTrainer(base_cls):  # type: ignore[misc, valid-type]
        soup_multipack: bool = True

        def _get_train_sampler(self, *args: Any, **kwargs: Any) -> Any:  # type: ignore[override]
            # v0.40.4 #65 — the live multipack path goes through
            # ``get_train_dataloader`` (below), which installs the
            # multipack ``MultipackBatchSampler`` directly as the
            # DataLoader's ``batch_sampler=``. This override stays as a
            # defensive no-op fallback that delegates to the base
            # implementation: HF Trainer's DataLoader iterates a
            # ``Sampler[int]`` (scalar indices) — returning a multipack
            # ``list[list[int]]`` here would be a shape mismatch if any
            # eval / prediction loop ever bypasses
            # ``get_train_dataloader`` and calls this directly.
            # ``*args/**kwargs`` accept the HF >=4.41 signature
            # (``train_dataset`` passed positionally).
            return super()._get_train_sampler(*args, **kwargs)

        def get_train_dataloader(self):  # type: ignore[override]
            """Return a DataLoader whose batch_sampler is multipack-aware.

            v0.40.4 #65 — when multipack state has been attached via
            :func:`attach_multipack_state`, build a flat-yield
            ``MultipackBatchSampler`` (``real_batches=False`` — yields
            ``list[int]`` per packed sequence, which is the contract HF
            ``DataLoader.batch_sampler`` expects). When state is missing,
            delegate to ``super().get_train_dataloader()`` so the subclass
            is safe to instantiate even when multipack is disabled at
            runtime (matches the ``_get_train_sampler`` fallback policy).
            """
            lengths = getattr(self, _LENGTHS_ATTR, None)
            max_seq = getattr(self, _MAX_SEQ_ATTR, None)
            batch_size = getattr(self, _BATCH_SIZE_ATTR, None)
            seed = getattr(self, _SEED_ATTR, 0)
            # `attach_multipack_state` rejects non-positive values up front;
            # the explicit None check below preserves the "state never
            # attached" path while letting an empty list (lengths=[]) fall
            # through to the same fallback (the sampler would reject it).
            if (
                lengths is None or not lengths
                or max_seq is None or batch_size is None
            ):
                return super().get_train_dataloader()

            from torch.utils.data import DataLoader

            train_dataset = getattr(self, "train_dataset", None)
            if train_dataset is None:
                # Defensive — HF Trainer always sets this before train()
                # invokes get_train_dataloader, but a user calling the
                # method directly might trip the unset case.
                return super().get_train_dataloader()

            args = getattr(self, "args", None)
            drop_last = False
            num_workers = 0
            pin_memory = False
            if args is not None:
                drop_last = bool(getattr(args, "dataloader_drop_last", False))
                num_workers = getattr(args, "dataloader_num_workers", 0) or 0
                pin_memory = bool(getattr(args, "dataloader_pin_memory", False))

            sampler = MultipackBatchSampler(
                lengths=list(lengths),
                batch_max_len=int(max_seq),
                batch_size=int(batch_size),
                real_batches=False,  # yield list[int] per pack — DataLoader-compatible
                seed=int(seed),
                drop_last=drop_last,
            )

            data_collator = getattr(self, "data_collator", None)
            return DataLoader(
                train_dataset,
                batch_sampler=sampler,
                collate_fn=data_collator,
                num_workers=num_workers,
                pin_memory=pin_memory,
            )

    MultipackTrainer.__name__ = f"Multipack{base_cls.__name__}"
    MultipackTrainer.__qualname__ = MultipackTrainer.__name__
    return MultipackTrainer


def attach_multipack_state(
    trainer: Any,
    *,
    lengths: Sequence[int],
    max_seq_len: int,
    batch_size: int,
    seed: int = 0,
) -> None:
    """Stash the multipack sampler config on ``trainer``."""
    if isinstance(max_seq_len, bool) or not isinstance(max_seq_len, int):
        raise TypeError("max_seq_len must be int")
    if isinstance(batch_size, bool) or not isinstance(batch_size, int):
        raise TypeError("batch_size must be int")
    if isinstance(seed, bool) or not isinstance(seed, int):
        raise TypeError("seed must be int")
    if max_seq_len <= 0:
        raise ValueError(f"max_seq_len must be > 0, got {max_seq_len}")
    if batch_size <= 0:
        raise ValueError(f"batch_size must be > 0, got {batch_size}")
    lengths_list = list(lengths)
    if not lengths_list:
        raise ValueError(
            "lengths must not be empty — empty dataset cannot drive "
            "MultipackBatchSampler. Check the tokenisation pipeline."
        )
    setattr(trainer, _LENGTHS_ATTR, lengths_list)
    setattr(trainer, _MAX_SEQ_ATTR, max_seq_len)
    setattr(trainer, _BATCH_SIZE_ATTR, batch_size)
    setattr(trainer, _SEED_ATTR, seed)