feat(trainer): multipack live HF Trainer DataLoader override (v0.40.4 Part B)

Closes #65 (deferred from v0.40.3).

make_multipack_trainer_class adds a get_train_dataloader override
that builds a MultipackBatchSampler(real_batches=False) — yields a
flat list[int] per packed sequence, which is the contract HF
DataLoader.batch_sampler expects — and installs it via
DataLoader(..., batch_sampler=sampler, collate_fn=self.data_collator,
num_workers=args.dataloader_num_workers,
pin_memory=args.dataloader_pin_memory). drop_last is forwarded from
TrainingArguments.dataloader_drop_last.

Falls back to super().get_train_dataloader() when state was never
attached OR when train_dataset is unset — defence-in-depth so the
subclass remains safe to instantiate even when multipack is later
disabled.

The state-presence guard switched from falsy (`not max_seq`) to
explicit `is None` (plus `not lengths` for empty-list defence) —
attach_multipack_state already rejects non-positive ints, so the
falsy guard would only mask configurator bugs.

_get_train_sampler override stays as a defensive no-op fallback that
ALWAYS delegates to super (review-fix from v0.40.4 code-review:
returning a multipack list[list[int]] from this method would cause a
shape mismatch if any HF eval / prediction loop bypasses
get_train_dataloader and calls _get_train_sampler directly).

SFT and Pretrain trainer wrappers now invoke
make_multipack_trainer_class(SFTTrainer) and attach_multipack_state(...)
when multipack: true. The v0.40.3 yellow advisory + standard-sampler
fallback is gone. Architecture allowlist
(validate_multipack_architecture) still gates at build time.

tests/test_v0403_part_b.py: TestSftAndPretrainWiringDeferred renamed
to TestSftAndPretrainWiringLive; the deferred-state test
(_get_train_sampler returns MultipackBatchSampler when state is set)
is replaced by the live-state test (_get_train_sampler always
delegates to super even with state attached).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

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

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__}"

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:

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