diff --git a/soup_cli/trainer/sft.py b/soup_cli/trainer/sft.py index 00d862c..4f29eeb 100644 --- a/soup_cli/trainer/sft.py +++ b/soup_cli/trainer/sft.py @@ -306,22 +306,39 @@ class SFTTrainerWrapper: "packed docs cannot attend across boundaries" ) - # v0.40.3: #65 multipack live wiring DEFERRED to v0.40.4 — adversarial - # review surfaced that HF Trainer's `_get_train_sampler` returns a - # `Sampler[int]` (DataLoader iterates ints), but `MultipackBatchSampler` - # yields `list[list[int]]`. A `get_train_dataloader` override (with - # `batch_sampler=...`) is required and lands next patch. The - # multipack_trainer.py helpers remain in the codebase as a stub used - # by unit tests — the schema gate keeps `multipack: true` from - # silently no-opping (rejected at config-load for unsupported tasks). - if getattr(tcfg, "multipack", False): - console.print( - "[yellow]multipack: live HF Trainer wiring is deferred to " - "v0.40.4 (DataLoader sampler-vs-batch-sampler mismatch in HF " - "Trainer). Falling back to the standard sampler for this " - "run.[/]" + # v0.40.4 #65 — multipack live wiring. ``make_multipack_trainer_class`` + # mixes a ``get_train_dataloader`` override into the SFTTrainer MRO + # that returns a DataLoader whose ``batch_sampler`` is the FFD + # bin-packing :class:`MultipackBatchSampler`. The factory is cached + # so two ``multipack: true`` runs against the same base class share + # the same subclass. + use_multipack = bool(getattr(tcfg, "multipack", False)) + if use_multipack: + from soup_cli.utils.multipack_sampler import ( + validate_multipack_architecture, ) - self.trainer = SFTTrainer(**trainer_kwargs) + from soup_cli.utils.multipack_trainer import ( + attach_multipack_state, + detect_arch_name, + lengths_from_dataset, + make_multipack_trainer_class, + ) + + arch = detect_arch_name(self.model) + if arch: + validate_multipack_architecture(arch) + trainer_cls = make_multipack_trainer_class(SFTTrainer) + self.trainer = trainer_cls(**trainer_kwargs) + attach_multipack_state( + self.trainer, + lengths=lengths_from_dataset(train_ds), + max_seq_len=cfg.data.max_length, + batch_size=batch_size, + seed=getattr(tcfg, "seed", 0) or 0, + ) + console.print("[green]Multipack FFD bin-packing sampler enabled[/]") + else: + self.trainer = SFTTrainer(**trainer_kwargs) self._output_dir = str(output_dir) self._batch_size = batch_size diff --git a/soup_cli/utils/multipack_trainer.py b/soup_cli/utils/multipack_trainer.py index 8b4fd20..73b1d01 100644 --- a/soup_cli/utils/multipack_trainer.py +++ b/soup_cli/utils/multipack_trainer.py @@ -125,15 +125,16 @@ def make_multipack_trainer_class(base_cls: type) -> type: 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. + 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 @@ -141,21 +142,80 @@ def make_multipack_trainer_class(base_cls: type) -> 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). + # 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) - if not lengths or not max_seq or not batch_size: - return super()._get_train_sampler(*args, **kwargs) - return MultipackBatchSampler( + # `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=True, + real_batches=False, # yield list[int] per pack — DataLoader-compatible seed=int(seed), - drop_last=False, + 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__}" diff --git a/tests/test_v0403_part_b.py b/tests/test_v0403_part_b.py index 049fb35..c5ee328 100644 --- a/tests/test_v0403_part_b.py +++ b/tests/test_v0403_part_b.py @@ -132,7 +132,13 @@ class TestMakeMultipackTrainerClass: assert result[1] == ("some_dataset",) assert result[2] == {"flag": True} - def test_returns_multipack_sampler_when_state_set(self): + def test_get_train_sampler_delegates_to_super_even_with_state(self): + # v0.40.4 #65 — the multipack path goes through + # ``get_train_dataloader`` (real_batches=False, batch_sampler= on + # DataLoader). ``_get_train_sampler`` stays as a defensive no-op + # fallback that ALWAYS delegates to super, even when state is + # attached, so eval / prediction loops that call it directly get + # the correct ``Sampler[int]`` shape (no nested-list shape mismatch). class Base: def __init__(self): pass @@ -149,11 +155,7 @@ class TestMakeMultipackTrainerClass: batch_size=2, seed=42, ) - sampler = instance._get_train_sampler() - # Must be a MultipackBatchSampler (not the string default). - from soup_cli.utils.multipack_sampler import MultipackBatchSampler - - assert isinstance(sampler, MultipackBatchSampler) + assert instance._get_train_sampler() == "default-sampler" class TestAttachMultipackState: @@ -217,25 +219,25 @@ class TestAttachMultipackState: ) -class TestSftAndPretrainWiringDeferred: - """v0.40.3 deferred #65 live wiring after the adversarial review surfaced - a HF Trainer DataLoader sampler-vs-batch-sampler shape mismatch. The SFT - and Pretrain wrappers print a yellow advisory and fall back to the - standard sampler when ``multipack: true``. Live wiring lands in v0.40.4. +class TestSftAndPretrainWiringLive: + """v0.40.4 #65 — live multipack HF Trainer wiring landed. The SFT and + Pretrain wrappers now instantiate the multipack subclass via + ``make_multipack_trainer_class(SFTTrainer)`` and call + ``attach_multipack_state`` when ``multipack: true``. """ - def test_sft_emits_deferred_advisory(self): + def test_sft_wires_live_factory(self): text = Path("soup_cli/trainer/sft.py").read_text(encoding="utf-8") - # Case-insensitive — comment uses DEFERRED, console string uses deferred. - assert "v0.40.4" in text and "deferred" in text.lower() - # The active wiring (factory call) MUST NOT appear in the SFT path - # — the multipack subclass is not built or instantiated. - assert "make_multipack_trainer_class(SFTTrainer)" not in text + assert "v0.40.4" in text + assert "make_multipack_trainer_class(SFTTrainer)" in text + # The deferred advisory must be GONE (no fallback in v0.40.4+). + assert "live HF Trainer wiring is deferred to" not in text - def test_pretrain_emits_deferred_advisory(self): + def test_pretrain_wires_live_factory(self): text = Path("soup_cli/trainer/pretrain.py").read_text(encoding="utf-8") - assert "v0.40.4" in text and "deferred" in text.lower() - assert "make_multipack_trainer_class(SFTTrainer)" not in text + assert "v0.40.4" in text + assert "make_multipack_trainer_class(SFTTrainer)" in text + assert "live HF Trainer wiring is deferred to" not in text class TestSamplerRespectsArchitectureAllowlist: diff --git a/tests/test_v0404_part_b.py b/tests/test_v0404_part_b.py new file mode 100644 index 0000000..7b2701f --- /dev/null +++ b/tests/test_v0404_part_b.py @@ -0,0 +1,179 @@ +"""Tests for v0.40.4 Part B — #65 multipack live HF Trainer wiring. + +The v0.40.3 release shipped the ``make_multipack_trainer_class`` factory +plus ``attach_multipack_state``, ``lengths_from_dataset``, and +``detect_arch_name``, but did not wire them into the SFT / Pretrain +trainer wrappers — adversarial review surfaced that HF Trainer's +DataLoader expects ``Sampler[int]`` while ``MultipackBatchSampler`` +yields ``list[list[int]]``. v0.40.4 fixes that by adding a +``get_train_dataloader`` override that installs the sampler as the +DataLoader's ``batch_sampler=`` kwarg. +""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from soup_cli.utils.multipack_trainer import ( + attach_multipack_state, + make_multipack_trainer_class, +) + + +class TestGetTrainDataloaderOverrideExists: + """The new override is part of the dynamically-generated subclass.""" + + def test_subclass_has_method(self): + class Base: + def get_train_dataloader(self): + return "default-dl" + + sub = make_multipack_trainer_class(Base) + # The override is defined on the subclass itself. + assert "get_train_dataloader" in sub.__dict__ + + def test_falls_back_when_state_missing(self): + class Base: + def __init__(self): + pass + + def get_train_dataloader(self): + return "fallback-dl" + + sub = make_multipack_trainer_class(Base) + instance = sub() + # No state attached → defer to super(). + assert instance.get_train_dataloader() == "fallback-dl" + + def test_falls_back_when_train_dataset_missing(self): + class Base: + def __init__(self): + pass + + def get_train_dataloader(self): + return "fallback-no-ds" + + sub = make_multipack_trainer_class(Base) + instance = sub() + # State is set, but ``train_dataset`` attr is missing. + attach_multipack_state( + instance, lengths=[3, 4, 5], max_seq_len=64, batch_size=2, + ) + assert instance.get_train_dataloader() == "fallback-no-ds" + + +class TestGetTrainDataloaderReturnsMultipackDataloader: + """When state is attached and ``train_dataset`` is set, the override + returns a ``DataLoader`` whose ``batch_sampler`` is the multipack + sampler (NOT the standard scalar-Sampler from HF Trainer). + """ + + def test_returns_multipack_batch_sampler(self): + try: + from torch.utils.data import DataLoader, Dataset + except ImportError: + pytest.skip("torch not installed") + + class TinyDataset(Dataset): + def __init__(self, n): + self.n = n + + def __len__(self): + return self.n + + def __getitem__(self, idx): + return {"x": idx} + + class Base: + def __init__(self): + self.train_dataset = TinyDataset(10) + self.data_collator = None + self.args = MagicMock( + dataloader_num_workers=0, dataloader_pin_memory=False, + ) + + def get_train_dataloader(self): + return "should-not-be-called" + + sub = make_multipack_trainer_class(Base) + instance = sub() + attach_multipack_state( + instance, + lengths=[10, 20, 30, 40, 50, 60, 70, 80, 90, 100], + max_seq_len=128, + batch_size=2, + seed=42, + ) + + dl = instance.get_train_dataloader() + assert isinstance(dl, DataLoader) + # The batch_sampler is a MultipackBatchSampler with real_batches=False. + from soup_cli.utils.multipack_sampler import MultipackBatchSampler + + assert isinstance(dl.batch_sampler, MultipackBatchSampler) + # Sanity: real_batches=False → yields list[int] per pack + # (NOT list[list[int]]). + first_pack = next(iter(dl.batch_sampler)) + assert isinstance(first_pack, list) + assert all(isinstance(x, int) for x in first_pack) + + +class TestGetTrainDataloaderForwardsDropLast: + """v0.40.4 H3 — `args.dataloader_drop_last` must reach the + MultipackBatchSampler; v0.40.4 first-cut hardcoded `drop_last=False`. + """ + + def test_drop_last_forwarded_to_sampler_constructor(self): + # Source-level proof: the override reads ``dataloader_drop_last`` + # from ``self.args`` and passes it as ``drop_last=`` to + # MultipackBatchSampler. (Live-spy patching is hard because the + # factory function captures the symbol via free-variable closure + # at definition time.) + text = Path("soup_cli/utils/multipack_trainer.py").read_text( + encoding="utf-8", + ) + # The override block reads dataloader_drop_last from args. + assert 'getattr(args, "dataloader_drop_last"' in text + # And the sampler constructor receives it as `drop_last=drop_last`. + assert "drop_last=drop_last" in text + + +class TestSftLiveWiring: + """Source-level proof that sft.py instantiates the multipack subclass.""" + + def test_sft_instantiates_subclass(self): + text = Path("soup_cli/trainer/sft.py").read_text(encoding="utf-8") + # The v0.40.3 yellow advisory string is GONE. + assert "live HF Trainer wiring is deferred" not in text + # The factory is invoked with SFTTrainer as the base. + assert "make_multipack_trainer_class(SFTTrainer)" in text + # State is attached. + assert "attach_multipack_state(" in text + # Architecture allowlist is consulted. + assert "validate_multipack_architecture" in text + + def test_pretrain_instantiates_subclass(self): + text = Path("soup_cli/trainer/pretrain.py").read_text(encoding="utf-8") + assert "live HF Trainer wiring is deferred" not in text + assert "make_multipack_trainer_class(SFTTrainer)" in text + assert "attach_multipack_state(" in text + + +class TestRealTrainerSubclassHasOverride: + """If transformers is installed, mix the override into the real + Trainer MRO and confirm the new ``get_train_dataloader`` method + shadows the parent. + """ + + def test_real_trainer_get_train_dataloader_in_dict(self): + try: + from transformers import Trainer + except ImportError: + pytest.skip("transformers not installed") + sub = make_multipack_trainer_class(Trainer) + # Our override is on the subclass itself, not inherited. + assert "get_train_dataloader" in sub.__dict__ + assert sub.__dict__["get_train_dataloader"] is not Trainer.get_train_dataloader