soup/soup_cli/utils/v028_features.py

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

import contextlib
from pathlib import Path
from typing import TYPE_CHECKING, Any, Iterator, Optional

if TYPE_CHECKING:
    from rich.console import Console

    from soup_cli.config.schema import TrainingConfig


def apply_v028_speed_memory(
    *,
    model: Any,
    tcfg: "TrainingConfig",
    base_model: str,
    console: Optional["Console"] = None,
    device: str = "cpu",
    backend: str = "transformers",
) -> 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":
        recipe = getattr(tcfg, "fp8_recipe", "tensorwise")
        try:
            from soup_cli.utils.fp8 import apply_fp8_training
            ok = bool(apply_fp8_training(model, recipe=recipe))
        except Exception:  # noqa: BLE001
            ok = False
        applied["fp8"] = ok
        if ok:
            _say(f"FP8 training enabled (Float8Linear, recipe={recipe})")
        else:
            _say(
                "FP8 training: torchao.float8 unavailable or no "
                "compatible linears", style="yellow",
            )

    # --- Kernel auto-compose -------------------------------------------------
    if getattr(tcfg, "kernel_auto_compose", False):
        picked_name = _bench_and_pick_kernel(
            model=model, device=device, backend=backend,
        )
        if picked_name is None:
            _say(
                "Kernel auto-compose: benchmarking unavailable on this host",
                style="yellow",
            )
        else:
            applied["kernel_auto_compose"] = True
            _say(f"Kernel auto-compose picked: {picked_name}")

    return applied


def _bench_and_pick_kernel(
    *, model: Any, device: str, backend: str,
) -> Optional[str]:
    """v0.35.0 #45 — benchmark candidate kernel combos and return the
    picked combo's name. Returns ``None`` on any benchmark / pick failure
    so the caller can degrade gracefully (no kernel_auto_compose flag set).
    """
    try:
        from soup_cli.utils.kernel_picker import (
            benchmark_kernel_combos,
            enumerate_kernel_combos,
            pick_best_kernel,
        )
        candidates = enumerate_kernel_combos(backend=backend, device=device)
        # On CPU / unsloth / mlx the candidate list is just [baseline]; no
        # benchmark needed — picker would raise on all-None times. Skip.
        if len(candidates) <= 1:
            return None
        timed = benchmark_kernel_combos(
            model=model, candidates=candidates, device=device,
        )
        picked = pick_best_kernel(timed)
        # Picker returns either a dict (current shape) or an object with
        # ``.name`` (legacy / namespace shape) — handle both defensively.
        if isinstance(picked, dict):
            return str(picked.get("name", "unknown"))
        return str(getattr(picked, "name", "unknown"))
    except Exception:  # noqa: BLE001 — picker is best-effort
        return None


@contextlib.contextmanager
def activation_offloading_context(
    tcfg: "TrainingConfig", output_dir: str,
) -> Iterator[None]:
    """Wrap a trainer's ``trainer.train()`` call with activation offloading.

    Centralises the cwd-containment guard for ``activation_offloading="disk"``
    (defence-in-depth — caller's ``cfg.output`` is also validated upstream)
    and the ``offload_context`` setup. ``None`` and ``"cpu"`` modes pass
    through to ``offload_context`` directly.

    Raises ``ValueError`` when ``activation_offloading="disk"`` and
    ``output_dir`` is outside the current working directory.
    """
    from soup_cli.utils.activation_offload import offload_context
    from soup_cli.utils.paths import is_under_cwd

    save_dir: Optional[str] = None
    mode = getattr(tcfg, "activation_offloading", None)
    if mode == "disk":
        if not is_under_cwd(output_dir):
            # Reduce to basename so $HOME / absolute paths don't leak (matches
            # the v0.34.0 crash.py policy).
            import os as _os

            raise ValueError(
                "activation_offloading='disk' requires the training output "
                "dir to be under the current working directory; got "
                f"basename={_os.path.basename(output_dir)!r}"
            )
        save_dir = str(Path(output_dir) / "_activation_offload")
    with offload_context(mode, save_dir=save_dir):
        yield


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. v0.35.0 (#60) extends coverage from {sft, dpo, pretrain} to
    every transformer-backend trainer.
    """
    return task in {
        "sft",
        "dpo",
        "pretrain",
        "grpo",
        "kto",
        "orpo",
        "simpo",
        "ipo",
        "bco",
        "preference",
        "ppo",
        "reward_model",
        "embedding",
    }


def warn_unsupported_features(
    tcfg: "TrainingConfig", 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. v0.35.0 #60 expanded coverage
    to every transformer-backend trainer; this helper now only fires for
    truly-unsupported tasks (e.g. a hypothetical future task or an MLX
    backend route).
    """
    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 wired for "
        f"task={task!r}. Flags will be silently ignored. See "
        "supports_v028_features() for the current supported task list."
    )