From 306d12c925c9bea18b753161f02d610885d5745c Mon Sep 17 00:00:00 2001
From: Joseph Magly <1159087+jmagly@users.noreply.github.com>
Date: Fri, 21 Aug 2026 18:39:07 -0400
Subject: [PATCH] fix(abliteration): propagate offload integrity failures

---
 obliteratus/abliterate.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py
index bc2fcd8..d77921f 100644
--- a/obliteratus/abliterate.py
+++ b/obliteratus/abliterate.py
@@ -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