From 2bb203ee4779592c6569c10ed74cdaeac38e18b2 Mon Sep 17 00:00:00 2001 From: joninco Date: Sat, 7 Mar 2026 09:10:45 -0500 Subject: [PATCH] fix(ara): store captured I/O tensors on CPU for multi-GPU robustness (#214) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- src/heretic/model.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/heretic/model.py b/src/heretic/model.py index 898d4df..a3556df 100644 --- a/src/heretic/model.py +++ b/src/heretic/model.py @@ -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