255 lines
8.9 KiB
Python
255 lines
8.9 KiB
Python
"""
|
|
Device, dtype, and VRAM utility functions for model loading.
|
|
Shared utilities for both GGUF and HuggingFace model loaders.
|
|
"""
|
|
|
|
import logging
|
|
|
|
import torch
|
|
|
|
# Try to import ComfyUI model management for device detection
|
|
try:
|
|
from comfy import model_management
|
|
|
|
COMFYUI_AVAILABLE = True
|
|
except ImportError:
|
|
COMFYUI_AVAILABLE = False
|
|
logging.warning("ComfyUI model_management not available, using basic device detection")
|
|
|
|
|
|
def calculate_optimal_gpu_layers(model_size_gb: float, available_vram_gb: float, reserve_gb: float = 2.0, context_size: int = 4096) -> int:
|
|
"""
|
|
Calculate optimal number of GPU layers based on available VRAM.
|
|
|
|
Args:
|
|
model_size_gb: Total model size in GB
|
|
available_vram_gb: Available VRAM in GB
|
|
reserve_gb: VRAM to reserve for other operations (default: 2.0 GB)
|
|
context_size: Context window size (used to estimate KV cache size)
|
|
|
|
Returns:
|
|
Number of layers to load on GPU (-1 for all, 0 for CPU only, or specific count)
|
|
"""
|
|
if available_vram_gb <= 0:
|
|
return 0 # CPU only
|
|
|
|
# 1. Reserve VRAM for system + KV cache
|
|
# Rough KV cache estimate: ~1-2GB for 4096 context, scales linearly with context
|
|
# Formula: 2 * 2 * layers * hidden_size * context_len (conservative estimate)
|
|
# Using a simpler heuristic: 0.5GB per 4096 tokens of context
|
|
kv_cache_gb = (context_size / 4096.0) * 0.5
|
|
effective_reserve = reserve_gb + kv_cache_gb
|
|
|
|
usable_vram = available_vram_gb - effective_reserve
|
|
if usable_vram <= 0:
|
|
return 0
|
|
|
|
if model_size_gb <= usable_vram:
|
|
return -1 # Fits fully
|
|
|
|
# 2. Estimate per-layer size dynamically
|
|
# Assume ~40-80 layers for 7B-14B models.
|
|
# If we don't know layer count, estimating based on file size is better than fixed 100MB
|
|
# Qwen2.5-14B has 48 layers. Qwen2.5-7B has 28 layers.
|
|
# Average layers approx 40 for safe estimation.
|
|
estimated_layers = 40
|
|
layer_size_gb = model_size_gb / estimated_layers
|
|
|
|
# 3. Calculate layers that fit
|
|
max_layers = int(usable_vram / layer_size_gb)
|
|
return max(0, max_layers)
|
|
|
|
|
|
def get_available_vram() -> float:
|
|
"""
|
|
Get available VRAM in GB using ComfyUI's memory management if available.
|
|
|
|
Returns:
|
|
Available VRAM in GB
|
|
"""
|
|
if COMFYUI_AVAILABLE:
|
|
try:
|
|
dev = model_management.get_torch_device()
|
|
free_memory = model_management.get_free_memory(dev)
|
|
return free_memory / (1024**3)
|
|
except Exception as e:
|
|
logging.warning(f"Failed to get VRAM from ComfyUI: {e}")
|
|
|
|
if torch.cuda.is_available():
|
|
try:
|
|
free_memory, _ = torch.cuda.mem_get_info(0)
|
|
return free_memory / (1024**3)
|
|
except Exception as e:
|
|
logging.warning(f"Failed to get VRAM from torch: {e}")
|
|
|
|
return 0.0
|
|
|
|
|
|
def get_device_for_model(device_preference: str = "auto") -> str:
|
|
"""
|
|
Get device string for model loading.
|
|
|
|
Args:
|
|
device_preference: "auto", "cuda", or "cpu"
|
|
|
|
Returns:
|
|
Device string ("cuda" or "cpu")
|
|
"""
|
|
if device_preference == "cpu":
|
|
return "cpu"
|
|
|
|
if device_preference == "cuda" or device_preference == "auto":
|
|
if COMFYUI_AVAILABLE:
|
|
try:
|
|
dev = model_management.get_torch_device()
|
|
if hasattr(dev, "type"):
|
|
if dev.type == "cuda":
|
|
return "cuda"
|
|
elif dev.type == "cpu":
|
|
return "cpu"
|
|
except Exception:
|
|
pass
|
|
|
|
# Fallback to torch detection
|
|
if torch.cuda.is_available():
|
|
return "cuda"
|
|
|
|
return "cpu"
|
|
|
|
|
|
def get_dtype_for_model(dtype_preference: str = "auto", device: str = "cuda") -> torch.dtype:
|
|
"""
|
|
Get dtype for HuggingFace model loading.
|
|
|
|
Args:
|
|
dtype_preference: "auto", "fp16", "bf16", "fp32", or "fp8"
|
|
device: Device string ("cuda" or "cpu")
|
|
|
|
Returns:
|
|
torch.dtype
|
|
"""
|
|
if dtype_preference == "fp32":
|
|
return torch.float32
|
|
|
|
if dtype_preference == "fp16":
|
|
return torch.float16
|
|
|
|
if dtype_preference == "bf16":
|
|
if device == "cuda" and torch.cuda.is_bf16_supported():
|
|
return torch.bfloat16
|
|
return torch.float16
|
|
|
|
if dtype_preference == "fp8":
|
|
# FP8 is CUDA-only and requires specific hardware support
|
|
if device == "cpu":
|
|
logging.debug("FP8 is not supported on CPU, falling back to fp16")
|
|
return torch.float16
|
|
|
|
# FP8 requires CUDA and specific PyTorch builds with FP8 support
|
|
if not torch.cuda.is_available():
|
|
logging.debug("FP8 requires CUDA, but CUDA is not available, falling back to fp16")
|
|
return torch.float16
|
|
|
|
if COMFYUI_AVAILABLE:
|
|
try:
|
|
float8_types = model_management.get_supported_float8_types()
|
|
if float8_types:
|
|
selected_type = float8_types[0]
|
|
# Verify FP8 storage is actually available by testing set_default_dtype
|
|
# This is what transformers will do internally, so we catch the error early
|
|
try:
|
|
original_dtype = torch.get_default_dtype()
|
|
torch.set_default_dtype(selected_type)
|
|
torch.set_default_dtype(original_dtype)
|
|
type_name = str(selected_type).replace("torch.", "")
|
|
logging.debug(f"Using FP8: {type_name}")
|
|
return selected_type
|
|
except (TypeError, RuntimeError) as e:
|
|
# FP8 dtype exists but storage is not available
|
|
logging.debug(f"FP8 storage not available: {e}, falling back to fp16")
|
|
return torch.float16
|
|
except Exception as e:
|
|
logging.debug(f"Failed to get FP8 types: {e}, falling back to fp16")
|
|
logging.debug("FP8 not available, falling back to fp16")
|
|
return torch.float16
|
|
|
|
if COMFYUI_AVAILABLE and device == "cuda":
|
|
try:
|
|
model_management.get_torch_device()
|
|
if torch.cuda.is_bf16_supported():
|
|
return torch.bfloat16
|
|
return torch.float16
|
|
except Exception:
|
|
pass
|
|
|
|
# Default: fp16 for CUDA, fp32 for CPU
|
|
if device == "cuda":
|
|
return torch.float16
|
|
return torch.float32
|
|
|
|
|
|
def load_model(
|
|
model_format: str,
|
|
model_path: str,
|
|
tokenizer_path: str | None = None,
|
|
n_gpu_layers: int | str = -1,
|
|
dtype: str | torch.dtype = "auto",
|
|
device_preference: str = "auto",
|
|
max_memory: dict | None = None,
|
|
attn_implementation: str = "auto",
|
|
context_size: int = 4096,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Unified model loader that supports both GGUF and HuggingFace formats.
|
|
|
|
Args:
|
|
model_format: "gguf" or "huggingface"
|
|
model_path: Path to model file (GGUF) or directory (HuggingFace)
|
|
tokenizer_path: Path to tokenizer directory (auto-detected if None)
|
|
n_gpu_layers: For GGUF: "auto" (calculate optimal), -1 for all, 0 for CPU only, or 1+ for specific count.
|
|
If value exceeds model's total layers, all layers will be used (no error).
|
|
dtype: For HuggingFace: "auto", "fp16", "bf16", "fp32", "fp8", or torch.dtype
|
|
device_preference: "auto", "cuda", or "cpu"
|
|
max_memory: For HuggingFace: Dict for device_map memory limits
|
|
attn_implementation: For HuggingFace: "auto", "flash_attention_2", or "sdpa"
|
|
context_size: For GGUF: Context window size
|
|
**kwargs: Additional arguments passed to model wrapper initialization
|
|
|
|
Returns:
|
|
Model wrapper instance (GGUFModelWrapper or HuggingFaceModelWrapper)
|
|
"""
|
|
import os
|
|
|
|
from ..model_manager import auto_detect_tokenizer_path
|
|
|
|
if model_format.lower() == "gguf":
|
|
from .gguf_loader import GGUFModelWrapper
|
|
|
|
if tokenizer_path is None:
|
|
# Ensure model_path is absolute and normalized before auto-detection
|
|
if not os.path.isabs(model_path):
|
|
model_path = os.path.abspath(model_path)
|
|
model_path = os.path.normpath(model_path)
|
|
tokenizer_path = auto_detect_tokenizer_path(model_path, model_format="gguf")
|
|
logging.info(f"Auto-detected tokenizer path for GGUF model: {tokenizer_path}")
|
|
return GGUFModelWrapper(
|
|
model_path, tokenizer_path, n_gpu_layers=n_gpu_layers, device_preference=device_preference, context_size=context_size, **kwargs
|
|
)
|
|
|
|
elif model_format.lower() == "huggingface":
|
|
from .huggingface_loader import HuggingFaceModelWrapper
|
|
|
|
return HuggingFaceModelWrapper(
|
|
model_path,
|
|
tokenizer_path,
|
|
dtype=dtype,
|
|
device_preference=device_preference,
|
|
max_memory=max_memory,
|
|
attn_implementation=attn_implementation,
|
|
**kwargs,
|
|
)
|
|
|
|
else:
|
|
raise ValueError(f"Unsupported model_format: {model_format}. Supported formats: 'gguf', 'huggingface'")
|