dragonscale 0.1.0__py3-none-any.whl

dragonscale/__init__.py

@@ -0,0 +1,39 @@
+"""DragonScale: unit-row weights and free-scale analysis.
+
+import dragonscale as ds
+
+report = ds.analysis.free_scales(model, x)   # which scales the loss cannot see
+session = ds.apply(model, "unit_gain")       # every Linear / Conv / Embedding, same function
+optimizer = torch.optim.AdamW(model.parameters())  # gains are module parameters
+hooks = ds.optim.attach(optimizer, session)  # tangent steps, rows kept at the radius
+...
+session.bake()                               # stock modules with the unit-row weights
+"""
+
+from dragonscale import analysis, functional, layer_types, monitor, nn, optim, spec
+from dragonscale.patch import Session, apply
+from dragonscale.plan import PRESETS, LayerPlan, Plan, Rule, make_plan, preset, validate
+from dragonscale.spec import RowSpec, defaults
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "PRESETS",
+    "LayerPlan",
+    "Plan",
+    "RowSpec",
+    "Rule",
+    "Session",
+    "analysis",
+    "apply",
+    "defaults",
+    "functional",
+    "layer_types",
+    "make_plan",
+    "monitor",
+    "nn",
+    "optim",
+    "preset",
+    "spec",
+    "validate",
+]

dragonscale/analysis/__init__.py

@@ -0,0 +1,29 @@
+"""Free-scale analysis: degree propagation over a real forward, the free directions it leaves,
+the scale-Jacobian spectrum, the plan that removes the dangerous ones, and scale moves.
+
+    report = ds.analysis.free_scales(model, x)
+    print(report.describe())
+    plan = ds.analysis.suggest(model, x, level="absorb")
+    session = ds.apply(model, ds.analysis.absorb_plan(model, plan, x))
+"""
+
+from dragonscale.analysis.absorb import absorb, absorb_plan
+from dragonscale.analysis.jacobian import ScaleJacobian, scale_jacobian
+from dragonscale.analysis.report import Direction, FreeScales, free_scales
+from dragonscale.analysis.rules import register_degree_rule
+from dragonscale.analysis.suggest import suggest
+from dragonscale.analysis.trace import DegreeTrace, trace_degrees
+
+__all__ = [
+    "DegreeTrace",
+    "Direction",
+    "FreeScales",
+    "ScaleJacobian",
+    "absorb",
+    "absorb_plan",
+    "free_scales",
+    "register_degree_rule",
+    "scale_jacobian",
+    "suggest",
+    "trace_degrees",
+]

dragonscale/analysis/absorb.py

@@ -0,0 +1,229 @@
+"""`absorb`: move a source's per-channel scale into its consumers' input columns, keeping the
+function; `absorb_plan` makes a plan's moves and absorbs its `init="absorb"` layers.
+
+The source is a module with a layer type (its rows are rescaled to the radius) or a norm with an
+affine (its weight is divided by its magnitude). Consumers are given, or found by
+`find_consumers`: the modules with a layer type whose first input carries the source's degree,
+taken in call order until the move is a free direction. The move must be free on the given inputs,
+and the output must not change.
+"""
+
+from collections.abc import Sequence
+from dataclasses import replace
+
+import torch
+from torch import Tensor, nn
+
+from dragonscale.analysis.report import FreeScales, solve
+from dragonscale.analysis.trace import DegreeTrace, trace_degrees
+from dragonscale.functional.rows import row_norms
+from dragonscale.layer_types import layer_type_of
+from dragonscale.plan.layer_plan import Plan
+from dragonscale.spec.rows import RowSpec
+from dragonscale_common.problems import raise_problems
+from dragonscale_kernels.layout import as_rows
+from dragonscale_kernels.ops import FLOOR
+
+
+def _source(
+    module: nn.Module, spec: RowSpec | None
+) -> tuple[Tensor, Tensor | None, RowSpec | None]:
+    """The source weight, its bias, and its row spec (None for a norm). A layer's spec must index
+    rows by output channel: rows along `out_axis`, or the layer type's own default."""
+    layer_type = layer_type_of(module)
+    if layer_type is not None:
+        default = layer_type.spec(module)
+        spec = spec or default
+        plain = spec.view is None and tuple(spec.row_dims) == (layer_type.out_axis,)
+        same = (tuple(spec.row_dims), spec.view) == (tuple(default.row_dims), default.view)
+        if not (plain or same):
+            raise ValueError(f"{layer_type.name}: absorb needs rows along the output channels")
+        bias = getattr(module, layer_type.bias, None) if layer_type.bias else None
+        return getattr(module, layer_type.weight), bias, spec
+    weight = getattr(module, "weight", None)
+    if not isinstance(weight, Tensor) or weight.dim() != 1:
+        raise ValueError(f"{type(module).__name__} is neither a layer type nor an affine norm")
+    return weight, getattr(module, "bias", None), None
+
+
+def _outputs(model: nn.Module, inputs, kwargs) -> list[Tensor]:
+    out = model(*inputs, **kwargs)
+    flat = (
+        out.values()
+        if isinstance(out, dict)
+        else (out if isinstance(out, tuple | list) else (out,))
+    )
+    return [t.detach().clone() for t in flat if isinstance(t, Tensor) and t.is_floating_point()]
+
+
+def _input_weight(module: nn.Module, channels: int) -> tuple[Tensor, int] | None:
+    """The weight of a module with an input-channel axis of size `channels`, and that axis."""
+    layer_type = layer_type_of(module)
+    if layer_type is None or layer_type.in_axis is None:
+        return None
+    weight = getattr(module, layer_type.weight)
+    return (weight, layer_type.in_axis) if weight.shape[layer_type.in_axis] == channels else None
+
+
+def _coeffs(params: dict[int, str], weight: Tensor, bias, consumers) -> dict[str, float]:
+    """Source weight (and bias) up, every consumer weight down."""
+    coeffs = {params[id(weight)]: 1.0}
+    if isinstance(bias, Tensor):
+        coeffs[params[id(bias)]] = 1.0
+    for cw, _ in consumers:
+        coeffs[params[id(cw)]] = coeffs.get(params[id(cw)], 0.0) - 1.0
+    return coeffs
+
+
+def _tied(model: nn.Module, weight: Tensor) -> list[str]:
+    """Every module that holds `weight` directly."""
+    return [p for p, m in model.named_modules() if any(q is weight for q in m.parameters(False))]
+
+
+def find_consumers(
+    model: nn.Module, source: str, trace: DegreeTrace, report: FreeScales
+) -> list[str] | None:
+    """The consumers of `source` in call order, added until the move is free; None when there is
+    no such set, the source weight is tied, or its rows are not along the output channels."""
+    modules = dict(model.named_modules())
+    params = {id(p): n for n, p in model.named_parameters()}
+    try:
+        weight, bias, spec = _source(modules[source], None)
+    except ValueError:
+        return None
+    if len(_tied(model, weight)) > 1:
+        return None
+    channels = weight.numel() if spec is None else spec.rows(weight.shape)
+    k = trace.names.index(params[id(weight)])
+    chosen: list[str] = []
+    for path, degree in trace.module_inputs.items():
+        if path == source or abs(degree[k].item()) < 1e-9:
+            continue
+        if _input_weight(modules[path], channels) is None:
+            continue
+        chosen.append(path)
+        consumers = [_input_weight(modules[p], channels) for p in chosen]
+        if report.contains(_coeffs(params, weight, bias, consumers)):
+            return chosen
+    return None
+
+
+def _restore_modes(modes: list[tuple[nn.Module, bool]]) -> None:
+    """Each module back to its training mode, parents first."""
+    for module, training in modes:
+        if module.training != training:
+            module.train(training)
+
+
+def absorb(
+    model: nn.Module,
+    source: str,
+    *inputs,
+    into: Sequence[str] | None = None,
+    spec: RowSpec | None = None,
+    rtol: float = 1e-4,
+    **kwargs,
+) -> float:
+    """Move `source`'s channel scale into `into` (None: found); returns the largest relative
+    output change. Runs the model in eval mode with TF32 off, and restores both after."""
+    modes = [(m, m.training) for m in model.modules()]
+    tf32 = torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32
+    model.eval()
+    torch.backends.cuda.matmul.allow_tf32 = torch.backends.cudnn.allow_tf32 = False
+    try:
+        return _absorb(model, source, inputs, into, spec, rtol, kwargs)
+    finally:
+        torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32 = tf32
+        _restore_modes(modes)
+
+
+@torch.no_grad()
+def _absorb(
+    model: nn.Module,
+    source: str,
+    inputs: tuple,
+    into: Sequence[str] | None,
+    spec: RowSpec | None,
+    rtol: float,
+    kwargs: dict,
+) -> float:
+    modules = dict(model.named_modules())
+    params = {id(p): n for n, p in model.named_parameters()}
+    weight, bias, spec = _source(modules[source], spec)
+    holders = _tied(model, weight)
+    if len(holders) > 1:
+        raise ValueError(f"{source}: its weight is tied to {holders}; absorb does not move it")
+    scale = weight.abs() if spec is None else row_norms(weight, spec) / spec.radius_of(weight.shape)
+    channels = scale.numel()
+
+    # Consumers: given, or found from the degree trace
+    trace = trace_degrees(model, *inputs, **kwargs)
+    report = solve(trace)
+    if into is None:
+        into = find_consumers(model, source, trace, report)
+        if into is None:
+            raise ValueError(f"{source}: no set of consumers makes a free direction; pass `into`")
+    consumers, problems = [], []
+    for path in into:
+        found = _input_weight(modules[path], channels)
+        if found is None:
+            problems.append(f"{path}: no input-channel axis of size {channels}")
+        else:
+            consumers.append(found)
+    raise_problems(problems)
+    if not report.contains(_coeffs(params, weight, bias, consumers)):
+        raise ValueError(f"absorbing {source} into {list(into)} is not a free direction")
+
+    before = _outputs(model, inputs, kwargs)
+    saved = [
+        weight.clone(),
+        None if bias is None else bias.clone(),
+        *(cw.clone() for cw, _ in consumers),
+    ]
+    factor = torch.where(scale > FLOOR, scale, torch.ones_like(scale))
+
+    # Source channels divided by their scale, consumer columns multiplied by it
+    if spec is None:
+        weight.div_(factor.to(weight.dtype))
+    else:
+        view = spec.view_of(weight.shape)
+        rows = weight.reshape(view) / as_rows(factor.to(weight.dtype), view, spec.row_dims)
+        weight.copy_(rows.reshape(weight.shape))
+    if isinstance(bias, Tensor):
+        bias.div_(factor.to(bias.dtype))
+    for cw, axis in consumers:
+        shape = [1] * cw.dim()
+        shape[axis] = -1
+        cw.mul_(factor.to(cw.dtype).reshape(shape))
+
+    # The output must not change; undo otherwise
+    after = _outputs(model, inputs, kwargs)
+    change = max(
+        ((a - b).abs().max() / b.abs().max().clamp_min(1e-30)).item()
+        for a, b in zip(after, before, strict=True)
+    )
+    if change > rtol:
+        for target, value in zip([weight, bias, *(cw for cw, _ in consumers)], saved, strict=True):
+            if value is not None:
+                target.copy_(value)
+        raise ValueError(f"absorbing {source} into {list(into)} changed the output by {change:.2e}")
+    return change
+
+
+def absorb_plan(model: nn.Module, plan: Plan, *inputs, rtol: float = 1e-4, **kwargs) -> Plan:
+    """Make the plan's moves in order, then absorb every other `init="absorb"` layer; returns the
+    plan without moves, ready for `ds.apply`. A moved norm's affine weight is frozen."""
+    modules = dict(model.named_modules())
+    layers = {layer.path: layer for layer in plan}
+    moved = set()
+    for source, into in plan.moves:
+        layer = layers.get(source)
+        spec = None if layer is None else layer.spec
+        absorb(model, source, *inputs, into=into, spec=spec, rtol=rtol, **kwargs)
+        if layer_type_of(modules[source]) is None:
+            modules[source].weight.requires_grad_(False)
+        moved.add(source)
+    for layer in plan:
+        if layer.init == "absorb" and layer.tied_to is None and layer.path not in moved:
+            absorb(model, layer.path, *inputs, spec=layer.spec, rtol=rtol, **kwargs)
+    return replace(plan, moves=())

dragonscale/analysis/jacobian.py

@@ -0,0 +1,59 @@
+"""`scale_jacobian`: `J[m, k] = <theta_k, grad_theta_k (rho_m . f(x))>` over random output
+projections; `J s` is how fast the output moves along the log-scale direction `s`."""
+
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor, nn
+from torch.utils._pytree import tree_leaves
+
+from dragonscale.analysis.trace import parameter_groups
+
+
+@dataclass(frozen=True)
+class ScaleJacobian:
+    """`matrix` is `[projections, groups]`; its singular values and right singular vectors."""
+
+    names: list[str]
+    matrix: Tensor
+    singular: Tensor
+    directions: Tensor
+
+    def pinning(self, s: Tensor) -> float:
+        """`|J s| / (|J| |s|)`: 0 for a free direction, small for a weakly pinned one."""
+        scale = self.matrix.norm().item() * s.norm().item()
+        return (self.matrix @ s).norm().item() / max(scale, 1e-300)
+
+
+def scale_jacobian(
+    model: nn.Module, *inputs, projections: int = 32, seed: int = 0, **kwargs
+) -> ScaleJacobian:
+    """The scale Jacobian of `model(*inputs, **kwargs)` over `projections` random projections;
+    the column of a parameter without `requires_grad` is 0."""
+    names, params = parameter_groups(model)
+    trained = [p for p in params if p.requires_grad]
+    with torch.enable_grad():
+        out = model(*inputs, **kwargs)
+        flat = torch.cat(
+            [
+                t.reshape(-1)
+                for t in tree_leaves(out)
+                if isinstance(t, Tensor) and t.is_floating_point()
+            ]
+        )
+        generator = torch.Generator(device="cpu").manual_seed(seed)
+        rows = []
+        for _ in range(projections):
+            rho = torch.randn(flat.shape, generator=generator, dtype=flat.dtype).to(flat.device)
+            grads = torch.autograd.grad(
+                (flat * rho).sum(), trained, retain_graph=True, allow_unused=True
+            )
+            by_param = {id(p): g for p, g in zip(trained, grads, strict=True)}
+            row = [
+                0.0 if by_param.get(id(p)) is None else (by_param[id(p)] * p).sum().item()
+                for p in params
+            ]
+            rows.append(row)
+    matrix = torch.tensor(rows, dtype=torch.float64)
+    _, singular, vh = torch.linalg.svd(matrix, full_matrices=False)
+    return ScaleJacobian(names, matrix, singular, vh)

dragonscale/analysis/report.py

@@ -0,0 +1,208 @@
+"""`free_scales`: the null space of a degree trace's constraints, as named directions, each
+labelled parameter-only, activation-moving, or dangerous (activation-moving with a matrix in it).
+
+A weakly pinned direction is free under every constraint except anchors: only a fixed-scale addend
+holds it, and once the scaled term outgrows that addend the direction is nearly free."""
+
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor, nn
+
+from dragonscale.analysis.trace import DegreeTrace, trace_degrees
+
+TOL = 1e-9
+
+
+@dataclass(frozen=True)
+class Direction:
+    """One free direction: the log-scale coefficient of every group in it, and its matrices
+    (`fast_ends`)."""
+
+    coeffs: dict[str, float]
+    moves_activations: bool
+    fast_ends: tuple[str, ...]
+    anchors: tuple[str, ...] = ()
+
+    @property
+    def dangerous(self) -> bool:
+        return self.moves_activations and bool(self.fast_ends)
+
+    def describe(self) -> str:
+        terms = " ".join(f"{c:+g}*{name}" for name, c in self.coeffs.items())
+        label = "dangerous" if self.dangerous else (
+            "activation-moving" if self.moves_activations else "parameter-only"
+        )  # fmt: skip
+        held = f" (held by: {'; '.join(self.anchors)})" if self.anchors else ""
+        return f"{label}: {terms}{held}"
+
+
+@dataclass(frozen=True)
+class FreeScales:
+    """Every free direction of one model on one input, in reduced row-echelon form, the weakly
+    pinned directions with the anchors that hold them (`weak`), and the architecture notes a weight
+    patch cannot address (`notes`)."""
+
+    names: list[str]
+    shapes: list[tuple[int, ...]]
+    directions: list[Direction]
+    basis: Tensor
+    constraints: int
+    unknown: dict[str, int]
+    notes: list[str]
+    weak: list[Direction]
+
+    @property
+    def dangerous(self) -> list[Direction]:
+        return [d for d in self.directions if d.dangerous]
+
+    def vector(self, coeffs: dict[str, float]) -> Tensor:
+        """A direction over every group from `{name: coefficient}`."""
+        out = torch.zeros(len(self.names), dtype=torch.float64)
+        for name, c in coeffs.items():
+            out[self.names.index(name)] = c
+        return out
+
+    def contains(self, coeffs: dict[str, float]) -> bool:
+        """Whether the direction `{name: coefficient}` is free."""
+        s = self.vector(coeffs)
+        if self.basis.numel() == 0:
+            return bool(s.abs().max() < TOL)
+        residual = s - self.basis @ torch.linalg.lstsq(self.basis, s[:, None]).solution[:, 0]
+        return bool(residual.norm() <= 1e-7 * max(s.norm().item(), 1.0))
+
+    def describe(self) -> str:
+        lines = [
+            f"{len(self.directions)} free direction(s), {len(self.dangerous)} dangerous, "
+            f"{self.constraints} constraints over {len(self.names)} groups"
+        ]
+        alone = [d for d in self.directions if len(d.coeffs) == 1 and not d.moves_activations]
+        lines += [f"  {d.describe()}" for d in self.directions if d not in alone]
+        if alone:
+            names = ", ".join(next(iter(d.coeffs)) for d in alone)
+            lines.append(f"  parameter-only, one group each ({len(alone)}): {names}")
+        if self.weak:
+            lines.append(f"{len(self.weak)} weakly pinned direction(s), held only by anchors:")
+            lines += [f"  {d.describe()}" for d in self.weak]
+        if self.unknown:
+            lines.append(f"  ops without a rule (treated as nonlinear): {dict(self.unknown)}")
+        if self.notes:
+            lines.append("architecture:")
+            lines += [f"  {note}" for note in self.notes]
+        return "\n".join(lines)
+
+
+def null_space(a: Tensor) -> Tensor:
+    """An orthonormal basis of `{s : a s = 0}`, `[groups, free]`."""
+    gram = a.T @ a
+    values, vectors = torch.linalg.eigh(gram)
+    scale = max(values.abs().max().item(), 1.0)
+    return vectors[:, values <= 1e-10 * scale]
+
+
+def echelon(basis: Tensor) -> Tensor:
+    """The rows of `basis.T` in reduced row-echelon form: one readable vector per direction."""
+    m = basis.T.clone()
+    rows, cols = m.shape
+    pivot_row = 0
+    for col in range(cols):
+        if pivot_row == rows:
+            break
+        best = pivot_row + m[pivot_row:, col].abs().argmax().item()
+        if m[best, col].abs() < 1e-8:
+            continue
+        m[[pivot_row, best]] = m[[best, pivot_row]]
+        m[pivot_row] /= m[pivot_row, col].clone()
+        for r in range(rows):
+            if r != pivot_row:
+                m[r] -= m[r, col] * m[pivot_row]
+        pivot_row += 1
+    m[m.abs() < 1e-9] = 0.0
+    return m
+
+
+def directions_of(result: DegreeTrace, basis: Tensor) -> list[Direction]:
+    """Readable directions spanning `basis`, labelled, each with the anchors it moves against."""
+    degrees = result.degrees()
+    anchors = result.anchor_matrix()
+    out = []
+    for row in echelon(basis):
+        coeffs = {n: round(c, 6) for n, c in zip(result.names, row.tolist(), strict=True) if c}
+        moves = bool((degrees @ row).abs().max() > TOL)
+        fast_ends = tuple(n for n in coeffs if len(result.shapes[result.names.index(n)]) >= 2)
+        held = []
+        if result.anchors:
+            for source, value in zip(result.anchor_sources, (anchors @ row).tolist(), strict=True):
+                if abs(value) > TOL and source not in held:
+                    held.append(source)
+        out.append(Direction(coeffs, moves, fast_ends, tuple(held)))
+    return out
+
+
+def complement(inner: Tensor, outer: Tensor) -> Tensor:
+    """An orthonormal basis of the part of span(`outer`) orthogonal to span(`inner`)."""
+    if outer.numel() == 0:
+        return outer
+    rest = outer - inner @ (inner.T @ outer) if inner.numel() else outer
+    u, sigma, _ = torch.linalg.svd(rest, full_matrices=False)
+    return u[:, sigma > 1e-8 * max(sigma.max().item(), 1.0)] if sigma.numel() else rest
+
+
+def architecture_notes(
+    result: DegreeTrace, basis: Tensor, directions: list[Direction], weak: list[Direction]
+) -> list[str]:
+    """Attention calls whose q and k trade scale freely, parameters of two or more modules that
+    scale the residual stream together (an activation-moving direction with every coefficient
+    positive), and dangerous weakly pinned directions."""
+    notes = []
+    if basis.numel():
+        free_qk = [
+            i for i, (dq, _) in enumerate(result.attention) if (basis.T @ dq).abs().max() > TOL
+        ]
+        if free_qk:
+            notes.append(
+                f"attention calls {free_qk}: q and k trade scale freely; "
+                "normalize q and k (QK-norm) with a gain on one side only"
+            )
+    for d in directions:
+        modules = {name.rpartition(".")[0] for name in d.coeffs}
+        if d.moves_activations and len(modules) >= 2 and all(c > 0 for c in d.coeffs.values()):
+            notes.append(
+                f"{', '.join(d.coeffs)} scale the residual stream together; "
+                "give the model input a fixed scale (an embedding: mode 'unit')"
+            )
+    for d in weak:
+        if d.dangerous:
+            notes.append(
+                f"{', '.join(d.coeffs)} are held only by a fixed-scale addend "
+                f"({'; '.join(d.anchors)}); once they outgrow it this scale is nearly free. Let "
+                "the output read this stream directly, without a norm in between, so the loss "
+                "sets its scale"
+            )
+    return notes
+
+
+def solve(result: DegreeTrace) -> FreeScales:
+    """The free directions of a degree trace, and its weakly pinned ones."""
+    hard = result.matrix()
+    basis = null_space(torch.cat([hard, result.anchor_matrix()]))
+    weak_basis = complement(basis, null_space(hard))
+    directions = [
+        Direction(d.coeffs, d.moves_activations, d.fast_ends) for d in directions_of(result, basis)
+    ]
+    weak = directions_of(result, weak_basis) if weak_basis.numel() else []
+    return FreeScales(
+        result.names,
+        result.shapes,
+        directions,
+        basis,
+        len(result.constraints) + len(result.anchors),
+        dict(result.unknown),
+        architecture_notes(result, basis, directions, weak),
+        weak,
+    )
+
+
+def free_scales(model: nn.Module, *inputs, **kwargs) -> FreeScales:
+    """Trace `model(*inputs, **kwargs)` once and return its free directions."""
+    return solve(trace_degrees(model, *inputs, **kwargs))