Upload tr3/heal/llama2-13b-base/tensor_formats.py with huggingface_hub

tr3/heal/llama2-13b-base/tensor_formats.py

@@ -0,0 +1,768 @@
+from __future__ import annotations
+
+"""General N-core tensor formats for a linear layer: tensor train (MPO) and tensor ring.
+
+WHY THIS EXISTS. The 2-core and 3-core results say the tensorisation step, not the healing
+step, is where the leverage is: at equal parameters the 3-core beat the 2-core AFTER healing
+at both budgets tried (4.72B: 38.94 vs ~36.5 full MMLU at 60M tokens; 5.74B: 44.91 vs 41.54
+probe), and it got there on roughly a quarter of the tokens. That invites two orthogonal
+questions this module exists to answer: how many cores, and which topology.
+
+TENSOR TRAIN vs TENSOR RING. A TT/MPO pins the two boundary bonds to 1 -- the chain has open
+ends. A tensor ring closes the chain: core k has shape (r_{k-1}, d_k, r_k) with r_N == r_0,
+and the weight is the TRACE of the cyclic product. TT is exactly the TR special case r_0 = 1,
+so at equal ranks TR is strictly more expressive. It is also strictly harder to fit: TT-SVD
+is a sequence of Eckart-Young-optimal truncations, while a ring has no such closed form, so
+TR is fitted here by TT initialisation followed by alternating least squares. That asymmetry
+matters when reading results -- a TR loss could mean "the format is worse" OR "the fit did
+not converge", so `fit_tr` reports its per-sweep error curve and refuses to return a fit
+worse than its own TT initialisation.
+
+INDEX CONVENTION (shared with mpo2.py / mpo3.py). Output mode k travels with input mode k:
+
+    R(W)[d_1, ..., d_N] = W[(o_1..o_N), (i_1..i_N)],   d_k = (o_k, i_k) flattened
+
+Every mode must be > 1. A mode of 1 collapses a site and silently reduces the core count,
+which would make an "N-core" result actually be an (N-1)-core one.
+
+WHITENING is diagonal and factorised across the N input modes, as in mpo3.py: the whitening
+is then a mode-local rescaling that survives the fit exactly and un-whitens by division.
+"""
+
+import torch
+
+
+# --------------------------------------------------------------------------- rearrangement
+def rearrange_n(W: torch.Tensor, osplit, isplit) -> torch.Tensor:
+    """W[(o_1..o_N),(i_1..i_N)] -> R[d_1,...,d_N] with d_k = o_k*i_k."""
+    n = len(osplit)
+    if len(isplit) != n:
+        raise ValueError("output and input must be split into the same number of modes")
+    perm = [j for k in range(n) for j in (k, n + k)]
+    return (W.reshape(*osplit, *isplit).permute(*perm)
+             .reshape(*[o * i for o, i in zip(osplit, isplit)]))
+
+
+def inv_rearrange_n(R: torch.Tensor, osplit, isplit) -> torch.Tensor:
+    """Inverse of `rearrange_n`."""
+    n = len(osplit)
+    inter = [x for k in range(n) for x in (osplit[k], isplit[k])]
+    perm = list(range(0, 2 * n, 2)) + list(range(1, 2 * n, 2))
+    o = 1
+    i = 1
+    for k in range(n):
+        o *= osplit[k]
+        i *= isplit[k]
+    return R.reshape(*inter).permute(*perm).reshape(o, i)
+
+
+def _apply_scales(R: torch.Tensor, osplit, isplit, scales, invert: bool):
+    """Multiply (or divide) R by sqrt(d_k[i_k]) along each input mode, in place of a dense L."""
+    n = len(osplit)
+    shape = [x for k in range(n) for x in (osplit[k], isplit[k])]
+    T = R.reshape(*shape)
+    for k in range(n):
+        s = scales[k].clamp_min(1e-12).sqrt()
+        view = [1] * (2 * n)
+        view[2 * k + 1] = isplit[k]
+        T = T / s.view(*view) if invert else T * s.view(*view)
+    return T.reshape(R.shape)
+
+
+# --------------------------------------------------------------------------- contraction
+def contract_tt(cores) -> torch.Tensor:
+    """Contract TT cores [(1,d1,r1), (r1,d2,r2), ..., (r_{N-1},dN,1)] -> R[d_1,...,d_N]."""
+    out = cores[0]
+    for g in cores[1:]:
+        a, b = out.shape[-1], g.shape[0]
+        if a != b:
+            raise ValueError(f"TT bond mismatch {a} vs {b}")
+        out = torch.tensordot(out, g, dims=([out.dim() - 1], [0]))
+    if out.shape[0] != 1 or out.shape[-1] != 1:
+        raise ValueError("TT boundary bonds must be 1")
+    return out.reshape(*out.shape[1:-1])
+
+
+def contract_tr(cores, chunk: int = 0) -> torch.Tensor:
+    """Contract TR cores [(r0,d1,r1), ..., (r_{N-1},dN,r0)] -> R[d_1,...,d_N].
+
+    The obvious route -- chain every tensordot, then take the diagonal over the two ring
+    indices -- materialises (r_0, d_1..d_N, r_0), i.e. r_0^2 times the size of the ANSWER.
+    On one LLaMA attention layer at ring rank 118 that is 934 GB, so it does not merely run
+    slowly, it cannot run at all. Folding the last contraction and the trace into a single
+    einsum keeps the peak at r_0 * prod(d_1..d_{N-1}) * r_{N-1}, and chunking the first mode
+    bounds it further for the wide MLP layers.
+    """
+    n = len(cores)
+    if cores[0].shape[0] != cores[-1].shape[-1]:
+        raise ValueError(f"ring does not close: {cores[0].shape[0]} vs {cores[-1].shape[-1]}")
+    d0 = cores[0].shape[1]
+    if chunk and d0 > chunk:
+        parts = []
+        for lo in range(0, d0, chunk):
+            sub = [cores[0][:, lo:min(lo + chunk, d0), :]] + list(cores[1:])
+            parts.append(contract_tr(sub, chunk=0))
+        return torch.cat(parts, dim=0)
+    out = cores[0]
+    for g in cores[1:-1]:
+        out = torch.tensordot(out, g, dims=([out.dim() - 1], [0]))
+    return torch.einsum("i...k,kdi->...d", out, cores[-1])
+
+
+# --------------------------------------------------------------------------- fitting
+def fit_tt(R: torch.Tensor, ranks):
+    """TT-SVD with the given internal ranks. Singular values split symmetrically at each bond,
+    matching fit_mpo2/fit_mpo3's gauge so cores stay on comparable scales (fp16-safe and
+    kinder to a blockwise quantizer)."""
+    dims = list(R.shape)
+    n = len(dims)
+    if len(ranks) != n - 1:
+        raise ValueError(f"{n} sites need {n-1} internal ranks, got {len(ranks)}")
+    cores, rest, rprev = [], R.reshape(dims[0], -1), 1
+    for k in range(n - 1):
+        M = rest.reshape(rprev * dims[k], -1)
+        U, S, Vh = torch.linalg.svd(M, full_matrices=False)
+        r = min(int(ranks[k]), S.shape[0])
+        q = S[:r].clamp_min(0).sqrt()
+        cores.append((U[:, :r] * q).reshape(rprev, dims[k], r))
+        rest = q[:, None] * Vh[:r]
+        rprev = r
+    cores.append(rest.reshape(rprev, dims[-1], 1))
+    return cores
+
+
+def tt_to_tr(cores):
+    """Reinterpret TT cores as TR cores with r_0 = 1 (an exact, lossless re-labelling)."""
+    return [c.clone() for c in cores]
+
+
+def _subchain(cores, k):
+    """Contract every core except k, in cyclic order, -> E[r_k, D, r_{k-1}] with D = prod_{j!=k} d_j."""
+    n = len(cores)
+    order = [(k + 1 + t) % n for t in range(n - 1)]
+    out = cores[order[0]]
+    for j in order[1:]:
+        out = torch.tensordot(out, cores[j], dims=([out.dim() - 1], [0]))
+    return out
+
+
+def _subchain_chunk(cores, k, lo, hi):
+    """`_subchain` with the FIRST following mode restricted to [lo, hi). The restriction rides
+    along the contraction untouched, so concatenating the chunks reproduces `_subchain`."""
+    n = len(cores)
+    order = [(k + 1 + t) % n for t in range(n - 1)]
+    out = cores[order[0]][:, lo:hi, :]
+    for j in order[1:]:
+        out = torch.tensordot(out, cores[j], dims=([out.dim() - 1], [0]))
+    return out
+
+
+def rank_of(cores, k):
+    """r_{k-1} * r_k for core k -- the width of that core's least-squares system."""
+    n = len(cores)
+    return cores[k].shape[0] * cores[k].shape[2]
+
+
+def _solve_maybe_cpu(AtA, AtW):
+    """Solve AtA X = AtW, falling back to the CPU when the GPU cannot hold the factorisation.
+
+    The normal-equation matrix is (r_{k-1} r_k) square, i.e. it grows as r^4. On down_proj at
+    ring rank 219 that is 47961^2 = 9.2 GB, and the factorisation needs a second copy -- no
+    amount of chunking the data helps, because this matrix does not depend on the data size.
+    Cholesky is tried first (the Gram is symmetric positive definite after the ridge) and
+    needs only one triangular factor; if even that does not fit, the solve moves to the host.
+    Slow, but a bounded slowdown beats an unbounded failure on the layers that matter most.
+    """
+    try:
+        L = torch.linalg.cholesky(AtA)
+        return torch.cholesky_solve(AtW, L)
+    except torch.OutOfMemoryError:
+        torch.cuda.empty_cache()
+    except Exception:
+        pass
+    try:
+        return torch.linalg.solve(AtA, AtW)
+    except torch.OutOfMemoryError:
+        torch.cuda.empty_cache()
+    dev = AtA.device
+    X = torch.linalg.solve(AtA.cpu().double(), AtW.cpu().double())
+    return X.to(device=dev, dtype=AtW.dtype)
+
+
+def tr_rel_err(cores, R: torch.Tensor, chunk_elems: int = 250_000_000) -> float:
+    """||contract_tr(cores) - R||_F / ||R||_F, accumulated in chunks.
+
+    Materialising the reconstruction to subtract it needs the reconstruction AND the einsum's
+    intermediate at once; on down_proj at ring rank 155 that intermediate is (155, 64, 688,
+    256) = 7 GB even with the first mode already chunked. Accumulating the squared error one
+    slice at a time keeps only one slice live, and the chunk width is derived from an element
+    budget rather than a hard-coded constant -- a constant is exactly what silently failed to
+    engage when a mode happened to equal it.
+    """
+    n = len(cores)
+    d0 = cores[0].shape[1]
+    mid = 1
+    for k in range(1, n - 1):
+        mid *= cores[k].shape[1]
+    per_unit = max(1, mid * cores[0].shape[0] * cores[-1].shape[0])
+    step = max(1, min(d0, int(chunk_elems // per_unit)))
+    acc = 0.0
+    for lo in range(0, d0, step):
+        hi = min(lo + step, d0)
+        sub = [cores[0][:, lo:hi, :]] + list(cores[1:])
+        diff = contract_tr(sub) - R[lo:hi]
+        acc += float(diff.pow(2).sum())
+        del sub, diff
+    return (acc ** 0.5) / float(R.norm())
+
+
+def fit_tr(R: torch.Tensor, ranks, sweeps: int = 30, tol: float = 1e-8, verbose: bool = False,
+           seed: int = 1234, restarts: int = 2, solver: str = "lstsq",
+           chunk: int = 0, chunk_elems: int = 250_000_000):
+    """Fit a tensor ring by TT initialisation + alternating least squares.
+
+    ranks = [r_0, r_1, ..., r_{N-1}] with r_N == r_0. The TT init uses r_0 = 1 (a legal ring),
+    then each ALS sweep re-solves every core in closed form given the others; the ring's extra
+    boundary rank is opened up by padding the initialisation. Because ALS on a ring is not
+    guaranteed to improve monotonically from an arbitrary start, the best iterate seen is kept
+    and the routine REFUSES to return anything worse than its own TT initialisation -- without
+    that guard a fitting failure would be indistinguishable from the format being worse, which
+    is precisely the comparison this module exists to make.
+    """
+    n = len(R.shape)
+    if len(ranks) != n:
+        raise ValueError(f"a ring over {n} sites needs {n} ranks (r_0..r_{{N-1}})")
+    dims = list(R.shape)
+    rnorm = float(R.norm())
+    rnorm2 = float(R.pow(2).sum())
+    last_sq = None
+    tt = fit_tt(R, ranks[1:])
+    base_err = float((contract_tt(tt) - R).norm() / R.norm())
+    r0 = int(ranks[0])
+    gen = torch.Generator(device="cpu").manual_seed(seed)
+
+    def padded_init():
+        """TT solution embedded in the ring, with the OPENED bond seeded on BOTH ends.
+
+        Seeding only one end is not enough and is a silent failure: core 0's least-squares
+        environment is built from cores 1..N-1, whose ring index comes from core N-1's
+        trailing bond. If that stays zero, the environment vanishes on exactly the components
+        the ring adds, every ALS solve returns zero there, and the fit never leaves the TT
+        solution -- it looks like "the ring buys nothing" rather than "the ring was never
+        initialised".
+        """
+        out = []
+        for k, c in enumerate(tt):
+            a, d, b = c.shape
+            na = r0 if k == 0 else a
+            nb = r0 if k == n - 1 else b
+            g = torch.zeros(na, d, nb, dtype=R.dtype, device=R.device)
+            g[:a, :, :b] = c
+            eps = 0.05 * float(c.norm()) / max(1, c.numel()) ** 0.5
+            if k == 0 and na > a:
+                g[a:, :, :] = eps * torch.randn(na - a, d, nb, generator=gen,
+                                                dtype=R.dtype).to(R.device)
+            if k == n - 1 and nb > b:
+                g[:, :, b:] = eps * torch.randn(na, d, nb - b, generator=gen,
+                                                dtype=R.dtype).to(R.device)
+            out.append(g)
+        return out
+
+    def random_init(g=None):
+        """i.i.d. cores whose ring contraction has the right variance.
+
+        The scale is derived, not guessed. For cores with i.i.d. N(0, s^2) entries the trace
+        of the cyclic product is a sum of prod_k r_k independent products of N Gaussians, so
+        Var(W) ~= (prod_k r_k) * s^(2N). Matching that to the target mean square gives
+        s = (E[R^2] / prod_k r_k)^(1/(2N)). Getting this wrong by a constant factor per core
+        compounds to that factor^N in the product, which for N=4,5 puts the initial iterate
+        many orders of magnitude away from the data and leaves ALS crawling.
+        """
+        rr = [int(ranks[k]) for k in range(n)]
+        prod = 1.0
+        for x in rr:
+            prod *= x
+        ms = float(R.pow(2).mean())
+        s = (ms / prod) ** (1.0 / (2 * n))
+        gg = g if g is not None else gen
+        return [s * torch.randn(rr[k], dims[k], rr[(k + 1) % n], generator=gg,
+                                dtype=R.dtype).to(R.device) for k in range(n)]
+
+    # ALS on a ring has genuine local minima -- a single start that stalls would be reported
+    # as "the ring is not better", which is exactly the confound this module must avoid. So
+    # several starts are run and the best kept; restarts are cheap next to the SVDs anyway.
+    inits = [("tt_padded", padded_init)]
+    for j in range(restarts):
+        g = torch.Generator(device="cpu").manual_seed(seed + 977 * (j + 1))
+        inits.append((f"random{j}", lambda g=g: random_init(g)))
+    best, best_err, hists, winner = tt_to_tr(tt), base_err, {}, "tt_init"
+    for init_name, init_fn in inits:
+        cores, ihist = init_fn(), []
+        for it in range(sweeps):
+            for k in range(n):
+                # The subchain leaves every d index separate and in CYCLIC order starting at
+                # k+1. R's remaining modes must be flattened in that SAME order -- movedim
+                # gives natural order, which for interior k is a different permutation and
+                # would solve the least-squares problem against mismatched columns. It still
+                # returns a finite "solution", just to the wrong system.
+                # The chunk decision is made from SHAPES ONLY and BEFORE anything large is
+                # built. An earlier version computed the full subchain first and then decided
+                # whether to chunk, so the unchunked allocation always happened anyway and the
+                # chunked branch saved nothing -- the memory-bounded path existed but could
+                # never take effect.
+                # Auto-size the chunk from a memory budget instead of comparing against a
+                # fixed number. A constant threshold silently does nothing when the mode
+                # happens to equal it -- which is exactly what let a 4-site fit try to
+                # materialise a 134 GB subchain (mode 64, chunk 64, "64 > 64" is false).
+                m_ = (k + 1) % n
+                other = 1
+                for j in range(n):
+                    if j != k:
+                        other *= dims[j]
+                per_unit = max(1, (other // max(1, dims[m_])) * rank_of(cores, k))
+                auto = max(1, min(dims[m_], int(chunk_elems // per_unit)))
+                use_chunk = auto < dims[m_]
+                if solver == "normal" and use_chunk:
+                    # Memory-bounded normal equations. Materialising the subchain costs
+                    # r_k * (prod_{j!=k} d_j) * r_{k-1} floats -- on down_proj at these ranks
+                    # that is (155, 688, 256, 155) = 17 GB, so the unchunked path cannot run
+                    # on the MLP layers at all. The Gram and right-hand side are sums over the
+                    # flattened "others" index, so they accumulate exactly over slices of the
+                    # first subchain mode; this is the same arithmetic, not an approximation.
+                    m = m_
+                    d_next = dims[m]
+                    rk_, rkm1_ = None, None
+                    AtA = AtW = None
+                    Rp = R.permute(*([k] + [(k + 1 + t) % n for t in range(n - 1)]))
+                    for lo in range(0, d_next, auto):
+                        hi = min(lo + auto, d_next)
+                        Ec = _subchain_chunk(cores, k, lo, hi)
+                        rk_, rkm1_ = Ec.shape[0], Ec.shape[-1]
+                        Ec = Ec.reshape(rk_, -1, rkm1_)
+                        Ac = Ec.permute(1, 2, 0).reshape(Ec.shape[1], rkm1_ * rk_)
+                        Wc = Rp[:, lo:hi].reshape(dims[k], -1)
+                        g = Ac.transpose(0, 1) @ Ac
+                        h = Ac.transpose(0, 1) @ Wc.transpose(0, 1)
+                        AtA = g if AtA is None else AtA + g
+                        AtW = h if AtW is None else AtW + h
+                        del Ec, Ac, Wc, g, h
+                    lam = 1e-8 * float(AtA.diagonal().abs().mean().clamp_min(1e-30))
+                    AtA_r = AtA + lam * torch.eye(AtA.shape[0], dtype=AtA.dtype,
+                                                  device=AtA.device)
+                    G = _solve_maybe_cpu(AtA_r, AtW)
+                    del AtA_r
+                    # ||R - A G||^2 = ||R||^2 - 2<AtW, G> + <G, AtA G>, all from quantities
+                    # already formed. Computing it by contracting the whole ring instead would
+                    # cost one full O(prod d_k) tensor per sweep per core.
+                    sq = rnorm2 - 2.0 * float((AtW * G).sum()) + float((G * (AtA @ G)).sum())
+                    last_sq = max(sq, 0.0)
+                    rkm1, rk = rkm1_, rk_
+                else:
+                    # ONE branch builds A/Wk and ONE consumes them. Splitting the decision
+                    # across two separate if/else chains left the non-chunked solve reachable
+                    # after the chunked path had run, where `A` is whatever the PREVIOUS core
+                    # left behind -- a silent wrong answer rather than an error.
+                    E = _subchain(cores, k)               # (r_k, d_{k+1},..,d_{k-1}, r_{k-1})
+                    rk, rkm1 = E.shape[0], E.shape[-1]
+                    E = E.reshape(rk, -1, rkm1)
+                    D = E.shape[1]
+                    A = E.permute(1, 2, 0).reshape(D, rkm1 * rk)
+                    perm = [k] + [(k + 1 + t) % n for t in range(n - 1)]
+                    # R's remaining modes must be flattened in the subchain's CYCLIC order --
+                    # movedim gives natural order, which for interior k solves the
+                    # least-squares problem against mismatched columns and still returns a
+                    # finite "solution", just to the wrong system.
+                    Wk = R.permute(*perm).reshape(dims[k], D)
+                    del E
+                    if solver == "normal":
+                        AtA = A.transpose(0, 1) @ A
+                        lam = 1e-8 * float(AtA.diagonal().abs().mean().clamp_min(1e-30))
+                        AtA += lam * torch.eye(AtA.shape[0], dtype=A.dtype, device=A.device)
+                        G = _solve_maybe_cpu(AtA, A.transpose(0, 1) @ Wk.transpose(0, 1))
+                        del AtA
+                    else:
+                        G = torch.linalg.lstsq(A, Wk.transpose(0, 1)).solution
+                    last_sq = None
+                    del A, Wk
+                cores[k] = (G.transpose(0, 1).reshape(dims[k], rkm1, rk)
+                             .permute(1, 0, 2).contiguous())
+            err = (float(last_sq) ** 0.5 / rnorm if last_sq is not None
+                   else float((contract_tr(cores, chunk=chunk) - R).norm() / rnorm))
+            ihist.append(round(err, 8))
+            if verbose:
+                print(f"    TR-ALS[{init_name}] sweep {it+1}: rel err {err:.6f}", flush=True)
+            if err < best_err - 1e-12:
+                best, best_err, winner = [c.clone() for c in cores], err, init_name
+            if it and abs(ihist[-2] - err) < tol * max(err, 1e-12):
+                break
+        hists[init_name] = ihist
+    # The algebraic residual is a difference of nearly equal large numbers; at a relative
+    # error of ~1e-2 it has already lost most of its fp32 significance. It is fine for
+    # steering ALS but not for the number that gets reported, so the winner is re-measured
+    # once by explicit contraction.
+    best_err = tr_rel_err(best, R)
+    if best_err > base_err:
+        return tt_to_tr(tt), base_err, {"fell_back_to_tt": True, "tt_err": base_err,
+                                        "winner": "tt_init", "als_hist": hists}
+    return best, best_err, {"fell_back_to_tt": winner == "tt_init", "tt_err": base_err,
+                            "winner": winner, "als_hist": hists}
+
+
+# --------------------------------------------------------------------------- bookkeeping
+def tt_params(dims, ranks) -> int:
+    r = [1] + [int(x) for x in ranks] + [1]
+    return sum(r[k] * dims[k] * r[k + 1] for k in range(len(dims)))
+
+
+def tr_params(dims, ranks) -> int:
+    r = [int(x) for x in ranks]
+    n = len(dims)
+    return sum(r[k] * dims[k] * r[(k + 1) % n] for k in range(n))
+
+
+def tt_rank_ceilings(dims):
+    """r_k <= min(prod_{j<=k} d_j, prod_{j>k} d_j): above that the truncation is vacuous."""
+    n, out = len(dims), []
+    for k in range(n - 1):
+        left = 1
+        for j in range(k + 1):
+            left *= dims[j]
+        right = 1
+        for j in range(k + 1, n):
+            right *= dims[j]
+        out.append(min(left, right))
+    return out
+
+
+def ranks_for_budget(dims, budget: int, ring: bool, alpha: float = 1.0):
+    """Largest ranks fitting `budget`, shaped by r_k proportional to ceiling_k ** alpha.
+
+    Bisection rather than a closed form: the count is linear in each rank but bilinear in
+    adjacent pairs, and cyclic for a ring, so there is no single quadratic to solve.
+
+    ALPHA EXISTS BECAUSE EQUAL FILL FRACTION IS WRONG WHEN THE CEILINGS ARE LOPSIDED. For a
+    4-site split of 4096 the ceilings are (64, 4096, 64): alpha=1 spends the budget on the
+    huge middle bond and returns ranks like (36, 2327, 36), strangling the two bonds that
+    actually gate information flow through the chain. alpha=0 gives every bond the same rank
+    instead. Neither is right a priori, so the caller is expected to MEASURE across alpha --
+    the reconstruction error, not the heuristic, decides.
+    """
+    ceil_tt = tt_rank_ceilings(dims)
+    ceils = ([min(ceil_tt)] + ceil_tt) if ring else ceil_tt
+    fn = tr_params if ring else tt_params
+    mx = max(ceils)
+    shape = [(c / mx) ** alpha for c in ceils]
+
+    def build(t):
+        return [max(1, min(int(round(t * mx * s)), c)) for s, c in zip(shape, ceils)]
+
+    lo, hi = 0.0, 1.0
+    if fn(dims, build(1.0)) <= budget:
+        return build(1.0)
+    for _ in range(60):
+        mid = (lo + hi) / 2
+        if fn(dims, build(mid)) > budget:
+            hi = mid
+        else:
+            lo = mid
+    r = build(lo)
+    while fn(dims, r) > budget and max(r) > 1:
+        r[r.index(max(r))] -= 1
+    return r
+
+
+def balanced_split(n: int, k: int):
+    """Most balanced k-way factorisation of n with every factor >= 2 (spread-minimising)."""
+    best = None
+    def rec(rem, parts):
+        nonlocal best
+        if len(parts) == k - 1:
+            if rem >= 2:
+                cand = parts + [rem]
+                key = max(cand) / min(cand)
+                if best is None or key < best[0]:
+                    best = (key, cand)
+            return
+        d = 2
+        while d * d <= rem or d <= rem:
+            if d > rem:
+                break
+            if rem % d == 0:
+                rec(rem // d, parts + [d])
+            d += 1
+    rec(n, [])
+    if best is None:
+        raise ValueError(f"no {k}-way factorisation of {n} with all factors >= 2")
+    return sorted(best[1])
+
+
+def self_test(device: str = "cpu") -> None:
+    torch.manual_seed(0)
+    dt = torch.float64
+    for osp, isp in [((4, 3), (2, 5)), ((4, 3, 2), (2, 5, 3)), ((2, 3, 2, 2), (2, 2, 3, 2))]:
+        n = len(osp)
+        out = 1
+        inn = 1
+        for x in osp:
+            out *= x
+        for x in isp:
+            inn *= x
+        W = torch.randn(out, inn, dtype=dt, device=device)
+        R = rearrange_n(W, osp, isp)
+        dims = [o * i for o, i in zip(osp, isp)]
+        assert list(R.shape) == dims, (R.shape, dims)
+        assert torch.equal(inv_rearrange_n(R, osp, isp), W), f"{n}-site rearrange not invertible"
+        # index identity, not just the round trip: a wrong-but-self-inverse permutation would
+        # pass the round trip and scramble every weight silently
+        for _ in range(40):
+            oi = [torch.randint(0, x, (1,)).item() for x in osp]
+            ii = [torch.randint(0, x, (1,)).item() for x in isp]
+            ridx = tuple(oi[k] * isp[k] + ii[k] for k in range(n))
+            orow = 0
+            for k in range(n):
+                orow = orow * osp[k] + oi[k]
+            ocol = 0
+            for k in range(n):
+                ocol = ocol * isp[k] + ii[k]
+            assert R[ridx] == W[orow, ocol], f"{n}-site rearrange index map wrong"
+
+        ceils = tt_rank_ceilings(dims)
+        cores = fit_tt(R, ceils)
+        e = float((contract_tt(cores) - R).norm() / R.norm())
+        assert e < 1e-10, f"{n}-site full-rank TT is not exact ({e:.3e})"
+        assert tt_params(dims, ceils) == sum(c.numel() for c in cores), "tt_params disagrees"
+
+        if n >= 3:
+            trunc = [max(1, c // 2) for c in ceils]
+            cores = fit_tt(R, trunc)
+            assert tt_params(dims, trunc) == sum(c.numel() for c in cores)
+            for a in (0.0, 0.5, 1.0):
+                rr = ranks_for_budget(dims, tt_params(dims, trunc), ring=False, alpha=a)
+                assert tt_params(dims, rr) <= tt_params(dims, trunc), \
+                    f"TT budget solver overran at alpha={a}"
+            # a ring at the SAME parameter budget must not do worse than the TT it starts from
+            budget = tt_params(dims, trunc)
+            rring = ranks_for_budget(dims, budget, ring=True)
+            assert tr_params(dims, rring) <= budget, "TR budget solver overran"
+            _c, err_tr, info = fit_tr(R, rring, sweeps=6)
+            assert err_tr <= info["tt_err"] + 1e-12, "fit_tr returned worse than its own TT init"
+            assert abs(float((contract_tr(_c) - R).norm() / R.norm()) - err_tr) < 1e-8, \
+                "reported TR error does not match the returned cores"
+
+    ring_chk = [torch.randn(3, 5, 4, dtype=dt), torch.randn(4, 6, 2, dtype=dt),
+                torch.randn(2, 7, 3, dtype=dt)]
+    full = contract_tr(ring_chk)
+    assert torch.allclose(full, contract_tr(ring_chk, chunk=2), atol=1e-12), \
+        "chunked contract_tr disagrees with the unchunked one"
+    ref = torch.einsum("iaj,jbk,kci->abc", *ring_chk)
+    assert torch.allclose(full, ref, atol=1e-12), "contract_tr is not the cyclic trace"
+
+    # a ring with r_0 > 1 must be able to represent things a TT of the same internal ranks cannot
+    dims = [6, 6, 6]
+    torch.manual_seed(1)
+    ring = [torch.randn(3, 6, 3, dtype=dt) for _ in range(3)]
+    Rr = contract_tr(ring)
+    tt_same = fit_tt(Rr, [3, 3])
+    e_tt = float((contract_tt(tt_same) - Rr).norm() / Rr.norm())
+    assert e_tt > 1e-3, ("a rank-(3,3,3) ring was reproduced by a rank-(3,3) TT; the ring "
+                         f"contraction is probably not closing the loop (err {e_tt:.3e})")
+    for slv in ("lstsq", "normal"):
+        got, err, info = fit_tr(Rr, [3, 3, 3], sweeps=200, restarts=3, solver=slv)
+        assert err < 1e-4, (f"TR-ALS[{slv}] failed to recover an exact rank-3 ring "
+                            f"(err {err:.3e}); winner={info['winner']} "
+                            f"tt_init={info['tt_err']:.3e}")
+    got, err, info = fit_tr(Rr, [3, 3, 3], sweeps=200, restarts=3)
+    # chunking changes only how the Gram is accumulated, so it must reach the same optimum
+    _g, err_ch, _i = fit_tr(Rr, [3, 3, 3], sweeps=200, restarts=3, solver="normal", chunk=2)
+    assert err_ch < 1e-4, f"chunked TR-ALS failed to recover an exact ring (err {err_ch:.3e})"
+    assert err < 1e-4, (f"TR-ALS failed to recover an exact rank-3 ring (err {err:.3e}); "
+                        f"winner={info['winner']} tt_init={info['tt_err']:.3e}")
+
+    assert balanced_split(4096, 3) == [16, 16, 16], balanced_split(4096, 3)
+    assert balanced_split(4096, 2) == [64, 64]
+    b = balanced_split(11008, 3)
+    assert 43 in b and b[0] * b[1] * b[2] == 11008, b
+    print("tensor_formats self-test PASSED (N-site rearrange exact & index-checked, full-rank "
+          "TT exact, budget solvers within budget, TR >= its TT init, TR strictly beats TT on "
+          "a true ring, balanced splits correct)")
+
+
+if __name__ == "__main__":
+    self_test()
+
+
+# --------------------------------------------------------------------------- layers
+class TRLinear(torch.nn.Module):
+    """Tensor-ring linear layer: W = Tr[G_0 G_1 ... G_{N-1}] after the shared rearrangement.
+
+    NOT a LowRankLinear subclass, unlike MPO2Linear/MPO3Linear. Those two could pretend to be
+    a pair of factors because a train has open ends; a ring genuinely has N 3-D cores and no
+    A/B roles to borrow, so anything keyed on `.A`/`.B` (the blockwise quantizer, the bond
+    rebalancer) has to be taught about `.cores` explicitly rather than silently doing the
+    wrong thing to a subclass it thinks it understands.
+
+    RECONSTRUCTION IS CHUNKED BY DEFAULT. Merging any two cores of a 3-site ring costs
+    r*d_i*d_j*r entries -- on down_proj at ring rank 142 that is 7.1 GB in bf16, per layer,
+    per forward. Chunking the first mode bounds it; the cost is more, smaller matmuls.
+    """
+
+    def __init__(self, cores, osplit, isplit, chunk_elems: int = 500_000_000,
+                 bias=None):
+        super().__init__()
+        n = len(cores)
+        if len(osplit) != n or len(isplit) != n:
+            raise ValueError(f"{n} cores need {n} output and {n} input modes")
+        if min(list(osplit) + list(isplit)) < 2:
+            raise ValueError(f"degenerate split o={tuple(osplit)} i={tuple(isplit)}: a mode of "
+                             f"1 collapses a site, making this an (N-1)-site network")
+        for k in range(n):
+            if cores[k].shape[1] != osplit[k] * isplit[k]:
+                raise ValueError(f"core {k} has mode {cores[k].shape[1]}, expected "
+                                 f"{osplit[k]}*{isplit[k]}")
+            if cores[k].shape[2] != cores[(k + 1) % n].shape[0]:
+                raise ValueError(f"ring does not close between cores {k} and {(k+1)%n}")
+        self.cores = torch.nn.ParameterList([torch.nn.Parameter(c) for c in cores])
+        self.osplit, self.isplit = tuple(int(x) for x in osplit), tuple(int(x) for x in isplit)
+        o = i = 1
+        for k in range(n):
+            o *= self.osplit[k]
+            i *= self.isplit[k]
+        self.out_features, self.in_features = o, i
+        self.ranks = tuple(int(c.shape[0]) for c in cores)
+        # Elements per unit of the FIRST mode in contract_tr's peak intermediate, which is
+        # (r_0, c, d_1, ..., d_{n-2}, r_{n-1}) -- the LAST core's mode is contracted away by
+        # the closing einsum and must not be counted. Including it over-estimated the cost by
+        # d_{n-1} (256x on an MLP layer), which drove the chunk to 1 and would have made every
+        # MLP forward 256 sequential contractions instead of 8.
+        mid = 1
+        for k in range(1, n - 1):
+            mid *= cores[k].shape[1]
+        per_unit = max(1, mid * self.ranks[0] * self.ranks[-1])
+        self._chunk = max(1, min(cores[0].shape[1], int(chunk_elems // per_unit)))
+        # non-persistent: never wants to show up in a checkpoint, and materialize_() is an
+        # inference-only opt-in (see below), so a freshly-loaded layer must start uncached.
+        self.register_buffer("_cached_weight", None, persistent=False)
+        # LLaMA projections have no bias, but Qwen2's q/k/v do. Dropping it on
+        # replacement is SILENT -- the shapes still match, the parameter-count gate
+        # still balances (achieved excludes bias, and so would the rebuilt model), and
+        # only the accuracy moves. So the bias is carried, not discarded.
+        self.bias = None if bias is None else torch.nn.Parameter(bias)
+
+    def weight(self) -> torch.Tensor:
+        if self._cached_weight is not None:
+            return self._cached_weight
+        cs = list(self.cores)
+        R = contract_tr(cs, chunk=self._chunk if self._chunk < cs[0].shape[1] else 0)
+        return inv_rearrange_n(R, self.osplit, self.isplit)
+
+    @torch.no_grad()
+    def materialize_(self):
+        """Cache the dense weight once. INFERENCE ONLY -- gradients cannot flow through the
+        cached tensor back to the cores, so this must never be called before/during training.
+
+        forward() calls weight() on every call, which redoes the full chunked ring
+        contraction (two tensor contractions + a trace per core) from scratch. That is fine
+        for a single forward pass, but autoregressive generation calls forward() once per
+        layer per GENERATED TOKEN with the cores completely unchanged between calls. Measured
+        on tr3_ring4b04_h1's GSM8K eval before this existed: 209s/item, ETA 73h for 1319
+        items -- every one of ~200 generated tokens re-contracted all 224 ring layers.
+        A train layer (LowRankLinear/MPO2Linear) never had this problem: F.linear(F.linear(x,
+        B), A) applies the two factors straight to the activation and never materializes W at
+        all, so there was nothing to cache there.
+        """
+        self._cached_weight = self.weight().detach().clone()
+
+    def unmaterialize_(self):
+        self._cached_weight = None
+
+    def forward(self, x):
+        b = None if self.bias is None else self.bias.to(x.dtype)
+        return torch.nn.functional.linear(x, self.weight().to(x.dtype), b)
+
+    def extra_repr(self) -> str:
+        return (f"o={self.osplit}, i={self.isplit}, ranks={self.ranks}, "
+                f"chunk={self._chunk}, materialized={self._cached_weight is not None}")
+
+
+class SparseAugTRLinear(TRLinear):
+    """TRLinear plus a sparse additive residual. Same int64-index contract as the MPO layers:
+    `sp_idx` must NEVER be cast to a floating dtype -- a `.half()` on a 45M-scale index
+    silently rounds it to a different, valid-looking position."""
+
+    def __init__(self, cores, osplit, isplit, sp_idx, sp_val,
+                 chunk_elems: int = 500_000_000, bias=None):
+        super().__init__(cores, osplit, isplit, chunk_elems, bias)
+        if sp_idx.numel() != sp_val.numel():
+            raise ValueError(f"{sp_idx.numel()} indices for {sp_val.numel()} values")
+        n = self.out_features * self.in_features
+        if sp_idx.numel() and (int(sp_idx.max()) >= n or int(sp_idx.min()) < 0):
+            raise ValueError("sparse index out of range for this layer")
+        self.register_buffer("sp_idx", sp_idx.to(torch.int64))
+        self.sp_val = torch.nn.Parameter(sp_val)
+
+    def weight(self) -> torch.Tensor:
+        # must check the cache here too, not just in TRLinear.weight(): materialize_() caches
+        # THIS class's weight() (ring + sparse, since self.weight() dispatches polymorphically
+        # from inside materialize_()), so without this check a later call would go
+        # super().weight() -> see the cache -> return the already-sparse-augmented tensor ->
+        # and then add the sparse residual a second time on top of it.
+        if self._cached_weight is not None:
+            return self._cached_weight
+        W = super().weight()
+        if self.sp_idx.numel():
+            W = W.reshape(-1).index_add(0, self.sp_idx, self.sp_val.to(W.dtype)).view_as(W)
+        return W
+
+    def extra_repr(self) -> str:
+        return super().extra_repr() + f", nnz={self.sp_idx.numel()}"
+
+
+def layer_self_test() -> None:
+    torch.manual_seed(0)
+    osp, isp = (4, 3, 2), (2, 5, 3)
+    dims = [o * i for o, i in zip(osp, isp)]
+    cores = [torch.randn(3, dims[0], 4), torch.randn(4, dims[1], 2), torch.randn(2, dims[2], 3)]
+    lay = TRLinear([c.clone() for c in cores], osp, isp, chunk_elems=10 ** 9)
+    ref = inv_rearrange_n(contract_tr(cores), osp, isp)
+    assert torch.allclose(lay.weight(), ref, atol=1e-5), "TRLinear.weight() != contract+inverse"
+    # chunking must not change the answer
+    lay_c = TRLinear([c.clone() for c in cores], osp, isp, chunk_elems=1)
+    assert lay_c._chunk < dims[0], "chunk did not engage"
+    assert torch.allclose(lay_c.weight(), ref, atol=1e-5), "chunked forward differs"
+    x = torch.randn(7, lay.in_features)
+    assert torch.allclose(lay(x), x @ ref.T, atol=1e-4), "forward != x W^T"
+    idx = torch.tensor([0, 5, lay.out_features * lay.in_features - 1], dtype=torch.int64)
+    val = torch.tensor([1.0, -2.0, 3.0])
+    sl = SparseAugTRLinear([c.clone() for c in cores], osp, isp, idx, val, chunk_elems=10 ** 9)
+    d = (sl.weight() - ref).reshape(-1)
+    assert torch.allclose(d[idx], val, atol=1e-5), "sparse residual not applied at sp_idx"
+    assert abs(float(d.abs().sum()) - float(val.abs().sum())) < 1e-4, "residual leaked"
+    try:
+        TRLinear([torch.zeros(2, 4, 2), torch.zeros(2, 3, 2), torch.zeros(2, 6, 2)],
+                 (4, 1, 2), (1, 3, 3))
+    except ValueError:
+        pass
+    else:
+        raise AssertionError("a degenerate split (mode = 1) was accepted")
+    # gradients must reach every core
+    lay.zero_grad()
+    lay(x).sum().backward()
+    for k, c in enumerate(lay.cores):
+        assert c.grad is not None and float(c.grad.abs().sum()) > 0, f"core {k} got no gradient"
+    # materialize_() must be a pure inference-speed cache: same weight, same forward output,
+    # and -- the one bug this is actually guarding against -- the sparse residual must not
+    # get added twice once SparseAugTRLinear.weight() is short-circuited by the cache.
+    pre = lay.weight().clone()
+    lay.materialize_()
+    assert lay._cached_weight is not None, "materialize_() did not populate the cache"
+    assert torch.allclose(lay.weight(), pre, atol=1e-6), "materialize_() changed weight()"
+    assert torch.allclose(lay(x), x @ ref.T, atol=1e-4), "forward != x W^T after materialize_()"
+    lay.unmaterialize_()
+    assert lay._cached_weight is None, "unmaterialize_() left the cache populated"
+    sl_ref = sl.weight().clone()
+    sl.materialize_()
+    assert torch.allclose(sl.weight(), sl_ref, atol=1e-6), \
+        "materialized SparseAugTRLinear.weight() differs -- sparse residual likely double-added"
+    print("TRLinear self-test PASSED (weight matches contraction, chunking exact, forward "
+          "correct, sparse localised, degenerate split refused, all cores differentiable, "
+          "materialize_() cache exact and idempotent)")