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
Test Suite Guide
Whenever we change any code-logic related to src/heretic/model.py or config.toml (e.g. row_normalization, full_normalization_lora_rank, winsorization_quantile, etc) which can affect a model's reproduciblity; Use these tests which are designed to verify that those changes does not affect reproducibility, unless they are meant to (like when we'll integrate ARA branch in future).
How to test
- Choose any model from tiny-random org which provides tiny models useful for debugging.
Example: tiny-random/minicpm5.
Note
It is highly recommended to use a model which does not have a
special_tokens_map.jsonfile in the repo. Because those files are almost always wrong intiny-random/*models compared to the original model.
- Clone that model repository using Git and generate the SHA256 hashes using
sha256sum:
On Linux:
sha256sum -b * > ../SHA256SUMS.LABEL
On Windows:
sha256sum * | Out-File -Encoding utf8NoBOM ../SHA256SUMS.LABEL
Tip
On windows,
sha256sumis generally pre-installed by Git for windows.
Verify with:
Get-Command sha256sum`
Expected:
CommandType Name Version Source
----------- ---- ------- ------
Application sha256sum.exe 0.0.0.0 C:\Program Files\Git\usr\bin\sha256sum...
Note
You must use Windows Powershell
v7.Xnot the core which isv5.1. This is required for-Encoding utf8NoBOMto work.See Differences between Windows PowerShell 5.1 and PowerShell 7.x documentation.
Where LABEL describes the type of system you are running the tests on.
Example:
SHA256SUMS.windows(For windows)SHA256SUMS.ci(For GitHub CI)SHA256SUMS.linux(For linux)
- Run the tests with:
uv run run_tests.py
The output hashes should FAIL against the Valid hashes in SHA256SUMS file of the test model you added. This is expected since Heretic changes the model. Without Step 2, the test model's folder will simply be ignored because it will not have a hash SUMS file to compare against.
- After that go to the output
TEST_MODEL_DIR/modelfolder and re-generate the Actual hashes based on the system you are using.
cd TEST_MODEL_DIR/model
sha256sum -b * > ../SHA256SUMS.LABEL # or use windows command.
- Re-run the tests with:
uv run run_tests.py
This time the tests should PASS because we added the new hashes which are expected to be reproduced on the same system.
- After that push the
SHA256SUMS.LABELfiles and wait for GitHub CI actions to run those tests.
Since PyTorch does not guarantee exact cross-system reproducibility regardless of configuration, multiple valid hashes can be provided for each output file. The above update must be performed for each TEST_MODEL_DIR and on each type of system.
For this, copy the Actual hash value for each mismatched unidentical file into a SHA256SUMS.ci file.
- After that push the
SHA256SUMS.cifiles and wait for GitHub CI actions to re-run those tests.
This time the tests should PASS because we added the new hashes which are expected to be reproduced on CI.