mirror of
https://github.com/p-e-w/heretic.git
synced 2026-09-14 08:06:14 -07:00
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user