new tasks, model pricing, auto quantisation, etc...

packages/node/meshnet_node/model_backend.py

@@ -7,7 +7,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Literal
 
-Quantization = Literal["bfloat16", "int8", "nf4"]
+Quantization = Literal["auto", "bfloat16", "int8", "nf4"]
 
 
 class ModelBackendError(RuntimeError):
@@ -31,14 +31,14 @@ class TensorPayload:
 
 
 def validate_quantization(value: str) -> Quantization:
-    if value not in {"bfloat16", "int8", "nf4"}:
-        raise ValueError("quantization must be one of: bfloat16, int8, nf4")
+    if value not in {"auto", "bfloat16", "int8", "nf4"}:
+        raise ValueError("quantization must be one of: auto, bfloat16, int8, nf4")
     return value  # type: ignore[return-value]
 
 
 def build_quantization_config(quantization: Quantization) -> Any | None:
     """Return a transformers BitsAndBytesConfig for quantized weights."""
-    if quantization == "bfloat16":
+    if quantization in {"auto", "bfloat16"}:
         return None
     try:
         import torch
@@ -65,7 +65,7 @@ class TorchModelShard:
         model_id: str,
         shard_start: int,
         shard_end: int,
-        quantization: Quantization = "bfloat16",
+        quantization: Quantization = "auto",
         cache_dir: Path | None = None,
     ) -> None:
         if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
@@ -77,7 +77,7 @@ class TorchModelShard:
 
         try:
             import torch
-            from transformers import AutoModelForCausalLM, AutoTokenizer
+            from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
         except ModuleNotFoundError as exc:
             raise MissingModelDependencyError(
                 "real model backend requires torch, transformers, safetensors, accelerate, and bitsandbytes"
@@ -85,17 +85,27 @@ class TorchModelShard:
 
         self.torch = torch
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        quant_config = build_quantization_config(quantization)
+        quant_config, dtype, uses_quantized_weights = _model_load_plan(
+            AutoConfig,
+            model_id,
+            quantization,
+            torch,
+            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 else None,
+            }
+            if quant_config is not None:
+                load_kwargs["quantization_config"] = quant_config
             self.model = AutoModelForCausalLM.from_pretrained(
                 model_id,
-                quantization_config=quant_config,
-                device_map="auto" if quant_config is not None else None,
-                dtype=torch.bfloat16,
-                low_cpu_mem_usage=True,
-                cache_dir=str(cache_dir) if cache_dir is not None else None,
+                **load_kwargs,
             )
-            if quant_config is None:
+            if not uses_quantized_weights:
                 self.model.to(self.device)
         except Exception as exc:
             if _looks_like_oom(exc):
@@ -340,12 +350,71 @@ def load_torch_shard(
     model_id: str,
     shard_start: int,
     shard_end: int,
-    quantization: Quantization = "bfloat16",
+    quantization: Quantization = "auto",
     cache_dir: Path | None = None,
 ) -> TorchModelShard:
     return TorchModelShard(model_id, shard_start, shard_end, quantization, cache_dir)
 
 
+def _model_load_plan(
+    auto_config: Any,
+    model_id: str,
+    quantization: Quantization,
+    torch: Any,
+    cache_dir: Path | None = None,
+) -> tuple[Any | None, Any, bool]:
+    """Return (explicit quant config, dtype, uses quantized weights)."""
+    if quantization != "auto":
+        quant_config = build_quantization_config(quantization)
+        return quant_config, torch.bfloat16, quant_config is not None
+
+    cfg = auto_config.from_pretrained(
+        model_id,
+        cache_dir=str(cache_dir) if cache_dir is not None else None,
+    )
+    if _native_quantization_config(cfg) is not None:
+        return None, _native_torch_dtype(cfg, torch), True
+    return None, _native_torch_dtype(cfg, torch), False
+
+
+def _config_candidates(cfg: Any) -> list[Any]:
+    candidates = [cfg]
+    get_text_config = getattr(cfg, "get_text_config", None)
+    if callable(get_text_config):
+        try:
+            candidates.append(get_text_config())
+        except Exception:
+            pass
+    text_config = getattr(cfg, "text_config", None)
+    if text_config is not None:
+        candidates.append(text_config)
+    return candidates
+
+
+def _native_quantization_config(cfg: Any) -> Any | None:
+    for candidate in _config_candidates(cfg):
+        quant_config = getattr(candidate, "quantization_config", None)
+        if quant_config:
+            return quant_config
+    return None
+
+
+def _native_torch_dtype(cfg: Any, torch: Any) -> Any:
+    for candidate in _config_candidates(cfg):
+        for attr in ("dtype", "torch_dtype"):
+            dtype = getattr(candidate, attr, None)
+            if dtype is None:
+                continue
+            if isinstance(dtype, str):
+                dtype_name = dtype.removeprefix("torch.")
+                dtype_value = getattr(torch, dtype_name, None)
+                if dtype_value is not None:
+                    return dtype_value
+            else:
+                return dtype
+    return torch.bfloat16
+
+
 def _model_layers(model: Any) -> Any:
     if hasattr(model, "model") and hasattr(model.model, "layers"):
         return model.model.layers