From 44f7842c9c6896cc2e3895624f2a0cea77fa0ea1 Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Thu, 27 Aug 2026 20:12:43 -0400 Subject: [PATCH 1/2] fix(qwen): preserve validated runtime on checkpoint reload --- README.md | 4 +++ app.py | 30 +++++++++++++++-- ci/test-risk-map.json | 1 + obliteratus/abliterate.py | 8 ++++- obliteratus/models/loader.py | 53 +++++++++++++++++++++++-------- tests/test_app_model_lifecycle.py | 15 +++++++++ tests/test_qwen35_contracts.py | 4 +++ 7 files changed, 99 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 7201fb5..1e2da22 100644 --- a/README.md +++ b/README.md @@ -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: diff --git a/app.py b/app.py index aad4e32..d6e8d6c 100644 --- a/app.py +++ b/app.py @@ -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: diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index 5f0651b..1661207 100644 --- a/ci/test-risk-map.json +++ b/ci/test-risk-map.json @@ -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", diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index de87d40..99e98d1 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -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, diff --git a/obliteratus/models/loader.py b/obliteratus/models/loader.py index 27e8244..4103813 100644 --- a/obliteratus/models/loader.py +++ b/obliteratus/models/loader.py @@ -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 diff --git a/tests/test_app_model_lifecycle.py b/tests/test_app_model_lifecycle.py index 1b8cbc0..f36d6f7 100644 --- a/tests/test_app_model_lifecycle.py +++ b/tests/test_app_model_lifecycle.py @@ -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, diff --git a/tests/test_qwen35_contracts.py b/tests/test_qwen35_contracts.py index 27e57e7..8b91e6f 100644 --- a/tests/test_qwen35_contracts.py +++ b/tests/test_qwen35_contracts.py @@ -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"]) From 691a5da746226ddc630ec1e1d419aa54e3cec7cb Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Thu, 27 Aug 2026 20:17:09 -0400 Subject: [PATCH 2/2] fix(qwen): recognize saved text derivatives --- obliteratus/models/loader.py | 8 ++++++-- tests/test_app_model_lifecycle.py | 10 ++++++---- tests/test_loader_boundaries.py | 15 +++++++++++++++ 3 files changed, 27 insertions(+), 6 deletions(-) diff --git a/obliteratus/models/loader.py b/obliteratus/models/loader.py index 4103813..c6eb0c4 100644 --- a/obliteratus/models/loader.py +++ b/obliteratus/models/loader.py @@ -666,7 +666,7 @@ def qwen_hybrid_runtime_overrides( 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": + if getattr(config, "model_type", "") not in {"qwen3_5", "qwen3_5_text"}: return {} import transformers @@ -876,7 +876,11 @@ def load_model( if task == "classification": config.num_labels = num_labels load_kwargs["config"] = config - is_qwen_hybrid = task == "causal_lm" and getattr(config, "model_type", "") == "qwen3_5" + is_qwen_hybrid = task == "causal_lm" and getattr( + config, + "model_type", + "", + ) in {"qwen3_5", "qwen3_5_text"} qwen_device_index = None if is_qwen_hybrid: qwen_overrides = qwen_hybrid_runtime_overrides( diff --git a/tests/test_app_model_lifecycle.py b/tests/test_app_model_lifecycle.py index f36d6f7..408a103 100644 --- a/tests/test_app_model_lifecycle.py +++ b/tests/test_app_model_lifecycle.py @@ -113,15 +113,17 @@ 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.AutoConfig.from_pretrained.return_value = SimpleNamespace( + model_type="qwen3_5_text", ) +loader._require_qwen_hybrid_kernels = Mock() +loader._estimate_model_memory_gb = Mock(return_value=54.0) +loader._qwen_single_device = Mock(return_value=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 = {} +app.AutoConfig.from_pretrained.return_value = SimpleNamespace(model_type="gpt2") quantization = object() loaded_model, loaded_tokenizer = app._reload_local_checkpoint( diff --git a/tests/test_loader_boundaries.py b/tests/test_loader_boundaries.py index 831f713..d63140d 100644 --- a/tests/test_loader_boundaries.py +++ b/tests/test_loader_boundaries.py @@ -573,6 +573,21 @@ def test_qwen_single_device_uses_per_device_free_memory(monkeypatch): assert loader._qwen_single_device(216.0, "4bit") == 1 +def test_qwen_saved_text_derivative_reuses_single_device_runtime(monkeypatch): + config = _config(model_type="qwen3_5_text") + monkeypatch.setattr(loader, "_require_qwen_hybrid_kernels", Mock()) + monkeypatch.setattr(loader, "_estimate_model_memory_gb", Mock(return_value=54.0)) + monkeypatch.setattr(loader, "_qwen_single_device", Mock(return_value=2)) + + assert loader.qwen_hybrid_runtime_overrides( + config, + torch.bfloat16, + ) == { + "attn_implementation": "sdpa", + "device_map": {"": 2}, + } + + def test_model_handle_metadata_snapshot_restore_summary_and_cleanup(tmp_path): model = _model() nested = SimpleNamespace(