flux2distill/surgery.py

"""Model surgery: capture single-block I/O, select blocks (SVD-energy), build student.

Selection criterion (plan.md S4): for each single block compute the residual delta
(output - input) on a calibration batch, take its singular values, and score the block
by the *absolute tail energy a rank-r surrogate cannot capture*: cost_i = sum(s[r:]**2).
Keep full the `keep_single` costliest blocks (high-rank / high-magnitude residual);
replace the rest with rank-r surrogates (cheapest to approximate). This accounts for
both residual magnitude and effective rank, unlike a scale-invariant energy ratio.
"""
from __future__ import annotations

import torch

from .surrogate import LowRankResidualSurrogate, make_surrogate


@torch.no_grad()
def capture_single_block_io(pipe, prompts, num_inference_steps=4, guidance_scale=1.0,
                            height=512, width=512, seed=0, max_tokens_per_block=12000):
    """Run the teacher pipeline and capture (X, Delta) for every single block via hooks.

    X     = hidden_states entering the block (concatenated [text,image] stream, dim d).
    Delta = block_output - X  (the residual the block adds).
    Returns dict[int] -> {"X": (M,d) cpu fp32, "Delta": (M,d) cpu fp32}.
    """
    tf = pipe.transformer
    singles = tf.single_transformer_blocks
    store = {i: {"X": [], "Delta": [], "n": 0} for i in range(len(singles))}
    handles = []

    def make_hook(i):
        def hook(module, args, kwargs, output):
            x = kwargs.get("hidden_states", args[0] if args else None)
            out = output[0] if isinstance(output, tuple) else output
            if store[i]["n"] >= max_tokens_per_block:
                return
            xf = x.detach().reshape(-1, x.shape[-1]).float()
            of = out.detach().reshape(-1, out.shape[-1]).float()
            take = min(xf.shape[0], max_tokens_per_block - store[i]["n"])
            store[i]["X"].append(xf[:take].cpu())
            store[i]["Delta"].append((of[:take] - xf[:take]).cpu())
            store[i]["n"] += take
        return hook

    for i, blk in enumerate(singles):
        handles.append(blk.register_forward_hook(make_hook(i), with_kwargs=True))

    try:
        gen = torch.Generator(device=pipe.device).manual_seed(seed)
        pipe(prompt=list(prompts), num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale, height=height, width=width,
             generator=gen, output_type="latent")
    finally:
        for h in handles:
            h.remove()

    io = {}
    for i in store:
        io[i] = {
            "X": torch.cat(store[i]["X"], 0),
            "Delta": torch.cat(store[i]["Delta"], 0),
        }
    return io


class _IdentityBlock(torch.nn.Module):
    """Pass-through replacement matching the single-block contract (for ablation)."""
    def forward(self, hidden_states, encoder_hidden_states=None, temb_mod=None,
                image_rotary_emb=None, joint_attention_kwargs=None,
                split_hidden_states=False, text_seq_len=None):
        if encoder_hidden_states is not None:
            text_seq_len = encoder_hidden_states.shape[1]
            hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
        if split_hidden_states:
            return hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:]
        return hidden_states


@torch.no_grad()
def importance_by_ablation(pipe, prompts, num_inference_steps=4, guidance_scale=1.0,
                           height=512, width=512, seed=0):
    """Leave-one-out importance: relative change in the final latent when each single
    block is replaced by identity. Higher = more important (keep); lower = safe to drop."""
    tf = pipe.transformer
    singles = tf.single_transformer_blocks
    n = len(singles)

    def run():
        gen = torch.Generator(device=pipe.device).manual_seed(seed)
        out = pipe(prompt=list(prompts), num_inference_steps=num_inference_steps,
                   guidance_scale=guidance_scale, height=height, width=width,
                   generator=gen, output_type="latent")
        return out.images.float()

    base = run()
    base_norm = base.norm() + 1e-8
    imp = {}
    for i in range(n):
        orig = singles[i]
        singles[i] = _IdentityBlock()
        lat = run()
        singles[i] = orig
        imp[i] = float((lat - base).norm() / base_norm)
    return imp


def select_blocks_by_importance(importance, drop_k):
    """Return (keep_idx, drop_idx) dropping the `drop_k` least-important single blocks."""
    order = sorted(importance, key=lambda i: importance[i])  # ascending importance
    drop_idx = sorted(order[:drop_k])
    keep_idx = sorted(order[drop_k:])
    return keep_idx, drop_idx


def select_blocks_svd_energy(io, rank, keep_single):
    """Return (keep_idx, surrogate_idx, stats). Surrogate cheapest-to-approximate blocks."""
    stats = []
    for i in sorted(io):
        D = io[i]["Delta"].double()
        s = torch.linalg.svdvals(D)            # (min(M,d),)
        total = (s ** 2).sum()
        tail = (s[rank:] ** 2).sum() if s.shape[0] > rank else torch.tensor(0.0, dtype=s.dtype)
        captured_ratio = float((1 - tail / (total + 1e-12)))
        stats.append({
            "block": i,
            "tail_energy": float(tail),                 # absolute cost of rank-r approx
            "captured_ratio": captured_ratio,           # fraction of energy in top-r
            "delta_rms": float(D.pow(2).mean().sqrt()),
        })
    # keep the `keep_single` blocks with the LARGEST tail energy (hardest to approximate)
    by_cost = sorted(stats, key=lambda d: d["tail_energy"], reverse=True)
    keep_idx = sorted(d["block"] for d in by_cost[:keep_single])
    surrogate_idx = sorted(d["block"] for d in by_cost[keep_single:])
    return keep_idx, surrogate_idx, stats


def attach_surrogates(transformer, surrogate_idx, kind="lowrank", rank=512, act="gelu",
                      heads=4, head_dim=128, conv_kernel=5, ffn_hidden=1024, ffn_idx=None,
                      device=None, dtype=None):
    """Replace the given single blocks with fresh surrogates of `kind` (lowrank |
    linear_attention). `ffn_idx`: subset of surrogate_idx that get an FFN (others are light);
    None => all get FFN. Used to (re)build the student structure before load_state_dict."""
    d = transformer.config.num_attention_heads * transformer.config.attention_head_dim
    ref = next(transformer.parameters())
    device = device or ref.device
    dtype = dtype or ref.dtype
    ffn_set = None if ffn_idx is None else set(ffn_idx)
    for i in surrogate_idx:
        use_ffn = True if ffn_set is None else (i in ffn_set)
        sur = make_surrogate(kind, d, rank=rank, act=act, heads=heads, head_dim=head_dim,
                             conv_kernel=conv_kernel, ffn_hidden=ffn_hidden, use_ffn=use_ffn)
        transformer.single_transformer_blocks[i] = sur.to(device=device, dtype=dtype)
    return transformer


@torch.no_grad()
def capture_block_io_seq(pipe, prompts, drop_idx, num_inference_steps=4, guidance_scale=1.0,
                         height=512, width=512, seed=0, max_seqs=16):
    """Sequence-preserving capture for warm-start: for each dropped single block, store the
    input X and teacher output Y as full (S, N, d) sequences (positions intact, so RoPE
    aligns), plus the shared image_rotary_emb captured once. Returns (io, rotary)."""
    tf = pipe.transformer
    singles = tf.single_transformer_blocks
    store = {i: {"X": [], "Y": [], "n": 0} for i in drop_idx}
    rotary = {}
    handles = []

    def make_hook(i):
        def hook(module, args, kwargs, output):
            if store[i]["n"] >= max_seqs:
                return
            x = kwargs.get("hidden_states", args[0] if args else None)
            y = output[0] if isinstance(output, tuple) else output
            store[i]["X"].append(x.detach().float().cpu())
            store[i]["Y"].append(y.detach().float().cpu())
            store[i]["n"] += x.shape[0]
            if "rotary" not in rotary and kwargs.get("image_rotary_emb") is not None:
                cos, sin = kwargs["image_rotary_emb"]
                rotary["rotary"] = (cos.detach().float().cpu(), sin.detach().float().cpu())
        return hook

    for i in drop_idx:
        handles.append(singles[i].register_forward_hook(make_hook(i), with_kwargs=True))
    try:
        gen = torch.Generator(device=pipe.device).manual_seed(seed)
        pipe(prompt=list(prompts), num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale, height=height, width=width,
             generator=gen, output_type="latent")
    finally:
        for h in handles:
            h.remove()

    io = {i: {"X": torch.cat(store[i]["X"], 0)[:max_seqs],
              "Y": torch.cat(store[i]["Y"], 0)[:max_seqs]} for i in drop_idx}
    return io, rotary["rotary"]


def build_student(transformer, surrogate_idx, io, rank=512, act="gelu",
                  ridge=1e-2, near_linear_scale=6.0, device="cpu"):
    """Replace the chosen single blocks with lstsq-initialized surrogates, in place.

    Returns dict[int] -> relative reconstruction error of each surrogate's warm start.
    """
    d = transformer.config.num_attention_heads * transformer.config.attention_head_dim
    errs = {}
    ref_param = next(transformer.parameters())
    for i in surrogate_idx:
        sur = LowRankResidualSurrogate(d, rank=rank, act=act)
        rel = sur.load_lstsq_init(io[i]["X"], io[i]["Delta"],
                                  ridge=ridge, near_linear_scale=near_linear_scale)
        sur = sur.to(device=ref_param.device, dtype=ref_param.dtype)
        transformer.single_transformer_blocks[i] = sur
        errs[i] = rel
    return errs