feat(abliteration): support safe offloaded weight surgery

This commit is contained in:
Joseph Magly
2026-08-21 18:36:18 -04:00
parent c65bedfb62
commit d1c0ee664d
6 changed files with 1113 additions and 51 deletions
+1
View File
@@ -131,6 +131,7 @@
"coverage_paths": [
"obliteratus/device.py",
"obliteratus/models/loader.py",
"obliteratus/models/offload_surgery.py",
"obliteratus/models/quant_dequant.py"
]
},
+16 -1
View File
@@ -50,6 +50,7 @@
],
"paths": [
"obliteratus/abliterate.py",
"obliteratus/models/offload_surgery.py",
"obliteratus/persistence_contracts.py",
"obliteratus/auto_obliterate.py",
"obliteratus/bayesian_optimizer.py",
@@ -74,8 +75,10 @@
"tests/test_sweep_contracts.py",
"tests/test_tourney_contracts.py",
"tests/test_checkpoint_atomicity.py",
"tests/test_offload_surgery.py",
"tests/test_persistence_contracts.py",
"tests/test_persistence_pipeline.py"
"tests/test_persistence_pipeline.py",
"tests/conditional/test_cuda_runtime.py"
]
},
{
@@ -458,6 +461,18 @@
"cuda-runtime"
]
},
{
"path": "obliteratus/models/offload_surgery.py",
"risk_class": "mixed-runtime",
"risk": "transactional logical-parameter mutation across live, CPU-offloaded, disk-backed, and CUDA execution layouts",
"required_tests": [
"tests/test_offload_surgery.py",
"tests/conditional/test_cuda_runtime.py"
],
"conditional_gates": [
"cuda-runtime"
]
},
{
"path": "obliteratus/mlx_backend.py",
"risk_class": "mixed-runtime",
+301 -50
View File
@@ -66,6 +66,13 @@ def _has_fused_quant_scale(container: nn.Module, name: str) -> bool:
dev.configure_cuda_alloc()
from obliteratus.models.loader import ModelHandle, load_model # noqa: E402
from obliteratus.models.offload_surgery import ( # noqa: E402
LogicalParameterTransaction,
UnsupportedOffloadLayoutError,
logical_module_device,
resolve_logical_parameter,
validate_offloaded_parameters,
)
from obliteratus.analysis.numerical_contracts import ( # noqa: E402
orthogonalize_subspace_rows,
project_weight_against_direction,
@@ -3572,6 +3579,71 @@ class AbliterationPipeline:
total_neurons_masked = 0
total_sae_projections = 0
# Resolve every selected layer's meta-resident parameter before any
# surgery. Unknown or unsupported Accelerate layouts therefore fail
# closed without leaving an earlier layer partially modified.
offloaded_selected = False
for layer_index in self._strong_layers:
layer = layers[layer_index]
meta_parameters = [
(name, parameter)
for name, parameter in layer.named_parameters()
if parameter.device.type == "meta"
]
if meta_parameters:
offloaded_selected = True
fused_names = [
name for name, parameter in meta_parameters if parameter.dim() > 2
]
if fused_names:
raise UnsupportedOffloadLayoutError(
"offloaded fused expert tensors are not yet a supported "
f"surgery layout: {', '.join(fused_names)}"
)
validate_offloaded_parameters(layer)
if offloaded_selected:
unsupported_modes = [
name
for enabled, name in (
(getattr(self, "use_lora_ablation", False), "LoRA ablation"),
(getattr(self, "attention_head_surgery", False), "attention-head surgery"),
(getattr(self, "safety_neuron_masking", False), "safety-neuron masking"),
(getattr(self, "expert_transplant", False), "expert transplant"),
)
if enabled
]
if unsupported_modes:
raise UnsupportedOffloadLayoutError(
"Accelerate-offloaded surgery does not yet support: "
+ ", ".join(unsupported_modes)
)
if self.handle:
model = self.handle.model
for head_name in ("lm_head", "embed_out", "output"):
head = getattr(model, head_name, None)
weight = getattr(head, "weight", None)
if weight is not None and weight.device.type == "meta":
resolve_logical_parameter(head, search_roots=(model,))
if self.project_embeddings:
for embedding_path in (
"model.embed_tokens",
"model.language_model.embed_tokens",
"transformer.wte",
"model.embed_in",
"gpt_neox.embed_in",
):
embedding = model
for path_component in embedding_path.split("."):
embedding = getattr(embedding, path_component, None)
if embedding is None:
break
weight = getattr(embedding, "weight", None)
if weight is not None and weight.device.type == "meta":
resolve_logical_parameter(embedding, search_roots=(model,))
break
# ── Bayesian optimization pre-pass ─────────────────────────────
# When enabled, run Optuna TPE to find optimal per-layer regularization
# before the standard projection loop. The found values override the
@@ -3724,7 +3796,7 @@ class AbliterationPipeline:
for idx in self._strong_layers:
subspace = self.refusal_subspaces[idx]
device = next(layers[idx].parameters()).device
device = logical_module_device(layers[idx])
# Layer-adaptive regularization: scale projection per-layer
layer_reg = self.regularization
@@ -3827,9 +3899,15 @@ class AbliterationPipeline:
norm_preserve=dir_norm_preserve,
regularization=attn_reg,
projection_row_fraction=self.projection_row_fraction,
offload_roots=(layers[idx],),
)
if self.project_biases:
count += self._project_bias(attn, d, _attn_names)
count += self._project_bias(
attn,
d,
_attn_names,
offload_roots=(layers[idx],),
)
# Additional head surgery: second-pass precision targeting
# on the top safety heads to remove residual refusal signal.
@@ -3870,6 +3948,7 @@ class AbliterationPipeline:
norm_preserve=dir_norm_preserve,
regularization=mlp_reg,
projection_row_fraction=self.projection_row_fraction,
offload_roots=(layers[idx],),
)
if ffn_count == 0:
# MoE path
@@ -3884,6 +3963,7 @@ class AbliterationPipeline:
norm_preserve=dir_norm_preserve,
regularization=mlp_reg,
project_biases=self.project_biases,
offload_roots=(layers[idx],),
)
elif self.invert_refusal and idx in self._expert_safety_scores:
# Selective MoE inversion: router reflected, safety
@@ -3892,6 +3972,7 @@ class AbliterationPipeline:
ffn, d, idx,
norm_preserve=dir_norm_preserve,
project_biases=self.project_biases,
offload_roots=(layers[idx],),
)
else:
ffn_count = self._project_moe_experts(
@@ -3900,10 +3981,16 @@ class AbliterationPipeline:
regularization=mlp_reg,
project_biases=self.project_biases,
projection_row_fraction=self.projection_row_fraction,
offload_roots=(layers[idx],),
)
elif self.projection_target == "output":
if self.project_biases:
ffn_count += self._project_bias(ffn, d, _FFN_OUT_NAMES)
ffn_count += self._project_bias(
ffn,
d,
_FFN_OUT_NAMES,
offload_roots=(layers[idx],),
)
else:
# Dense model: also project FFN input projections
# (up_proj, gate_proj carry refusal signal too)
@@ -3912,10 +3999,14 @@ class AbliterationPipeline:
norm_preserve=dir_norm_preserve,
regularization=mlp_reg,
projection_row_fraction=self.projection_row_fraction,
offload_roots=(layers[idx],),
)
if self.project_biases:
ffn_count += self._project_bias(
ffn, d, _FFN_OUT_NAMES + _FFN_IN_NAMES,
ffn,
d,
_FFN_OUT_NAMES + _FFN_IN_NAMES,
offload_roots=(layers[idx],),
)
# Safety-neuron masking (applied after projection for
@@ -4023,6 +4114,7 @@ class AbliterationPipeline:
norm_preserve=self.norm_preserve,
regularization=sae_reg,
projection_row_fraction=self.projection_row_fraction,
offload_roots=(layers[idx],),
)
except (AttributeError, RuntimeError):
pass
@@ -4033,6 +4125,7 @@ class AbliterationPipeline:
norm_preserve=self.norm_preserve,
regularization=sae_reg,
projection_row_fraction=self.projection_row_fraction,
offload_roots=(layers[idx],),
)
if fc == 0:
fc = self._project_moe_experts(
@@ -4041,6 +4134,7 @@ class AbliterationPipeline:
regularization=sae_reg,
project_biases=False,
projection_row_fraction=self.projection_row_fraction,
offload_roots=(layers[idx],),
)
sae_count += fc
except (AttributeError, RuntimeError):
@@ -4107,7 +4201,7 @@ class AbliterationPipeline:
model = self.handle.model
if last_strong in self.refusal_subspaces:
subspace = self.refusal_subspaces[last_strong]
lm_device = self._get_model_device(model)
lm_device = logical_module_device(model)
# Pre-transfer subspace and resolve lm_head module once
subspace_on_device = subspace.to(lm_device)
lm_head_name = None
@@ -4131,9 +4225,11 @@ class AbliterationPipeline:
and lm_head_obj is not None
and hasattr(lm_head_obj, "weight")
)
lm_original_norm = 0.0
if lm_multi_dir:
lm_original_norm = lm_head_obj.weight.data.norm().item()
lm_saved_norms = (
self._capture_layer_weight_norms(lm_head_obj)
if lm_multi_dir
else {}
)
for dir_idx in range(subspace_on_device.shape[0]):
d = subspace_on_device[dir_idx].unsqueeze(-1)
lm_head_count += self._project_out_advanced(
@@ -4144,14 +4240,8 @@ class AbliterationPipeline:
)
del d
# Restore lm_head norm once after all directions
if lm_multi_dir and lm_original_norm > 0 and lm_head_obj is not None:
new_norm = lm_head_obj.weight.data.norm().item()
if new_norm > 0 and not math.isnan(new_norm) and not math.isinf(new_norm):
ratio = lm_original_norm / new_norm
if ratio > _MAX_NORM_RATIO:
ratio = _MAX_NORM_RATIO
if abs(ratio - 1.0) > 1e-6:
lm_head_obj.weight.data.mul_(ratio)
if lm_multi_dir and lm_saved_norms and lm_head_obj is not None:
self._restore_layer_weight_norms(lm_head_obj, lm_saved_norms)
del subspace_on_device
if lm_head_count > 0:
total_modified += lm_head_count
@@ -4207,6 +4297,7 @@ class AbliterationPipeline:
norm_preserve=True, # always norm-preserve embeds
regularization=emb_reg,
projection_row_fraction=self.projection_row_fraction,
offload_roots=(model,),
)
break
del d
@@ -4930,7 +5021,13 @@ class AbliterationPipeline:
if identity in seen:
continue
param_name = f"{module_name}.weight" if module_name else "weight"
data, _requires_replacement = AbliterationPipeline._dequantize_weight(module)
if weight is not None and weight.device.type == "meta":
data = resolve_logical_parameter(
module,
search_roots=(layer,),
).tensor
else:
data, _requires_replacement = AbliterationPipeline._dequantize_weight(module)
norms[param_name] = data.float().norm().item()
seen.add(identity)
return norms
@@ -4974,7 +5071,16 @@ class AbliterationPipeline:
continue
original_norm = saved_norms[param_name]
data, requires_replacement = AbliterationPipeline._dequantize_weight(module)
transaction = None
if weight is not None and weight.device.type == "meta":
transaction = resolve_logical_parameter(
module,
search_roots=(layer,),
)
data = transaction.tensor
requires_replacement = False
else:
data, requires_replacement = AbliterationPipeline._dequantize_weight(module)
ratio = norm_restoration_ratio(
original_norm,
data.float().norm().item(),
@@ -4985,7 +5091,9 @@ class AbliterationPipeline:
continue
with torch.no_grad():
if storage_kind == "integer":
if transaction is not None:
transaction.commit(data.mul(ratio))
elif storage_kind == "integer":
# Integer storage has no scale/zero-point contract. Casting
# back can erase the update, so materialize logical float
# values while preserving Parameter identity and ties.
@@ -5020,6 +5128,7 @@ class AbliterationPipeline:
norm_preserve: bool = False,
regularization: float = 0.0,
projection_row_fraction: float = 1.0,
offload_roots: tuple[nn.Module, ...] = (),
) -> int:
"""Advanced projection with norm preservation and regularization.
@@ -5039,16 +5148,45 @@ class AbliterationPipeline:
in-place operations on packed NF4 storage are silent no-ops.
"""
count = 0
candidates: list[
tuple[nn.Module, LogicalParameterTransaction | None, torch.Tensor, bool]
] = []
for name in candidate_names:
proj = getattr(module, name, None)
if proj is None or not hasattr(proj, "weight"):
continue
W, is_quantized = AbliterationPipeline._dequantize_weight(proj)
if proj.weight.device.type == "meta":
transaction = resolve_logical_parameter(
proj,
search_roots=(module, *offload_roots),
)
candidates.append((proj, transaction, transaction.tensor, False))
else:
W, is_quantized = AbliterationPipeline._dequantize_weight(proj)
transaction = (
None
if is_quantized
else resolve_logical_parameter(proj)
)
candidates.append((proj, transaction, W, is_quantized))
# Resolve every backing store before the first mutation. An unsupported
# later candidate therefore cannot leave an earlier projection changed.
updates: list[
tuple[
nn.Module,
LogicalParameterTransaction | None,
torch.Tensor,
bool,
torch.Tensor,
]
] = []
for proj, transaction, W, is_quantized in candidates:
result = project_weight_against_direction(
W,
direction,
direction.to(W.device),
norm_preserve=norm_preserve,
regularization=regularization,
projection_row_fraction=projection_row_fraction,
@@ -5057,11 +5195,23 @@ class AbliterationPipeline:
if not result.projected:
continue
W.copy_(result.weight)
if is_quantized:
AbliterationPipeline._replace_quantized_weight(proj, W)
updates.append((proj, transaction, W, is_quantized, result.weight))
count += 1
committed: list[LogicalParameterTransaction] = []
try:
for proj, transaction, W, is_quantized, updated in updates:
if transaction is not None:
transaction.commit(updated)
committed.append(transaction)
else:
W.copy_(updated)
if transaction is None and is_quantized:
AbliterationPipeline._replace_quantized_weight(proj, W)
count += 1
except Exception:
for transaction in reversed(committed):
transaction.rollback()
raise
return count
@@ -5070,6 +5220,7 @@ class AbliterationPipeline:
module: nn.Module,
direction: torch.Tensor,
candidate_names: list[str],
offload_roots: tuple[nn.Module, ...] = (),
) -> int:
"""Project the refusal direction out of bias terms.
@@ -5081,6 +5232,9 @@ class AbliterationPipeline:
Gabliteration, grimjim) do not project biases.
"""
count = 0
candidates: list[
tuple[LogicalParameterTransaction | None, torch.Tensor]
] = []
for name in candidate_names:
proj = getattr(module, name, None)
if proj is None or not hasattr(proj, "bias"):
@@ -5088,16 +5242,40 @@ class AbliterationPipeline:
if proj.bias is None:
continue
b = proj.bias.data
if proj.bias.device.type == "meta":
transaction = resolve_logical_parameter(
proj,
"bias",
search_roots=(module, *offload_roots),
)
candidates.append((transaction, transaction.tensor))
else:
transaction = resolve_logical_parameter(proj, "bias")
candidates.append((transaction, transaction.tensor))
updates: list[tuple[LogicalParameterTransaction, torch.Tensor]] = []
for transaction, b in candidates:
d = direction.to(device=b.device, dtype=b.dtype).squeeze() # (hidden_dim,)
if b.shape[0] == d.shape[0]:
# Bias is (out_features,) = (hidden_dim,) for output projections
component = (b @ d).unsqueeze(0) * d # scalar * direction
proj.bias.data = b - component.squeeze()
count += 1
updated = b - component.squeeze()
assert transaction is not None
updates.append((transaction, updated))
# else: dimension mismatch — expected for GQA k/v projections,
# fused QKV (c_attn), and MoE routers. Skip silently.
committed: list[LogicalParameterTransaction] = []
try:
for transaction, updated in updates:
transaction.commit(updated)
committed.append(transaction)
count += 1
except Exception:
for transaction in reversed(committed):
transaction.rollback()
raise
return count
@staticmethod
@@ -5246,7 +5424,10 @@ class AbliterationPipeline:
return 0
@staticmethod
def _stabilize_router_weights(ffn_module: nn.Module):
def _stabilize_router_weights(
ffn_module: nn.Module,
offload_roots: tuple[nn.Module, ...] = (),
):
"""Clamp router weights after projection to prevent extreme routing.
After projecting the refusal direction from router weights, modified
@@ -5260,11 +5441,23 @@ class AbliterationPipeline:
for rname in _ROUTER_NAMES:
gate = getattr(ffn_module, rname, None)
if gate is not None and hasattr(gate, "weight"):
W = gate.weight.data
transaction = None
if gate.weight.device.type == "meta":
transaction = resolve_logical_parameter(
gate,
search_roots=(ffn_module, *offload_roots),
)
W = transaction.tensor
else:
W = gate.weight.data
std = W.std()
if std > 0:
mean = W.mean()
gate.weight.data = W.clamp(mean - 3 * std, mean + 3 * std)
stabilized = W.clamp(mean - 3 * std, mean + 3 * std)
if transaction is not None:
transaction.commit(stabilized)
else:
W.copy_(stabilized)
return
# Auto-detect fallback
if getattr(ffn_module, "experts", None) is not None:
@@ -5273,12 +5466,25 @@ class AbliterationPipeline:
continue
if not hasattr(child, "weight"):
continue
W = child.weight
if W.shape[0] < 512 and W.shape[0] != W.shape[-1]:
weight = child.weight
if weight.shape[0] < 512 and weight.shape[0] != weight.shape[-1]:
transaction = None
if weight.device.type == "meta":
transaction = resolve_logical_parameter(
child,
search_roots=(ffn_module, *offload_roots),
)
W = transaction.tensor
else:
W = weight.data
std = W.data.std()
if std > 0:
mean = W.data.mean()
child.weight.data = W.data.clamp(mean - 3 * std, mean + 3 * std)
stabilized = W.data.clamp(mean - 3 * std, mean + 3 * std)
if transaction is not None:
transaction.commit(stabilized)
else:
W.copy_(stabilized)
return
@staticmethod
@@ -5289,6 +5495,7 @@ class AbliterationPipeline:
regularization: float = 0.0,
project_biases: bool = False,
projection_row_fraction: float = 1.0,
offload_roots: tuple[nn.Module, ...] = (),
) -> int:
"""Project refusal direction from all MoE components.
@@ -5329,10 +5536,11 @@ class AbliterationPipeline:
norm_preserve=norm_preserve,
regularization=regularization,
projection_row_fraction=projection_row_fraction,
offload_roots=offload_roots,
)
if project_biases:
count += AbliterationPipeline._project_bias(
ffn_module, direction, [rname],
ffn_module, direction, [rname], offload_roots,
)
router_found = True
break # only one router per MoE block
@@ -5363,10 +5571,11 @@ class AbliterationPipeline:
norm_preserve=norm_preserve,
regularization=regularization,
projection_row_fraction=projection_row_fraction,
offload_roots=offload_roots,
)
if project_biases:
count += AbliterationPipeline._project_bias(
ffn_module, direction, [child_name],
ffn_module, direction, [child_name], offload_roots,
)
router_found = True
break
@@ -5386,6 +5595,7 @@ class AbliterationPipeline:
norm_preserve=norm_preserve,
regularization=regularization,
projection_row_fraction=projection_row_fraction,
offload_roots=offload_roots,
)
# Input projections
count += AbliterationPipeline._project_out_advanced(
@@ -5393,13 +5603,14 @@ class AbliterationPipeline:
norm_preserve=norm_preserve,
regularization=regularization,
projection_row_fraction=projection_row_fraction,
offload_roots=offload_roots,
)
if project_biases:
count += AbliterationPipeline._project_bias(
shared, direction, _FFN_OUT_NAMES,
shared, direction, _FFN_OUT_NAMES, offload_roots,
)
count += AbliterationPipeline._project_bias(
shared, direction, _FFN_IN_NAMES,
shared, direction, _FFN_IN_NAMES, offload_roots,
)
break
@@ -5411,6 +5622,7 @@ class AbliterationPipeline:
norm_preserve=norm_preserve,
regularization=regularization,
projection_row_fraction=projection_row_fraction,
offload_roots=offload_roots,
)
# ── Routed expert projection ──────────────────────────────────────
@@ -5449,6 +5661,7 @@ class AbliterationPipeline:
norm_preserve=norm_preserve,
regularization=regularization,
projection_row_fraction=projection_row_fraction,
offload_roots=offload_roots,
)
# Input projections (up_proj, gate_proj, w1, w3, etc.)
expert_count += AbliterationPipeline._project_out_advanced(
@@ -5456,13 +5669,14 @@ class AbliterationPipeline:
norm_preserve=norm_preserve,
regularization=regularization,
projection_row_fraction=projection_row_fraction,
offload_roots=offload_roots,
)
if project_biases:
expert_count += AbliterationPipeline._project_bias(
expert, direction, _FFN_OUT_NAMES,
expert, direction, _FFN_OUT_NAMES, offload_roots,
)
expert_count += AbliterationPipeline._project_bias(
expert, direction, _FFN_IN_NAMES,
expert, direction, _FFN_IN_NAMES, offload_roots,
)
count += expert_count
@@ -5470,7 +5684,9 @@ class AbliterationPipeline:
# Stabilize router weights after projection to prevent extreme logits
# that cause CUDA illegal memory access during generation.
if count > 0:
AbliterationPipeline._stabilize_router_weights(ffn_module)
AbliterationPipeline._stabilize_router_weights(
ffn_module, offload_roots
)
return count
@@ -5481,6 +5697,7 @@ class AbliterationPipeline:
layer_idx: int,
norm_preserve: bool = False,
project_biases: bool = False,
offload_roots: tuple[nn.Module, ...] = (),
) -> int:
"""MoE excision with selective inversion (refusal reflection).
@@ -5531,9 +5748,12 @@ class AbliterationPipeline:
ffn_module, direction, [rname],
norm_preserve=norm_preserve,
regularization=router_reg,
offload_roots=offload_roots,
)
if project_biases:
count += self._project_bias(ffn_module, direction, [rname])
count += self._project_bias(
ffn_module, direction, [rname], offload_roots
)
break
# Router auto-detection fallback
@@ -5550,6 +5770,7 @@ class AbliterationPipeline:
ffn_module, direction, [child_name],
norm_preserve=norm_preserve,
regularization=router_reg,
offload_roots=offload_roots,
)
break
@@ -5563,9 +5784,15 @@ class AbliterationPipeline:
shared, direction, _FFN_OUT_NAMES + _FFN_IN_NAMES,
norm_preserve=norm_preserve,
regularization=reflect_reg,
offload_roots=offload_roots,
)
if project_biases:
count += self._project_bias(shared, direction, _FFN_OUT_NAMES + _FFN_IN_NAMES)
count += self._project_bias(
shared,
direction,
_FFN_OUT_NAMES + _FFN_IN_NAMES,
offload_roots,
)
break
# ── Routed experts: selective inversion ───────────────────────
@@ -5581,9 +5808,15 @@ class AbliterationPipeline:
expert, direction, _FFN_OUT_NAMES + _FFN_IN_NAMES,
norm_preserve=norm_preserve,
regularization=reg,
offload_roots=offload_roots,
)
if project_biases:
count += self._project_bias(expert, direction, _FFN_OUT_NAMES + _FFN_IN_NAMES)
count += self._project_bias(
expert,
direction,
_FFN_OUT_NAMES + _FFN_IN_NAMES,
offload_roots,
)
else:
# Fused 3D: per-expert differentiation via per-slice processing.
# Safety experts get reflected, capability experts get standard removal.
@@ -5609,7 +5842,7 @@ class AbliterationPipeline:
# Stabilize router weights after reflection to prevent extreme logits
# that cause CUDA illegal memory access during generation.
if count > 0:
self._stabilize_router_weights(ffn_module)
self._stabilize_router_weights(ffn_module, offload_roots)
return count
@@ -5621,6 +5854,7 @@ class AbliterationPipeline:
norm_preserve: bool = False,
regularization: float = 0.0,
project_biases: bool = False,
offload_roots: tuple[nn.Module, ...] = (),
) -> int:
"""Expert-Granular Abliteration: per-expert direction projection.
@@ -5648,9 +5882,12 @@ class AbliterationPipeline:
ffn_module, direction, [rname],
norm_preserve=norm_preserve,
regularization=regularization,
offload_roots=offload_roots,
)
if project_biases:
count += self._project_bias(ffn_module, direction, [rname])
count += self._project_bias(
ffn_module, direction, [rname], offload_roots
)
router_found = True
break
if not router_found:
@@ -5662,6 +5899,7 @@ class AbliterationPipeline:
ffn_module, direction, [child_name],
norm_preserve=norm_preserve,
regularization=regularization,
offload_roots=offload_roots,
)
break
@@ -5673,16 +5911,22 @@ class AbliterationPipeline:
count += self._project_out_advanced(
shared, direction, _FFN_OUT_NAMES + _FFN_IN_NAMES,
norm_preserve=norm_preserve, regularization=regularization,
offload_roots=offload_roots,
)
if project_biases:
count += self._project_bias(shared, direction, _FFN_OUT_NAMES + _FFN_IN_NAMES)
count += self._project_bias(
shared,
direction,
_FFN_OUT_NAMES + _FFN_IN_NAMES,
offload_roots,
)
break
# ── Routed experts: per-expert directions ──
experts = getattr(ffn_module, "experts", None)
if experts is None:
if count > 0:
self._stabilize_router_weights(ffn_module)
self._stabilize_router_weights(ffn_module, offload_roots)
return count
expert_count = 0
@@ -5699,14 +5943,21 @@ class AbliterationPipeline:
expert, ed, _FFN_OUT_NAMES,
norm_preserve=norm_preserve,
regularization=regularization,
offload_roots=offload_roots,
)
expert_count += self._project_out_advanced(
expert, ed, _FFN_IN_NAMES,
norm_preserve=norm_preserve,
regularization=regularization,
offload_roots=offload_roots,
)
if project_biases:
expert_count += self._project_bias(expert, ed, _FFN_OUT_NAMES + _FFN_IN_NAMES)
expert_count += self._project_bias(
expert,
ed,
_FFN_OUT_NAMES + _FFN_IN_NAMES,
offload_roots,
)
else:
# Fused 3D: process per-expert with individual directions
expert_count += self._project_fused_3d_granular(
@@ -5726,7 +5977,7 @@ class AbliterationPipeline:
count += expert_count
if count > 0:
self._stabilize_router_weights(ffn_module)
self._stabilize_router_weights(ffn_module, offload_roots)
return count
@staticmethod
+376
View File
@@ -0,0 +1,376 @@
"""Transactional access to Accelerate-offloaded logical parameters.
Accelerate represents an offloaded parameter with a live ``meta`` tensor and
keeps the logical value in an ``AlignDevicesHook.weights_map``. Surgery must
therefore update that backing map without leaving a materialized live tensor or
silently depending on an unknown hook layout.
This module is the only OBLITERATUS compatibility boundary allowed to write to
Accelerate backing stores. Callers receive a detached logical tensor and commit
the completed update atomically; exceptions before commit leave both the live
module and the backing store unchanged.
"""
from __future__ import annotations
from collections.abc import Mapping, MutableMapping, Sequence
from dataclasses import dataclass, field
from importlib.metadata import PackageNotFoundError, version
from typing import Any
import torch
import torch.nn as nn
from obliteratus.runtime_contracts import classify_weight_storage
class OffloadSurgeryError(RuntimeError):
"""Base error for unsupported or inconsistent offloaded surgery state."""
class UnsupportedOffloadLayoutError(OffloadSurgeryError):
"""Accelerate exposed a layout that this compatibility adapter cannot write."""
class UnsupportedOffloadedQuantizationError(OffloadSurgeryError):
"""An offloaded packed or quantized parameter cannot be updated safely."""
def _accelerate_release() -> tuple[int, int, int]:
try:
raw = version("accelerate").split("+", 1)[0]
except PackageNotFoundError as exc: # pragma: no cover - dependency contract
raise UnsupportedOffloadLayoutError(
"Accelerate is required to resolve a meta-resident parameter"
) from exc
numeric: list[int] = []
for part in raw.split(".")[:3]:
digits = "".join(character for character in part if character.isdigit())
numeric.append(int(digits or 0))
while len(numeric) < 3:
numeric.append(0)
release = tuple(numeric)
if not ((release[0] == 0 and release[1] >= 24) or release[0] == 1):
raise UnsupportedOffloadLayoutError(
f"Accelerate {raw} is outside the tested offload-surgery range "
"(>=0.24,<2); refusing compatibility access"
)
return release
def _align_hooks(hook: Any) -> list[Any]:
"""Return supported offloading hooks from direct or sequential layouts."""
from accelerate.hooks import AlignDevicesHook, SequentialHook
if isinstance(hook, AlignDevicesHook):
return [hook] if hook.offload else []
if isinstance(hook, SequentialHook):
hooks: list[Any] = []
for child in hook.hooks:
hooks.extend(_align_hooks(child))
return hooks
return []
def _relative_module_path(owner: nn.Module, target: nn.Module) -> str | None:
for name, candidate in owner.named_modules():
if candidate is target:
return name
return None
@dataclass(frozen=True)
class _WriteTarget:
"""A validated readable/writable view of one backing-store key."""
container: Any
key: str
loader: bool = False
def read(self) -> torch.Tensor:
value = self.container[self.key]
if not isinstance(value, torch.Tensor):
raise UnsupportedOffloadLayoutError(
f"Accelerate backing value {self.key!r} is not a tensor"
)
return value
def write(self, value: torch.Tensor) -> None:
if self.loader:
self.container.state_dict[self.key] = value
if self.key not in self.container.all_keys:
self.container.all_keys.append(self.key)
return
self.container[self.key] = value
def _write_target(weights_map: Mapping[str, torch.Tensor], key: str) -> _WriteTarget:
"""Resolve public mappings and the two version-gated Accelerate wrappers."""
from accelerate.utils.offload import OffloadedWeightsLoader, PrefixedDataset
if isinstance(weights_map, PrefixedDataset):
return _write_target(weights_map.dataset, f"{weights_map.prefix}{key}")
if isinstance(weights_map, OffloadedWeightsLoader):
return _WriteTarget(weights_map, key, loader=True)
if isinstance(weights_map, MutableMapping):
return _WriteTarget(weights_map, key)
raise UnsupportedOffloadLayoutError(
"Accelerate weights_map is read-only and has no supported writable "
f"backing store ({type(weights_map).__module__}.{type(weights_map).__name__})"
)
def _same_storage(left: torch.Tensor, right: torch.Tensor) -> bool:
if left is right:
return True
if left.device.type == "meta" or right.device.type == "meta":
return False
return (
left.shape == right.shape
and left.dtype == right.dtype
and left.data_ptr() == right.data_ptr()
)
def _alias_targets(primary: _WriteTarget, source: torch.Tensor) -> list[_WriteTarget]:
"""Find inexpensive in-memory aliases so tied backing values stay tied."""
targets = [primary]
if primary.loader:
candidates = primary.container.state_dict.items()
elif isinstance(primary.container, MutableMapping):
candidates = primary.container.items()
else: # pragma: no cover - construction guarantees one of the above
candidates = ()
for key, value in candidates:
if key == primary.key or not isinstance(value, torch.Tensor):
continue
if _same_storage(value, source):
targets.append(_WriteTarget(primary.container, key, loader=primary.loader))
return targets
@dataclass
class LogicalParameterTransaction:
"""A detached logical parameter with an atomic commit operation."""
module: nn.Module
parameter_name: str
tensor: torch.Tensor
offloaded: bool
backing_key: str | None = None
_targets: list[_WriteTarget] = field(default_factory=list, repr=False)
_parameter_identity: int = field(init=False, repr=False)
_committed: bool = field(default=False, init=False, repr=False)
_rollback_values: list[tuple[_WriteTarget | None, torch.Tensor]] = field(
default_factory=list,
init=False,
repr=False,
)
def __post_init__(self) -> None:
self._parameter_identity = id(getattr(self.module, self.parameter_name))
def commit(self, updated: torch.Tensor) -> None:
"""Commit a complete logical value without exposing partial mutation."""
parameter = getattr(self.module, self.parameter_name)
if id(parameter) != self._parameter_identity:
raise OffloadSurgeryError(
f"{self.parameter_name} identity changed during logical surgery"
)
if updated.shape != self.tensor.shape:
raise OffloadSurgeryError(
f"logical update changed shape from {tuple(self.tensor.shape)} "
f"to {tuple(updated.shape)}"
)
if not updated.is_floating_point():
raise UnsupportedOffloadedQuantizationError(
"logical surgery requires a floating-point result"
)
if not self.offloaded:
original = parameter.data.detach().clone()
with torch.no_grad():
parameter.data.copy_(updated.to(parameter.device, parameter.dtype))
self._rollback_values = [(None, original)]
self._committed = True
return
if parameter.device.type != "meta":
raise OffloadSurgeryError(
"offloaded parameter was unexpectedly materialized before commit"
)
committed = updated.detach().to(device="cpu", dtype=self.tensor.dtype).clone()
originals = [(target, target.read()) for target in self._targets]
try:
for target, _original in originals:
target.write(committed)
except Exception as exc:
for target, original in reversed(originals):
try:
target.write(original)
except Exception:
pass
raise OffloadSurgeryError(
f"failed to commit logical parameter {self.backing_key!r}; "
"the backing store was rolled back"
) from exc
if parameter.device.type != "meta" or id(parameter) != self._parameter_identity:
for target, original in reversed(originals):
target.write(original)
raise OffloadSurgeryError(
"offload invariant changed during commit; backing update rolled back"
)
self._rollback_values = originals
self._committed = True
def rollback(self) -> None:
"""Restore the value captured by the most recent successful commit."""
if not self._committed:
return
parameter = getattr(self.module, self.parameter_name)
if id(parameter) != self._parameter_identity:
raise OffloadSurgeryError(
f"cannot roll back changed {self.parameter_name} identity"
)
if self.offloaded:
for target, original in reversed(self._rollback_values):
if target is not None:
target.write(original)
if parameter.device.type != "meta":
raise OffloadSurgeryError(
"offloaded parameter was materialized during rollback"
)
else:
original = self._rollback_values[0][1]
with torch.no_grad():
parameter.data.copy_(original)
self._rollback_values = []
self._committed = False
def resolve_logical_parameter(
module: nn.Module,
parameter_name: str = "weight",
*,
search_roots: Sequence[nn.Module] = (),
) -> LogicalParameterTransaction:
"""Resolve a live or Accelerate-offloaded parameter for safe surgery.
``search_roots`` supplies known ancestors such as the decoder layer. This
is required when Accelerate attaches one ``place_submodules`` hook to a
parent and stores keys such as ``self_attn.o_proj.weight``.
"""
parameter = getattr(module, parameter_name, None)
if parameter is None or not isinstance(parameter, (nn.Parameter, torch.Tensor)):
raise OffloadSurgeryError(
f"module has no tensor parameter named {parameter_name!r}"
)
if parameter.device.type != "meta":
return LogicalParameterTransaction(
module=module,
parameter_name=parameter_name,
tensor=parameter.data,
offloaded=False,
)
_accelerate_release()
roots: list[nn.Module] = []
for root in (module, *search_roots):
if all(root is not existing for existing in roots):
roots.append(root)
candidates: list[tuple[int, nn.Module, str, Any]] = []
inspected_hooks = 0
for root in roots:
for _owner_name, owner in root.named_modules():
relative = _relative_module_path(owner, module)
if relative is None:
continue
hook = getattr(owner, "_hf_hook", None)
if hook is None:
continue
align_hooks = _align_hooks(hook)
inspected_hooks += len(align_hooks)
key = f"{relative}.{parameter_name}" if relative else parameter_name
for align_hook in align_hooks:
candidates.append((relative.count("."), owner, key, align_hook))
# A hook closest to the target has the shortest relative key and is the
# least ambiguous source of truth.
candidates.sort(key=lambda candidate: candidate[0])
attempted: list[str] = []
for _depth, _owner, key, hook in candidates:
weights_map = getattr(hook, "weights_map", None)
if weights_map is None:
continue
attempted.append(key)
try:
source = weights_map[key]
except (KeyError, IndexError):
continue
if not isinstance(source, torch.Tensor):
raise UnsupportedOffloadLayoutError(
f"Accelerate backing value {key!r} is not a tensor"
)
storage_kind = classify_weight_storage(
module_class_name=module.__class__.__name__,
parameter_class_name=parameter.__class__.__name__,
has_quant_state=hasattr(parameter, "quant_state"),
data_is_floating_point=source.is_floating_point(),
)
if storage_kind != "float":
raise UnsupportedOffloadedQuantizationError(
f"offloaded {storage_kind} parameter {key!r} cannot be "
"safely updated before supported quantized write-back exists"
)
primary = _write_target(weights_map, key)
return LogicalParameterTransaction(
module=module,
parameter_name=parameter_name,
tensor=source,
offloaded=True,
backing_key=key,
_targets=_alias_targets(primary, source),
)
attempted_text = ", ".join(dict.fromkeys(attempted)) or parameter_name
if inspected_hooks == 0:
raise UnsupportedOffloadLayoutError(
f"meta-resident parameter {parameter_name!r} has no supported "
"offloading AlignDevicesHook in the supplied module roots"
)
raise UnsupportedOffloadLayoutError(
f"no authoritative Accelerate backing value found for meta parameter; "
f"attempted keys: {attempted_text}"
)
def validate_offloaded_parameters(root: nn.Module) -> None:
"""Fail closed on every meta parameter before a surgery pass mutates data."""
seen: set[int] = set()
for _module_name, module in root.named_modules():
for parameter_name, parameter in module.named_parameters(recurse=False):
if parameter.device.type != "meta" or id(parameter) in seen:
continue
resolve_logical_parameter(module, parameter_name, search_roots=(root,))
seen.add(id(parameter))
def logical_module_device(root: nn.Module) -> torch.device:
"""Return a usable compute device even when a layer starts on ``meta``."""
for _module_name, module in root.named_modules():
for parameter_name, parameter in module.named_parameters(recurse=False):
if parameter.device.type != "meta":
return parameter.device
return resolve_logical_parameter(
module,
parameter_name,
search_roots=(root,),
).tensor.device
raise OffloadSurgeryError("module has no parameters from which to select a device")
+28
View File
@@ -4,8 +4,11 @@ from __future__ import annotations
import pytest
import torch
import torch.nn as nn
from accelerate.hooks import AlignDevicesHook, add_hook_to_module
from obliteratus import device
from obliteratus.abliterate import AbliterationPipeline
pytestmark = pytest.mark.gpu
@@ -23,6 +26,31 @@ def test_cuda_discovery_dtype_placement_and_operation():
assert torch.isfinite(result).all()
def test_cuda_execution_observes_offloaded_surgery_and_restores_meta_state():
if not torch.cuda.is_available():
pytest.skip("requires a CUDA runner; set ENABLE_CUDA_GATE and attach the cuda label")
module = nn.Module()
module.proj = nn.Linear(4, 4, bias=False)
original = module.proj.weight.detach().clone()
hook = AlignDevicesHook(execution_device="cuda", offload=True)
add_hook_to_module(module.proj, hook)
count = AbliterationPipeline._project_out_advanced(
module,
torch.tensor([[1.0], [0.0], [0.0], [0.0]], device="cuda"),
["proj"],
)
output = module.proj(torch.ones(1, 4, device="cuda"))
expected = original.clone()
expected[:, 0] = 0
assert count == 1
assert output.device.type == "cuda"
assert module.proj.weight.device.type == "meta"
torch.testing.assert_close(hook.weights_map["weight"], expected)
def test_bitsandbytes_quantization_operation():
if not torch.cuda.is_available():
pytest.skip("requires a CUDA runner; set ENABLE_CUDA_GATE and attach the cuda label")
+391
View File
@@ -0,0 +1,391 @@
"""CPU-only contracts for transactional Accelerate offload surgery."""
from __future__ import annotations
from collections.abc import MutableMapping
import pytest
import torch
import torch.nn as nn
from accelerate.big_modeling import disk_offload
from accelerate.hooks import (
AlignDevicesHook,
ModelHook,
SequentialHook,
add_hook_to_module,
)
from accelerate.utils.modeling import get_state_dict_offloaded_model
import obliteratus.models.offload_surgery as offload_surgery
from obliteratus.abliterate import AbliterationPipeline
from obliteratus.models.offload_surgery import (
OffloadSurgeryError,
UnsupportedOffloadLayoutError,
UnsupportedOffloadedQuantizationError,
logical_module_device,
resolve_logical_parameter,
validate_offloaded_parameters,
)
def _direct_offload(module: nn.Module) -> AlignDevicesHook:
hook = AlignDevicesHook(execution_device="cpu", offload=True)
add_hook_to_module(module, hook)
return hook
def test_direct_align_hook_updates_backing_weight_and_preserves_meta_state():
module = nn.Linear(3, 2, bias=False)
original = module.weight.detach().clone()
hook = _direct_offload(module)
parameter_identity = id(module.weight)
transaction = resolve_logical_parameter(module)
transaction.commit(transaction.tensor + 2)
assert module.weight.device.type == "meta"
assert id(module.weight) == parameter_identity
torch.testing.assert_close(hook.weights_map["weight"], original + 2)
def test_sequential_hook_finds_nested_align_devices_hook():
module = nn.Linear(3, 2, bias=False)
original = module.weight.detach().clone()
align = AlignDevicesHook(execution_device="cpu", offload=True)
add_hook_to_module(module, SequentialHook(ModelHook(), align))
transaction = resolve_logical_parameter(module)
transaction.commit(transaction.tensor - 1)
assert module.weight.device.type == "meta"
torch.testing.assert_close(align.weights_map["weight"], original - 1)
def test_parent_prefixed_weight_map_resolves_nested_parameter():
block = nn.Module()
block.self_attn = nn.Module()
block.self_attn.o_proj = nn.Linear(3, 2, bias=False)
original = block.self_attn.o_proj.weight.detach().clone()
hook = AlignDevicesHook(
execution_device="cpu",
offload=True,
place_submodules=True,
)
add_hook_to_module(block, hook)
transaction = resolve_logical_parameter(
block.self_attn.o_proj,
search_roots=(block,),
)
assert transaction.backing_key == "self_attn.o_proj.weight"
transaction.commit(transaction.tensor * 0.5)
assert block.self_attn.o_proj.weight.device.type == "meta"
torch.testing.assert_close(
hook.weights_map["self_attn.o_proj.weight"],
original * 0.5,
)
def test_disk_backed_prefixed_loader_uses_updated_authoritative_value(tmp_path):
model = nn.Sequential(nn.Linear(3, 2, bias=False))
original = model[0].weight.detach().clone()
disk_offload(model, tmp_path, execution_device=torch.device("cpu"))
transaction = resolve_logical_parameter(model[0], search_roots=(model,))
transaction.commit(transaction.tensor + 3)
assert model[0].weight.device.type == "meta"
hook = model[0]._hf_hook
torch.testing.assert_close(hook.weights_map["weight"], original + 3)
gathered = get_state_dict_offloaded_model(model)
assert all(tensor.device.type != "meta" for tensor in gathered.values())
torch.testing.assert_close(gathered["0.weight"], original + 3)
def test_save_reload_observes_surgery_instead_of_stale_backing_weight():
model = nn.Sequential(nn.Linear(3, 2, bias=False))
original = model[0].weight.detach().clone()
_direct_offload(model[0])
transaction = resolve_logical_parameter(model[0], search_roots=(model,))
transaction.commit(transaction.tensor - 4)
state_dict = get_state_dict_offloaded_model(model)
reloaded = nn.Sequential(nn.Linear(3, 2, bias=False))
reloaded.load_state_dict(state_dict)
torch.testing.assert_close(reloaded[0].weight, original - 4)
assert reloaded[0].weight.device.type == "cpu"
def test_tied_backing_values_retain_identity_and_consistent_values():
block = nn.Module()
block.first = nn.Linear(3, 3, bias=False)
block.second = nn.Linear(3, 3, bias=False)
shared = block.first.weight.detach().clone()
shared_meta = nn.Parameter(torch.empty_like(shared, device="meta"))
block.first.weight = shared_meta
block.second.weight = shared_meta
hook = AlignDevicesHook(
execution_device="cpu",
offload=True,
weights_map={"first.weight": shared, "second.weight": shared},
place_submodules=True,
)
block._hf_hook = hook
live_identity = id(block.first.weight)
transaction = resolve_logical_parameter(block.first, search_roots=(block,))
transaction.commit(transaction.tensor + 1)
first = hook.weights_map["first.weight"]
second = hook.weights_map["second.weight"]
assert first is second
torch.testing.assert_close(first, shared + 1)
assert block.first.weight is block.second.weight
assert id(block.first.weight) == live_identity
class _FailOnceAliasMap(MutableMapping[str, torch.Tensor]):
def __init__(self, shared: torch.Tensor):
self.values = {"first.weight": shared, "second.weight": shared}
self.failed = False
def __getitem__(self, key: str) -> torch.Tensor:
return self.values[key]
def __setitem__(self, key: str, value: torch.Tensor) -> None:
if key == "second.weight" and not self.failed:
self.failed = True
raise OSError("simulated backing-store failure")
self.values[key] = value
def __delitem__(self, key: str) -> None:
del self.values[key]
def __iter__(self):
return iter(self.values)
def __len__(self) -> int:
return len(self.values)
def test_commit_failure_rolls_back_all_backing_aliases_and_leaves_meta_state():
block = nn.Module()
block.first = nn.Linear(3, 3, bias=False)
original = block.first.weight.detach().clone()
block.first.weight = nn.Parameter(torch.empty_like(original, device="meta"))
backing = _FailOnceAliasMap(original)
block._hf_hook = AlignDevicesHook(
execution_device="cpu",
offload=True,
weights_map=backing,
place_submodules=True,
)
transaction = resolve_logical_parameter(block.first, search_roots=(block,))
with pytest.raises(OffloadSurgeryError, match="rolled back"):
transaction.commit(transaction.tensor + 5)
torch.testing.assert_close(backing["first.weight"], original)
torch.testing.assert_close(backing["second.weight"], original)
assert block.first.weight.device.type == "meta"
def test_quantized_offloaded_weight_fails_before_any_mutation():
module = nn.Linear(3, 2, bias=False, device="meta")
backing = {"weight": torch.ones(2, 3, dtype=torch.uint8)}
module._hf_hook = AlignDevicesHook(
execution_device="cpu",
offload=True,
weights_map=backing,
)
with pytest.raises(UnsupportedOffloadedQuantizationError):
resolve_logical_parameter(module)
assert module.weight.device.type == "meta"
assert backing["weight"].dtype == torch.uint8
assert torch.equal(backing["weight"], torch.ones(2, 3, dtype=torch.uint8))
def test_read_only_unknown_mapping_fails_during_resolution():
module = nn.Linear(3, 2, bias=False, device="meta")
class MappingOnly:
def __getitem__(self, key):
return torch.ones(2, 3)
def __iter__(self):
return iter(("weight",))
def __len__(self):
return 1
# Registering through Mapping is deliberately avoided: this object models
# a future wrapper with read access but no supported write contract.
hook = AlignDevicesHook(execution_device="cpu", offload=True)
hook.weights_map = MappingOnly()
module._hf_hook = hook
with pytest.raises(UnsupportedOffloadLayoutError, match="read-only"):
resolve_logical_parameter(module)
def test_validation_resolves_all_meta_parameters_before_surgery():
block = nn.Module()
block.proj = nn.Linear(3, 2, bias=False)
hook = AlignDevicesHook(
execution_device="cpu",
offload=True,
place_submodules=True,
)
add_hook_to_module(block, hook)
validate_offloaded_parameters(block)
hook.weights_map = {}
with pytest.raises(UnsupportedOffloadLayoutError, match="attempted keys"):
validate_offloaded_parameters(block)
def test_advanced_projection_updates_offloaded_logical_weight():
module = nn.Module()
module.proj = nn.Linear(3, 2, bias=False)
original = module.proj.weight.detach().clone()
hook = _direct_offload(module.proj)
direction = torch.tensor([[1.0], [0.0], [0.0]])
count = AbliterationPipeline._project_out_advanced(
module,
direction,
["proj"],
)
assert count == 1
assert module.proj.weight.device.type == "meta"
expected = original.clone()
expected[:, 0] = 0
torch.testing.assert_close(hook.weights_map["weight"], expected)
def test_advanced_projection_resolves_all_candidates_before_mutation():
block = nn.Module()
block.first = nn.Linear(3, 2, bias=False)
first_original = block.first.weight.detach().clone()
first_hook = _direct_offload(block.first)
block.second = nn.Linear(3, 2, bias=False, device="meta")
block.second._hf_hook = AlignDevicesHook(
execution_device="cpu",
offload=True,
weights_map={"weight": torch.ones(2, 3, dtype=torch.uint8)},
)
with pytest.raises(UnsupportedOffloadedQuantizationError):
AbliterationPipeline._project_out_advanced(
block,
torch.tensor([[1.0], [0.0], [0.0]]),
["first", "second"],
)
torch.testing.assert_close(first_hook.weights_map["weight"], first_original)
assert block.first.weight.device.type == "meta"
def test_logical_module_device_uses_authoritative_backing_device():
module = nn.Linear(3, 2, bias=False)
_direct_offload(module)
assert logical_module_device(module) == torch.device("cpu")
def test_unknown_accelerate_major_version_fails_closed(monkeypatch):
module = nn.Linear(3, 2, bias=False)
original = module.weight.detach().clone()
hook = _direct_offload(module)
monkeypatch.setattr(offload_surgery, "version", lambda _name: "2.0.0")
with pytest.raises(UnsupportedOffloadLayoutError, match="outside the tested"):
resolve_logical_parameter(module)
assert module.weight.device.type == "meta"
torch.testing.assert_close(hook.weights_map["weight"], original)
def test_bias_projection_updates_parent_prefixed_backing_value():
block = nn.Module()
block.proj = nn.Linear(3, 3, bias=True)
original = block.proj.bias.detach().clone()
hook = AlignDevicesHook(
execution_device="cpu",
offload=True,
place_submodules=True,
)
add_hook_to_module(block, hook)
count = AbliterationPipeline._project_bias(
block,
torch.tensor([[1.0], [0.0], [0.0]]),
["proj"],
offload_roots=(block,),
)
assert count == 1
assert block.proj.bias.device.type == "meta"
expected = original.clone()
expected[0] = 0
torch.testing.assert_close(hook.weights_map["proj.bias"], expected)
class _FailSecondParameterMap(MutableMapping[str, torch.Tensor]):
def __init__(self, first: torch.Tensor, second: torch.Tensor):
self.values = {"first.weight": first, "second.weight": second}
self.failed = False
def __getitem__(self, key: str) -> torch.Tensor:
return self.values[key]
def __setitem__(self, key: str, value: torch.Tensor) -> None:
if key == "second.weight" and not self.failed:
self.failed = True
raise OSError("simulated second-parameter failure")
self.values[key] = value
def __delitem__(self, key: str) -> None:
del self.values[key]
def __iter__(self):
return iter(self.values)
def __len__(self) -> int:
return len(self.values)
def test_multi_parameter_projection_rolls_back_prior_commits():
block = nn.Module()
block.first = nn.Linear(3, 2, bias=False)
block.second = nn.Linear(3, 2, bias=False)
first = block.first.weight.detach().clone()
second = block.second.weight.detach().clone()
block.first.weight = nn.Parameter(torch.empty_like(first, device="meta"))
block.second.weight = nn.Parameter(torch.empty_like(second, device="meta"))
backing = _FailSecondParameterMap(first, second)
block._hf_hook = AlignDevicesHook(
execution_device="cpu",
offload=True,
weights_map=backing,
place_submodules=True,
)
with pytest.raises(OffloadSurgeryError, match="rolled back"):
AbliterationPipeline._project_out_advanced(
block,
torch.tensor([[1.0], [0.0], [0.0]]),
["first", "second"],
offload_roots=(block,),
)
torch.testing.assert_close(backing["first.weight"], first)
torch.testing.assert_close(backing["second.weight"], second)
assert block.first.weight.device.type == "meta"
assert block.second.weight.device.type == "meta"