From 55d1b9312cd07c506ac177dd81eb9048b63fefcd Mon Sep 17 00:00:00 2001 From: Alpamys Date: Mon, 27 Apr 2026 18:55:50 +0500 Subject: [PATCH] feat(speed,memory): v0.28.0 features go multi-trainer (v0.33.0 Part C) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- soup_cli/config/schema.py | 34 ++-- soup_cli/data/collators.py | 113 ++++++++++++ soup_cli/trainer/dpo.py | 8 +- soup_cli/trainer/pretrain.py | 8 +- soup_cli/utils/gradient_ckpt.py | 83 +++++++++ soup_cli/utils/v028_features.py | 126 +++++++++++++ tests/test_part_c.py | 305 ++++++++++++++++++++++++++++++++ tests/test_training_speed.py | 19 +- 8 files changed, 665 insertions(+), 31 deletions(-) create mode 100644 soup_cli/data/collators.py create mode 100644 soup_cli/utils/v028_features.py create mode 100644 tests/test_part_c.py diff --git a/soup_cli/config/schema.py b/soup_cli/config/schema.py index 530baa7..512f8b8 100644 --- a/soup_cli/config/schema.py +++ b/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 diff --git a/soup_cli/data/collators.py b/soup_cli/data/collators.py new file mode 100644 index 0000000..2df1aec --- /dev/null +++ b/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 diff --git a/soup_cli/trainer/dpo.py b/soup_cli/trainer/dpo.py index ec7f22b..306056b 100644 --- a/soup_cli/trainer/dpo.py +++ b/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 diff --git a/soup_cli/trainer/pretrain.py b/soup_cli/trainer/pretrain.py index 4e702ac..c5c4fe6 100644 --- a/soup_cli/trainer/pretrain.py +++ b/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 diff --git a/soup_cli/utils/gradient_ckpt.py b/soup_cli/utils/gradient_ckpt.py index e93cc64..1b34f50 100644 --- a/soup_cli/utils/gradient_ckpt.py +++ b/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.` (LLaMA), `transformer.h.` (GPT-2), + # `model.decoder.layers.` (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: diff --git a/soup_cli/utils/v028_features.py b/soup_cli/utils/v028_features.py new file mode 100644 index 0000000..362581f --- /dev/null +++ b/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." + ) diff --git a/tests/test_part_c.py b/tests/test_part_c.py new file mode 100644 index 0000000..5f89e30 --- /dev/null +++ b/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.`` + 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 diff --git a/tests/test_training_speed.py b/tests/test_training_speed.py index 84d6ee0..abd5303 100644 --- a/tests/test_training_speed.py +++ b/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: