dash QOL

packages/node/meshnet_node/model_backend.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import base64
 from dataclasses import dataclass
+import json
 from pathlib import Path
 from typing import Any, Literal
 
@@ -22,6 +23,10 @@ class InsufficientVRAMError(ModelBackendError):
     """Raised when a requested shard cannot fit in available CUDA memory."""
 
 
+class PartialModelLoadUnsupported(ModelBackendError):
+    """Raised when a shard cannot be materialized from a local snapshot subset."""
+
+
 @dataclass(frozen=True)
 class TensorPayload:
     body: bytes
@@ -94,20 +99,39 @@ class TorchModelShard:
             None if load_source != model_id else cache_dir,
         )
         try:
-            load_kwargs = {
-                "device_map": "auto" if uses_quantized_weights else None,
-                "dtype": dtype,
-                "low_cpu_mem_usage": True,
-                "cache_dir": str(cache_dir) if cache_dir is not None and load_source == model_id else None,
-            }
-            if quant_config is not None:
-                load_kwargs["quantization_config"] = quant_config
-            self.model = AutoModelForCausalLM.from_pretrained(
+            total_layers_hint = _total_layers_for_local_snapshot(AutoConfig, load_source)
+            if _should_partial_materialize_shard(
                 load_source,
-                **load_kwargs,
-            )
-            if not uses_quantized_weights:
-                self.model.to(self.device)
+                shard_start,
+                shard_end,
+                total_layers_hint=total_layers_hint,
+                uses_quantized_weights=uses_quantized_weights,
+            ):
+                self.model = _load_partial_model_from_snapshot(
+                    AutoConfig,
+                    AutoModelForCausalLM,
+                    torch,
+                    load_source,
+                    shard_start,
+                    shard_end,
+                    dtype,
+                    self.device,
+                )
+            else:
+                load_kwargs = {
+                    "device_map": "auto" if uses_quantized_weights else None,
+                    "dtype": dtype,
+                    "low_cpu_mem_usage": True,
+                    "cache_dir": str(cache_dir) if cache_dir is not None and load_source == model_id else None,
+                }
+                if quant_config is not None:
+                    load_kwargs["quantization_config"] = quant_config
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    load_source,
+                    **load_kwargs,
+                )
+                if not uses_quantized_weights:
+                    self.model.to(self.device)
         except Exception as exc:
             if _looks_like_oom(exc):
                 raise InsufficientVRAMError(
@@ -357,6 +381,135 @@ def load_torch_shard(
     return TorchModelShard(model_id, shard_start, shard_end, quantization, cache_dir)
 
 
+def _total_layers_for_local_snapshot(auto_config: Any, load_source: str) -> int | None:
+    snapshot_dir = Path(load_source)
+    if not (snapshot_dir / "config.json").exists():
+        return None
+    from .model_catalog import layers_from_config
+
+    try:
+        cfg = auto_config.from_pretrained(str(snapshot_dir))
+    except Exception:
+        return None
+    return layers_from_config(cfg)
+
+
+def _should_partial_materialize_shard(
+    load_source: str,
+    shard_start: int,
+    shard_end: int,
+    *,
+    total_layers_hint: int | None,
+    uses_quantized_weights: bool,
+) -> bool:
+    if uses_quantized_weights:
+        return False
+    snapshot_dir = Path(load_source)
+    if not snapshot_dir.exists() or not (snapshot_dir / "config.json").exists():
+        return False
+    if not (snapshot_dir / "model.safetensors.index.json").exists():
+        return False
+    if total_layers_hint is None:
+        return False
+    return not (shard_start == 0 and shard_end >= total_layers_hint - 1)
+
+
+def _load_partial_model_from_snapshot(
+    auto_config: Any,
+    auto_model_for_causal_lm: Any,
+    torch: Any,
+    load_source: str,
+    shard_start: int,
+    shard_end: int,
+    dtype: Any,
+    device: Any,
+    *,
+    init_empty_weights_fn: Any | None = None,
+    set_tensor_fn: Any | None = None,
+    safe_open_fn: Any | None = None,
+) -> Any:
+    from .model_catalog import layers_from_config
+    from .safetensors_selection import (
+        INDEX_FILENAME,
+        select_tensor_names_for_layers_from_index,
+    )
+
+    if init_empty_weights_fn is None:
+        from accelerate import init_empty_weights as init_empty_weights_fn
+    if set_tensor_fn is None:
+        from accelerate.utils import set_module_tensor_to_device as set_tensor_fn
+    if safe_open_fn is None:
+        from safetensors import safe_open as safe_open_fn
+
+    snapshot_dir = Path(load_source)
+    cfg = auto_config.from_pretrained(str(snapshot_dir))
+    total_layers = layers_from_config(cfg)
+    if total_layers is None:
+        raise PartialModelLoadUnsupported(
+            f"could not determine num_hidden_layers for local snapshot {snapshot_dir}"
+        )
+    if shard_end >= total_layers:
+        raise ValueError(
+            f"shard_end {shard_end} exceeds last layer index {total_layers - 1}"
+        )
+
+    index_path = snapshot_dir / INDEX_FILENAME
+    try:
+        index = json.loads(index_path.read_text(encoding="utf-8"))
+    except FileNotFoundError as exc:
+        raise PartialModelLoadUnsupported(
+            f"missing SafeTensors index for partial load: {index_path}"
+        ) from exc
+    weight_map = index.get("weight_map")
+    if not isinstance(weight_map, dict):
+        raise PartialModelLoadUnsupported(f"{INDEX_FILENAME} must contain a weight_map object")
+
+    tensor_names = select_tensor_names_for_layers_from_index(
+        weight_map,
+        shard_start,
+        shard_end,
+        total_layers=total_layers,
+    )
+    if not tensor_names:
+        raise PartialModelLoadUnsupported(
+            f"no checkpoint tensors matched layers {shard_start}-{shard_end} in {snapshot_dir}"
+        )
+
+    with init_empty_weights_fn():
+        model = auto_model_for_causal_lm.from_config(cfg, torch_dtype=dtype)
+    tie_weights = getattr(model, "tie_weights", None)
+    if callable(tie_weights):
+        tie_weights()
+
+    tensors_by_file: dict[str, list[str]] = {}
+    for tensor_name in sorted(tensor_names):
+        rel_file = weight_map.get(tensor_name)
+        if not isinstance(rel_file, str):
+            continue
+        tensors_by_file.setdefault(rel_file, []).append(tensor_name)
+
+    for rel_file, names in tensors_by_file.items():
+        checkpoint_file = snapshot_dir / rel_file
+        if not checkpoint_file.exists():
+            raise PartialModelLoadUnsupported(
+                f"checkpoint file advertised in {INDEX_FILENAME} is missing: {checkpoint_file}"
+            )
+        with safe_open_fn(str(checkpoint_file), framework="pt", device="cpu") as handle:
+            for tensor_name in names:
+                set_tensor_fn(
+                    model,
+                    tensor_name,
+                    device,
+                    value=handle.get_tensor(tensor_name),
+                    dtype=dtype,
+                )
+
+    for module in _active_modules_for_shard(model, shard_start, shard_end):
+        if hasattr(module, "to"):
+            module.to(device)
+    return model
+
+
 def _model_load_plan(
     auto_config: Any,
     model_id: str,
@@ -442,6 +595,37 @@ def _position_embeddings(model: Any) -> Any | None:
     return None
 
 
+def _rotary_embedding_module(model: Any) -> Any | None:
+    if hasattr(model, "model") and hasattr(model.model, "rotary_emb"):
+        return model.model.rotary_emb
+    if hasattr(model, "transformer") and hasattr(model.transformer, "rotary_emb"):
+        return model.transformer.rotary_emb
+    return None
+
+
+def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> list[Any]:
+    active: list[Any] = []
+
+    def add(module: Any | None) -> None:
+        if module is None:
+            return
+        if any(existing is module for existing in active):
+            return
+        active.append(module)
+
+    if shard_start == 0:
+        add(_embed_tokens(model))
+        add(_position_embeddings(model))
+    add(_rotary_embedding_module(model))
+    for layer in _model_layers(model)[shard_start:shard_end + 1]:
+        add(layer)
+    total_layers = len(_model_layers(model))
+    if shard_end >= total_layers - 1:
+        add(_final_norm(model))
+        add(getattr(model, "lm_head", None))
+    return active
+
+
 def _final_norm(model: Any) -> Any | None:
     if hasattr(model, "model") and hasattr(model.model, "norm"):
         return model.model.norm
@@ -485,11 +669,7 @@ def _rotary_position_embeddings(model: Any, hidden_states: Any, position_ids: An
     """Return model-level rotary embeddings required by newer HF decoder layers."""
     if position_ids is None:
         return None
-    rotary = None
-    if hasattr(model, "model") and hasattr(model.model, "rotary_emb"):
-        rotary = model.model.rotary_emb
-    elif hasattr(model, "transformer") and hasattr(model.transformer, "rotary_emb"):
-        rotary = model.transformer.rotary_emb
+    rotary = _rotary_embedding_module(model)
     if rotary is None:
         return None
     return rotary(hidden_states, position_ids)