mirror of
https://github.com/p-e-w/heretic.git
synced 2026-11-06 05:16:39 +00:00
Varun Balani
a2e626fd03
* perf(ara): compute closure outputs via low-rank update instead of dense weight ARA's modify_model rebuilt the full dense effective weight (W_base + B@A) and ran a full input @ W_eff.T on every L-BFGS closure evaluation, which dominates the per-trial cost. Since input @ W_base.T equals the captured module output (minus bias) and is constant across closures, and the LoRA update is low-rank, the closure output is input_output_base + (input @ A.T) @ B.T. This avoids both the dense [d_out x d_in] weight materialization and the full dense matmul; only the rank-r term is recomputed. With preserve_row_magnitudes, row normalization becomes a per-output-row diagonal rescale applied to the outputs, so the identity holds there too. Algebraically exact (rel err ~1e-5); not bit-identical due to changed float summation order. modify-step speedup ~1.4x (gemma-2-2b/H100) to ~3.2x (Qwen2.5-7B/H100), scaling with layer size; final-model Pareto front is preserved. * test: assert low-rank ARA closure equals the dense formulation Adds an equivalence test that runs both the original dense closure (build W_eff, multiply) and the low-rank reformulation on the same inputs and asserts they match to float tolerance, with and without preserve_row_magnitudes. This guards the reformulation's correctness directly, rather than relying on bit-identical output hashes (which necessarily change under any FLOP-reordering optimization). * simplify: assume bias-free linear modules in ARA closure The dense formulation already matches the raw weight product (input @ W.T) against the captured output, so a biased module was a pre-existing apples-to-oranges comparison. Rather than reconstruct the bias-free base from the captured output, assert the modules are bias-free (true for the o_proj / down_proj modules ARA targets across essentially all current LLMs). The captured output is then exactly input @ W_base.T, so the closure uses it directly and the separate *_base variables are gone. No behavior change for the bias-free modules ARA runs on. * fix(ara): handle linear modules with a bias term The captured module output is input @ W_base.T + bias, but ARA only modifies the weight. Instead of asserting that no bias is present (which fails on gpt-oss), subtract the bias to get the bias-free product, add the low-rank update and apply the row-norm rescale to that, then add the bias back. This also fixes the original dense closure, which compared input @ W_eff.T against a bias-inclusive output for such models. Bias-free modules are unchanged. Extend the equivalence test to cover modules with a bias, with and without row normalization. * test: update ARA output hashes for the low-rank closure Only the two ARA tests (gpt-oss, seed-oss) change. Hashes taken from CI on ubuntu-latest: gpt-oss is identical on Python 3.10-3.13; seed-oss has one hash for 3.10-3.12 (SHA256SUMS.ci) and another for 3.13 (SHA256SUMS.ci2). * fix(ara): compute the base product in float32 instead of reusing the captured output The low-rank closure used the hook's captured module output as the constant input @ W_base.T term. That output was computed in the model's dtype (e.g. bfloat16), which is about 3e-3 relative off from the float32 product the original dense closure used, so the optimizer saw a different objective. Compute input @ W_base.T once per module from the float32 weight instead. This matches the dense closure to float32 rounding (about 1e-6 relative on gemma-3-270m) while keeping the low-rank speedup. * test: update ARA output hashes for the float32 base product * ci: update hash for deepseek-v4-ara test * ci: update hash for deepseek-v4-ara test