Skip to content

perf(ara): compute closure outputs via low-rank update instead of dense weight - #499

Open
varunbln wants to merge 11 commits into
p-e-w:masterfrom
varunbln:perf/ara-lowrank-closure
Open

varunbln wants to merge 11 commits into
p-e-w:masterfrom
varunbln:perf/ara-lowrank-closure

Conversation

@varunbln

@varunbln varunbln commented Oct 16, 2026 •

Copy link
Copy Markdown
Contributor

hi again

ARA's modify_model rebuilds the full dense effective weight and runs a full input @ W_eff.T on 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.T is exactly the captured module output (minus bias). It is constant across closures, so it is precomputed once.
  • Only the rank-r term (input @ A.T) @ B.T is recomputed. This avoids both the dense [d_out x d_in] weight materialization and the full dense matmul.
  • With 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)

Hardware Model modify step end-to-end
M3 Pro (MPS) gemma-2-2b-it 2.45x 1.88x
H100 NVL gemma-2-2b-it 1.36x 1.25x
H100 NVL Qwen2.5-7B-Instruct 3.23x 2.46x

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. modify is 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:

  • baseline: (0, 4.37) (1, 0.218) (2, 0.086) (10, 0.072)
  • optimized: (0, 1.87) (1, 0.123) (2, 0.101) (3, 0.068)

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 okay

Testing

  • ruff format --check, ruff check --extend-select I, and ty check pass.
  • Standalone numerical check: reformulated vs dense closure output matches to about 1e-5, with preserve_row_magnitudes both on and off.
  • End to end on gemma-2-2b-it and Qwen2.5-7B-Instruct (CUDA), and gemma-2-2b-it (MPS).

Note: transformers 5.18+ requires torch>=2.6. This is unrelated to the change. I am flagging it from the CUDA runs.

…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.
Comment thread src/heretic/modifiers/ara.py Outdated
else:
bias = bias.float().to(lora_A.device)
good_output_base = good_output - bias
bad_output_base = bad_output - bias

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@varunbln varunbln Oct 17, 2026 •

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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!

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@p-e-w

p-e-w commented Oct 17, 2026

Copy link
Copy Markdown
Owner

Thanks! Those are impressive performance numbers, but I do wonder why you chose ancient models for the benchmark?

also, since ARA is marked reproducible = True, do you want this to be behind a flag?

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.

Note: transformers 5.18+ requires torch>=2.6. This is unrelated to the change. I am flagging it from the CUDA runs.

Thanks, I will update the docs for the 2.0 release.

@p-e-w

p-e-w commented Oct 17, 2026

Copy link
Copy Markdown
Owner

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).
@varunbln

varunbln commented Oct 17, 2026 •

Copy link
Copy Markdown
Contributor Author

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:

  • stock 381.0s -> low-rank 76.8s = 4.96x
  • end-to-end (modify + score) 502.2s -> 175.4s = 2.86x

The win scales with layer size, since the eliminated matmul is proportional to d_in * d_out. On the same hardware the end-to-end numbers go 2b: 1.25x, 7B: 2.46x, 27B: 2.86x, so modern/larger models benefit more, not less.

@varunbln

Copy link
Copy Markdown
Contributor Author

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 tests/test_ara_lowrank_equivalence.py, which runs the original dense closure and the low-rank closure on the same inputs and asserts they match to tolerance, both with and without preserve_row_magnitudes. That verifies the reformulation is correct by construction, which for this kind of change is arguably a stronger guarantee than a frozen output hash (the hash only tells you something changed, not whether the change is correct). The model hashes can then be re-baselined once, with that test guarding the path.

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.

Comment thread src/heretic/modifiers/ara.py Outdated
# 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

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But then why do we need these variables at all?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.
@varunbln
varunbln force-pushed the perf/ara-lowrank-closure branch from b9f0adf to 49e37f3 Compare October 19, 2026 08:45
@p-e-w

p-e-w commented Oct 19, 2026

Copy link
Copy Markdown
Owner

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).
@varunbln

Copy link
Copy Markdown
Contributor Author

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:

  • Bias handling. The assert bias is None I added failed on the gpt-oss test model, which has biased linear layers. I removed it and now handle bias properly instead: subtract the bias from the captured output, add the low-rank update and apply the row-norm rescale to that, then add the bias back. Bias-free modules are unchanged. This also fixes the old dense closure, which compared a bias-free product against a bias-inclusive output for those models. The equivalence test now covers biased layers too.
  • Test hashes. Only the two ARA tests change. I took the new hashes from the CI run: gpt-oss/SHA256SUMS.linux (same hash on 3.10 to 3.13), and seed-oss/SHA256SUMS.ci (3.10 to 3.12) and .ci2 (3.13).

I didn't touch the .windows files or seed-oss/SHA256SUMS.linux, since I don't have those machines to regenerate them. They'll need updating on your side.

@p-e-w

p-e-w commented Oct 20, 2026

Copy link
Copy Markdown
Owner

The assert bias is None I added failed on the gpt-oss test model, which has biased linear layers.

👁️ 👁️ 👁️

@p-e-w

p-e-w commented Oct 20, 2026

Copy link
Copy Markdown
Owner

I'm unable to reproduce the Pareto front after this change. The results are completely different, which shouldn't be the case.

Command

heretic --seed 555 unsloth/gemma-3-270m-it

master

   [Trial  37] Refusals: 0/100, KL divergence: 1.0440
   [Trial   4] Refusals: 1/100, KL divergence: 0.9828
   [Trial  55] Refusals: 2/100, KL divergence: 0.5610
   [Trial   2] Refusals: 4/100, KL divergence: 0.3704
   [Trial   8] Refusals: 6/100, KL divergence: 0.2821
   [Trial  20] Refusals: 7/100, KL divergence: 0.0979
   [Trial  96] Refusals: 8/100, KL divergence: 0.0713
   [Trial  34] Refusals: 12/100, KL divergence: 0.0566
   [Trial  87] Refusals: 16/100, KL divergence: 0.0365
   [Trial  86] Refusals: 19/100, KL divergence: 0.0338
   [Trial  90] Refusals: 23/100, KL divergence: 0.0286

This PR

   [Trial  27] Refusals: 0/100, KL divergence: 2.2300
   [Trial  37] Refusals: 3/100, KL divergence: 0.1463
   [Trial  86] Refusals: 6/100, KL divergence: 0.0804
   [Trial  57] Refusals: 9/100, KL divergence: 0.0689
   [Trial  83] Refusals: 10/100, KL divergence: 0.0674
   [Trial  68] Refusals: 12/100, KL divergence: 0.0558
   [Trial  96] Refusals: 13/100, KL divergence: 0.0527
   [Trial   8] Refusals: 16/100, KL divergence: 0.0299
   [Trial  88] Refusals: 17/100, KL divergence: 0.0297
   [Trial  69] Refusals: 18/100, KL divergence: 0.0164
   [Trial  62] Refusals: 29/100, KL divergence: 0.0105
   [Trial  97] Refusals: 30/100, KL divergence: 0.0091
   [Trial 100] Refusals: 37/100, KL divergence: 0.0089
   [Trial  40] Refusals: 41/100, KL divergence: 0.0076
   [Trial  80] Refusals: 44/100, KL divergence: 0.0076
   [Trial  42] Refusals: 47/100, KL divergence: 0.0056
   [Trial  66] Refusals: 48/100, KL divergence: 0.0049
   [Trial  90] Refusals: 94/100, KL divergence: 0.0012
   [Trial  17] Refusals: 97/100, KL divergence: 0.0005
   [Trial  32] Refusals: 98/100, KL divergence: 0.0000
   [Trial  43] Refusals: 98/100, KL divergence: 0.0000
   [Trial  61] Refusals: 98/100, KL divergence: 0.0000

I can confirm a 30% end-to-end speedup, which is great, but obviously, the Pareto fronts should be identical, at least approximately.

@p-e-w

p-e-w commented Oct 21, 2026

Copy link
Copy Markdown
Owner

It appears that unfortunately, this PR is completely broken. Same command as above (heretic --seed 555 unsloth/gemma-3-270m-it), I get these for the first trial:

master

Running trial 1 of 10...
* Parameters:
  * start_layer_index = 7
  * end_layer_index = 10
  * preserve_good_behavior_weight = 0.9445
  * steer_bad_behavior_weight = 0.0557
  * overcorrect_relative_weight = 0.7556
  * neighbor_count = 3
* Resetting model...
* Modifying model using Arbitrary-Rank Ablation (ARA)...
* Evaluating...
  * Refusals: 52/100
  * KL divergence: 2.5757

This PR

Running trial 1 of 10...
* Parameters:
  * start_layer_index = 7
  * end_layer_index = 10
  * preserve_good_behavior_weight = 0.9445
  * steer_bad_behavior_weight = 0.0557
  * overcorrect_relative_weight = 0.7556
  * neighbor_count = 3
* Resetting model...
* Modifying model using Arbitrary-Rank Ablation (ARA)...
* Evaluating...
  * Refusals: 17/100
  * KL divergence: 2.7618

The 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.
@varunbln

Copy link
Copy Markdown
Contributor Author

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 input @ W_base.T. That tensor is bf16, while master recomputes the product in float32 on every closure call. On gemma-3-270m layer 8 o_proj, the captured base differs from the fp32 product by a relative 3.0e-3. That is enough to steer L-BFGS somewhere else.

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, --seed 555, 2 trials, gemma-3-270m-it). The reference is this branch with only ara.py reverted to master:

variant trial 1 trial 2
master ARA 11/100, KL 2.6550 8/100, KL 0.4119
master ARA, second run identical identical
previous PR head 15/100, KL 2.8367 16/100, KL 0.2900
fixed PR 14/100, KL 2.5708 8/100, KL 0.3679

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 --seed 555 on the fixed branch and see whether the front looks comparable to master on your hardware? If it does not, I am fine with closing this.

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.

@p-e-w

p-e-w commented Oct 22, 2026

Copy link
Copy Markdown
Owner

Nope, it's more complicated than that.

master

? Which trial do you want to use? (Use arrow keys)
 » [Trial  37] Refusals: 0/100, KL divergence: 1.0440
   [Trial   4] Refusals: 1/100, KL divergence: 0.9828
   [Trial  55] Refusals: 2/100, KL divergence: 0.5610
   [Trial   2] Refusals: 4/100, KL divergence: 0.3704
   [Trial   8] Refusals: 6/100, KL divergence: 0.2821
   [Trial  20] Refusals: 7/100, KL divergence: 0.0979
   [Trial  96] Refusals: 8/100, KL divergence: 0.0713
   [Trial  34] Refusals: 12/100, KL divergence: 0.0566
   [Trial  87] Refusals: 16/100, KL divergence: 0.0365
   [Trial  86] Refusals: 19/100, KL divergence: 0.0338
   [Trial  90] Refusals: 23/100, KL divergence: 0.0286
   Run additional trials
   Exit program

This PR (updated)

? Which trial do you want to use? (Use arrow keys)
 » [Trial  45] Refusals: 0/100, KL divergence: 0.9720
   [Trial  65] Refusals: 1/100, KL divergence: 0.4958
   [Trial   2] Refusals: 2/100, KL divergence: 0.4071
   [Trial  99] Refusals: 4/100, KL divergence: 0.3680
   [Trial   8] Refusals: 6/100, KL divergence: 0.2846
   [Trial  20] Refusals: 7/100, KL divergence: 0.1191
   Run additional trials
   Exit program

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?

@p-e-w

p-e-w commented Oct 23, 2026

Copy link
Copy Markdown
Owner

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.

@varunbln

Copy link
Copy Markdown
Contributor Author

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. n_startup_trials is 30, so trials 1 to 30 come from the seeded random sampler and have the same parameters in every variant. From trial 31 on, TPE conditions on earlier results, so any change in any result changes every later parameter. Fronts built mostly from trials after 30 will never line up. Trials 1 to 30 can be compared one to one.

Test 1: perturb master itself. I made two changes to master's ara.py that are mathematically a no-op:

  • p64: compute input @ W_eff.T in float64 and cast back. This is strictly more accurate.
  • ulp: multiply the inputs by 1 + 2**-23, i.e. one float32 ulp.

Then I ran --seed 555 --n-trials 30 on gemma-3-270m-it (Mac, MPS) and compared the 30 startup trials to master:

variant vs master identical trials mean abs Δrefusals max Δrefusals median rel ΔKL
master, second run 30/30 0 0 0%
master + float64 0/30 8.5 55 9%
master + 1 ulp 0/30 5.0 27 9%
this PR 0/30 6.3 31 9%

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 o_proj I/O (400 harmless, 400 harmful prompts), used your trial 1 parameters, and ran the same L-BFGS settings (5 outer steps, 20 inner, history 10, strong Wolfe) on CPU in float32.

  • Master and the low rank objective agree to 7 digits at random points and at a trained point. Same function.
  • With the inputs scaled by 1 + k ulps for k = 0..7, master's final loss lands in one of two basins: -0.123 -0.122 -0.121 -0.291 -0.125 -0.121 -0.292 -0.291. The low rank version gives -0.291 -0.293 -0.121 -0.292 -0.293 -0.126 -0.121 -0.292. Same two basins, picked effectively at random by rounding.
  • One ulp changes master's B @ A by 55% (relative norm) after the first outer step. The kNN term is a likely amplifier: topk makes a discrete choice of neighbours, and after one ulp only 84% of bad rows keep the same 3 nearest good neighbours.

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.

@p-e-w

p-e-w commented Oct 23, 2026

Copy link
Copy Markdown
Owner

There are no basins, it's more complicated than that. I'll post the full explanation in the next hour.

@p-e-w

p-e-w commented Oct 23, 2026

Copy link
Copy Markdown
Owner

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.

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants