core.decompose
Transform a standard state dict into a flat weightless theta bundle with analytic decomposition, and verify the result.
core/decompose.py is model-agnostic. It operates on plain numpy.ndarray values keyed by the standard MAX/HuggingFace state-dict layout and emits the flat weightless key layout that model wrappers remap to their nested names.
Families
Two module constants define which projections are decomposed:
FAMILIES: all seven,q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj.ATTN_FAMILIES: the attention subsetq/k/v/o, used whentie_families="attn".
decompose_standard_state_dict_to_weightless
The top-level entry point.
decompose_standard_state_dict_to_weightless(
state_dict, # Mapping[str, array-like]
rank=32, # LoRA rank for the residual SVD
shared_blocks=1, # number of contiguous layer bands (G)
) -> dict[str, np.ndarray]It groups projection weights by family, computes each band’s W0 as the mean of its layers, and decomposes every layer into s, U, V. Non-projection weights pass through unchanged; it drops a tied lm_head.weight, which MAX materializes from the embeddings at runtime.
The returned dict uses these keys:
W0.{fam} # G=1
W0.{band}.{fam} # G>1
layers.{i}.{fam}.s # [out_dim]
layers.{i}.{fam}.U # [out_dim, rank]
layers.{i}.{fam}.V # [in_dim, rank]decompose_projection_weight
Decomposes a single layer’s projection given the shared block:
decompose_projection_weight(weight, w0, rank) -> (s, U, V)s is the per-output-row least-squares scale against w0; U and V come from a truncated SVD of the residual. When rank == 0, U and V have zero columns (a pure scaled shared block). When the residual’s effective rank is smaller than the requested rank (for example, k_proj/v_proj in GQA models where out_dim < rank), U and V are padded up to the full rank with zeros, so every family produces consistently shaped tensors.
Verifying quality
Two helpers check how well theta reproduces the original weight:
w = reconstruct_weight(s, w0, U, V) # diag(s) @ W0 + U @ V.T
err = relative_frobenius_error(original, w) # ||orig - w||_F / ||orig||_FDetecting pre-converted checkpoints
is_weightless_theta_dict(state_dict) returns True when a state dict already uses theta keys (either G=1 or banded), so the auto-decompose adapter can pass it through untouched. theta_bundle_bytes(theta) sums the byte size of all arrays in a bundle, which is useful for reporting the shipped size.