mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-09-21 17:00:50 +02:00
Pliny the YoungerandClaude Fable 5.1
89b63ce216
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 <noreply@anthropic.com>
639 lines
21 KiB
Python
639 lines
21 KiB
Python
"""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 _MAX_NORM_RATIO, 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"
|
|
|
|
|
|
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"
|
|
|