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
+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: