Upload 80 files

obliteratus/models/loader.py

@@ -6,7 +6,6 @@ import copy
 import logging
 import tempfile
 from dataclasses import dataclass, field
-from pathlib import Path
 from typing import Optional
 
 import torch
@@ -103,7 +102,7 @@ def _available_gpu_memory_gb() -> float:
     total = 0.0
     for i in range(torch.cuda.device_count()):
         props = torch.cuda.get_device_properties(i)
-        total += props.total_mem / (1024 ** 3)
+        total += props.total_memory / (1024 ** 3)
     return total
 
 
@@ -116,7 +115,7 @@ def load_model(
     num_labels: int = 2,
     quantization: str | None = None,
     offload_folder: str | None = None,
-    skip_snapshot: bool = False,
+    skip_snapshot: bool | None = None,
 ) -> ModelHandle:
     """Load a HuggingFace model and tokenizer, returning a ModelHandle.
 
@@ -130,7 +129,10 @@ def load_model(
         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: If True, skip the initial state dict snapshot to save memory.
+        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).
     """
     if task not in TASK_MODEL_MAP:
         raise ValueError(f"Unknown task {task!r}. Choose from {list(TASK_MODEL_MAP)}")
@@ -215,10 +217,15 @@ def load_model(
     )
 
     # Skip snapshot for large models to avoid doubling memory usage
-    if not skip_snapshot:
-        if est_gb > 0 and est_gb > gpu_gb * 0.5:
+    if skip_snapshot is True:
+        pass  # user explicitly opted out
+    elif skip_snapshot is False:
+        handle.snapshot()  # user explicitly forced snapshot
+    else:
+        # Auto-decide: skip only when GPU is present and model is tight on memory
+        if gpu_gb > 0 and est_gb > 0 and est_gb > gpu_gb * 0.5:
             logger.warning(
-                f"Skipping state dict snapshot to save memory "
+                f"Auto-skipping state dict snapshot to save memory "
                 f"(model ~{est_gb:.0f} GB vs GPU {gpu_gb:.0f} GB). "
                 f"Use skip_snapshot=False to force."
             )