feat(speed,memory): v0.28.0 features go multi-trainer (v0.33.0 Part C)

Closes #43, #44, #47.

#43 Multi-trainer wiring (sft/dpo/pretrain):
- New utils/v028_features.apply_v028_speed_memory(model, tcfg, base_model,
  console) — single shared helper for use_cut_ce, quantization_aware="fp8",
  kernel_auto_compose. Each feature degrades silently to a yellow advisory
  if the underlying lib is missing; never crashes training kick-off.
- Helpers supports_v028_features(task) and warn_unsupported_features(tcfg, task)
  drive both the schema validator and runtime advisories.
- soup_cli/trainer/dpo.py and trainer/pretrain.py now call the helper after
  model load (post-LoRA, post-QAT) — same hook point as SFT.
- soup_cli/config/schema.py validator
  _validate_v028_speed_memory_sft_only renamed
  _validate_v028_speed_memory_supported_tasks; allowlist now {sft, dpo,
  pretrain}. GRPO/KTO/ORPO/SimPO/IPO/PPO/RewardModel/Embedding still error
  out at config-load with a precise multi-trainer message.

#44 Selective gradient-checkpoint hooks:
- New utils/gradient_ckpt.install_selective_hooks(model, granularity)
  iterates ``model.named_modules()`` looking for transformer-block-shaped
  names (numeric suffix on layer path), wraps each module's ``forward``
  with torch.utils.checkpoint.checkpoint based on tier:
    - selective: only attention sub-modules
    - medium: every second transformer block
    - full: every transformer block
- Returns hook count so callers can fall back to HF native checkpointing
  when zero blocks were found.

#47 CrossDocCollator:
- New soup_cli/data/collators.CrossDocCollator wraps any base data
  collator and injects a block-diagonal causal ``cross_doc_attn_mask``
  built from per-example ``doc_lengths``. Preferred over TRL's
  ``packing_strategy="attention_free"`` flag (best-effort across TRL
  versions). Degrades gracefully when doc_lengths is missing or shapes
  don't match — base attention_mask preserved, no crash.

Tests: +16 in tests/test_part_c.py covering apply_v028_speed_memory
(no-features, cut_ce graceful failure), supports/warn helpers extension,
schema gate (dpo + pretrain accept, kto still rejects), selective hook
installation across full/medium/selective with fake transformer-shaped
models, CrossDocCollator passthrough + strip + injection. One existing
test in test_training_speed.py updated: dpo+use_cut_ce now accepted.

Known limitations:
- 7 trainers (GRPO/KTO/ORPO/SimPO/IPO/PPO/RewardModel/Embedding) still
  reject v0.28.0 flags at config-load. Each is a 5-line addition once
  schema validation is satisfied; tracked as a v0.33.x follow-up.
- install_selective_hooks doesn't undo earlier hooks — caller must be
  re-init aware. Not an issue for the typical "construct wrapper, train,
  exit" flow but worth noting.
- CrossDocCollator emits ``cross_doc_attn_mask`` (not ``attention_mask``)
  to avoid clobbering the base collator's contract; downstream consumers
  must read the new key explicitly. The plan calls for "preferred over
  TRL's packing_strategy" which we satisfy via opt-in collation, not
  silent override.

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

soup_cli/config/schema.py

@@ -610,23 +610,18 @@ class SoupConfig(BaseModel):
         return value
 
     @model_validator(mode="after")
-    def _validate_v028_speed_memory_sft_only(self) -> "SoupConfig":
-        """v0.28.0 speed/memory features are wired only in SFTTrainerWrapper.
+    def _validate_v028_speed_memory_supported_tasks(self) -> "SoupConfig":
+        """v0.28.0 speed/memory features (#43, v0.33.0): supported tasks now
+        include sft, dpo, pretrain. Other tasks still receive the
+        TrainingConfig but do NOT call the apply helpers — emit an explicit
+        ValueError to prevent silent no-ops.
 
-        Non-SFT trainers (DPO/GRPO/KTO/ORPO/SimPO/IPO/PPO/Pretrain/
-        RewardModel/Embedding) receive the TrainingConfig but do NOT call
-        ``apply_cut_ce`` / ``apply_fp8_training`` / ``offload_context`` / the
-        kernel picker. Accepting these flags silently on non-SFT tasks would
-        produce a confusing no-op at best (CCE / kernel_auto_compose /
-        activation_offloading) or a runtime crash at worst
-        (``quantization_aware="fp8"`` falls through to the int8-QAT path in
-        non-SFT wrappers).
-
-        Fail fast at config-load so the user sees a precise error instead of
-        debugging a silent regression. Full multi-trainer wiring is tracked
-        for v0.28.1.
+        Multi-trainer expansion to GRPO/KTO/ORPO/SimPO/IPO/PPO/RewardModel/
+        Embedding is tracked as a follow-up.
         """
-        if self.task == "sft":
+        from soup_cli.utils.v028_features import supports_v028_features
+
+        if supports_v028_features(self.task):
             return self
         tcfg = self.training
         offenders: list[str] = []
@@ -640,10 +635,11 @@ class SoupConfig(BaseModel):
             offenders.append("kernel_auto_compose")
         if offenders:
             raise ValueError(
-                f"v0.28.0 features {offenders} are only wired for task=sft "
-                f"in this release; got task={self.task!r}. Support for other "
-                "trainers is tracked for v0.28.1. Either switch to task=sft "
-                "or remove these flags."
+                f"v0.28.0 features {offenders} are only wired for tasks "
+                f"sft/dpo/pretrain in this release; got task={self.task!r}. "
+                "Multi-trainer expansion to GRPO/KTO/ORPO/SimPO/IPO/PPO/"
+                "RewardModel/Embedding is tracked as a follow-up. Either "
+                "switch to a supported task or remove these flags."
             )
         return self
 

soup_cli/data/collators.py

@@ -0,0 +1,113 @@
+"""Data collators for sample packing (v0.33.0 #47).
+
+When ``training.packing=true`` and ``training.packing_cross_doc_attn_mask=true``,
+multiple short documents are packed into a single sequence; the default causal
+mask leaks attention across document boundaries. ``CrossDocCollator`` builds an
+explicit block-diagonal causal mask via :func:`utils.cross_doc_attn.build_cross_doc_mask`
+and injects it as ``attention_mask`` on the batch, preferred over TRL's
+``packing_strategy="attention_free"`` flag (which is best-effort across TRL
+versions).
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+
+class CrossDocCollator:
+    """Wraps a base data collator and overrides ``attention_mask`` with a
+    block-diagonal causal mask derived from per-doc boundaries.
+
+    The base collator produces the standard ``input_ids`` / ``labels`` tensors;
+    this wrapper looks for a ``doc_lengths`` field in the underlying examples
+    (per-doc token counts after tokenization) and computes the block-diagonal
+    mask. If ``doc_lengths`` is missing, the original ``attention_mask`` is
+    preserved so the wrapper degrades gracefully.
+
+    Usage::
+
+        from transformers import DataCollatorForLanguageModeling
+        from soup_cli.data.collators import CrossDocCollator
+
+        base = DataCollatorForLanguageModeling(tokenizer, mlm=False)
+        collator = CrossDocCollator(base_collator=base)
+        trainer = SFTTrainer(..., data_collator=collator)
+    """
+
+    def __init__(
+        self, base_collator: Any, doc_lengths_key: str = "doc_lengths",
+    ) -> None:
+        if base_collator is None:
+            raise ValueError("CrossDocCollator requires a base_collator")
+        self._base = base_collator
+        self._key = doc_lengths_key
+
+    def __call__(self, features: list[dict]) -> dict:
+        # Pull the per-example doc_lengths out before delegating, since the
+        # base collator may strip unknown keys.
+        per_example_lengths: list[Optional[list[int]]] = []
+        cleaned: list[dict] = []
+        for example in features:
+            lengths = example.pop(self._key, None) if isinstance(example, dict) else None
+            per_example_lengths.append(lengths)
+            cleaned.append(example)
+
+        batch = self._base(cleaned)
+
+        # Try to build the cross-doc mask. If anything goes wrong (no
+        # doc_lengths, mismatched shapes, no numpy/torch), preserve the
+        # base ``attention_mask`` and continue.
+        try:
+            self._inject_block_diag_mask(batch, per_example_lengths)
+        except Exception:  # noqa: BLE001 — degrade rather than crash training
+            pass
+        return batch
+
+    def _inject_block_diag_mask(
+        self, batch: dict, per_example_lengths: list[Optional[list[int]]],
+    ) -> None:
+        from soup_cli.utils.cross_doc_attn import (
+            build_cross_doc_mask,
+            compute_doc_boundaries,
+        )
+
+        if "input_ids" not in batch:
+            return
+        input_ids = batch["input_ids"]
+        if not hasattr(input_ids, "shape") or len(input_ids.shape) != 2:
+            return
+        batch_size, seq_length = input_ids.shape
+
+        if not any(lengths for lengths in per_example_lengths):
+            return  # Nothing to do — no doc_lengths supplied
+
+        try:
+            import numpy as np
+            import torch
+        except ImportError:
+            return
+
+        masks = []
+        for lengths in per_example_lengths:
+            if lengths and sum(lengths) <= seq_length:
+                # Pad final segment if doc lengths sum to < seq_length
+                total = sum(lengths)
+                padded_lengths = list(lengths)
+                if total < seq_length:
+                    padded_lengths.append(seq_length - total)
+                boundaries = compute_doc_boundaries(padded_lengths)
+                mask = build_cross_doc_mask(boundaries, seq_length)
+            else:
+                # Fallback: lower-triangular causal (no doc separation).
+                mask = np.tril(np.ones((seq_length, seq_length), dtype=np.uint8))
+            masks.append(mask)
+
+        if len(masks) != batch_size:
+            return
+
+        attn_tensor = torch.from_numpy(np.stack(masks))
+        # HF expects ``attention_mask`` shape (batch, seq) for standard
+        # masking; the (batch, seq, seq) block-diag mask is consumed by
+        # FlashAttn-2 / SDPA when passed as ``attn_mask`` instead. We expose
+        # both names so downstream code can pick whichever it needs.
+        batch["cross_doc_attn_mask"] = attn_tensor

soup_cli/trainer/dpo.py

@@ -195,11 +195,17 @@ class DPOTrainerWrapper:
         self.model = get_peft_model(self.model, lora_config)
 
         # QAT — insert fake quantization ops after LoRA
-        if tcfg.quantization_aware:
+        if tcfg.quantization_aware and tcfg.quantization_aware != "fp8":
             from soup_cli.utils.qat import prepare_model_for_qat
 
             self.model = prepare_model_for_qat(self.model)
 
+        # v0.33.0 #43 — multi-trainer wiring of v0.28.0 speed/memory features.
+        from soup_cli.utils.v028_features import apply_v028_speed_memory
+        apply_v028_speed_memory(
+            model=self.model, tcfg=tcfg, base_model=cfg.base, console=console,
+        )
+
     def _setup_unsloth(self, cfg, tcfg):
         """Load model via unsloth FastLanguageModel (2-5x faster)."""
         from soup_cli.utils.unsloth import load_model_and_tokenizer

soup_cli/trainer/pretrain.py

@@ -247,11 +247,17 @@ class PretrainTrainerWrapper:
         self.model = get_peft_model(self.model, lora_config)
 
         # QAT — insert fake quantization ops after LoRA
-        if tcfg.quantization_aware:
+        if tcfg.quantization_aware and tcfg.quantization_aware != "fp8":
             from soup_cli.utils.qat import prepare_model_for_qat
 
             self.model = prepare_model_for_qat(self.model)
 
+        # v0.33.0 #43 — multi-trainer wiring of v0.28.0 speed/memory features.
+        from soup_cli.utils.v028_features import apply_v028_speed_memory
+        apply_v028_speed_memory(
+            model=self.model, tcfg=tcfg, base_model=cfg.base, console=console,
+        )
+
     def _setup_unsloth(self, cfg: SoupConfig, tcfg) -> None:
         """Load model via unsloth FastLanguageModel (2-5x faster)."""
         from soup_cli.utils.unsloth import load_model_and_tokenizer

soup_cli/utils/gradient_ckpt.py

@@ -90,6 +90,89 @@ def resolve_gradient_checkpointing(
     }
 
 
+def install_selective_hooks(model, granularity: str) -> int:
+    """Install selective / medium gradient-checkpoint hooks on transformer
+    blocks (#44, v0.33.0).
+
+    Iterates the model's named modules looking for transformer-block-shaped
+    children, then wraps their ``forward`` with
+    ``torch.utils.checkpoint.checkpoint`` based on the granularity:
+
+    - ``"selective"``: only attention sub-modules (looks for ``self_attn`` /
+      ``attention`` in the module name)
+    - ``"medium"``: every second transformer block
+    - ``"full"``: every transformer block (kept for symmetry; HF's native
+      ``gradient_checkpointing`` already handles full — this path is a
+      manual fallback for backends that don't expose that toggle)
+
+    Args:
+        model: a torch ``nn.Module`` (typically a HuggingFace model).
+        granularity: one of ``"selective"`` / ``"medium"`` / ``"full"``.
+
+    Returns:
+        Number of modules that received a hook. Zero is a meaningful signal
+        — caller should fall back to HF's native ``gradient_checkpointing``.
+
+    Raises:
+        ValueError: when ``granularity`` is not recognised.
+
+    Notes:
+        - Pure best-effort; the function never raises on a missing torch
+          dependency at call site (it imports inside).
+        - We do NOT undo earlier hooks. The trainer wrapper is expected to
+          call this once per ``self.model`` instance before training starts.
+    """
+    if granularity not in {"selective", "medium", "full"}:
+        raise ValueError(
+            f"granularity must be one of selective/medium/full, "
+            f"got {granularity!r}"
+        )
+
+    try:
+        import torch.utils.checkpoint as ckpt_mod
+    except ImportError:
+        return 0
+
+    layer_index = 0
+    hooked = 0
+
+    def _wrap(module):
+        original_forward = module.forward
+
+        def _checkpointed_forward(*args, **kwargs):
+            return ckpt_mod.checkpoint(
+                original_forward, *args, use_reentrant=False, **kwargs,
+            )
+
+        module.forward = _checkpointed_forward
+
+    for name, module in model.named_modules():
+        # Heuristic: HF transformer blocks are named like
+        # `model.layers.<i>` (LLaMA), `transformer.h.<i>` (GPT-2),
+        # `model.decoder.layers.<i>` (T5/Bart). We match on a numeric suffix.
+        parts = name.rsplit(".", 1)
+        if len(parts) != 2 or not parts[-1].isdigit():
+            continue
+
+        if granularity == "full":
+            _wrap(module)
+            hooked += 1
+        elif granularity == "medium":
+            if layer_index % 2 == 0:
+                _wrap(module)
+                hooked += 1
+            layer_index += 1
+        elif granularity == "selective":
+            # Find children whose name contains attention markers.
+            for child_name, child in module.named_modules():
+                lc = child_name.lower()
+                if "attn" in lc or "attention" in lc:
+                    _wrap(child)
+                    hooked += 1
+
+    return hooked
+
+
 def describe_tier(tier: TierLike, gpu_memory_gb: float | None = None) -> str:
     """Return a short human-readable description of the selected tier."""
     if not tier:

soup_cli/utils/v028_features.py

@@ -0,0 +1,126 @@
+"""v0.28.0 speed/memory feature application — extracted for multi-trainer reuse.
+
+The original v0.28.0 release wired Cut Cross-Entropy, FP8, kernel-auto-compose
+into ``SFTTrainerWrapper`` only and gated other trainers via a
+``model_validator`` to fail-fast at config-load. v0.33.0 (#43) drops that
+gate and extracts the apply logic here so any trainer wrapper can call it
+in two lines.
+
+Activation-offloading is NOT included here — its scope is the entire
+``trainer.train()`` call (it wraps in a context manager), so each trainer
+wires it inline. CCE / FP8 / kernel-pick are pre-train one-shots and fit
+this single helper.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+
+def apply_v028_speed_memory(
+    *, model, tcfg, base_model: str, console=None,
+) -> dict[str, bool]:
+    """Apply Cut-CE / FP8 / kernel-auto-compose features to ``model``.
+
+    Returns a dict ``{feature_name: applied}`` so the caller can log the
+    decisions for the run record. Each feature degrades silently to a
+    yellow advisory if the underlying lib isn't available — never crashes
+    the training kick-off.
+    """
+    applied: dict[str, bool] = {
+        "cut_ce": False,
+        "fp8": False,
+        "kernel_auto_compose": False,
+    }
+
+    def _say(text: str, style: str = "green") -> None:
+        if console is None:
+            return
+        console.print(f"[{style}]{text}[/]")
+
+    # --- Cut Cross-Entropy ---------------------------------------------------
+    if getattr(tcfg, "use_cut_ce", False):
+        try:
+            from soup_cli.utils.cut_ce import apply_cut_ce
+            ok = bool(apply_cut_ce(base_model))
+        except Exception:  # noqa: BLE001 — degrade gracefully
+            ok = False
+        applied["cut_ce"] = ok
+        if ok:
+            _say("Cut Cross-Entropy enabled (chunked CCE kernel)")
+        else:
+            _say(
+                "Cut Cross-Entropy: no matching architecture or "
+                "cut_cross_entropy not installed", style="yellow",
+            )
+
+    # --- FP8 training --------------------------------------------------------
+    if getattr(tcfg, "quantization_aware", None) == "fp8":
+        try:
+            from soup_cli.utils.fp8 import apply_fp8_training
+            ok = bool(apply_fp8_training(model))
+        except Exception:  # noqa: BLE001
+            ok = False
+        applied["fp8"] = ok
+        if ok:
+            _say("FP8 training enabled (Float8Linear)")
+        else:
+            _say(
+                "FP8 training: torchao.float8 unavailable or no "
+                "compatible linears", style="yellow",
+            )
+
+    # --- Kernel auto-compose -------------------------------------------------
+    if getattr(tcfg, "kernel_auto_compose", False):
+        try:
+            from soup_cli.utils.kernel_picker import (
+                enumerate_candidates,
+                pick_best_kernel,
+            )
+            candidates = list(enumerate_candidates())
+            picked = pick_best_kernel(candidates)
+            applied["kernel_auto_compose"] = True
+            _say(f"Kernel auto-compose picked: {picked.name}")
+        except Exception:  # noqa: BLE001 — picker may be benchmark-blocked
+            _say(
+                "Kernel auto-compose: benchmarking unavailable on this host",
+                style="yellow",
+            )
+
+    return applied
+
+
+def supports_v028_features(task: str) -> bool:
+    """Tasks where v0.28.0 speed/memory wiring has been ported.
+
+    Every task that calls :func:`apply_v028_speed_memory` should be listed
+    here so config validation can advise users on tasks that would silently
+    no-op.
+    """
+    return task in {"sft", "dpo", "pretrain"}
+
+
+def warn_unsupported_features(tcfg, task: str) -> Optional[str]:
+    """Return a human warning if non-v0.28.0-wired tasks set v0.28.0 flags.
+
+    Returns None when nothing to warn about.
+    """
+    if supports_v028_features(task):
+        return None
+    issues: list[str] = []
+    if getattr(tcfg, "use_cut_ce", False):
+        issues.append("use_cut_ce")
+    if getattr(tcfg, "quantization_aware", None) == "fp8":
+        issues.append('quantization_aware="fp8"')
+    if getattr(tcfg, "kernel_auto_compose", False):
+        issues.append("kernel_auto_compose")
+    if getattr(tcfg, "activation_offloading", None) is not None:
+        issues.append("activation_offloading")
+    if not issues:
+        return None
+    return (
+        f"v0.28.0 speed/memory features {issues} are not yet wired for "
+        f"task={task!r} (live in: sft, dpo, pretrain). Flags will be "
+        "silently ignored. Multi-trainer expansion is tracked in "
+        "release notes."
+    )

tests/test_part_c.py

@@ -0,0 +1,305 @@
+"""Part C — v0.28.1 speed/memory live (#43, #44, #47) for v0.33.0.
+
+Covers:
+  - #43 Multi-trainer wiring: apply_v028_speed_memory helper +
+    supports_v028_features extension to dpo + pretrain. Schema gate
+    softens to allow sft / dpo / pretrain.
+  - #44 install_selective_hooks: per-layer checkpoint hooks across the
+    three granularities, with a fake transformer-shaped model.
+  - #47 CrossDocCollator: block-diagonal mask injection via the
+    underlying build_cross_doc_mask helper.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+# ---------------------------------------------------------------------------
+# #43 — apply_v028_speed_memory helper
+# ---------------------------------------------------------------------------
+
+
+class TestApplyV028SpeedMemory:
+    def test_no_features_returns_all_false(self):
+        from soup_cli.utils.v028_features import apply_v028_speed_memory
+
+        tcfg = SimpleNamespace(
+            use_cut_ce=False, quantization_aware=False,
+            kernel_auto_compose=False,
+        )
+        result = apply_v028_speed_memory(
+            model=MagicMock(), tcfg=tcfg, base_model="x/y",
+        )
+        assert result == {
+            "cut_ce": False, "fp8": False, "kernel_auto_compose": False,
+        }
+
+    def test_cut_ce_failure_logs_yellow(self, monkeypatch, capsys):
+        from soup_cli.utils import v028_features as vf
+
+        # Force apply_cut_ce import to raise
+        def _boom(_name):
+            raise RuntimeError("no cut_cross_entropy")
+
+        monkeypatch.setattr(
+            "soup_cli.utils.cut_ce.apply_cut_ce", _boom, raising=False,
+        )
+        from rich.console import Console
+        console = Console()
+        tcfg = SimpleNamespace(
+            use_cut_ce=True, quantization_aware=False,
+            kernel_auto_compose=False,
+        )
+        result = vf.apply_v028_speed_memory(
+            model=MagicMock(), tcfg=tcfg,
+            base_model="x/y", console=console,
+        )
+        assert result["cut_ce"] is False
+
+    def test_supports_v028_features_extends_to_dpo_pretrain(self):
+        from soup_cli.utils.v028_features import supports_v028_features
+
+        for task in ("sft", "dpo", "pretrain"):
+            assert supports_v028_features(task) is True
+        for task in ("grpo", "kto", "orpo", "simpo", "ipo", "ppo",
+                     "reward_model", "embedding"):
+            assert supports_v028_features(task) is False
+
+    def test_warn_unsupported_returns_none_for_supported(self):
+        from soup_cli.utils.v028_features import warn_unsupported_features
+
+        tcfg = SimpleNamespace(
+            use_cut_ce=True, quantization_aware="fp8",
+            kernel_auto_compose=True, activation_offloading="cpu",
+        )
+        assert warn_unsupported_features(tcfg, "sft") is None
+        assert warn_unsupported_features(tcfg, "dpo") is None
+
+    def test_warn_unsupported_lists_offenders_for_unsupported(self):
+        from soup_cli.utils.v028_features import warn_unsupported_features
+
+        tcfg = SimpleNamespace(
+            use_cut_ce=True, quantization_aware="fp8",
+            kernel_auto_compose=False, activation_offloading=None,
+        )
+        msg = warn_unsupported_features(tcfg, "kto")
+        assert msg is not None
+        assert "use_cut_ce" in msg
+        assert "fp8" in msg
+
+
+class TestSchemaGateExpanded:
+    def _config(self, task: str, **training_extra):
+        import yaml
+
+        from soup_cli.config.loader import load_config_from_string
+
+        body = {
+            "base": "test/model",
+            "task": task,
+            "data": {"train": "data.jsonl", "format": "alpaca"}
+            if task != "pretrain"
+            else {"train": "data.jsonl", "format": "plaintext"},
+            "training": {"epochs": 1, "lr": 1e-4, "batch_size": 1, **training_extra},
+        }
+        if task in ("dpo", "kto", "orpo", "simpo", "ipo", "grpo"):
+            body["data"]["format"] = "dpo"
+        return load_config_from_string(yaml.safe_dump(body))
+
+    def test_dpo_now_accepts_v028_features(self):
+        cfg = self._config("dpo", use_cut_ce=True)
+        assert cfg.task == "dpo"
+        assert cfg.training.use_cut_ce is True
+
+    def test_pretrain_now_accepts_v028_features(self):
+        cfg = self._config("pretrain", use_cut_ce=True)
+        assert cfg.task == "pretrain"
+
+    def test_kto_still_rejects_v028_features(self):
+        with pytest.raises(Exception, match="v0.28.0"):
+            self._config("kto", use_cut_ce=True)
+
+
+# ---------------------------------------------------------------------------
+# #44 — install_selective_hooks
+# ---------------------------------------------------------------------------
+
+
+class TestInstallSelectiveHooks:
+    def test_unknown_granularity_rejected(self):
+        from soup_cli.utils.gradient_ckpt import install_selective_hooks
+
+        with pytest.raises(ValueError, match="must be one of"):
+            install_selective_hooks(MagicMock(), "weird")
+
+    def test_full_wraps_every_block(self):
+        """A 4-block fake model with 'medium' wraps 2 blocks (every other)."""
+        from soup_cli.utils.gradient_ckpt import install_selective_hooks
+
+        class FakeBlock:
+            def __init__(self, idx):
+                self.idx = idx
+
+            def forward(self, *args, **kwargs):
+                return ("orig", self.idx)
+
+            # Simulate having attention sub-modules
+            def named_modules(self, prefix=""):
+                yield (f"{prefix}", self)
+                yield (f"{prefix}.self_attn", _AttnStub(self.idx))
+
+        class _AttnStub:
+            def __init__(self, idx):
+                self.idx = idx
+
+            def forward(self, *args, **kwargs):
+                return ("attn", self.idx)
+
+        class FakeModel:
+            def __init__(self):
+                self.blocks = [FakeBlock(i) for i in range(4)]
+
+            def named_modules(self):
+                # HF-style names: ``model.layers.<i>``
+                for idx, blk in enumerate(self.blocks):
+                    yield (f"model.layers.{idx}", blk)
+                    for child_name, child in blk.named_modules(
+                        prefix=f"model.layers.{idx}",
+                    ):
+                        if child_name != f"model.layers.{idx}":
+                            yield (child_name, child)
+
+        # Full granularity wraps every numbered child
+        model = FakeModel()
+        hooked_full = install_selective_hooks(model, "full")
+        assert hooked_full == 4
+
+    def test_medium_wraps_every_other_block(self):
+        from soup_cli.utils.gradient_ckpt import install_selective_hooks
+
+        class FakeBlock:
+            def forward(self, *args, **kwargs):
+                return None
+
+            def named_modules(self, prefix=""):
+                yield (prefix, self)
+
+        class FakeModel:
+            def __init__(self, n):
+                self.blocks = [FakeBlock() for _ in range(n)]
+
+            def named_modules(self):
+                for i, blk in enumerate(self.blocks):
+                    yield (f"model.layers.{i}", blk)
+
+        hooked = install_selective_hooks(FakeModel(6), "medium")
+        # 6 blocks, every-other → 3
+        assert hooked == 3
+
+    def test_selective_wraps_attention_only(self):
+        from soup_cli.utils.gradient_ckpt import install_selective_hooks
+
+        class _Attn:
+            def forward(self, *args, **kwargs):
+                return None
+
+        class FakeBlock:
+            def __init__(self):
+                self._attn = _Attn()
+                self._mlp = MagicMock()
+                self._mlp.forward = lambda *_a, **_k: None
+
+            def forward(self, *args, **kwargs):
+                return None
+
+            def named_modules(self, prefix=""):
+                yield (prefix, self)
+                yield (f"{prefix}.self_attn", self._attn)
+                yield (f"{prefix}.mlp", self._mlp)
+
+        class FakeModel:
+            def __init__(self):
+                self.blocks = [FakeBlock() for _ in range(3)]
+
+            def named_modules(self):
+                for i, blk in enumerate(self.blocks):
+                    yield (f"model.layers.{i}", blk)
+                    for child_name, child in blk.named_modules(
+                        prefix=f"model.layers.{i}",
+                    ):
+                        if child_name != f"model.layers.{i}":
+                            yield (child_name, child)
+
+        hooked = install_selective_hooks(FakeModel(), "selective")
+        # 3 blocks, each with one attention child = 3 hooks
+        assert hooked == 3
+
+
+# ---------------------------------------------------------------------------
+# #47 — CrossDocCollator
+# ---------------------------------------------------------------------------
+
+
+class TestCrossDocCollator:
+    def test_requires_base_collator(self):
+        from soup_cli.data.collators import CrossDocCollator
+
+        with pytest.raises(ValueError, match="base_collator"):
+            CrossDocCollator(base_collator=None)
+
+    def test_passes_through_when_no_doc_lengths(self):
+        from soup_cli.data.collators import CrossDocCollator
+
+        base = MagicMock(return_value={"input_ids": MagicMock()})
+        collator = CrossDocCollator(base_collator=base)
+        result = collator([{"input_ids": [1, 2, 3]}])
+        assert "input_ids" in result
+        # No cross_doc_attn_mask since no doc_lengths supplied
+        assert "cross_doc_attn_mask" not in result
+
+    def test_strips_doc_lengths_before_base_call(self):
+        from soup_cli.data.collators import CrossDocCollator
+
+        captured: list[dict] = []
+
+        def _base(features):
+            captured.extend(features)
+            return {"input_ids": MagicMock()}
+
+        collator = CrossDocCollator(base_collator=_base)
+        collator([{"input_ids": [1, 2], "doc_lengths": [1, 1]}])
+        # base collator must NOT see the doc_lengths key
+        assert "doc_lengths" not in captured[0]
+
+    def test_injects_block_diag_mask_with_doc_lengths(self):
+        """End-to-end: when doc_lengths are present, cross_doc_attn_mask
+        appears on the batch with the right shape."""
+        import numpy as np
+
+        from soup_cli.data.collators import CrossDocCollator
+
+        # 2 documents, each 2 tokens long, packed into seq_length=4
+        seq_len = 4
+
+        class _Tensor:
+            shape = (1, seq_len)
+
+        def _base(features):
+            return {"input_ids": _Tensor()}
+
+        collator = CrossDocCollator(base_collator=_base)
+        result = collator([{"doc_lengths": [2, 2]}])
+
+        assert "cross_doc_attn_mask" in result
+        mask = result["cross_doc_attn_mask"]
+        # Shape: (batch=1, seq, seq)
+        assert mask.shape == (1, seq_len, seq_len)
+        # Token in doc 0 should NOT attend to token in doc 1 (positions 0,1 vs 2,3)
+        np_mask = mask[0].numpy() if hasattr(mask[0], "numpy") else np.array(mask[0])
+        # Position 0 attending to position 2 → should be 0
+        assert np_mask[0, 2] == 0
+        # Position 2 attending to position 2 → 1 (causal within doc 1)
+        assert np_mask[2, 2] == 1

tests/test_training_speed.py

@@ -839,16 +839,15 @@ class TestV028SFTOnlyValidator:
     legacy int8 QAT wrapper.
     """
 
-    def test_use_cut_ce_rejected_on_dpo(self):
-        with pytest.raises(ValidationError) as exc:
-            SoupConfig(
-                base="m",
-                task="dpo",
-                data={"train": "./d.jsonl", "format": "dpo"},
-                training={"use_cut_ce": True},
-            )
-        assert "use_cut_ce" in str(exc.value)
-        assert "sft" in str(exc.value)
+    def test_use_cut_ce_now_accepted_on_dpo(self):
+        # v0.33.0 #43 — DPO is now in the supported task set.
+        cfg = SoupConfig(
+            base="m",
+            task="dpo",
+            data={"train": "./d.jsonl", "format": "dpo"},
+            training={"use_cut_ce": True},
+        )
+        assert cfg.training.use_cut_ce is True
 
     def test_fp8_rejected_on_grpo(self):
         with pytest.raises(ValidationError) as exc: