fix(ara): store captured I/O tensors on CPU for multi-GPU robustness (#214)

Extends d79a443 — that commit correctly moves I/O tensors to the weight
matrix's device before L-BFGS optimization, but the captured tensors
remain on their original GPU between trials. When reset_model() reloads
the model, device assignments can change, leaving orphaned tensors on
GPUs that now need that VRAM for the reloaded weights.

Moving to CPU at capture time in get_module_io ensures:
- Zero VRAM wasted on stale device assignments between trials
- Clean CPU→target transfer regardless of how devices shuffle on reload
- No overhead on single-GPU (.cpu() is a no-op when already on CPU,
  and .to(device) in ara_abliterate handles the final placement)
This commit is contained in:
joninco
2026-03-07 09:10:45 -05:00
committed by GitHub
parent d79a443e6f
commit 2bb203ee47
+4 -2
View File
@@ -803,8 +803,10 @@ class Model:
# inputs[0] and outputs have shape (prompt, position, component),
# so this extracts the input/output at the end of each prompt.
input = inputs[0][:, -1, :].detach().clone()
output = outputs[:, -1, :].detach().clone()
# Move to CPU to decouple from device assignments, which can
# change between model reloads in multi-GPU configurations.
input = inputs[0][:, -1, :].detach().clone().cpu()
output = outputs[:, -1, :].detach().clone().cpu()
# The modules associated with a component (e.g. expert MLPs)
# are not necessarily invoked in order, nor are all of them