fix(qwen): preserve validated runtime on checkpoint reload

This commit is contained in:
Joseph Magly
2026-08-27 20:12:43 -04:00
parent fff7d4a24b
commit 44f7842c9c
7 changed files with 99 additions and 16 deletions
+4
View File
@@ -557,6 +557,10 @@ obliteratus obliterate Qwen/Qwen3.8-27B --dtype float16 --quantization 4bit
> PyTorch/CUDA-compatible kernel build with `pip install -e ".[qwen-hybrid]"`.
> The pristine and post-edit quality gates remain mandatory; a failed pristine
> checkpoint is never modified.
> Saved-checkpoint Chat reloads use the same validated Qwen hybrid runtime
> contract as the surgery load: compatible FLA/causal-conv1d kernels, SDPA,
> and the complete text model on one CUDA device. Generic multi-GPU
> `device_map="auto"` sharding is rejected for reloads as well as initial loads.
Install the optional backend before selecting a bitsandbytes mode:
+28 -2
View File
@@ -78,7 +78,7 @@ from obliteratus.model_load_settings import (
resolve_model_load_settings,
)
from obliteratus.persistence_contracts import validate_reloadable_checkpoint
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
_gpu_lifecycle = from_environment()
@@ -156,7 +156,33 @@ def _load_model_to_device(
if local_files_only:
kwargs["local_files_only"] = True
if dev.supports_device_map_auto():
config_kwargs = {
key: value
for key, value in {
"trust_remote_code": trust_remote_code or None,
"token": token,
"local_files_only": local_files_only or None,
}.items()
if value is not None
}
config = AutoConfig.from_pretrained(pretrained_path, **config_kwargs)
from obliteratus.models.loader import qwen_hybrid_runtime_overrides
quantization = None
if quantization_config is not None:
if getattr(quantization_config, "load_in_4bit", False):
quantization = "4bit"
elif getattr(quantization_config, "load_in_8bit", False):
quantization = "8bit"
qwen_overrides = qwen_hybrid_runtime_overrides(
config,
torch_dtype or torch.float32,
quantization,
)
if qwen_overrides:
kwargs["config"] = config
kwargs.update(qwen_overrides)
elif dev.supports_device_map_auto():
kwargs["device_map"] = "auto"
try:
+1
View File
@@ -34,6 +34,7 @@
],
"required_tests": [
"tests/test_app_vram.py",
"tests/test_app_model_lifecycle.py",
"tests/test_app_model_load_settings.py",
"tests/test_model_load_settings.py",
"tests/test_cli.py",
+7 -1
View File
@@ -4143,6 +4143,7 @@ class AbliterationPipeline:
f"compound amplification (directions are not re-extracted)"
)
effective_passes = 1
self._effective_refinement_passes = effective_passes
# Track previous directions for cosine-similarity early-exit
_prev_directions: dict[int, torch.Tensor] = {}
@@ -7382,7 +7383,12 @@ class AbliterationPipeline:
"direction_method": self.direction_method,
"norm_preserve": self.norm_preserve,
"regularization": self.regularization,
"refinement_passes": self.refinement_passes,
"refinement_passes": getattr(
self,
"_effective_refinement_passes",
self.refinement_passes,
),
"requested_refinement_passes": self.refinement_passes,
"project_biases": self.project_biases,
"use_chat_template": self.use_chat_template,
"use_whitened_svd": self.use_whitened_svd,
+40 -13
View File
@@ -655,6 +655,37 @@ def _qwen_single_device(estimate_gb: float, quantization: str | None) -> int:
)
def qwen_hybrid_runtime_overrides(
config,
torch_dtype: torch.dtype,
quantization: str | None = None,
) -> dict:
"""Return the mandatory runtime placement for a Qwen3.5 hybrid model.
Initial surgery loads and later local-checkpoint reloads must share this
contract. Generic Accelerate sharding is not valid for the recurrent
DeltaNet execution path, even when aggregate VRAM is sufficient.
"""
if getattr(config, "model_type", "") != "qwen3_5":
return {}
import transformers
version_match = re.match(r"^(\d+)\.(\d+)", transformers.__version__)
version = tuple(map(int, version_match.groups())) if version_match else (0, 0)
if version < (5, 8):
raise RuntimeError(
"Qwen3.8 requires transformers>=5.8 for the validated hybrid runtime"
)
_require_qwen_hybrid_kernels()
estimate_gb = _estimate_model_memory_gb(config, torch_dtype)
device_index = _qwen_single_device(estimate_gb, quantization)
return {
"attn_implementation": "sdpa",
"device_map": {"": device_index},
}
def _hf_token() -> str | None:
"""Resolve the optional Hugging Face read token."""
return resolve_secret("HF_TOKEN")
@@ -848,17 +879,13 @@ def load_model(
is_qwen_hybrid = task == "causal_lm" and getattr(config, "model_type", "") == "qwen3_5"
qwen_device_index = None
if is_qwen_hybrid:
import transformers
version_match = re.match(r"^(\d+)\.(\d+)", transformers.__version__)
version = tuple(map(int, version_match.groups())) if version_match else (0, 0)
if version < (5, 8):
raise RuntimeError(
"Qwen3.8 requires transformers>=5.8 for the validated hybrid runtime"
)
_require_qwen_hybrid_kernels()
qwen_device_index = _qwen_single_device(est_gb, quantization)
load_kwargs["attn_implementation"] = "sdpa"
qwen_overrides = qwen_hybrid_runtime_overrides(
config,
torch_dtype,
quantization,
)
load_kwargs.update(qwen_overrides)
qwen_device_index = qwen_overrides["device_map"][""]
logger.info(
"Loading Qwen3.8 through AutoModelForCausalLM as an explicit text-only "
"derivative on cuda:%d; vision and MTP checkpoint tensors are not part "
@@ -908,8 +935,8 @@ def load_model(
load_kwargs["device_map"] = "auto"
if qwen_device_index is not None:
# Generic layer sharding is not a validated execution path for the
# recurrent DeltaNet state. Keep the complete text model on one GPU.
# Preserve the validated override after generic quantization/device
# policy has run. The helper is also used by local Chat reloads.
load_kwargs["device_map"] = {"": qwen_device_index}
# Offload support: provide a folder for disk offloading when GPU memory is insufficient
+15
View File
@@ -82,9 +82,11 @@ import pathlib
import sys
import threading
import time
from types import SimpleNamespace
from unittest.mock import Mock
import app
from obliteratus.models import loader
root = pathlib.Path(sys.argv[1])
checkpoint = root / "completed"
@@ -103,11 +105,24 @@ app.AutoModelForCausalLM.from_pretrained = Mock(
side_effect=lambda source, **kwargs: (calls.append((source, kwargs)), model)[1]
)
app.AutoTokenizer.from_pretrained = Mock(return_value=tokenizer)
app.AutoConfig.from_pretrained = Mock(
return_value=SimpleNamespace(model_type="gpt2"),
)
app.dev.supports_device_map_auto = lambda: True
app._load_model_to_device(checkpoint, local_files_only=True)
assert calls[0][0] == checkpoint
assert calls[0][1]["local_files_only"] is True
loader.qwen_hybrid_runtime_overrides = Mock(
return_value={"attn_implementation": "sdpa", "device_map": {"": 2}},
)
app._load_model_to_device(checkpoint, local_files_only=True)
assert calls[-1][1]["attn_implementation"] == "sdpa"
assert calls[-1][1]["device_map"] == {"": 2}
assert calls[-1][1]["device_map"] != "auto"
loader.qwen_hybrid_runtime_overrides.reset_mock(return_value=True)
loader.qwen_hybrid_runtime_overrides.return_value = {}
quantization = object()
loaded_model, loaded_tokenizer = app._reload_local_checkpoint(
checkpoint,
+4
View File
@@ -216,6 +216,10 @@ def test_qwen_excise_route_does_not_touch_inputs_gates_or_lm_head():
torch.equal(after[name], before[name])
for name in before.keys() - allowed
)
assert pipeline._effective_refinement_passes == 1
method_config = pipeline._build_metadata()["method_config"]
assert method_config["refinement_passes"] == 1
assert method_config["requested_refinement_passes"] == 2
@pytest.mark.parametrize("architecture", ["qwen3_5_text", "qwen3_5_moe"])