OBLITERATUS/obliteratus/models/loader.py

"""Load HuggingFace models and wrap them for ablation."""

from __future__ import annotations

import logging
import os
import tempfile
from dataclasses import dataclass, field
from typing import Optional

import sys as _sys

import torch
from obliteratus import device as dev
from obliteratus.models import quant_dequant as qd
from obliteratus.runtime_contracts import (
    effective_model_memory_gb,
    quantized_model_fits_gpu,
    resolve_model_load_policy,
    should_snapshot_model,
    validate_model_load_request,
)
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    PreTrainedModel,
    PreTrainedTokenizerBase,
)

try:
    from transformers import AutoModelForImageTextToText
except (ImportError, AttributeError):
    AutoModelForImageTextToText = None

logger = logging.getLogger(__name__)

_QUANTIZATION_STATE_COMPONENTS = frozenset(
    {
        "absmax",
        "nested_absmax",
        "nested_quant_map",
        "quant_map",
        "quant_state",
    }
)


# ---------------------------------------------------------------------------
# Compat shims for transformers ≥5.0 breaking changes.
#
# Many HuggingFace model repos ship custom modeling code (loaded via
# trust_remote_code=True) that imports symbols from their pre-5.x locations.
# We monkey-patch the old module paths so loading works without downgrading.
#
# Every section is wrapped in try/except so a failure in one shim never
# breaks unrelated functionality.  Patches are purely additive — we never
# remove attributes that already exist.
# ---------------------------------------------------------------------------

# ── 1. utils.generic → utils.output_capturing ──────────────────────
# OutputRecorder, check_model_inputs, _CAN_RECORD_REGISTRY moved.
# Affected: MiniMax-M2.x, DeepSeek-V3
try:
    import transformers.utils.generic as _tfu_generic
    try:
        from transformers.utils import output_capturing as _oc
        for _old, _new in [
            ("OutputRecorder", "OutputRecorder"),
            ("check_model_inputs", "capture_outputs"),
            ("_CAN_RECORD_REGISTRY", "_CAN_RECORD_REGISTRY"),
        ]:
            if not hasattr(_tfu_generic, _old) and hasattr(_oc, _new):
                setattr(_tfu_generic, _old, getattr(_oc, _new))
    except ImportError:
        pass
except Exception:
    pass

# ── 2. utils.generic.working_or_temp_dir ───────────────────────────
# Removed in 5.x.  Trivial contextmanager replacement.
# Affected: GLM-4 / ChatGLM custom code
try:
    import transformers.utils.generic as _tfu_generic  # noqa: F811 – may already be imported
    if not hasattr(_tfu_generic, "working_or_temp_dir"):
        import contextlib as _ctxlib
        import tempfile as _tmpmod

        @_ctxlib.contextmanager
        def _working_or_temp_dir(working_dir=None):
            if working_dir is not None:
                yield working_dir
            else:
                with _tmpmod.TemporaryDirectory() as tmp:
                    yield tmp

        _tfu_generic.working_or_temp_dir = _working_or_temp_dir
except Exception:
    pass

# ── 3. utils.import_utils: removed availability checks ─────────────
# is_torch_fx_available   → removed (torch.fx always present in torch≥2.0)
# is_tf_available         → removed (TF backend dropped in v5)
# is_flax_available       → removed (Flax backend dropped in v5)
# is_safetensors_available→ removed (safetensors is now mandatory)
# Affected: various model repos that defensively check backends
try:
    import transformers.utils.import_utils as _tfu_imports
    _import_shims = {
        "is_torch_fx_available": lambda: True,
        "is_tf_available": lambda: False,
        "is_flax_available": lambda: False,
        "is_safetensors_available": lambda: True,
    }
    for _name, _fn in _import_shims.items():
        if not hasattr(_tfu_imports, _name):
            setattr(_tfu_imports, _name, _fn)
    # Also patch the top-level transformers.utils re-export so both
    # ``from transformers.utils import is_tf_available`` and
    # ``from transformers.utils.import_utils import is_tf_available`` work.
    try:
        import transformers.utils as _tu
        for _name, _fn in _import_shims.items():
            if not hasattr(_tu, _name):
                setattr(_tu, _name, _fn)
    except Exception:
        pass
except Exception:
    pass

# ── 4. pytorch_utils: removed version-check constants ──────────────
# ``is_torch_greater_or_equal_than_X_Y`` constants removed in v4.48+.
# Affected: DeepSeek-V3/R1/V2-Lite, MiniCPM3, older custom code
try:
    import transformers.pytorch_utils as _pt_utils
    # transformers ≥5.0 requires torch ≥2.0, so every historical gate is True.
    for _ver in [
        "is_torch_greater_or_equal_than_2_4",
        "is_torch_greater_or_equal_than_2_3",
        "is_torch_greater_or_equal_than_2_2",
        "is_torch_greater_or_equal_than_2_1",
        "is_torch_greater_or_equal_than_2_0",
        "is_torch_greater_or_equal_than_1_13",
        "is_torch_greater_or_equal_than_1_12",
        "is_torch_greater_or_equal_than_1_11",
        "is_torch_greater_or_equal_than_1_10",
        "is_torch_greater_or_equal_than_1_9",
        "is_torch_greater_or_equal_than_1_8",
        "is_torch_greater_or_equal_than_1_6",
    ]:
        if not hasattr(_pt_utils, _ver):
            setattr(_pt_utils, _ver, True)
except Exception:
    pass

# ── 5. generation_utils module → transformers.generation ────────────
# Entire module removed; old custom code does
#   ``from transformers.generation_utils import GenerationMixin``
# Affected: older generation-customising model repos
try:
    import transformers.generation_utils  # noqa: F401 – already exists
except ModuleNotFoundError:
    try:
        import transformers.generation as _gen
        _sys.modules["transformers.generation_utils"] = _gen
    except Exception:
        pass

# ── 6. deepspeed module → transformers.integrations.deepspeed ───────
# Affected: model repos with DeepSpeed training code
try:
    import transformers.deepspeed  # noqa: F401 – already exists
except ModuleNotFoundError:
    try:
        import transformers.integrations.deepspeed as _ds
        _sys.modules["transformers.deepspeed"] = _ds
    except Exception:
        pass

# ── 7. DynamicCache.get_max_length → get_max_cache_shape ───────────
# Removed in v4.49+.  DeepSeek-V3/R1 custom code calls .get_max_length().
try:
    from transformers.cache_utils import DynamicCache as _DC
    if not hasattr(_DC, "get_max_length") and hasattr(_DC, "get_max_cache_shape"):
        _DC.get_max_length = _DC.get_max_cache_shape
except Exception:
    pass

# ── 8. LogitsWarper → LogitsProcessor ──────────────────────────────
# LogitsWarper removed in v5.0 (deprecated v4.46).  Drop-in alias.
# Affected: MiniCPM-o custom code
# NOTE: submodule patch runs here; top-level ``transformers.LogitsWarper``
# is deferred to _apply_deferred_shims() because the _LazyModule may reset
# its __dict__ during initial import.
try:
    import transformers.generation.logits_process as _lp_mod
    if not hasattr(_lp_mod, "LogitsWarper"):
        from transformers.generation.logits_process import LogitsProcessor as _LP
        _lp_mod.LogitsWarper = _LP
except Exception:
    pass

# ── 9. processing_utils._validate_images_text_input_order ──────────
# Removed in v5.0rc3.  Kimi-VL custom code imports it.
try:
    import transformers.processing_utils as _proc
    if not hasattr(_proc, "_validate_images_text_input_order"):
        def _validate_images_text_input_order(images=None, text=None, **kw):
            return images, text
        _proc._validate_images_text_input_order = _validate_images_text_input_order
except Exception:
    pass

# ── 10. TF/Flax weight constants (removed with TF backend) ─────────
try:
    import transformers.utils as _tu  # noqa: F811
    for _cname, _cval in [
        ("TF_WEIGHTS_NAME", "tf_model.h5"),
        ("TF2_WEIGHTS_NAME", "tf_model.h5"),
    ]:
        if not hasattr(_tu, _cname):
            setattr(_tu, _cname, _cval)
except Exception:
    pass

# ── 11. file_utils.cached_path → huggingface_hub fallback ──────────
# Removed in v4.22.  Very old model repos use it for file download.
def _cached_path_shim(url_or_filename, cache_dir=None, **kwargs):
    """Best-effort replacement for the removed legacy cached-path helper."""
    if os.path.exists(str(url_or_filename)):
        return str(url_or_filename)
    try:
        from huggingface_hub import hf_hub_download

        parts = str(url_or_filename).rsplit("/", 1)
        if len(parts) == 2:
            return hf_hub_download(
                repo_id=parts[0],
                filename=parts[1],
                cache_dir=cache_dir,
            )
    except Exception:
        pass
    return str(url_or_filename)


try:
    import transformers.file_utils as _fu
    if not hasattr(_fu, "cached_path"):
        _fu.cached_path = _cached_path_shim
except Exception:
    pass


# ── 12. PreTrainedModel.all_tied_weights_keys ──────────────────────
# Referenced by the accelerate device-map integration in transformers 5.x,
# but older remote-code model classes never define it.  Aggregate the
# legacy ``_tied_weights_keys`` (list of names or dict) across modules into
# the {param_name: target_name} mapping accelerate expects.
# Affected: Nemotron Omni (NemotronH_Nano_Omni_Reasoning_V3)
def _install_all_tied_weights_keys_compat(model_cls) -> bool:
    """Install the legacy tied-weight adapter only when the class lacks it."""
    if hasattr(model_cls, "all_tied_weights_keys"):
        return False

    def _get_all_tied_weights_keys(self):
        cached = self.__dict__.get("_all_tied_weights_keys_store")
        if cached is not None:
            return cached
        expander = getattr(self, "get_expanded_tied_weights_keys", None)
        if callable(expander):
            try:
                expanded = expander(all_submodels=True)
            except Exception:
                # Legacy remote implementations may still expose the modern
                # method name while keeping an incompatible list contract.
                pass
            else:
                self.__dict__["_all_tied_weights_keys_store"] = expanded
                return expanded
        tied: dict = {}
        for module in self.modules():
            keys = getattr(module, "_tied_weights_keys", None)
            if not keys:
                continue
            if isinstance(keys, dict):
                candidates = keys
            else:  # legacy list of patterns/names — one tied group
                names = list(keys)
                candidates = {name: names[0] for name in names}
            for name, target in candidates.items():
                # Only expose keys that actually resolve on this model —
                # multimodal wrappers inherit "lm_head.weight" from their
                # LM submodule's config but have no top-level lm_head.
                try:
                    self.get_submodule(name.rsplit(".", 1)[0])
                except AttributeError:
                    continue
                tied.setdefault(name, target)
        return tied

    def _set_all_tied_weights_keys(self, value):
        self.__dict__["_all_tied_weights_keys_store"] = value

    model_cls.all_tied_weights_keys = property(
        _get_all_tied_weights_keys, _set_all_tied_weights_keys,
    )
    return True


try:
    from transformers import PreTrainedModel as _PTM

    _install_all_tied_weights_keys_compat(_PTM)
except Exception:
    pass


def _force_eager_config(config) -> None:
    """Recursively replace unavailable Flash Attention 2 selections."""
    from transformers import PretrainedConfig

    if getattr(config, "_attn_implementation", None) == "flash_attention_2":
        config._attn_implementation = "eager"
    for subconfig in vars(config).values():
        if isinstance(subconfig, PretrainedConfig):
            _force_eager_config(subconfig)


def _config_uses_flash_attention_2(config) -> bool:
    """Return whether this config or a nested Transformers config selects FA2."""
    from transformers import PretrainedConfig

    if getattr(config, "_attn_implementation", None) == "flash_attention_2":
        return True
    return any(
        _config_uses_flash_attention_2(subconfig)
        for subconfig in vars(config).values()
        if isinstance(subconfig, PretrainedConfig)
    )


# ── Deferred shims ──────────────────────────────────────────────────
# Some patches must wait until the _LazyModule has fully initialized
# (it replaces its __dict__ during bootstrap).  We apply these once,
# lazily, the first time load_model() is called.
_DEFERRED_SHIMS_APPLIED = False


def _apply_deferred_shims():
    global _DEFERRED_SHIMS_APPLIED
    if _DEFERRED_SHIMS_APPLIED:
        return
    _DEFERRED_SHIMS_APPLIED = True

    tf_mod = _sys.modules.get("transformers")
    if tf_mod is None:
        return

    # LogitsWarper → LogitsProcessor on the top-level transformers namespace
    try:
        if not hasattr(tf_mod, "LogitsWarper"):
            from transformers.generation.logits_process import LogitsProcessor
            tf_mod.__dict__["LogitsWarper"] = LogitsProcessor
            if hasattr(tf_mod, "_objects"):
                tf_mod._objects["LogitsWarper"] = LogitsProcessor
    except Exception:
        pass

    # is_tf_available / is_flax_available / is_safetensors_available
    # on the top-level namespace (complements shim 3 which patches submodules)
    try:
        for name, val in [
            ("is_tf_available", lambda: False),
            ("is_flax_available", lambda: False),
            ("is_safetensors_available", lambda: True),
        ]:
            if not hasattr(tf_mod, name):
                tf_mod.__dict__[name] = val
                if hasattr(tf_mod, "_objects"):
                    tf_mod._objects[name] = val
    except Exception:
        pass


TASK_MODEL_MAP = {
    "causal_lm": AutoModelForCausalLM,
    "classification": AutoModelForSequenceClassification,
}


_IMAGE_TEXT_MODEL_TYPES = {
    "gemma4_unified",
    "gemma4",
}


def _select_model_class(task: str, config: AutoConfig):
    """Return the HF AutoModel class appropriate for a task/config pair."""
    if task not in TASK_MODEL_MAP:
        raise ValueError(f"Unknown task {task!r}. Choose from {list(TASK_MODEL_MAP)}")

    model_type = getattr(config, "model_type", "")
    architectures = tuple(getattr(config, "architectures", None) or ())
    is_image_text = (
        model_type in _IMAGE_TEXT_MODEL_TYPES
        or any("ForConditionalGeneration" in arch for arch in architectures)
        and any("Gemma4" in arch for arch in architectures)
    )
    if task == "causal_lm" and is_image_text:
        if AutoModelForImageTextToText is None:
            raise RuntimeError(
                "AutoModelForImageTextToText is required for Gemma 4 unified models. "
                "Upgrade transformers to a version that provides it."
            )
        return AutoModelForImageTextToText

    return TASK_MODEL_MAP[task]


@dataclass
class ModelHandle:
    """Wrapper around a HF model + tokenizer with metadata useful for ablation."""

    model: PreTrainedModel
    tokenizer: PreTrainedTokenizerBase
    config: AutoConfig
    model_name: str
    task: str
    architecture: str = ""
    num_layers: int = 0
    num_heads: int = 0
    hidden_size: int = 0
    intermediate_size: int = 0
    _original_state: Optional[dict] = field(default=None, repr=False)
    _offload_dir: Optional[str] = field(default=None, repr=False)
    _owns_offload_dir: bool = field(default=False, repr=False)

    def __post_init__(self):
        cfg = self.config
        self.architecture = cfg.model_type
        # For composite configs (e.g. VL models like Qwen3.5), the text model
        # attributes live under a nested text_config.  Fall through to it when
        # the top-level config doesn't have the standard attributes.
        text_cfg = getattr(cfg, "text_config", None)
        self.num_layers = getattr(cfg, "num_hidden_layers", 0) or (
            getattr(text_cfg, "num_hidden_layers", 0) if text_cfg else 0
        )
        self.num_heads = getattr(cfg, "num_attention_heads", 0) or (
            getattr(text_cfg, "num_attention_heads", 0) if text_cfg else 0
        )
        self.hidden_size = getattr(cfg, "hidden_size", 0) or (
            getattr(text_cfg, "hidden_size", 0) if text_cfg else 0
        )
        self.intermediate_size = getattr(cfg, "intermediate_size", 0) or (
            getattr(text_cfg, "intermediate_size", 0) if text_cfg else 0
        )

    def snapshot(self):
        """Save a copy of the model state dict so we can restore after ablation.

        Tensors are moved to CPU to avoid doubling GPU memory usage on
        multi-GPU (device_map) setups.
        """
        self._original_state = {k: v.cpu().clone() for k, v in self.model.state_dict().items()}

    def restore(self):
        """Restore the model to the snapshot state.

        Moves CPU-saved tensors back to each parameter's current device.
        """
        if self._original_state is None:
            raise RuntimeError("No snapshot to restore — call .snapshot() first.")
        # Map each key to the device where the model currently holds it
        current_state = self.model.state_dict()
        restored = {}
        for k, v in self._original_state.items():
            target = current_state[k].device if k in current_state else None
            restored[k] = v.to(target) if target is not None else v
        incompatible = self.model.load_state_dict(restored, strict=False)
        missing = [
            key
            for key in incompatible.missing_keys
            if not _is_quantization_state_key(key)
        ]
        unexpected = [
            key
            for key in incompatible.unexpected_keys
            if not _is_quantization_state_key(key)
        ]
        if missing or unexpected:
            details = []
            if missing:
                details.append(f"missing keys: {missing}")
            if unexpected:
                details.append(f"unexpected keys: {unexpected}")
            raise RuntimeError(f"Snapshot restore was incomplete ({'; '.join(details)})")

    def cleanup(self):
        """Remove temporary offload directory if one was auto-created."""
        if self._offload_dir is not None and self._owns_offload_dir:
            import shutil
            try:
                shutil.rmtree(self._offload_dir, ignore_errors=True)
            except Exception:
                pass
        self._offload_dir = None
        self._owns_offload_dir = False

    def __del__(self):
        self.cleanup()

    def summary(self) -> dict:
        return {
            "model_name": self.model_name,
            "architecture": self.architecture,
            "task": self.task,
            "num_layers": self.num_layers,
            "num_heads": self.num_heads,
            "hidden_size": self.hidden_size,
            "intermediate_size": self.intermediate_size,
            "total_params": sum(p.numel() for p in self.model.parameters()),
        }


def _estimate_model_memory_gb(config: AutoConfig, dtype: torch.dtype) -> float:
    """Rough estimate of model weight memory in GB."""
    # Estimate total params from config.  For composite configs (VL models),
    # fall through to text_config when top-level attributes are missing.
    text_cfg = getattr(config, "text_config", None)
    hidden = getattr(config, "hidden_size", 0) or (
        getattr(text_cfg, "hidden_size", 0) if text_cfg else 0
    )
    n_layers = getattr(config, "num_hidden_layers", 0) or (
        getattr(text_cfg, "num_hidden_layers", 0) if text_cfg else 0
    )
    intermediate = getattr(config, "intermediate_size", 0) or (
        getattr(text_cfg, "intermediate_size", hidden * 4) if text_cfg else hidden * 4
    )
    vocab = getattr(config, "vocab_size", 0) or (
        getattr(text_cfg, "vocab_size", 0) if text_cfg else 0
    )

    if hidden == 0 or n_layers == 0:
        return 0.0

    # For MoE models, the FFN is replicated per expert
    num_experts = getattr(config, "num_local_experts", None) or getattr(config, "num_experts", 1)

    # Per layer: attn (4 * hidden^2) + ffn (3 * hidden * intermediate * num_experts) + norms
    per_layer = 4 * hidden * hidden + num_experts * 3 * hidden * intermediate
    # Embedding + LM head
    embedding = 2 * vocab * hidden
    total_params = per_layer * n_layers + embedding

    bytes_per_param = {torch.float32: 4, torch.float16: 2, torch.bfloat16: 2}.get(dtype, 2)
    return total_params * bytes_per_param / (1024 ** 3)


def _is_quantization_state_key(key: str) -> bool:
    """Return whether a state-dict key is bitsandbytes quantization metadata."""
    return any(
        component in _QUANTIZATION_STATE_COMPONENTS
        or component.startswith("bitsandbytes__")
        for component in key.split(".")
    )


def _effective_model_memory_gb(estimate_gb: float, quantization: str | None) -> float:
    """Adjust a full-precision weight estimate for runtime quantization."""
    return effective_model_memory_gb(estimate_gb, quantization)


def _bounded_max_memory(
    gpu_memory_utilization: float | None = None,
) -> dict[int | str, str]:
    """Build Accelerate memory limits with inference and host headroom."""
    max_memory: dict[int | str, str] = {}
    for index in range(dev.device_count()):
        total = torch.cuda.get_device_properties(index).total_memory
        if gpu_memory_utilization is None:
            reserve = max(int(total * 0.15), 2 * 1024 ** 3)
            usable = total - reserve
        else:
            usable = int(total * gpu_memory_utilization)
        max_memory[index] = f"{usable // (1024 ** 2)}MiB"
    total_ram, _ = dev._system_memory_gb()
    cpu_budget_gb = int(total_ram * 0.85)
    max_memory["cpu"] = f"{max(cpu_budget_gb, 4)}GiB"
    return max_memory


def _available_gpu_memory_gb() -> float:
    """Return free accelerator memory in GB (CUDA, MPS, or 0 for CPU)."""
    return dev.get_total_free_gb()


def _hf_token() -> str | None:
    """Return the HF_TOKEN from environment, or None."""
    return os.environ.get("HF_TOKEN") or None


def load_model(
    model_name: str,
    task: str = "causal_lm",
    device: str = "auto",
    dtype: str = "float32",
    trust_remote_code: bool = False,
    num_labels: int = 2,
    quantization: str | None = None,
    offload_folder: str | None = None,
    skip_snapshot: bool | None = None,
    gpu_memory_utilization: float | None = None,
    revision: str | None = None,
    local_files_only: bool = False,
) -> ModelHandle:
    """Load a HuggingFace model and tokenizer, returning a ModelHandle.

    Args:
        model_name: HuggingFace model identifier (e.g. "gpt2", "meta-llama/Llama-2-7b-hf").
        task: One of "causal_lm", "classification".
        device: Torch device string. "auto" uses accelerate's device_map.
        dtype: Weight dtype — "float32", "float16", "bfloat16".
        trust_remote_code: Whether to trust remote code from the Hub.
        num_labels: Number of labels for classification tasks.
        quantization: None, "4bit", or "8bit". Requires bitsandbytes.
        offload_folder: Directory for disk offloading when model exceeds GPU memory.
            If None and offloading is needed, a temp directory is created automatically.
        skip_snapshot: Controls initial state dict snapshot.
            None (default): auto-decide based on GPU memory headroom.
            True: always skip (saves memory).
            False: always snapshot (force even for large models).
        gpu_memory_utilization: Optional GPU VRAM fraction in ``(0, 1]``. When
            omitted, reserves 15% or 2 GiB per GPU, whichever is larger.
        revision: Optional Hub branch, tag, or commit passed to every loader.
        local_files_only: Refuse network access and use only locally cached files.
    """
    _apply_deferred_shims()

    validate_model_load_request(
        model_name,
        task,
        quantization,
        dtype,
        valid_tasks=TASK_MODEL_MAP,
    )
    if gpu_memory_utilization is not None and (
        isinstance(gpu_memory_utilization, bool)
        or not isinstance(gpu_memory_utilization, (int, float))
        or not 0.0 < float(gpu_memory_utilization) <= 1.0
    ):
        raise ValueError("gpu_memory_utilization must be a number in (0, 1]")
    if gpu_memory_utilization is not None:
        gpu_memory_utilization = float(gpu_memory_utilization)

    dtype_map = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}
    torch_dtype = dtype_map[dtype]
    resolved_device = dev.get_device(device)
    if dtype == "bfloat16" and not dev.supports_bfloat16(resolved_device):
        raise RuntimeError(
            f"bfloat16 is not supported on '{resolved_device}'. Use float16 or float32.",
        )

    token = _hf_token()
    hf_kwargs = {
        "trust_remote_code": trust_remote_code,
        "token": token,
        "revision": revision,
        "local_files_only": local_files_only,
    }

    try:
        config = AutoConfig.from_pretrained(model_name, **hf_kwargs)
    except PermissionError:
        fallback_cache = os.path.join(tempfile.gettempdir(), "hf_home", "hub")
        os.makedirs(fallback_cache, exist_ok=True)
        config = AutoConfig.from_pretrained(
            model_name, cache_dir=fallback_cache, **hf_kwargs,
        )
    except OSError as e:
        # Gated repo access denied — provide a clear, actionable error.
        err_msg = str(e)
        if "gated repo" in err_msg.lower() or "access to model" in err_msg.lower():
            raise RuntimeError(
                f"Access denied for gated model '{model_name}'.\n\n"
                f"This model requires you to:\n"
                f"  1. Accept the license at https://huggingface.co/{model_name}\n"
                f"  2. Set your HF_TOKEN: export HF_TOKEN=hf_...\n"
                f"     (or add it to your HF Space secrets)\n\n"
                f"Token {'is' if token else 'is NOT'} currently set."
            ) from e
        raise
    except (ValueError, KeyError) as e:
        # Unrecognized model_type — don't silently escalate trust_remote_code.
        # Provide a clear error with guidance instead.
        raise RuntimeError(
            f"Architecture '{model_name}' is not recognized by transformers "
            f"{__import__('transformers').__version__}. "
            f"Try: pip install --upgrade transformers\n"
            f"If this model requires custom code, pass trust_remote_code=True explicitly."
        ) from e

    # FP8 / NVFP4 checkpoints: dequantize to a plain float copy on disk,
    # then load through the normal float path.  Anything quantized that we
    # don't explicitly support fails loudly here, before any weight loads.
    quant_detection = qd.detect_quant_scheme(
        model_name,
        token=token,
        revision=revision,
        local_files_only=local_files_only,
    )
    _dequant_tmp = None
    _dequant_source = None
    if quant_detection.scheme is qd.QuantScheme.UNSUPPORTED:
        raise RuntimeError(
            f"Unsupported quantization in '{model_name}': {quant_detection.reason}"
        )
    if quant_detection.scheme is not qd.QuantScheme.NONE:
        logger.warning(
            "Quantized checkpoint detected (%s) — dequantizing to %s for "
            "surgery. Peak memory is the full %s model size; output is "
            "saved as %s.",
            quant_detection.scheme.value, dtype, dtype, dtype,
        )
        _dequant_source = model_name
        _dequant_tmp, _ = qd.materialize_dequantized_checkpoint(
            model_name,
            quant_detection,
            out_dtype=torch_dtype,
            token=token,
            revision=revision,
            local_files_only=local_files_only,
        )
        model_name = _dequant_tmp
        config = AutoConfig.from_pretrained(
            model_name, trust_remote_code=trust_remote_code, token=token,
        )

    # Memory estimation and warnings (skip for natively quantized models — estimate is wrong)
    native_quant = getattr(config, "quantization_config", None)
    load_policy = resolve_model_load_policy(
        device=device,
        resolved_device=resolved_device,
        dtype=dtype,
        quantization=quantization,
        has_native_quantization=native_quant is not None,
        device_map_auto_supported=dev.supports_device_map_auto(resolved_device),
        bitsandbytes_supported=dev.supports_bitsandbytes(resolved_device),
    )
    est_gb = _estimate_model_memory_gb(config, torch_dtype) if native_quant is None else 0.0
    gpu_gb = _available_gpu_memory_gb()
    if est_gb > 0 and gpu_gb > 0:
        logger.info(f"Estimated model size: {est_gb:.1f} GB | Available GPU: {gpu_gb:.1f} GB")
        if est_gb > gpu_gb * 0.9 and quantization is None:
            logger.warning(
                f"Model (~{est_gb:.0f} GB) may exceed GPU memory ({gpu_gb:.0f} GB). "
                f"Consider using quantization='4bit' or quantization='8bit'."
            )

    # Some repos pin attn_implementation="flash_attention_2" in config.json.
    # If flash-attn isn't installed, loading would fail — fall back to eager.
    # Multimodal wrappers pin it on nested sub-configs too (e.g. Nemotron
    # Omni's llm_config), so walk those as well.
    if _config_uses_flash_attention_2(config):
        try:
            from transformers.utils import is_flash_attn_2_available
            _fa2 = is_flash_attn_2_available()
        except Exception:
            _fa2 = False
        if not _fa2:
            logger.warning(
                "config.json pins flash_attention_2 but flash-attn is not "
                "installed — falling back to eager attention."
            )
            _force_eager_config(config)

    model_cls = _select_model_class(task, config)
    load_kwargs: dict = {
        "pretrained_model_name_or_path": model_name,
        "config": config,
        **hf_kwargs,
    }
    if load_policy.include_torch_dtype:
        load_kwargs["torch_dtype"] = torch_dtype
    if task == "classification":
        config.num_labels = num_labels
        load_kwargs["config"] = config

    # Quantization support (requires bitsandbytes)
    if load_policy.quantization_backend == "native":
        # Model ships with native quantization (e.g. Mxfp4Config) — don't layer BitsAndBytes
        # on top, and don't override its compute dtype with our torch_dtype
        logger.info(
            f"Model has native quantization ({type(native_quant).__name__}), "
            f"skipping BitsAndBytes and using model's native dtype"
        )
        load_kwargs["device_map"] = "auto"
    elif load_policy.quantization_backend == "bitsandbytes":
        try:
            import bitsandbytes  # noqa: F401
        except ImportError:
            raise RuntimeError(
                f"Quantization '{quantization}' requires bitsandbytes: "
                f"pip install -U bitsandbytes>=0.46.1"
            )
        from transformers import BitsAndBytesConfig

        # Enable fp32 CPU offload so that models too large to fit entirely on
        # GPU (even quantized) can spill to CPU without crashing bitsandbytes.
        # This is critical for frontier MoE models (GLM-5 744B, DeepSeek-V3 685B,
        # Mistral Large 3 675B, etc.) on single-GPU setups.
        if quantization == "4bit":
            load_kwargs["quantization_config"] = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch_dtype,
                bnb_4bit_quant_type="nf4",
                llm_int8_enable_fp32_cpu_offload=True,
            )
        else:
            load_kwargs["quantization_config"] = BitsAndBytesConfig(
                load_in_8bit=True,
                llm_int8_enable_fp32_cpu_offload=True,
            )
        load_kwargs["device_map"] = "auto"

    # device_map="auto" is only reliable on CUDA (accelerate doesn't support MPS).
    if "device_map" not in load_kwargs and load_policy.use_device_map_auto:
        load_kwargs["device_map"] = "auto"

    # Offload support: provide a folder for disk offloading when GPU memory is insufficient
    _offload_dir = None
    _owns_offload_dir = False
    if load_kwargs.get("device_map") == "auto":
        if offload_folder:
            _offload_dir = offload_folder
            load_kwargs["offload_folder"] = offload_folder
        else:
            # Auto-create a temp offload dir so from_pretrained never crashes
            # when Accelerate needs disk offloading
            _offload_dir = tempfile.mkdtemp(prefix="obliteratus_offload_")
            _owns_offload_dir = True
            load_kwargs["offload_folder"] = _offload_dir
            logger.info(f"Auto-created offload folder: {_offload_dir}")

        effective_est_gb = _effective_model_memory_gb(est_gb, quantization)
        quantized_fit = (
            gpu_memory_utilization is None
            and quantized_model_fits_gpu(est_gb, quantization, gpu_gb)
        )
        if quantized_fit:
            logger.info(
                f"Quantized estimate ({effective_est_gb:.1f} GB) fits GPU "
                f"({gpu_gb:.0f} GB) — skipping max_memory constraint"
            )
        elif dev.is_cuda():
            max_memory = _bounded_max_memory(gpu_memory_utilization)
            load_kwargs["max_memory"] = max_memory
            logger.info(
                f"GPU memory budget: {', '.join(f'GPU{k}={v}' for k, v in max_memory.items() if k != 'cpu')}"
            )

    try:
        model = model_cls.from_pretrained(**load_kwargs)
    except PermissionError as e:
        # Cache dir (typically ~/.cache/huggingface) is not writable — common in
        # containers running as UID with no home dir. Retry with /tmp cache.
        logger.warning(
            "PermissionError loading model (%s). Retrying with cache_dir=/tmp/hf_home/hub", e
        )
        fallback_cache = os.path.join(tempfile.gettempdir(), "hf_home", "hub")
        os.makedirs(fallback_cache, exist_ok=True)
        load_kwargs["cache_dir"] = fallback_cache
        model = model_cls.from_pretrained(**load_kwargs)
    except OSError as e:
        err_msg = str(e)
        if "gated repo" in err_msg.lower() or "access to model" in err_msg.lower():
            raise RuntimeError(
                f"Access denied for gated model '{model_name}'.\n\n"
                f"This model requires you to:\n"
                f"  1. Accept the license at https://huggingface.co/{model_name}\n"
                f"  2. Set your HF_TOKEN: export HF_TOKEN=hf_...\n"
                f"     (or add it to your HF Space secrets)\n\n"
                f"Token {'is' if token else 'is NOT'} currently set."
            ) from e
        raise
    except (ValueError, KeyError) as e:
        err_msg = str(e)
        if "does not recognize this architecture" in err_msg or "model type" in err_msg:
            model_type = getattr(config, "model_type", "unknown")
            raise RuntimeError(
                f"Model architecture '{model_type}' is not supported by transformers "
                f"{__import__('transformers').__version__}. "
                f"Run: pip install --upgrade transformers\n"
                f"If this model was released very recently, it may require "
                f"pip install git+https://github.com/huggingface/transformers.git"
            ) from e
        raise
    finally:
        # from_pretrained fully materializes weights before returning. Clean up
        # on success and across every typed/retry failure path above.
        if _dequant_tmp is not None:
            import shutil as _shutil

            _shutil.rmtree(_dequant_tmp, ignore_errors=True)
            logger.info("Removed temporary dequantized checkpoint %s", _dequant_tmp)

    if load_policy.move_to_resolved_device:
        # Explicit devices and auto-selected MPS/CPU load on CPU before moving.
        model = model.to(resolved_device)

    model.eval()

    # Multimodal wrappers (e.g. Nemotron Omni) require pixel_values/audio in
    # forward() — unusable for text-only surgery.  Unwrap the language-model
    # submodule; the pipeline only ever edits the LM.
    if hasattr(model, "language_model"):
        import inspect as _inspect
        lm = getattr(model, "language_model", None)
        try:
            params = _inspect.signature(model.forward).parameters
            needs_media = any(
                p.default is _inspect.Parameter.empty
                and p.kind in (_inspect.Parameter.POSITIONAL_ONLY,
                               _inspect.Parameter.POSITIONAL_OR_KEYWORD)
                and p.name not in ("self", "input_ids")
                for p in params.values()
            )
        except (TypeError, ValueError):
            needs_media = False
        if needs_media and isinstance(lm, PreTrainedModel):
            logger.warning(
                "Multimodal wrapper %s requires media inputs in forward() — "
                "unwrapping language_model (%s) for text-only surgery.",
                type(model).__name__, type(lm).__name__,
            )
            model = lm
            config = model.config

    # transformers 5.x may call prepare_inputs_for_generation with
    # cache_position=None; older remote code (Nemotron-H) assumes a tensor.
    _pifg = getattr(model, "prepare_inputs_for_generation", None)
    if _pifg is not None:
        def _prepare_inputs_safe(input_ids, *args, cache_position=None, **kwargs):
            if cache_position is None:
                past_len = 0
                pkv = kwargs.get("past_key_values") or (args[0] if args else None)
                if pkv is not None:
                    try:
                        past_len = pkv.get_seq_length()
                    except AttributeError:
                        try:
                            past_len = pkv[0][0].shape[-2]
                        except Exception:
                            past_len = 0
                cache_position = torch.arange(
                    past_len, past_len + input_ids.shape[-1],
                    device=input_ids.device,
                )
            return _pifg(input_ids, *args, cache_position=cache_position, **kwargs)

        model.prepare_inputs_for_generation = _prepare_inputs_safe

    # transformers 5.x expects module._tied_weights_keys as a dict; older
    # remote code (Nemotron-H) defines it as a list.  Normalize: when word
    # embeddings aren't actually tied, nothing should be dropped at save.
    _tie_embeddings = bool(getattr(config, "tie_word_embeddings", False))
    for _m in model.modules():
        _twk = getattr(_m, "_tied_weights_keys", None)
        if isinstance(_twk, list):
            _m._tied_weights_keys = {k: k for k in _twk} if _tie_embeddings else {}

    # Free accelerator cache after loading
    dev.empty_cache()

    if _dequant_tmp is not None:
        model._obliteratus_dequantized_scheme = quant_detection.scheme.value

    try:
        tokenizer = AutoTokenizer.from_pretrained(
            _dequant_source or model_name,
            **hf_kwargs,
        )
    except PermissionError:
        fallback_cache = os.path.join(tempfile.gettempdir(), "hf_home", "hub")
        tokenizer = AutoTokenizer.from_pretrained(
            _dequant_source or model_name,
            cache_dir=fallback_cache,
            **hf_kwargs,
        )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    handle = ModelHandle(
        model=model,
        tokenizer=tokenizer,
        config=config,
        model_name=_dequant_source or model_name,
        task=task,
        _offload_dir=_offload_dir,
        _owns_offload_dir=_owns_offload_dir,
    )

    # Skip snapshot for large models to avoid doubling memory usage.
    remaining_gpu_gb = (
        dev.get_total_free_gb()
        if skip_snapshot is None and gpu_gb > 0 and native_quant is not None
        else gpu_gb
    )
    snapshot = should_snapshot_model(
        skip_snapshot=skip_snapshot,
        initial_gpu_free_gb=gpu_gb,
        remaining_gpu_free_gb=remaining_gpu_gb,
        has_native_quantization=native_quant is not None,
        estimate_gb=est_gb,
        quantization=quantization,
    )
    if snapshot:
        handle.snapshot()
    elif skip_snapshot is None and gpu_gb > 0 and native_quant is not None:
        logger.warning(
            f"Auto-skipping state dict snapshot for natively quantized model "
            f"(free GPU: {remaining_gpu_gb:.1f} GB / {gpu_gb:.1f} GB). "
            f"Use skip_snapshot=False to force."
        )
    elif skip_snapshot is None and gpu_gb > 0 and est_gb > 0:
        effective_gb = _effective_model_memory_gb(est_gb, quantization)
        logger.warning(
            f"Auto-skipping state dict snapshot to save memory "
            f"(model ~{effective_gb:.0f} GB vs GPU {gpu_gb:.0f} GB). "
            f"Use skip_snapshot=False to force."
        )

    return handle