mirror of https://github.com/razor-ai/soup.git
236 lines
8.1 KiB
Python
236 lines
8.1 KiB
Python
"""Per-module LR groups — v0.41.0 Part B.
|
|
|
|
Maps a list of (pattern, lr) entries onto a model's named parameters and
|
|
produces a list of optimizer parameter groups suitable for any
|
|
``torch.optim.Optimizer`` constructor.
|
|
|
|
Schema validation:
|
|
- ``lr_groups`` is a list of (pattern, lr) pairs (or dict alias).
|
|
- Capped at ``MAX_LR_GROUPS=32`` entries.
|
|
- Pattern: non-empty string, ≤256 chars, no null byte.
|
|
- LR: float in (0, 1.0].
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import re
|
|
from dataclasses import dataclass
|
|
from typing import Any, Iterable, List, Optional, Set, Tuple
|
|
|
|
MAX_LR_GROUPS = 32
|
|
_MAX_PATTERN_LEN = 256
|
|
_LR_LOWER_EXCLUSIVE = 0.0
|
|
_LR_UPPER_INCLUSIVE = 1.0
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LrGroup:
|
|
"""A single (regex-pattern, lr) override entry.
|
|
|
|
``pattern`` is matched against the parameter's *fully qualified* name
|
|
(e.g. ``model.layers.0.self_attn.q_proj.weight``) using ``re.search``.
|
|
The first matching group wins; remaining params fall through to the
|
|
base LR.
|
|
"""
|
|
|
|
pattern: str
|
|
lr: float
|
|
|
|
|
|
def parse_lr_groups(value: object) -> Optional[List[LrGroup]]:
|
|
"""Parse + validate the raw schema value.
|
|
|
|
Accepts:
|
|
- ``None`` → returns None (feature off).
|
|
- List of dicts ``[{pattern: str, lr: float}, ...]``.
|
|
- List of (pattern, lr) pairs (lists or tuples).
|
|
- Mapping ``{pattern: lr, ...}`` — order preserved (dict insertion).
|
|
|
|
Empty list/dict returns None (feature off, no-op).
|
|
"""
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, dict):
|
|
items: List[Tuple[Any, Any]] = list(value.items())
|
|
elif isinstance(value, list):
|
|
items = []
|
|
for entry in value:
|
|
if isinstance(entry, dict):
|
|
if set(entry.keys()) != {"pattern", "lr"}:
|
|
raise ValueError(
|
|
f"lr_groups dict entries must have exactly "
|
|
f"{{'pattern', 'lr'}} keys, got {sorted(entry.keys())}"
|
|
)
|
|
items.append((entry["pattern"], entry["lr"]))
|
|
elif isinstance(entry, (list, tuple)):
|
|
if len(entry) != 2:
|
|
raise ValueError(
|
|
"lr_groups pair entries must be (pattern, lr), "
|
|
f"got {len(entry)} elements"
|
|
)
|
|
items.append((entry[0], entry[1]))
|
|
else:
|
|
raise ValueError(
|
|
"lr_groups list entries must be dicts or (pattern, lr) "
|
|
f"pairs, got {type(entry).__name__}"
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
"lr_groups must be a list of (pattern, lr) pairs or a dict, "
|
|
f"got {type(value).__name__}"
|
|
)
|
|
if not items:
|
|
return None
|
|
if len(items) > MAX_LR_GROUPS:
|
|
raise ValueError(
|
|
f"lr_groups exceeds cap of {MAX_LR_GROUPS}, got {len(items)}"
|
|
)
|
|
seen: Set[str] = set()
|
|
out: List[LrGroup] = []
|
|
for raw_pattern, raw_lr in items:
|
|
pattern = _validate_pattern(raw_pattern)
|
|
if pattern in seen:
|
|
raise ValueError(
|
|
f"lr_groups contains duplicate pattern {pattern!r}"
|
|
)
|
|
seen.add(pattern)
|
|
lr_value = _validate_lr(raw_lr, pattern)
|
|
out.append(LrGroup(pattern=pattern, lr=lr_value))
|
|
return out
|
|
|
|
|
|
def _validate_pattern(raw: object) -> str:
|
|
if not isinstance(raw, str):
|
|
raise ValueError(
|
|
f"lr_groups pattern must be a string, got {type(raw).__name__}"
|
|
)
|
|
if not raw:
|
|
raise ValueError("lr_groups pattern must be non-empty")
|
|
if "\x00" in raw:
|
|
raise ValueError("lr_groups pattern must not contain null bytes")
|
|
if len(raw) > _MAX_PATTERN_LEN:
|
|
raise ValueError(
|
|
f"lr_groups pattern exceeds {_MAX_PATTERN_LEN} chars"
|
|
)
|
|
try:
|
|
compiled = re.compile(raw)
|
|
except re.error as exc:
|
|
raise ValueError(
|
|
f"lr_groups pattern {raw!r} is not a valid regex: {exc}"
|
|
) from None
|
|
# Best-effort ReDoS probe: a 256-char pattern compiled against a
|
|
# 128-char benign sample completes in microseconds for sane regexes.
|
|
# Catastrophic-backtracking patterns like ``(a+)+`` will hang on
|
|
# this synthetic input. We bound the work via signal-free timing
|
|
# (Python's re has no timeout pre-3.11). The probe is a sanity
|
|
# check, not a hard guarantee — the 256-char length cap above is
|
|
# the primary defence.
|
|
try:
|
|
compiled.search("a" * 128)
|
|
except re.error as exc: # pragma: no cover — runtime regex errors
|
|
raise ValueError(
|
|
f"lr_groups pattern {raw!r} failed runtime probe: {exc}"
|
|
) from None
|
|
return raw
|
|
|
|
|
|
def _validate_lr(raw: object, pattern: str) -> float:
|
|
# Accept str forms of floats — PyYAML parses ``1e-4`` (no dot) as a
|
|
# string in many versions; coercing here keeps the YAML surface friendly.
|
|
if isinstance(raw, str):
|
|
try:
|
|
raw = float(raw)
|
|
except (TypeError, ValueError):
|
|
raise ValueError(
|
|
f"lr_groups[{pattern!r}].lr must be a number, "
|
|
f"got string {raw!r}"
|
|
) from None
|
|
if isinstance(raw, bool) or not isinstance(raw, (int, float)):
|
|
raise ValueError(
|
|
f"lr_groups[{pattern!r}].lr must be a number, "
|
|
f"got {type(raw).__name__}"
|
|
)
|
|
lr_value = float(raw)
|
|
if not math.isfinite(lr_value):
|
|
raise ValueError(
|
|
f"lr_groups[{pattern!r}].lr must be finite, got {lr_value}"
|
|
)
|
|
if lr_value <= _LR_LOWER_EXCLUSIVE or lr_value > _LR_UPPER_INCLUSIVE:
|
|
raise ValueError(
|
|
f"lr_groups[{pattern!r}].lr={lr_value} must be in "
|
|
f"({_LR_LOWER_EXCLUSIVE}, {_LR_UPPER_INCLUSIVE}]"
|
|
)
|
|
return lr_value
|
|
|
|
|
|
def lr_groups_from_schema(
|
|
raw: Optional[List[dict]],
|
|
) -> Optional[List[LrGroup]]:
|
|
"""Convert ``TrainingConfig.lr_groups`` (List[Dict]) into
|
|
runtime ``List[LrGroup]`` for ``build_optimizer_param_groups``.
|
|
|
|
The schema stores the canonical ``[{pattern, lr}, ...]`` shape so YAML
|
|
round-trips cleanly via ``model_dump``; trainer code consumes typed
|
|
``LrGroup`` instances. Returns ``None`` for ``None`` / empty input.
|
|
"""
|
|
if not raw:
|
|
return None
|
|
return [LrGroup(pattern=entry["pattern"], lr=float(entry["lr"]))
|
|
for entry in raw]
|
|
|
|
|
|
def build_optimizer_param_groups(
|
|
named_params: Iterable[Tuple[str, Any]],
|
|
base_lr: float,
|
|
lr_groups: Optional[List[LrGroup]],
|
|
) -> List[dict]:
|
|
"""Map ``model.named_parameters()`` onto optimizer param groups.
|
|
|
|
Each parameter is assigned to the *first* matching ``lr_groups`` entry
|
|
(search order = list order). Unmatched parameters land in a final
|
|
base-LR group.
|
|
|
|
Returns a list of dicts ``[{params: [...], lr: float, name: str}, ...]``
|
|
suitable for ``torch.optim.AdamW(group_dicts, ...)``. Empty groups are
|
|
omitted from the output.
|
|
"""
|
|
if not isinstance(base_lr, (int, float)) or isinstance(base_lr, bool):
|
|
raise ValueError(
|
|
f"base_lr must be a number, got {type(base_lr).__name__}"
|
|
)
|
|
if base_lr <= 0:
|
|
raise ValueError(f"base_lr must be > 0, got {base_lr}")
|
|
|
|
materialised = list(named_params)
|
|
if lr_groups is None or not lr_groups:
|
|
return [{"params": [p for _, p in materialised], "lr": float(base_lr),
|
|
"name": "base"}]
|
|
|
|
compiled = [(g, re.compile(g.pattern)) for g in lr_groups]
|
|
buckets: List[List[Any]] = [[] for _ in compiled]
|
|
base_bucket: List[Any] = []
|
|
for pname, param in materialised:
|
|
for idx, (_, regex) in enumerate(compiled):
|
|
if regex.search(pname):
|
|
buckets[idx].append(param)
|
|
break
|
|
else:
|
|
base_bucket.append(param)
|
|
out: List[dict] = []
|
|
for (group, _), bucket in zip(compiled, buckets):
|
|
if not bucket:
|
|
continue
|
|
out.append({
|
|
"params": bucket,
|
|
"lr": float(group.lr),
|
|
"name": f"lr_group:{group.pattern}",
|
|
})
|
|
if base_bucket:
|
|
out.append({
|
|
"params": base_bucket,
|
|
"lr": float(base_lr),
|
|
"name": "base",
|
|
})
|
|
return out
|