Merge pull request #179 from jmagly/fix/168-qwen38-reload-contract

fix(qwen): preserve validated runtime on checkpoint reload

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:
 

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:

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

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,

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", "") not in {"qwen3_5", "qwen3_5_text"}:
+        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")
@@ -845,20 +876,20 @@ 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:
-        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 +939,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

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,26 @@ 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
 
+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"
+app.AutoConfig.from_pretrained.return_value = SimpleNamespace(model_type="gpt2")
+
 quantization = object()
 loaded_model, loaded_tokenizer = app._reload_local_checkpoint(
     checkpoint,

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(

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"])