diff --git a/ci/conditional-test-policy.json b/ci/conditional-test-policy.json index 043bde4..f5e5b9c 100644 --- a/ci/conditional-test-policy.json +++ b/ci/conditional-test-policy.json @@ -131,6 +131,7 @@ "coverage_paths": [ "obliteratus/device.py", "obliteratus/models/loader.py", + "obliteratus/models/offload_surgery.py", "obliteratus/models/quant_dequant.py" ] }, diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index 658f106..85c3fd5 100644 --- a/ci/test-risk-map.json +++ b/ci/test-risk-map.json @@ -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", diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index ca5322c..d77921f 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -66,6 +66,14 @@ 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, + OffloadSurgeryError, + 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 +3580,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 +3797,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 +3900,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. @@ -3848,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}). " @@ -3870,6 +3951,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 +3966,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 +3975,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 +3984,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 +4002,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 @@ -3936,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}). " @@ -4023,7 +4119,10 @@ class AbliterationPipeline: norm_preserve=self.norm_preserve, regularization=sae_reg, projection_row_fraction=self.projection_row_fraction, + offload_roots=(layers[idx],), ) + except OffloadSurgeryError: + raise except (AttributeError, RuntimeError): pass if sae_ffn is not None: @@ -4033,6 +4132,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,8 +4141,11 @@ class AbliterationPipeline: regularization=sae_reg, project_biases=False, projection_row_fraction=self.projection_row_fraction, + offload_roots=(layers[idx],), ) sae_count += fc + except OffloadSurgeryError: + raise except (AttributeError, RuntimeError): pass del sd @@ -4107,7 +4210,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 +4234,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 +4249,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 +4306,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 +5030,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 +5080,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 +5100,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 +5137,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 +5157,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 +5204,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 +5229,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 +5241,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 +5251,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 +5433,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 +5450,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 +5475,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 +5504,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 +5545,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 +5580,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 +5604,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 +5612,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 +5631,7 @@ class AbliterationPipeline: norm_preserve=norm_preserve, regularization=regularization, projection_row_fraction=projection_row_fraction, + offload_roots=offload_roots, ) # ── Routed expert projection ────────────────────────────────────── @@ -5449,6 +5670,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 +5678,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 +5693,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 +5706,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 +5757,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 +5779,7 @@ class AbliterationPipeline: ffn_module, direction, [child_name], norm_preserve=norm_preserve, regularization=router_reg, + offload_roots=offload_roots, ) break @@ -5563,9 +5793,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 +5817,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 +5851,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 +5863,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 +5891,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 +5908,7 @@ class AbliterationPipeline: ffn_module, direction, [child_name], norm_preserve=norm_preserve, regularization=regularization, + offload_roots=offload_roots, ) break @@ -5673,16 +5920,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 +5952,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 +5986,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 diff --git a/obliteratus/models/offload_surgery.py b/obliteratus/models/offload_surgery.py new file mode 100644 index 0000000..e09cdaf --- /dev/null +++ b/obliteratus/models/offload_surgery.py @@ -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") diff --git a/tests/conditional/test_cuda_runtime.py b/tests/conditional/test_cuda_runtime.py index 2ac4bee..0600c59 100644 --- a/tests/conditional/test_cuda_runtime.py +++ b/tests/conditional/test_cuda_runtime.py @@ -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") diff --git a/tests/test_offload_surgery.py b/tests/test_offload_surgery.py new file mode 100644 index 0000000..b002ea2 --- /dev/null +++ b/tests/test_offload_surgery.py @@ -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"