fix(abliteration): propagate offload integrity failures

This commit is contained in:
Joseph Magly
2026-08-21 18:39:07 -04:00
parent d1c0ee664d
commit 306d12c925
+9
View File
@@ -68,6 +68,7 @@ dev.configure_cuda_alloc()
from obliteratus.models.loader import ModelHandle, load_model # noqa: E402
from obliteratus.models.offload_surgery import ( # noqa: E402
LogicalParameterTransaction,
OffloadSurgeryError,
UnsupportedOffloadLayoutError,
logical_module_device,
resolve_logical_parameter,
@@ -3926,6 +3927,8 @@ class AbliterationPipeline:
norm_preserve=dir_norm_preserve,
regularization=0.0, # full removal of residual
)
except OffloadSurgeryError:
raise
except (AttributeError, RuntimeError) as e:
warnings.warn(
f"Layer {idx}: attention projection failed ({type(e).__name__}: {e}). "
@@ -4027,6 +4030,8 @@ class AbliterationPipeline:
total_neurons_masked += n_masked
count += ffn_count
except OffloadSurgeryError:
raise
except (AttributeError, RuntimeError) as e:
warnings.warn(
f"Layer {idx}: FFN projection failed ({type(e).__name__}: {e}). "
@@ -4116,6 +4121,8 @@ class AbliterationPipeline:
projection_row_fraction=self.projection_row_fraction,
offload_roots=(layers[idx],),
)
except OffloadSurgeryError:
raise
except (AttributeError, RuntimeError):
pass
if sae_ffn is not None:
@@ -4137,6 +4144,8 @@ class AbliterationPipeline:
offload_roots=(layers[idx],),
)
sae_count += fc
except OffloadSurgeryError:
raise
except (AttributeError, RuntimeError):
pass
del sd