diff --git a/src/soup_cli/commands/train.py b/src/soup_cli/commands/train.py index e8d3a4b..6eace57 100644 --- a/src/soup_cli/commands/train.py +++ b/src/soup_cli/commands/train.py @@ -24,6 +24,118 @@ if TYPE_CHECKING: # pragma: no cover - type hints only, no runtime import console = Console() +# Optimizers the analytical hardware-fit predictor understands (mirror of +# hardware_fit._VALID_OPTIMIZERS); an unknown optimizer maps to the +# highest-state default so the estimate stays conservative. +_HW_FIT_OPTIMIZERS = frozenset({ + "adamw_torch", "adamw_torch_fused", "adafactor", "sgd", + "adamw_bnb_8bit", "paged_adamw_8bit", "lion_8bit", + "lomo", "adalomo", "schedule_free_adamw", +}) + + +def _build_hardware_fit_input(cfg): + """Best-effort ``HardwareFitInput`` from a ``SoupConfig``. + + Returns ``None`` when the run is not statically predictable (batch_size + ``"auto"``, unknown model size, unsupported quant, out-of-range dims), in + which case the caller skips the gate rather than guess. + """ + from soup_cli.utils.gpu import model_size_from_name + from soup_cli.utils.hardware_fit import HardwareFitInput + + tcfg = cfg.training + bs = getattr(tcfg, "batch_size", None) + if not isinstance(bs, int) or isinstance(bs, bool): + return None # "auto" resolves later — can't predict yet + params_b = model_size_from_name(getattr(cfg, "base", "") or "") + if not isinstance(params_b, (int, float)) or params_b <= 0: + return None + seq_len = getattr(cfg.data, "max_length", None) + if not isinstance(seq_len, int) or isinstance(seq_len, bool): + return None + quant = {"none": "none", "4bit": "4bit", "8bit": "8bit"}.get( + str(getattr(tcfg, "quantization", "none") or "none") + ) + if quant is None: + return None + if quant == "4bit": + peft = "qlora" + elif ( + getattr(tcfg, "unfrozen_parameters", None) + or getattr(tcfg, "freeze_layers", None) + or getattr(tcfg, "freeze_ratio", None) + ): + peft = "full" # Spectrum / freeze-based full fine-tuning + else: + peft = "lora" + optimizer = str(getattr(tcfg, "optimizer", "adamw_torch") or "adamw_torch") + if optimizer not in _HW_FIT_OPTIMIZERS: + optimizer = "adamw_torch" + gc = bool(getattr(tcfg, "gradient_checkpointing", False)) + try: + return HardwareFitInput( + params_b=float(params_b), + seq_len=int(seq_len), + batch_size=int(bs), + optimizer=optimizer, + quant=quant, + peft=peft, + gradient_checkpointing=gc, + ) + except (ValueError, TypeError): + return None # dims out of the predictor's supported range + + +def _hardware_fit_preflight(cfg, gpu_info, *, allow_oom_attempt: bool) -> None: + """Refuse (or warn) before launch when the predicted peak VRAM won't fit. + + Skips silently on CPU / when VRAM is unknown / when the run isn't + statically predictable, so CI and small runs are unaffected. Honors the + documented ``--allow-oom-attempt`` opt-out. + """ + total_bytes = 0 + try: + total_bytes = int(gpu_info.get("memory_total_bytes", 0) or 0) + except (AttributeError, TypeError, ValueError): + return + if total_bytes <= 0: + return # no CUDA VRAM to predict against + inp = _build_hardware_fit_input(cfg) + if inp is None: + return + from soup_cli.utils.hardware_fit import VRAM_SAFETY_MARGIN, decide_hardware_fit + + report = decide_hardware_fit(inp, available_vram_gb=total_bytes / 1e9) + if report.ok: + return + b = report.breakdown + tail = ( + "[yellow]--allow-oom-attempt set: launching anyway.[/]" + if allow_oom_attempt + else "Reduce batch_size / max_length, enable gradient_checkpointing or " + "quantization, or pass [bold]--allow-oom-attempt[/] to try anyway." + ) + console.print( + Panel( + f"Predicted peak VRAM [bold]{report.peak_vram_gb:.1f} GB[/] " + f"(+{int(VRAM_SAFETY_MARGIN * 100)}% margin = " + f"{report.required_with_margin_gb:.1f} GB) exceeds " + f"{report.available_vram_gb:.1f} GB available.\n" + f"weights {b.weights_gb:.1f} | optim {b.optimizer_gb:.1f} | " + f"grads {b.gradients_gb:.1f} | activations {b.activations_gb:.1f} " + f"| overhead {b.overhead_gb:.1f} GB\n\n" + tail, + title=( + "[yellow]Hardware-fit warning[/]" + if allow_oom_attempt + else "[bold red]Hardware-fit gate[/]" + ), + border_style="yellow" if allow_oom_attempt else "red", + ) + ) + if not allow_oom_attempt: + raise typer.Exit(1) + def train( config: str = typer.Option( @@ -209,6 +321,14 @@ def train( "training steps. Output: /profiles/.trace.json" ), ), + allow_oom_attempt: bool = typer.Option( + False, + "--allow-oom-attempt", + help=( + "Bypass the analytical hardware-fit VRAM gate and launch even when " + "the run is predicted to run out of GPU memory (opt-out)." + ), + ), diagnose_gate: str = typer.Option( None, "--diagnose-gate", @@ -739,6 +859,8 @@ def train( script_args.extend(["--repro-receipt", repro_receipt]) if profile_run: script_args.append("--profile") + if allow_oom_attempt: + script_args.append("--allow-oom-attempt") if track_energy: script_args.append("--track-energy") if energy_country: @@ -797,6 +919,11 @@ def train( ) cfg.training.quantization = "none" + # Hardware-fit preflight: refuse (unless --allow-oom-attempt) when the + # analytical VRAM predictor says the run won't fit. Skips silently on CPU + # or when the config isn't statically predictable (e.g. batch_size='auto'). + _hardware_fit_preflight(cfg, gpu_info, allow_oom_attempt=allow_oom_attempt) + backend_label = cfg.backend if cfg.backend == "unsloth": backend_label = "unsloth [green](fast mode)[/]" diff --git a/src/soup_cli/utils/mod.py b/src/soup_cli/utils/mod.py index 0def212..f069b24 100644 --- a/src/soup_cli/utils/mod.py +++ b/src/soup_cli/utils/mod.py @@ -148,36 +148,119 @@ def _resolve_hidden_size(model: Any, layers: Any) -> int | None: return None +def _gather_dim1(tensor, idx): + """Gather ``tensor`` along its sequence dim (1) by ``idx`` ``[B, cap]``. + + Handles a leading batch of 1 (broadcast RoPE cos/sin) by expanding first, + and trailing feature dims of any rank. ``[B, T]`` -> ``[B, cap]``; + ``[B, T, D]`` -> ``[B, cap, D]``; ``[1, T, D]`` -> ``[B, cap, D]``. + """ + import torch + + bsz, cap = idx.shape + lead = tensor.shape[0] + extra = tuple(tensor.shape[2:]) + if lead == 1 and bsz > 1: + tensor = tensor.expand(bsz, *tensor.shape[1:]) + view = idx.reshape(bsz, cap, *([1] * len(extra))).expand(bsz, cap, *extra) + return torch.gather(tensor, 1, view) + + +def _gather_block_inputs(hidden_states, topk, args, kwargs): + """Gather the decoder-block inputs for the selected sub-sequence. + + Returns ``(sub_hidden, sub_kwargs)`` when the positional inputs can be + gathered SAFELY, else ``None`` (the caller then falls back to running the + full block + gate-blend). Conservative by design: any positional forward + args, a KV cache, ``use_cache``, or a non-4D-causal attention mask abort the + savings path so a mis-gathered mask can never silently corrupt attention. + """ + import torch + + if args: # older positional signature — don't risk mis-identifying args + return None + if kwargs.get("use_cache"): + return None + if ( + kwargs.get("past_key_value") is not None + or kwargs.get("past_key_values") is not None + ): + return None + + bsz, cap = topk.shape + seq_len = hidden_states.shape[1] + hidden = hidden_states.shape[-1] + new_kwargs = dict(kwargs) + + sub_hidden = _gather_dim1(hidden_states, topk) # [B, cap, H] + if sub_hidden.shape != (bsz, cap, hidden): + return None + + am = kwargs.get("attention_mask") + if am is not None: + # Only a square 4D causal mask [B, h, T, T] can be safely narrowed to + # [B, h, cap, cap] (gather query dim then key dim). Anything else aborts. + if am.dim() != 4 or am.shape[-1] != seq_len or am.shape[-2] != seq_len: + return None + heads = am.shape[1] + qi = topk.reshape(bsz, 1, cap, 1).expand(bsz, heads, cap, seq_len) + am_q = torch.gather(am, 2, qi) # [B, h, cap, T] + ki = topk.reshape(bsz, 1, 1, cap).expand(bsz, heads, cap, cap) + new_kwargs["attention_mask"] = torch.gather(am_q, 3, ki) # [B, h, cap, cap] + + pid = kwargs.get("position_ids") + if pid is not None: + if pid.dim() != 2 or pid.shape[-1] != seq_len: + return None + new_kwargs["position_ids"] = _gather_dim1(pid, topk) + + cache_pos = kwargs.get("cache_position") + if cache_pos is not None: + if cache_pos.dim() != 1 or cache_pos.shape[0] != seq_len: + return None + # [T] -> gather with the first row of indices (cache_position is shared) + new_kwargs["cache_position"] = cache_pos.index_select(0, topk[0]) + + pe = kwargs.get("position_embeddings") + if pe is not None: + if not (isinstance(pe, (tuple, list)) and len(pe) == 2): + return None + gathered = [] + for part in pe: + if ( + not hasattr(part, "dim") + or part.dim() != 3 + or part.shape[1] != seq_len + ): + return None + gathered.append(_gather_dim1(part, topk)) + new_kwargs["position_embeddings"] = tuple(gathered) + + return sub_hidden, new_kwargs + + def _make_mod_forward(original, router, capacity_factor: float): """Wrap a decoder-layer forward to route only the top-k tokens. The router scores every token; the top ``floor(T * capacity_factor)`` tokens - receive the block's residual update (gated by the router weight so the router - learns), the rest keep their input hidden state. Returns the original output - on any shape mismatch (best-effort — never crashes training). + are GATHERED into a shorter sub-sequence, the block runs on ONLY those + tokens (real compute saving), and the gated result is scattered back — + unselected tokens keep their input hidden state unchanged. When the block's + positional inputs can't be safely gathered, it falls back to running the + full block once and gate-blending (correct, but no saving). Never crashes + training (returns the original output on any shape mismatch). """ - def mod_forward(hidden_states, *args, **kwargs): + def _fallback(hidden_states, args, kwargs, router_logits, topk): import torch out = original(hidden_states, *args, **kwargs) new_hidden = out[0] if isinstance(out, tuple) else out if ( - not hasattr(hidden_states, "shape") - or len(hidden_states.shape) != 3 - or not hasattr(new_hidden, "shape") + not hasattr(new_hidden, "shape") or new_hidden.shape != hidden_states.shape ): return out - seq_len = hidden_states.shape[1] - try: - cap = mod_capacity(int(seq_len), capacity_factor) - except (TypeError, ValueError): - return out - if cap >= seq_len: - return out # no routing benefit - router_logits = router(hidden_states).squeeze(-1) # [B, T] - topk = torch.topk(router_logits, k=cap, dim=-1).indices # [B, cap] mask = torch.zeros_like(router_logits) mask.scatter_(1, topk, 1.0) weights = (torch.sigmoid(router_logits) * mask).unsqueeze(-1) @@ -186,6 +269,47 @@ def _make_mod_forward(original, router, capacity_factor: float): return (blended,) + tuple(out[1:]) return blended + def mod_forward(hidden_states, *args, **kwargs): + import torch + + if not hasattr(hidden_states, "shape") or len(hidden_states.shape) != 3: + return original(hidden_states, *args, **kwargs) + seq_len = hidden_states.shape[1] + try: + cap = mod_capacity(int(seq_len), capacity_factor) + except (TypeError, ValueError): + return original(hidden_states, *args, **kwargs) + if cap >= seq_len: + return original(hidden_states, *args, **kwargs) # no routing benefit + + router_logits = router(hidden_states).squeeze(-1) # [B, T] + # Sort the selected indices ascending so the gathered sub-sequence keeps + # causal order for its self-attention + RoPE positions. + topk = torch.topk(router_logits, k=cap, dim=-1).indices # [B, cap] + topk, _ = torch.sort(topk, dim=-1) + + gathered = _gather_block_inputs(hidden_states, topk, args, kwargs) + if gathered is None: + return _fallback(hidden_states, args, kwargs, router_logits, topk) + + sub_hidden, sub_kwargs = gathered + sub_out = original(sub_hidden, **sub_kwargs) # block on cap tokens only + sub_new = sub_out[0] if isinstance(sub_out, tuple) else sub_out + if not hasattr(sub_new, "shape") or sub_new.shape != sub_hidden.shape: + # Unexpected sub-block output — bail to the safe full path. + return _fallback(hidden_states, args, kwargs, router_logits, topk) + + # Learnable depth gate on the selected tokens. + sel_logits = torch.gather(router_logits, 1, topk) # [B, cap] + gate = torch.sigmoid(sel_logits).unsqueeze(-1) # [B, cap, 1] + sub_updated = sub_hidden + gate * (sub_new - sub_hidden) + + idx = topk.unsqueeze(-1).expand(-1, -1, hidden_states.shape[-1]) + result = hidden_states.scatter(1, idx, sub_updated) + if isinstance(sub_out, tuple): + return (result,) + tuple(sub_out[1:]) + return result + mod_forward.__name__ = "mod_forward" return mod_forward diff --git a/src/soup_cli/utils/reward_hack_control.py b/src/soup_cli/utils/reward_hack_control.py index d1ab0a5..4a7a745 100644 --- a/src/soup_cli/utils/reward_hack_control.py +++ b/src/soup_cli/utils/reward_hack_control.py @@ -155,10 +155,11 @@ def combine_signals(signals: Mapping[str, Any], names: Sequence[str]) -> float: def smooth_signal(new: float, window: Sequence[float], *, method: str) -> float: - """Smooth a scalar signal. ``none`` → new; ``ema`` → ``alpha·new + - (1-alpha)·prev`` with alpha = ``_EMA_ALPHA`` = 0.5 (prev = ``window[-1]``, - or ``new`` when the window is empty); ``median`` → median of - ``window + [new]``. + """Smooth a scalar signal. ``none`` → new; ``ema`` → windowed EMA (alpha = + ``_EMA_ALPHA`` = 0.5) folded over the RETAINED window (oldest→newest) then + the new sample, so a larger ``reward_hack_smoothing_window`` incorporates + more history; an empty window returns ``new`` and a 1-element window reduces + to ``alpha·new + (1-alpha)·prev``. ``median`` → median of ``window + [new]``. """ if method not in SMOOTHING_METHODS: raise ValueError( @@ -171,11 +172,14 @@ def smooth_signal(new: float, window: Sequence[float], *, method: str) -> float: if method == "ema": if not win: return fnew - # Standard recursive EMA: alpha weights the NEW sample against the - # previous value (window[-1]). Note: a true EMA is inherently - # window-size-independent, so `reward_hack_smoothing_window` only - # affects the `median` method — see the v0.71.26 known limitation. - return _EMA_ALPHA * fnew + (1.0 - _EMA_ALPHA) * win[-1] + # Windowed EMA: previously only ``win[-1]`` was read, so the retained + # window size had no effect. Fold alpha over the whole window so that + # ``reward_hack_smoothing_window`` genuinely bounds how much history the + # smoother incorporates. + ema = win[0] + for prev in win[1:]: + ema = _EMA_ALPHA * prev + (1.0 - _EMA_ALPHA) * ema + return _EMA_ALPHA * fnew + (1.0 - _EMA_ALPHA) * ema return float(statistics.median(win + [fnew])) diff --git a/tests/test_code_review_deferred.py b/tests/test_code_review_deferred.py new file mode 100644 index 0000000..c560e5e --- /dev/null +++ b/tests/test_code_review_deferred.py @@ -0,0 +1,177 @@ +"""Regression tests for the 3 items deferred in the first code-review pass: +windowed EMA smoothing, the hardware-fit OOM preflight, and real MoD +token-dropping. +""" + +from __future__ import annotations + +import pytest + +from soup_cli.config.loader import load_config_from_string + +# ─────────────────────────── EMA smoothing window ─────────────────────────── + + +def test_ema_smoothing_now_respects_window_size(): + from soup_cli.utils.reward_hack_control import smooth_signal + + # Windowed EMA folds alpha over the whole retained window, so a longer + # window (bounded by reward_hack_smoothing_window) changes the result. + assert smooth_signal(0.4, [0.1, 0.2], method="ema") == pytest.approx(0.275) + assert smooth_signal(0.4, [0.2], method="ema") == pytest.approx(0.3) + assert smooth_signal(1.0, [0.0], method="ema") != smooth_signal( + 1.0, [1.0, 0.0, 0.0], method="ema" + ) + + +# ─────────────────────────── hardware-fit OOM preflight ───────────────────── + +_FIT_YAML = """ +base: meta-llama/Llama-2-7b-hf +task: sft +data: + train: train.jsonl + max_length: 2048 +training: + batch_size: 8 + quantization: none +""" + +_AUTO_BS_YAML = """ +base: meta-llama/Llama-2-7b-hf +task: sft +data: + train: train.jsonl + max_length: 2048 +training: + batch_size: auto + quantization: none +""" + + +def test_build_hardware_fit_input_from_config(): + from soup_cli.commands.train import _build_hardware_fit_input + + inp = _build_hardware_fit_input(load_config_from_string(_FIT_YAML)) + assert inp is not None + assert inp.batch_size == 8 + assert inp.seq_len == 2048 + assert inp.params_b >= 6.0 # a 7B base + + # batch_size="auto" isn't statically predictable -> skip the gate. + assert _build_hardware_fit_input(load_config_from_string(_AUTO_BS_YAML)) is None + + +def test_hardware_fit_preflight_gate_and_optout(): + import typer + + from soup_cli.commands import train as train_mod + + cfg = load_config_from_string(_FIT_YAML) + + # 4 GB can't hold a 7B model -> refuse (exit) by default. + with pytest.raises(typer.Exit): + train_mod._hardware_fit_preflight( + cfg, {"memory_total_bytes": int(4e9)}, allow_oom_attempt=False + ) + + # --allow-oom-attempt -> warn, don't refuse. + train_mod._hardware_fit_preflight( + cfg, {"memory_total_bytes": int(4e9)}, allow_oom_attempt=True + ) + + # No detectable VRAM (CPU / CI) -> skip silently. + train_mod._hardware_fit_preflight( + cfg, {"memory_total_bytes": 0}, allow_oom_attempt=False + ) + + # Plenty of VRAM -> fits, no refuse. + train_mod._hardware_fit_preflight( + cfg, {"memory_total_bytes": int(500e9)}, allow_oom_attempt=False + ) + + +# ─────────────────────────── MoD real token-dropping ──────────────────────── + + +def _mod_pieces(hidden: int, capacity_factor: float): + import torch.nn as nn + + from soup_cli.utils.mod import _make_mod_forward + + class _RecordingLayer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.seqs: list = [] + self.kw: list = [] + self.lin = nn.Linear(hidden, hidden) + + def forward(self, hs, *args, **kwargs): + self.seqs.append(hs.shape[1]) + self.kw.append(kwargs) + return self.lin(hs) # per-token (no attention mixing) — plumbing test + + layer = _RecordingLayer() + router = nn.Linear(hidden, 1, bias=False) + fwd = _make_mod_forward(layer.forward, router, capacity_factor) + return layer, router, fwd + + +def test_mod_forward_gathers_and_saves_compute(): + torch = pytest.importorskip("torch") + + hidden, seq, batch = 4, 8, 2 + torch.manual_seed(0) + layer, router, fwd = _mod_pieces(hidden, 0.5) # cap = 4 + x = torch.randn(batch, seq, hidden) + + cap = 4 + topk = torch.topk(router(x).squeeze(-1), k=cap, dim=-1).indices + topk, _ = torch.sort(topk, dim=-1) + + out = fwd(x) + # The block ran on ONLY the cap tokens — the whole point (real savings). + assert layer.seqs == [cap] + assert out.shape == x.shape + # Unselected tokens pass through unchanged. + selected = torch.zeros(batch, seq, dtype=torch.bool) + selected.scatter_(1, topk, True) + for b in range(batch): + for t in range(seq): + if not selected[b, t]: + assert torch.allclose(out[b, t], x[b, t]) + + +def test_mod_forward_gathers_positional_inputs(): + torch = pytest.importorskip("torch") + + hidden, seq, batch, head_dim, heads = 4, 8, 2, 6, 1 + torch.manual_seed(1) + layer, router, fwd = _mod_pieces(hidden, 0.5) # cap = 4 + x = torch.randn(batch, seq, hidden) + cos = torch.randn(batch, seq, head_dim) + sin = torch.randn(batch, seq, head_dim) + attn = torch.zeros(batch, heads, seq, seq) + + fwd(x, position_embeddings=(cos, sin), attention_mask=attn) + kw = layer.kw[-1] + # RoPE + mask narrowed to the sub-sequence (cap), proving real savings with + # correctly-gathered positional inputs. + assert layer.seqs == [4] + assert kw["position_embeddings"][0].shape == (batch, 4, head_dim) + assert kw["attention_mask"].shape == (batch, heads, 4, 4) + + +def test_mod_forward_falls_back_on_positional_args(): + torch = pytest.importorskip("torch") + + hidden, seq, batch = 4, 8, 2 + torch.manual_seed(2) + layer, router, fwd = _mod_pieces(hidden, 0.5) + x = torch.randn(batch, seq, hidden) + + # A positional forward arg aborts the gather path -> full block (seq == T) + # + gate-blend fallback (correct, no savings). + out = fwd(x, object()) + assert layer.seqs == [seq] + assert out.shape == x.shape diff --git a/tests/test_code_review_medium_low.py b/tests/test_code_review_medium_low.py index c5b76f2..790ecd6 100644 --- a/tests/test_code_review_medium_low.py +++ b/tests/test_code_review_medium_low.py @@ -59,14 +59,17 @@ def test_tool_call_args_subset_penalizes_hallucinated_args(): assert "0.5 if not out_args else 0.5" not in src -def test_median_smoothing_uses_window(): +def test_ema_and_median_use_window_size(): from soup_cli.utils.reward_hack_control import smooth_signal - # `median` genuinely uses the retained window (smoothing_window has effect - # here). EMA is recursive/window-independent by design — see the v0.71.26 - # known-limitation note; its 2-tap form is asserted by test_v07126. + # Windowed EMA: a longer retained window folds in more history, so the + # result differs from the 1-element (2-tap) case — proving + # reward_hack_smoothing_window now has effect for EMA. + short = smooth_signal(1.0, [0.0], method="ema") + longer = smooth_signal(1.0, [1.0, 0.0, 0.0], method="ema") + assert short != longer + # median genuinely uses the retained window too. assert smooth_signal(10.0, [1.0, 2.0], method="median") == 2.0 - assert smooth_signal(1.0, [0.0], method="ema") == 0.5 def test_sse_metric_push_preserves_zero(): diff --git a/tests/test_v07126.py b/tests/test_v07126.py index 3de6fcb..9430191 100644 --- a/tests/test_v07126.py +++ b/tests/test_v07126.py @@ -343,8 +343,13 @@ class TestSmoothSignal: def test_ema(self): from soup_cli.utils.reward_hack_control import smooth_signal - # 0.5*prev + 0.5*new; prev = window[-1] = 0.2 - assert smooth_signal(0.4, [0.1, 0.2], method="ema") == pytest.approx(0.3) + # Windowed EMA folds alpha over the whole retained window (oldest first) + # then the new sample, so smoothing_window has effect: + # ema0 = 0.1 -> ema1 = 0.5*0.2 + 0.5*0.1 = 0.15 + # out = 0.5*0.4 + 0.5*0.15 = 0.275 + assert smooth_signal(0.4, [0.1, 0.2], method="ema") == pytest.approx(0.275) + # A 1-element window reduces to the 2-tap form (0.5*0.4 + 0.5*0.2 = 0.3). + assert smooth_signal(0.4, [0.2], method="ema") == pytest.approx(0.3) def test_ema_empty_window_returns_new(self): from soup_cli.utils.reward_hack_control import smooth_signal