mirror of https://github.com/razor-ai/soup.git
fix: close out the 3 deferred code-review items (EMA, hardware-fit, MoD)
These were documented as known limitations in the MEDIUM/LOW pass; now fixed.
1. reward_hack EMA smoothing window — smooth_signal("ema") folded only
window[-1], so reward_hack_smoothing_window had no effect. Now a windowed
EMA folds alpha over the whole retained window (oldest→newest) then the new
sample, so a larger window incorporates more history; a 1-element window
reduces to the old 2-tap form. Updated test_v07126 (0.3 → 0.275).
2. hardware_fit OOM gate wired into `soup train` — the analytical VRAM
predictor was never called despite its docstring. Added
_build_hardware_fit_input (SoupConfig → HardwareFitInput, best-effort;
None when not statically predictable, e.g. batch_size="auto") and
_hardware_fit_preflight, run after device detection. Refuses on predicted
OOM (peak × 1.1 > available) unless the documented --allow-oom-attempt
opt-out is passed; skips silently on CPU / unknown VRAM, and the flag is
threaded through the --gpus re-exec.
3. MoD real token-dropping — mod_forward ran the full block on ALL tokens then
masked (zero compute savings). Now the top-k tokens are gathered into a
shorter sub-sequence, the block runs on ONLY those tokens (real saving),
the gated result is scattered back, and unselected tokens pass through
unchanged. Positional inputs (RoPE cos/sin, 4D-causal attention_mask,
position_ids, cache_position) are gathered to the sub-sequence; any
unsafe-to-gather case (positional forward args, KV cache, non-4D mask)
falls back to the prior correct blend so attention can never be silently
mis-computed. Validated on CPU (gather/scatter/passthrough/savings +
fallback); the sub-sequence-attention numerics still warrant GPU validation
at scale.
Adds tests/test_code_review_deferred.py (7 tests). ruff clean; full suite
14867 passed / 120 skipped.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Alpamys
parent
defd3151cf
commit
cba51e153b
|
|
@ -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: <output>/profiles/<run_id>.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)[/]"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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]))
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue