From 89b63ce216c5a43f18933b0c20a40250caa61402 Mon Sep 17 00:00:00 2001 From: Pliny the Younger Date: Wed, 2 Sep 2026 11:54:14 -0400 Subject: [PATCH 1/2] feat(abliteration): update offloaded fused MoE expert tensors in place (#115) transformers 5 stores routed experts as fused 3D parameters (experts.gate_up_proj / experts.down_proj) for Mixtral, Qwen3-MoE, DeepSeek-V3, GLM-4 MoE, Llama 4, gpt-oss and OLMoE. When such a layer is CPU- or disk-offloaded by Accelerate, EXCISE failed closed with "offloaded fused expert tensors are not yet a supported surgery layout", so every frontier MoE checkpoint large enough to need offload could not be abliterated at all. Route the three fused projections and the fused bias projection through the existing LogicalParameterTransaction adapter: a meta-resident fused parameter is resolved to a private copy of its authoritative backing value, projected per expert, and committed atomically; the live parameter stays on meta and quantized or unknown backing layouts still fail before mutation. Thread offload_roots through the MoE dispatch call sites so parent-prefixed Accelerate hooks resolve, and drop the preflight rejection of 3D meta parameters (validate_offloaded_parameters already resolves them). Tests: fused 3D, bias, granular and selective-inversion projections on offloaded backing stores, bounded norm restoration, parent-prefixed hook resolution, quantized fail-closed, commit-failure rollback, and a full offline pipeline run on a disk-offloaded tiny Mixtral fixture that verifies the saved checkpoint carries the updated experts with no meta tensors. Co-Authored-By: Claude Fable 5.1 --- README.md | 4 + obliteratus/abliterate.py | 121 ++++++++++--- tests/fixtures/tiny_offline_model.py | 82 ++++++++- tests/test_offline_integration.py | 128 +++++++++++++- tests/test_offload_surgery.py | 249 ++++++++++++++++++++++++++- 5 files changed, 549 insertions(+), 35 deletions(-) diff --git a/README.md b/README.md index a921e33..f8fced1 100644 --- a/README.md +++ b/README.md @@ -519,6 +519,10 @@ When you have multiple GPUs, OBLITERATUS uses accelerate's `device_map="auto"` t This means multi-GPU sharding is a **memory solution, not a speed solution**. It lets you run models that don't fit on one GPU, but it won't make small models run faster. In fact, more GPUs can be *slower* due to inter-GPU data transfer overhead at layer boundaries. +### Offloaded layers and fused MoE experts + +When the combined GPU budget is still too small, accelerate offloads whole layers to CPU RAM or disk and leaves `meta` placeholders in the live model. OBLITERATUS performs surgery on those layers through the authoritative offload backing store instead of materializing them. Linear projections, bias projections, and fused 3D expert tensors are all handled by the same transactional path, so the fused `experts.gate_up_proj` / `experts.down_proj` layout that transformers 5 uses for MoE architectures (Mixtral, Qwen3-MoE, DeepSeek-V3, GLM-4 MoE, Llama 4, gpt-oss, OLMoE) is updated in place and the saved checkpoint is identical to what live surgery would produce. Quantized or unrecognized offload layouts fail closed before any weight is modified. Attention-head surgery, safety-neuron masking, expert transplant, and LoRA ablation are still rejected for offloaded layers. + ### Selecting GPUs Use `--gpus` to control which GPUs are used: diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index 503599e..9e68724 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -4150,25 +4150,17 @@ class AbliterationPipeline: # 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. + # closed without leaving an earlier layer partially modified. Fused + # 3D expert tensors (transformers 5 MoE layouts) are resolved like any + # other logical parameter and updated through the same transactional + # backing-store adapter. 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: + if any( + parameter.device.type == "meta" for parameter in layer.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: @@ -5815,6 +5807,32 @@ class AbliterationPipeline: raise return count + @staticmethod + def _resolve_fused_parameter( + container: nn.Module, + name: str, + offload_roots: tuple[nn.Module, ...], + ) -> tuple[torch.Tensor, LogicalParameterTransaction | None]: + """Return writable logical data for a fused expert parameter. + + A live parameter is returned as its own storage, so the existing + in-place projection path is unchanged. A meta-resident parameter is + resolved through the Accelerate offload adapter: the returned tensor + is a private copy of the authoritative backing value, so a failure + before ``transaction.commit`` leaves the backing store untouched and + the live parameter stays offloaded after the commit. Quantized or + unknown backing layouts fail closed inside the adapter. + """ + param = getattr(container, name) + if param.device.type != "meta": + return param.data, None + transaction = resolve_logical_parameter( + container, + name, + search_roots=(container, *offload_roots), + ) + return transaction.tensor.detach().clone(), transaction + @staticmethod def _project_fused_3d( container: nn.Module, @@ -5822,6 +5840,7 @@ class AbliterationPipeline: param_names: list[str], norm_preserve: bool, scale: float, + offload_roots: tuple[nn.Module, ...] = (), ) -> int: """Project refusal direction from fused 3D expert parameters. @@ -5848,7 +5867,14 @@ class AbliterationPipeline: f"dequantized this checkpoint to float. This is a bug; " f"please report it." ) - is_quantized = AbliterationPipeline._is_quantized_param(param) + transaction: LogicalParameterTransaction | None = None + if param.device.type == "meta": + data, transaction = AbliterationPipeline._resolve_fused_parameter( + container, name, offload_roots, + ) + is_quantized = False + else: + is_quantized = AbliterationPipeline._is_quantized_param(param) if is_quantized: try: import bitsandbytes as bnb @@ -5864,7 +5890,7 @@ class AbliterationPipeline: stacklevel=2, ) continue - else: + elif transaction is None: data = param.data # FP8/packed fused params are quantized storage — the loader # dequantizes them; raw upcasting here would corrupt scales. @@ -5917,6 +5943,10 @@ class AbliterationPipeline: count += 1 if count > 0: + if transaction is not None: + # Commit the completed logical update to the offload + # backing store; the live parameter remains on meta. + transaction.commit(data) # Write back (re-quantize if needed) if is_quantized: try: @@ -5944,19 +5974,29 @@ class AbliterationPipeline: container: nn.Module, direction: torch.Tensor, bias_names: list[str], + offload_roots: tuple[nn.Module, ...] = (), ) -> int: - """Project refusal direction from fused 2D expert biases.""" + """Project refusal direction from fused 2D expert biases. + + Offloaded (meta-resident) biases are updated through the transactional + offload adapter so the authoritative backing value changes while the + live parameter stays offloaded. + """ for bname in bias_names: bp = getattr(container, bname, None) if bp is None or not isinstance(bp, (nn.Parameter, torch.Tensor)): continue - b = bp.data + b, transaction = AbliterationPipeline._resolve_fused_parameter( + container, bname, offload_roots, + ) d_sq = direction.to(device=b.device, dtype=b.dtype).squeeze() if b.dim() == 2 and b.shape[-1] == d_sq.shape[0]: for ei in range(b.shape[0]): comp = (b[ei] @ d_sq) * d_sq b[ei].sub_(comp) del comp + if transaction is not None: + transaction.commit(b) return b.shape[0] return 0 @@ -6174,6 +6214,7 @@ class AbliterationPipeline: fused_out = AbliterationPipeline._project_fused_3d( experts, direction, ["down_proj", "w2"], norm_preserve=norm_preserve, scale=scale, + offload_roots=offload_roots, ) if fused_out > 0: expert_count += fused_out @@ -6181,10 +6222,12 @@ class AbliterationPipeline: expert_count += AbliterationPipeline._project_fused_3d( experts, direction, _FFN_IN_NAMES, norm_preserve=norm_preserve, scale=scale, + offload_roots=offload_roots, ) if project_biases: expert_count += AbliterationPipeline._project_fused_bias( experts, direction, ["down_proj_bias", "w2_bias"], + offload_roots=offload_roots, ) count += expert_count return count @@ -6363,6 +6406,7 @@ class AbliterationPipeline: reflect_scale=self.reflection_strength, remove_scale=1.0, norm_preserve=norm_preserve, + offload_roots=offload_roots, ) count += self._project_fused_3d_selective_inversion( experts, direction, _FFN_IN_NAMES, @@ -6370,10 +6414,12 @@ class AbliterationPipeline: reflect_scale=self.reflection_strength, remove_scale=1.0, norm_preserve=norm_preserve, + offload_roots=offload_roots, ) if project_biases: count += self._project_fused_bias( experts, direction, ["down_proj_bias", "w2_bias"], + offload_roots=offload_roots, ) # Stabilize router weights after reflection to prevent extreme logits @@ -6501,15 +6547,18 @@ class AbliterationPipeline: experts, direction, expert_dirs, ["down_proj", "w2"], norm_preserve=norm_preserve, scale=scale, + offload_roots=offload_roots, ) expert_count += self._project_fused_3d_granular( experts, direction, expert_dirs, _FFN_IN_NAMES, norm_preserve=norm_preserve, scale=scale, + offload_roots=offload_roots, ) if project_biases: expert_count += self._project_fused_bias( experts, direction, ["down_proj_bias", "w2_bias"], + offload_roots=offload_roots, ) count += expert_count @@ -6525,6 +6574,7 @@ class AbliterationPipeline: param_names: list[str], norm_preserve: bool, scale: float, + offload_roots: tuple[nn.Module, ...] = (), ) -> int: """Project fused 3D expert params with per-expert directions. @@ -6556,7 +6606,14 @@ class AbliterationPipeline: f"dequantized this checkpoint to float. This is a bug; " f"please report it." ) - is_quantized = AbliterationPipeline._is_quantized_param(param) + transaction: LogicalParameterTransaction | None = None + if param.device.type == "meta": + data, transaction = AbliterationPipeline._resolve_fused_parameter( + container, pname, offload_roots, + ) + is_quantized = False + else: + is_quantized = AbliterationPipeline._is_quantized_param(param) if is_quantized: try: import bitsandbytes as bnb @@ -6617,7 +6674,12 @@ class AbliterationPipeline: W.mul_(ratio) count += 1 - if is_quantized and count > 0: + if transaction is not None: + if count > 0: + # Commit the completed logical update to the offload + # backing store; the live parameter remains on meta. + transaction.commit(data) + elif is_quantized and count > 0: try: import bitsandbytes as bnb quantized, new_state = bnb.functional.quantize_4bit( @@ -6649,6 +6711,7 @@ class AbliterationPipeline: reflect_scale: float, remove_scale: float, norm_preserve: bool, + offload_roots: tuple[nn.Module, ...] = (), ) -> int: """Fused 3D projection with per-expert inversion differentiation. @@ -6682,7 +6745,14 @@ class AbliterationPipeline: f"dequantized this checkpoint to float. This is a bug; " f"please report it." ) - is_quantized = AbliterationPipeline._is_quantized_param(param) + transaction: LogicalParameterTransaction | None = None + if param.device.type == "meta": + data, transaction = AbliterationPipeline._resolve_fused_parameter( + container, pname, offload_roots, + ) + is_quantized = False + else: + is_quantized = AbliterationPipeline._is_quantized_param(param) if is_quantized: try: import bitsandbytes as bnb @@ -6740,7 +6810,12 @@ class AbliterationPipeline: W.mul_(ratio) count += 1 - if is_quantized and count > 0: + if transaction is not None: + if count > 0: + # Commit the completed logical update to the offload + # backing store; the live parameter remains on meta. + transaction.commit(data) + elif is_quantized and count > 0: try: import bitsandbytes as bnb quantized, new_state = bnb.functional.quantize_4bit( diff --git a/tests/fixtures/tiny_offline_model.py b/tests/fixtures/tiny_offline_model.py index c954485..119eedc 100644 --- a/tests/fixtures/tiny_offline_model.py +++ b/tests/fixtures/tiny_offline_model.py @@ -9,7 +9,13 @@ import torch from tokenizers import Tokenizer from tokenizers.models import WordLevel from tokenizers.pre_tokenizers import Whitespace -from transformers import GPT2Config, GPT2LMHeadModel, PreTrainedTokenizerFast +from transformers import ( + GPT2Config, + GPT2LMHeadModel, + MixtralConfig, + MixtralForCausalLM, + PreTrainedTokenizerFast, +) FIXTURE_SEED = 20260814 @@ -28,20 +34,24 @@ FIXTURE_VOCAB = { } +def _fixture_tokenizer() -> PreTrainedTokenizerFast: + tokenizer_backend = Tokenizer(WordLevel(FIXTURE_VOCAB, unk_token="")) + tokenizer_backend.pre_tokenizer = Whitespace() + return PreTrainedTokenizerFast( + tokenizer_object=tokenizer_backend, + pad_token="", + eos_token="", + unk_token="", + ) + + def build_tiny_offline_model(destination: Path) -> Path: """Create a tiny random-init GPT-2 model without downloads or caches.""" destination = Path(destination) destination.mkdir(parents=True, exist_ok=False) torch.manual_seed(FIXTURE_SEED) - tokenizer_backend = Tokenizer(WordLevel(FIXTURE_VOCAB, unk_token="")) - tokenizer_backend.pre_tokenizer = Whitespace() - tokenizer = PreTrainedTokenizerFast( - tokenizer_object=tokenizer_backend, - pad_token="", - eos_token="", - unk_token="", - ) + tokenizer = _fixture_tokenizer() config = GPT2Config( vocab_size=len(FIXTURE_VOCAB), n_positions=128, @@ -78,3 +88,57 @@ def build_tiny_offline_model(destination: Path) -> Path: encoding="utf-8", ) return destination + + +def build_tiny_offline_moe_model(destination: Path) -> Path: + """Create a tiny random-init Mixtral model with fused 3D expert tensors. + + transformers 5 stores routed experts as ``experts.gate_up_proj`` and + ``experts.down_proj`` parameters of shape ``(num_experts, ...)``, which is + the layout frontier MoE checkpoints expose to weight surgery. + """ + destination = Path(destination) + destination.mkdir(parents=True, exist_ok=False) + + torch.manual_seed(FIXTURE_SEED) + tokenizer = _fixture_tokenizer() + config = MixtralConfig( + vocab_size=len(FIXTURE_VOCAB), + hidden_size=16, + intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + num_local_experts=4, + num_experts_per_tok=2, + max_position_embeddings=2048, + bos_token_id=1, + eos_token_id=1, + pad_token_id=0, + ) + model = MixtralForCausalLM(config) + model.save_pretrained(destination, safe_serialization=True) + tokenizer.save_pretrained(destination) + + manifest = { + "fixture": "tiny-offline-mixtral-moe", + "provenance": "generated locally from configuration with random initialization", + "training_data": None, + "third_party_weights": None, + "license": "AGPL-3.0-only (part of the OBLITERATUS test suite)", + "seed": FIXTURE_SEED, + "architecture": { + "model_type": "mixtral", + "layers": 2, + "hidden_size": 16, + "attention_heads": 2, + "experts": 4, + "experts_per_token": 2, + "vocabulary_size": len(FIXTURE_VOCAB), + }, + } + (destination / "fixture-provenance.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return destination diff --git a/tests/test_offline_integration.py b/tests/test_offline_integration.py index 58e383d..35e0d61 100644 --- a/tests/test_offline_integration.py +++ b/tests/test_offline_integration.py @@ -12,13 +12,18 @@ import pytest import torch import yaml from datasets import Dataset -from transformers import AutoModelForCausalLM, AutoTokenizer +from accelerate import dispatch_model +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer from obliteratus.abliterate import AbliterationPipeline from obliteratus.config import DatasetConfig, ModelConfig, StrategyConfig, StudyConfig +from obliteratus.models.loader import ModelHandle from obliteratus.reporting.report import AblationReport from obliteratus.runner import run_study -from tests.fixtures.tiny_offline_model import build_tiny_offline_model +from tests.fixtures.tiny_offline_model import ( + build_tiny_offline_model, + build_tiny_offline_moe_model, +) pytestmark = [pytest.mark.cpu, pytest.mark.integration] @@ -239,6 +244,125 @@ def test_multidirection_pipeline_restores_tiny_model_layer_norms(tmp_path): assert metadata["method_config"]["norm_preserve"] is True +def _dispatch_decoder_layers_to_disk(model, offload_dir: Path) -> None: + """Offload every decoder layer to disk so its parameters become meta tensors.""" + device_map = {"lm_head": "cpu"} + for name, _child in model.model.named_children(): + if name == "layers": + for index in range(len(model.model.layers)): + device_map[f"model.layers.{index}"] = "disk" + else: + device_map[f"model.{name}"] = "cpu" + dispatch_model(model, device_map=device_map, offload_dir=str(offload_dir)) + + +def _run_moe_pipeline(source: Path, output: Path, load_model) -> Path: + import obliteratus.abliterate as abliterate_module + + previous = abliterate_module.load_model + abliterate_module.load_model = load_model + try: + pipeline = AbliterationPipeline( + model_name=str(source), + output_dir=str(output), + device="cpu", + dtype="float32", + method="advanced", + max_seq_length=8, + verify_sample_size=1, + refusal_max_tokens=1, + max_perplexity_increase=1000.0, + max_degenerate_fraction=1.0, + harmful_prompts=["harmful request", "harmful answer"], + harmless_prompts=["harmless request", "harmless answer"], + ) + return pipeline.run() + finally: + abliterate_module.load_model = previous + + +def _moe_load_model(offload_dir: Path | None, live: dict): + def load_model(model_name, task, **_kwargs): + config = AutoConfig.from_pretrained(model_name, local_files_only=True) + model = AutoModelForCausalLM.from_pretrained( + model_name, + config=config, + dtype=torch.float32, + local_files_only=True, + ) + if offload_dir is not None: + _dispatch_decoder_layers_to_disk(model, offload_dir) + live["model"] = model + return ModelHandle( + model=model, + tokenizer=AutoTokenizer.from_pretrained(model_name, local_files_only=True), + config=config, + model_name=model_name, + task=task, + ) + + return load_model + + +# transformers emits this advisory whenever ``hf_device_map`` contains cpu/disk +# entries. The pipeline materializes the complete state dict itself before +# ``save_pretrained`` (see ``_gather_state_dict``), so the advisory is expected +# here and unrelated warnings stay fatal. +@pytest.mark.filterwarnings( + "ignore:Attempting to save a model with offloaded modules.*:UserWarning", +) +def test_disk_offloaded_fused_moe_surgery_matches_live_surgery(tmp_path): + """Offload must be transparent: same checkpoint as live surgery, byte for byte.""" + source = build_tiny_offline_moe_model(tmp_path / "source") + offload_dir = tmp_path / "offload" + offload_dir.mkdir() + original = _state_dict(source) + + live_output = _run_moe_pipeline( + source, tmp_path / "live", _moe_load_model(None, {}), + ) + offloaded_live: dict[str, torch.nn.Module] = {} + offloaded_output = _run_moe_pipeline( + source, tmp_path / "offloaded", _moe_load_model(offload_dir, offloaded_live), + ) + + assert live_output == tmp_path / "live" + assert offloaded_output == tmp_path / "offloaded" + model = offloaded_live["model"] + expert_parameters = { + name: parameter + for name, parameter in model.named_parameters() + if ".experts." in name + } + assert expert_parameters, "fixture must expose fused expert parameters" + assert all(parameter.dim() == 3 for parameter in expert_parameters.values()) + # Surgery must leave the live model offloaded rather than materializing it. + assert all(parameter.device.type == "meta" for parameter in expert_parameters.values()) + + live_state = _state_dict(live_output) + offloaded_state = _state_dict(offloaded_output) + assert offloaded_state.keys() == live_state.keys() == original.keys() + assert all(tensor.device.type != "meta" for tensor in offloaded_state.values()) + assert all(torch.isfinite(tensor).all() for tensor in offloaded_state.values()) + changed_experts = [ + name + for name in expert_parameters + if not torch.equal(original[name], live_state[name]) + ] + assert changed_experts, "live surgery did not touch the fused expert tensors" + # Bitwise identical on macOS; assert_close with default float32 tolerances + # keeps the contract robust to BLAS threading differences on other hosts + # while still catching any real divergence (a wrong projection is ~1e-2). + for name in live_state: + torch.testing.assert_close( + offloaded_state[name], + live_state[name], + msg=lambda detail, name=name: ( + f"offloaded surgery diverged from live surgery on {name}: {detail}" + ), + ) + + def test_installed_wheel_cli_loads_local_model_without_repository_imports(tmp_path): source = build_tiny_offline_model(tmp_path / "source") isolated_workdir = tmp_path / "outside-repository" diff --git a/tests/test_offload_surgery.py b/tests/test_offload_surgery.py index b002ea2..0db5105 100644 --- a/tests/test_offload_surgery.py +++ b/tests/test_offload_surgery.py @@ -17,7 +17,7 @@ from accelerate.hooks import ( 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.abliterate import _MAX_NORM_RATIO, AbliterationPipeline from obliteratus.models.offload_surgery import ( OffloadSurgeryError, UnsupportedOffloadLayoutError, @@ -389,3 +389,250 @@ def test_multi_parameter_projection_rolls_back_prior_commits(): torch.testing.assert_close(backing["second.weight"], second) assert block.first.weight.device.type == "meta" assert block.second.weight.device.type == "meta" + + +class _FusedExperts(nn.Module): + """transformers 5 fused expert container (Mixtral / GPT-OSS parameter layout).""" + + def __init__(self, experts: int = 3, hidden: int = 4, intermediate: int = 6): + super().__init__() + torch.manual_seed(1234) + self.gate_up_proj = nn.Parameter(torch.randn(experts, 2 * intermediate, hidden)) + self.down_proj = nn.Parameter(torch.randn(experts, hidden, intermediate)) + self.down_proj_bias = nn.Parameter(torch.randn(experts, hidden)) + + +def _axis_direction(axis: int = 1, hidden: int = 4) -> torch.Tensor: + direction = torch.zeros(hidden, 1) + direction[axis, 0] = 1.0 + return direction + + +def test_fused_3d_projection_updates_offloaded_expert_backing_weights(): + experts = _FusedExperts() + original_down = experts.down_proj.detach().clone() + original_gate_up = experts.gate_up_proj.detach().clone() + hook = _direct_offload(experts) + direction = _axis_direction() + + count = AbliterationPipeline._project_fused_3d( + experts, + direction, + ["down_proj", "w2"], + norm_preserve=False, + scale=1.0, + ) + + assert count == 3 + assert experts.down_proj.device.type == "meta" + assert experts.gate_up_proj.device.type == "meta" + updated = hook.weights_map["down_proj"] + expected = original_down.clone() + expected[:, 1, :] = 0 # (experts, hidden, intermediate): hidden axis is dim 1 + torch.testing.assert_close(updated, expected) + torch.testing.assert_close(hook.weights_map["gate_up_proj"], original_gate_up) + + count = AbliterationPipeline._project_fused_3d( + experts, + direction, + ["up_proj", "gate_proj", "gate_up_proj"], + norm_preserve=False, + scale=1.0, + ) + + assert count == 3 + expected = original_gate_up.clone() + expected[:, :, 1] = 0 # (experts, 2*intermediate, hidden): hidden axis is dim 2 + torch.testing.assert_close(hook.weights_map["gate_up_proj"], expected) + + +def test_fused_3d_projection_preserves_norm_through_offload_backing(): + experts = _FusedExperts() + original = experts.down_proj.detach().clone() + hook = _direct_offload(experts) + + count = AbliterationPipeline._project_fused_3d( + experts, + _axis_direction(), + ["down_proj"], + norm_preserve=True, + scale=1.0, + ) + + assert count == 3 + updated = hook.weights_map["down_proj"] + projected = original.clone() + projected[:, 1, :] = 0 + for expert in range(3): + # Norm restoration is a bounded rescale of the projected slice: it + # recovers the original norm unless that needs more than the cap. + assert torch.count_nonzero(updated[expert, 1]) == 0 + ratio = min( + original[expert].norm().item() / projected[expert].norm().item(), + _MAX_NORM_RATIO, + ) + torch.testing.assert_close(updated[expert], projected[expert] * ratio) + + +def test_fused_bias_projection_updates_offloaded_backing_value(): + experts = _FusedExperts() + original = experts.down_proj_bias.detach().clone() + hook = _direct_offload(experts) + + count = AbliterationPipeline._project_fused_bias( + experts, + _axis_direction(), + ["down_proj_bias", "w2_bias"], + ) + + assert count == 3 + assert experts.down_proj_bias.device.type == "meta" + expected = original.clone() + expected[:, 1] = 0 + torch.testing.assert_close(hook.weights_map["down_proj_bias"], expected) + + +def test_granular_and_selective_fused_projection_update_offloaded_backing(): + experts = _FusedExperts() + original_down = experts.down_proj.detach().clone() + original_gate_up = experts.gate_up_proj.detach().clone() + hook = _direct_offload(experts) + shared = _axis_direction().squeeze() + per_expert = {0: _axis_direction(axis=0).squeeze()} + + count = AbliterationPipeline._project_fused_3d_granular( + experts, + shared, + per_expert, + ["down_proj"], + norm_preserve=False, + scale=1.0, + ) + + assert count == 3 + assert experts.down_proj.device.type == "meta" + granular = hook.weights_map["down_proj"] + assert torch.count_nonzero(granular[0, 0]) == 0 + torch.testing.assert_close(granular[0, 1], original_down[0, 1]) + for expert in (1, 2): + assert torch.count_nonzero(granular[expert, 1]) == 0 + torch.testing.assert_close(granular[expert, 0], original_down[expert, 0]) + + count = AbliterationPipeline._project_fused_3d_selective_inversion( + experts, + shared, + ["gate_up_proj"], + safety_indices={2}, + reflect_scale=2.0, + remove_scale=1.0, + norm_preserve=False, + ) + + assert count == 3 + assert experts.gate_up_proj.device.type == "meta" + selective = hook.weights_map["gate_up_proj"] + for expert in (0, 1): + assert torch.count_nonzero(selective[expert, :, 1]) == 0 + torch.testing.assert_close(selective[2, :, 1], -original_gate_up[2, :, 1]) + torch.testing.assert_close(selective[:, :, 0], original_gate_up[:, :, 0]) + + +def test_fused_projection_resolves_parent_prefixed_expert_backing(): + layer = nn.Module() + layer.mlp = nn.Module() + layer.mlp.experts = _FusedExperts() + original = layer.mlp.experts.down_proj.detach().clone() + hook = AlignDevicesHook(execution_device="cpu", offload=True, place_submodules=True) + add_hook_to_module(layer, hook) + + count = AbliterationPipeline._project_fused_3d( + layer.mlp.experts, + _axis_direction(), + ["down_proj"], + norm_preserve=False, + scale=1.0, + offload_roots=(layer,), + ) + + assert count == 3 + assert layer.mlp.experts.down_proj.device.type == "meta" + expected = original.clone() + expected[:, 1, :] = 0 + torch.testing.assert_close(hook.weights_map["mlp.experts.down_proj"], expected) + + +def test_offloaded_quantized_fused_expert_fails_before_mutation(): + experts = nn.Module() + experts.down_proj = nn.Parameter( + torch.empty(2, 4, 6, dtype=torch.uint8, device="meta"), + requires_grad=False, + ) + backing = {"down_proj": torch.ones(2, 4, 6, dtype=torch.uint8)} + experts._hf_hook = AlignDevicesHook( + execution_device="cpu", + offload=True, + weights_map=backing, + ) + + with pytest.raises(UnsupportedOffloadedQuantizationError): + AbliterationPipeline._project_fused_3d( + experts, + _axis_direction(), + ["down_proj"], + norm_preserve=False, + scale=1.0, + ) + + assert experts.down_proj.device.type == "meta" + assert torch.equal(backing["down_proj"], torch.ones(2, 4, 6, dtype=torch.uint8)) + + +class _ReadOnlyFusedMap(MutableMapping[str, torch.Tensor]): + def __init__(self, tensor: torch.Tensor): + self.values = {"down_proj": tensor} + + def __getitem__(self, key: str) -> torch.Tensor: + return self.values[key] + + def __setitem__(self, key: str, value: torch.Tensor) -> None: + raise OSError("simulated fused backing-store failure") + + 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_fused_projection_commit_failure_leaves_backing_and_meta_state_intact(): + torch.manual_seed(99) + original = torch.randn(2, 4, 6) + experts = nn.Module() + experts.down_proj = nn.Parameter( + torch.empty_like(original, device="meta"), + requires_grad=False, + ) + backing = _ReadOnlyFusedMap(original.clone()) + experts._hf_hook = AlignDevicesHook( + execution_device="cpu", + offload=True, + weights_map=backing, + ) + + with pytest.raises(OffloadSurgeryError, match="rolled back"): + AbliterationPipeline._project_fused_3d( + experts, + _axis_direction(), + ["down_proj"], + norm_preserve=False, + scale=1.0, + ) + + # The projection worked on a private copy, so the authoritative value is + # byte-identical and the live parameter is still offloaded. + torch.testing.assert_close(backing["down_proj"], original) + assert experts.down_proj.device.type == "meta" + From 131a4a8b55255ed33f8f36bee60cc67915a58044 Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Wed, 2 Sep 2026 14:32:29 -0400 Subject: [PATCH 2/2] docs: align fused offload claims with test contract --- README.md | 2 +- tests/test_offline_integration.py | 2 +- tests/test_offload_surgery.py | 1 - 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index f8fced1..cb3a7f7 100644 --- a/README.md +++ b/README.md @@ -521,7 +521,7 @@ This means multi-GPU sharding is a **memory solution, not a speed solution**. It ### Offloaded layers and fused MoE experts -When the combined GPU budget is still too small, accelerate offloads whole layers to CPU RAM or disk and leaves `meta` placeholders in the live model. OBLITERATUS performs surgery on those layers through the authoritative offload backing store instead of materializing them. Linear projections, bias projections, and fused 3D expert tensors are all handled by the same transactional path, so the fused `experts.gate_up_proj` / `experts.down_proj` layout that transformers 5 uses for MoE architectures (Mixtral, Qwen3-MoE, DeepSeek-V3, GLM-4 MoE, Llama 4, gpt-oss, OLMoE) is updated in place and the saved checkpoint is identical to what live surgery would produce. Quantized or unrecognized offload layouts fail closed before any weight is modified. Attention-head surgery, safety-neuron masking, expert transplant, and LoRA ablation are still rejected for offloaded layers. +When the combined GPU budget is still too small, accelerate offloads whole layers to CPU RAM or disk and leaves `meta` placeholders in the live model. OBLITERATUS performs surgery on those layers through the authoritative offload backing store instead of materializing them. Linear projections, bias projections, and fused 3D expert tensors are all handled by the same transactional path, so the fused `experts.gate_up_proj` / `experts.down_proj` layout that transformers 5 uses for MoE architectures (Mixtral, Qwen3-MoE, DeepSeek-V3, GLM-4 MoE, Llama 4, gpt-oss, OLMoE) is updated in place and the saved checkpoint matches what live surgery would produce within the configured floating-point tolerance. Quantized or unrecognized offload layouts fail closed before any weight is modified. Attention-head surgery, safety-neuron masking, expert transplant, and LoRA ablation are still rejected for offloaded layers. ### Selecting GPUs diff --git a/tests/test_offline_integration.py b/tests/test_offline_integration.py index 35e0d61..464fd2f 100644 --- a/tests/test_offline_integration.py +++ b/tests/test_offline_integration.py @@ -312,7 +312,7 @@ def _moe_load_model(offload_dir: Path | None, live: dict): "ignore:Attempting to save a model with offloaded modules.*:UserWarning", ) def test_disk_offloaded_fused_moe_surgery_matches_live_surgery(tmp_path): - """Offload must be transparent: same checkpoint as live surgery, byte for byte.""" + """Offload must be transparent: checkpoint tensors match live surgery.""" source = build_tiny_offline_moe_model(tmp_path / "source") offload_dir = tmp_path / "offload" offload_dir.mkdir() diff --git a/tests/test_offload_surgery.py b/tests/test_offload_surgery.py index 0db5105..fd6b02f 100644 --- a/tests/test_offload_surgery.py +++ b/tests/test_offload_surgery.py @@ -635,4 +635,3 @@ def test_fused_projection_commit_failure_leaves_backing_and_meta_state_intact(): # byte-identical and the live parameter is still offloaded. torch.testing.assert_close(backing["down_proj"], original) assert experts.down_proj.device.type == "meta" -