Repository navigation
Conversation
…se 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.
| else: | ||
| bias = bias.float().to(lora_A.device) | ||
| good_output_base = good_output - bias | ||
| bad_output_base = bad_output - bias |
There was a problem hiding this comment.
Shouldn't the captured output already account for the bias? It's captured by the hook from the output of the entire nn.Linear module, no?
There was a problem hiding this comment.
Right, the hook captures the full nn.Linear output, so it's input @ W.T + bias. But the original closure computes new_good_output = good_input @ W_eff.T with no bias term. To keep the reformulation exactly equivalent, good_output_base has to be the bias-free input @ W_base.T, so I subtract the captured bias to recover it. Then input @ W_base.T + (input @ A.T) @ B.T == input @ (W_base + B@A).T, which matches the original product. For bias-free o_proj/down_proj (Qwen, Gemma, Llama) the branch is a no-op; it only matters for modules that carry a bias, where skipping the subtraction would shift new_good_output by +bias and diverge from the dense path.
There was a problem hiding this comment.
That doesn't really make sense though. If the module contains a bias term then the implementation is simply broken (and was already broken before).
We can't correct for the bias in one place, then run an optimization that tries to match the output from an unbiased linear operator (bad_output_base/new_bad_output) with the original output that incorporates the bias (bad_output). This is comparing apples to oranges.
We either flat out assume that modules have no bias (which I believe is true for the vast majority of LLMs), or we have to make sure that the bias is either present everywhere or stripped out everywhere. As it stands, we are comparing apples to oranges.
There was a problem hiding this comment.
i think i agree? it's a pre-existing apples-to-oranges for bias modules. Just to be precise about what the subtraction did: the original already compares input @ W_eff.T (no bias) against the captured good_output (with bias), and my version passed that same captured good_output to the loss unchanged. The subtraction only reconstructed the bias-free input @ W_base.T so the low-rank term summed back to exactly input @ W_eff.T, i.e. it reproduced the original behavior (bias quirk included) rather than changing it.
But I agree assuming no bias is the clean call, since it holds for the o_proj/down_proj modules ARA targets across basically all current LLMs. Dropped the branch in b9f0adf, so good_output_base is just the captured output. It's a no-op for bias-free modules (identical weights/hashes), and the confusing half-correction is gone. Happy to add an assert module.base_layer.bias is None if you'd like the assumption enforced loudly.
There was a problem hiding this comment.
the original already compares input @ W_eff.T (no bias) against the captured good_output (with bias)
Correct, this was already relying on there being no bias.
To be clear, we never need to keep incorrect or suboptimal behavior. Heretic's reproducibility system is version-based for a reason. We are not shackled to the past (especially not with ARA, which isn't yet in any released version).
Happy to add an assert module.base_layer.bias is None if you'd like the assumption enforced loudly.
Yes please!
There was a problem hiding this comment.
Agreed. Added assert module.base_layer.bias is None so the no-bias assumption fails loudly instead of silently comparing a bias-free output against a bias-inclusive one.
|
Thanks! Those are impressive performance numbers, but I do wonder why you chose ancient models for the benchmark?
No. Reproducibility is always constrained to a particular Heretic release, and ARA isn't released in stable yet, so we can still change the implementation without breaking anything.
Thanks, I will update the docs for the 2.0 release. |
|
It's super unfortunate that this changes all test hashes though. The entire purpose of the tests is to feel confident that such changes don't break anything, and instead we have to make the tests satisfy the change instead of the other way round... |
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).
|
Fair point. Those were just what I had ungated and could fit in a small rented-GPU budget. Here is Qwen3.8-27B on an H100, same seed, 3 trials, modify step:
The win scales with layer size, since the eliminated matmul is proportional to |
|
You're right that this is unavoidable. Any FLOP-reordering optimization changes the summation order, so the weights differ at the float-rounding level (about 1e-5 relative) and bit-identical output isn't achievable. I pushed As a behavioral cross-check, a 60-trial Pareto run on gemma-2-2b produced an equivalent refusals/KL front to the dense path, so the final selected models come out equivalent in quality too. |
| # assumption the captured module output *is* input @ W_base.T, which | ||
| # is the constant term the closure adds the low-rank update to. | ||
| good_output_base = good_output | ||
| bad_output_base = bad_output |
There was a problem hiding this comment.
But then why do we need these variables at all?
There was a problem hiding this comment.
Right, once the bias handling was gone they were just aliases. Removed them: the closure now adds the low-rank update to the captured good_output/bad_output directly.
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.
b9f0adf to
49e37f3
Compare
|
Please update the test hashes so this can be merged. |
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.
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).
|
Updated, and CI is green on all four Python versions (run on my fork: https://github.com/varunbln/heretic/actions/runs/39424996627). The run here is waiting for workflow approval. Two changes:
I didn't touch the |
👁️ 👁️ 👁️ |
|
I'm unable to reproduce the Pareto front after this change. The results are completely different, which shouldn't be the case. CommandmasterThis PRI can confirm a 30% end-to-end speedup, which is great, but obviously, the Pareto fronts should be identical, at least approximately. |
|
It appears that unfortunately, this PR is completely broken. Same command as above ( masterThis PRThe parameters are identical, but the results (refusals and KLD) differ dramatically. The difference is way, way beyond what could reasonably be explained by floating point order shenanigans. |
…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.
|
You were right, and thanks for running it. There was a real bug. Cause. The low-rank closure reused the module output captured by the forward hook as the constant base term Fix (3359a7d). The base products are now computed once per module in float32 from the inputs, outside the closure. The bias is added after the row norm scaling, as before. Relative difference against the dense master computation is now 1.6e-6, down from 3.0e-3. The equivalence test builds the product from the input too, so it would have caught this. Same machine comparison (Mac, MPS,
One caveat. Unchanged master gives 52/100 for trial 1 on your machine and 11/100 on mine, so the optimizer amplifies small floating point differences a lot. Any change that reorders the matmuls will not reproduce your fronts bit for bit, even when it is numerically correct. I think the fair test is whether the Pareto front is about as good, not whether it is identical. The speedup is unchanged, since the fp32 base product is computed once per module and not per closure call. Could you rerun Two practical notes. The upstream CI run needs your approval. I updated the Linux hashes for the gpt-oss and seed-oss tests from a CI run on my fork. The Windows hashes would need regenerating on your side. |
Nope, it's more complicated than that.masterThis PR (updated)The Pareto fronts are nothing alike. But it gets even more mysterious: The Pareto front of this PR is also nothing like what it was before your latest change (see above). They all do reproduce with the same code though, so the process itself is deterministic. Could it really be that the optimization is so unstable that even the slightest perturbation of the computation completely changes the result? How can we test this? |
|
I have identified the root cause of this issue. Your PR is fine, the problem is fundamental to the current ARA formulation. I will open a PR with an explanation and a fix soon. |
|
Yes, I think it really is that unstable, and it is a property of master, not of this PR. Here is how I tested it. Why the fronts share no trials. Test 1: perturb master itself. I made two changes to master's
Then I ran
So master is deterministic, but a one-ulp change to its inputs moves results by about as much as this PR does. The PR sits between the two no-op perturbations. Test 2: what happens inside one module. I captured real layer 8
ARA optimizes each module separately, and a trial touches several modules, so each trial is a handful of coin flips between basins. That gives the scatter above, and with TPE feeding on it, entirely different fronts. What this means for the PR. I don't think any reordering of this computation can match master's fronts, including a more accurate one. The question is whether the front is as good on average, and a single seed can't answer that. If it would help, I can run several seeds of master vs PR and compare hypervolume. I'm also fine if you would rather not take a change whose effect can't be checked by exact reproduction. Separate from this PR. The bimodality looks worth fixing in ARA itself, since today it makes results depend on hardware (52/100 for you vs 11/100 for me on unchanged master, trial 1). In the runs I traced on that module, loss was still falling at the 5th outer step. More steps, or a few restarts keeping the lowest loss, might make ARA land in the better basin consistently. Happy to look into that in a separate issue if you're interested. |
|
There are no basins, it's more complicated than that. I'll post the full explanation in the next hour. |
|
See #522. I will do one final sanity test now, then merge this. |
| # so compute it once here (in float32, from the float32 weight) | ||
| # instead of on every closure evaluation. It must not be taken from | ||
| # the captured module output: that was computed in the model's dtype | ||
| # (e.g. bfloat16) and would change the objective. |
There was a problem hiding this comment.
But we're still using the captured good_output and bad_output as arguments to ara_loss below.
I don't think this logic makes much sense. Why would we suddenly recompute this one thing? This feels if anything like an inconsistency, and in particular, it is not the cause of the trajectory divergence (see #522). The trajectories remain divergent even in full precision.
Whatever we do here, using the captured output in one place while avoiding it in another is surely the wrong approach. Above all else, we should keep this change as simple as possible.
hi again
ARA's
modify_modelrebuilds the full dense effective weight and runs a fullinput @ W_eff.Ton every L-BFGS closure evaluation. That dense product dominates the per-trial cost, and it is avoidable because the LoRA update is low-rank.The optimization
input @ (W_base + B@A).T = input @ W_base.T + (input @ A.T) @ B.T.input @ W_base.Tis exactly the captured module output (minus bias). It is constant across closures, so it is precomputed once.rterm(input @ A.T) @ B.Tis recomputed. This avoids both the dense[d_out x d_in]weight materialization and the full dense matmul.preserve_row_magnitudes, row normalization is a per-output-row rescale. It is applied to the outputs via a diagonal factor (W_row_norms / ‖effective row‖), so the identity holds in that path too.Measured speedup (seed 0, L-BFGS startup trials, so params are identical per pair)
The win scales with layer size, since the dense matmul is proportional to
d_in·d_out. It is modest on a tiny model with a fast GPU, and large on realistically sized models even on an H100.modifyis about 80% of each trial on both MPS and H100, so the step speedup carries through to end-to-end.Final-model quality is preserved (gemma-2-2b, 60 trials per side, H100). Pareto front (refusals/100, KL), lower is better on both axes:
At matched refusal levels the optimized front is equal or better at the low-refusal end. The two TPE searches diverge but land on comparable fronts.
Reproducibility caveat. This changes floating-point summation order. Abliterated weights are numerically equivalent (relative error about 1e-5) but not bit-identical, and per-trial refusal counts shift slightly.
also, since ARA is marked
reproducible = True, do you want this to be behind a flag? lmk - since its a clear speedup i think making it the default is okayTesting
ruff format --check,ruff check --extend-select I, andty checkpass.preserve_row_magnitudesboth on and off.Note: transformers 5.18+ requires
torch>=2.6. This is unrelated to the change. I am flagging it from the CUDA runs.