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.3 ships this factory but **does not** wire it into the SFT /
       Pretrain trainer wrappers. Adversarial review surfaced that 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]]``. Live
       wiring requires a ``get_train_dataloader`` override (with the
       sampler installed as ``batch_sampler=`` on the underlying
       ``DataLoader``) and lands in v0.40.4. The factory remains in code
       as the stub end-point used by unit tests.
    """
    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]
            # *args/**kwargs accept newer HF signature
            # (transformers >=4.41 passes train_dataset as positional kwarg).
            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)
            if not lengths or not max_seq or not batch_size:
                return super()._get_train_sampler(*args, **kwargs)
            return MultipackBatchSampler(
                lengths=list(lengths),
                batch_max_len=int(max_seq),
                batch_size=int(batch_size),
                real_batches=True,
                seed=int(seed),
                drop_last=False,
            )

    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)