mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-30 06:30:37 +02:00
fix(abliteration): propagate offload integrity failures
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user