Merge pull request #196 from younger-plinius/feat/115-offloaded-fused-experts

feat(abliteration): update offloaded fused MoE expert tensors in place (#115)

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 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
 
 Use `--gpus` to control which GPUs are used:

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(

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="<unk>"))
+    tokenizer_backend.pre_tokenizer = Whitespace()
+    return PreTrainedTokenizerFast(
+        tokenizer_object=tokenizer_backend,
+        pad_token="<pad>",
+        eos_token="<eos>",
+        unk_token="<unk>",
+    )
+
+
 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="<unk>"))
-    tokenizer_backend.pre_tokenizer = Whitespace()
-    tokenizer = PreTrainedTokenizerFast(
-        tokenizer_object=tokenizer_backend,
-        pad_token="<pad>",
-        eos_token="<eos>",
-        unk_token="<unk>",
-    )
+    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

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: checkpoint tensors match live surgery."""
+    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"

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,249 @@ 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"