wassname/jsteer
1
0
Fork 0
mirror of https://github.com/wassname/jsteer.git synced 2026-09-09 11:25:03 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: _to_vector forces CPU fp32 -- vjp path returned cuda Vectors, crashing U4 step-2 cosine vs cpu reference

Browse Source 
Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
This commit is contained in:
wassname
2026-07-10 13:19:54 +08:00
co-authored by Claudypoo
parent 6f59986139
commit 261a79ebdd
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
jsteer/jacobian.py
+5 -2
View File
@@ -82,9 +82,12 @@ def _resolve_layers(layers, n_layers: int) -> list[int]:
def _to_vector(cfg: SteeringConfig, per_layer: dict[int, Tensor]) -> Vector:
"""Wrap unit directions as a steering-lite Vector (mean_diff's layout:
stacked["v"] with leading k=1 dim, shared empty)."""
stacked["v"] with leading k=1 dim, shared empty). Always CPU fp32 --
Jacobian.pullback is CPU-native but the VJP path accumulates on cuda;
without the .cpu() the two paths return device-inconsistent Vectors
(Claude: found by U4 step-2 crash, pueue 550)."""
shared = {l: {} for l in per_layer}
stacked = {l: {"v": _unit(v.float()).unsqueeze(0)} for l, v in per_layer.items()}
stacked = {l: {"v": _unit(v.float().cpu()).unsqueeze(0)} for l, v in per_layer.items()}
return Vector(cfg, shared, stacked)