mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-09-27 11:41:55 +02:00
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 <noreply@anthropic.com>
This commit is contained in:
Pliny the Younger
co-authored by
Claude Fable 5.1
Claude Fable 5.1
parent
e39f908832
commit
89b63ce216
Vendored
+73
-9
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user