239 lines
9.6 KiB
Python
239 lines
9.6 KiB
Python
"""
|
|
HuggingFace model loader using transformers with GPU layer management.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import time
|
|
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
from ..model_manager import auto_detect_tokenizer_path
|
|
from .device_utils import (
|
|
get_device_for_model,
|
|
get_dtype_for_model,
|
|
)
|
|
|
|
|
|
class HuggingFaceModelWrapper:
|
|
"""Wrapper for HuggingFace models using transformers with GPU layer management."""
|
|
|
|
def __init__(
|
|
self,
|
|
model_path: str,
|
|
tokenizer_path: str | None = None,
|
|
dtype: str | torch.dtype = "auto",
|
|
device_preference: str = "auto",
|
|
max_memory: dict | None = None,
|
|
attn_implementation: str = "auto",
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Initialize HuggingFace model wrapper.
|
|
|
|
Args:
|
|
model_path: Path to HuggingFace model directory
|
|
tokenizer_path: Optional separate tokenizer path (defaults to model_path)
|
|
dtype: "auto", "fp16", "bf16", "fp32", "fp8", or torch.dtype
|
|
device_preference: "auto", "cuda", or "cpu"
|
|
max_memory: Dict for device_map memory limits (e.g., {0: "10GB", "cpu": "30GB"})
|
|
attn_implementation: "auto", "flash_attention_2", or "sdpa"
|
|
**kwargs: Additional arguments for AutoModelForCausalLM.from_pretrained
|
|
"""
|
|
logging.info(f"Loading HuggingFace model: {os.path.basename(model_path)}")
|
|
load_start = time.time()
|
|
|
|
device = get_device_for_model(device_preference)
|
|
logging.debug(f"Device preference: {device_preference}, selected: {device}")
|
|
|
|
if isinstance(dtype, str):
|
|
if dtype == "auto" and device_preference == "auto":
|
|
torch_dtype = get_dtype_for_model("fp16", device)
|
|
else:
|
|
torch_dtype = get_dtype_for_model(dtype, device)
|
|
else:
|
|
torch_dtype = dtype
|
|
|
|
logging.debug(f"Using dtype: {torch_dtype}")
|
|
|
|
if attn_implementation == "auto":
|
|
try:
|
|
import flash_attn # noqa: F401
|
|
|
|
if device == "cuda" and torch.cuda.is_available():
|
|
attn_impl = "flash_attention_2"
|
|
logging.debug("Auto-detected attention: flash_attention_2")
|
|
else:
|
|
attn_impl = "sdpa"
|
|
logging.debug("Auto-detected attention: sdpa")
|
|
except ImportError:
|
|
attn_impl = "sdpa"
|
|
logging.debug("Auto-detected attention: sdpa")
|
|
else:
|
|
attn_impl = attn_implementation
|
|
if attn_impl == "flash_attention_2" and (device != "cuda" or not torch.cuda.is_available()):
|
|
logging.warning("flash_attention_2 requires NVIDIA CUDA, but CUDA is not available. Falling back to sdpa.")
|
|
attn_impl = "sdpa"
|
|
logging.debug(f"Using attention: {attn_impl}")
|
|
|
|
load_kwargs = {
|
|
"dtype": torch_dtype,
|
|
"trust_remote_code": True,
|
|
"attn_implementation": attn_impl,
|
|
}
|
|
|
|
if device_preference == "cpu":
|
|
load_kwargs["device_map"] = {"": "cpu"}
|
|
logging.debug("Using CPU only")
|
|
elif device == "cuda":
|
|
load_kwargs["device_map"] = "auto"
|
|
if max_memory is not None:
|
|
load_kwargs["max_memory"] = max_memory
|
|
logging.debug("Using CUDA with device_map='auto' and max_memory limits")
|
|
else:
|
|
logging.debug("Using CUDA with device_map='auto' (automatic GPU/CPU split)")
|
|
else:
|
|
load_kwargs["device_map"] = {"": "cpu"}
|
|
logging.debug("Using CPU only")
|
|
|
|
quantization_enabled = False
|
|
if "quantization_config" in kwargs:
|
|
quantization_enabled = True
|
|
logging.debug("8-bit quantization enabled")
|
|
|
|
load_kwargs.update(kwargs)
|
|
|
|
if quantization_enabled:
|
|
logging.info("Quantizing model to 8-bit (this may take a few minutes for large models)...")
|
|
|
|
if tokenizer_path is None:
|
|
if not os.path.isabs(model_path):
|
|
model_path = os.path.abspath(model_path)
|
|
model_path = os.path.normpath(model_path)
|
|
if model_path.endswith(".gguf"):
|
|
raise ValueError(
|
|
f"Invalid model path for HuggingFace format: {model_path}\nHuggingFace models must be directories, not .gguf files.\n"
|
|
)
|
|
tokenizer_path = auto_detect_tokenizer_path(model_path, model_format="huggingface")
|
|
logging.debug(f"Auto-detected tokenizer path: {tokenizer_path}")
|
|
|
|
tokenizer_load_path = tokenizer_path if tokenizer_path else model_path
|
|
|
|
self.model = AutoModelForCausalLM.from_pretrained(model_path, **load_kwargs)
|
|
|
|
load_time = time.time() - load_start
|
|
logging.info(f"HuggingFace model loaded in {load_time:.2f}s")
|
|
|
|
logging.debug(f"Loading tokenizer from: {tokenizer_load_path}")
|
|
tokenizer_start = time.time()
|
|
tokenizer_path_abs = os.path.abspath(tokenizer_load_path) if not os.path.isabs(tokenizer_load_path) else tokenizer_load_path
|
|
|
|
try:
|
|
from huggingface_hub.utils import HFValidationError
|
|
except ImportError:
|
|
HFValidationError = Exception
|
|
|
|
if os.path.exists(tokenizer_path_abs) and os.path.isdir(tokenizer_path_abs):
|
|
try:
|
|
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path_abs, trust_remote_code=True, local_files_only=True)
|
|
except (HFValidationError, Exception) as e:
|
|
if isinstance(e, HFValidationError) or "Repo id must use alphanumeric" in str(e):
|
|
logging.warning(f"Tokenizer path validation error, retrying with explicit local path handling: {e}")
|
|
else:
|
|
logging.warning(f"Failed to load tokenizer with local_files_only=True: {e}")
|
|
logging.warning("Trying without local_files_only...")
|
|
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path_abs, trust_remote_code=True)
|
|
else:
|
|
try:
|
|
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_load_path, trust_remote_code=True)
|
|
except (HFValidationError, Exception) as e:
|
|
if (isinstance(e, HFValidationError) or "Repo id must use alphanumeric" in str(e)) and (
|
|
os.path.sep in tokenizer_load_path
|
|
or (os.name == "nt" and len(tokenizer_load_path) > 1 and tokenizer_load_path[1] == ":")
|
|
):
|
|
raise ValueError(
|
|
f"Tokenizer path does not exist: {tokenizer_path_abs}\n"
|
|
f"Please ensure the tokenizer directory exists at this location."
|
|
) from e
|
|
raise
|
|
tokenizer_time = time.time() - tokenizer_start
|
|
logging.debug(f"Tokenizer loaded in {tokenizer_time:.2f}s")
|
|
|
|
# Fix for "The attention mask and the pad token id were not set" warning
|
|
if self.tokenizer.pad_token is None:
|
|
self.tokenizer.pad_token = self.tokenizer.eos_token
|
|
logging.debug("Tokenizer pad_token was None, set to eos_token")
|
|
|
|
if hasattr(self.model, "device"):
|
|
self.device = self.model.device
|
|
elif hasattr(self.model, "hf_device_map"):
|
|
self.device = torch.device(device) if device == "cuda" else torch.device("cpu")
|
|
else:
|
|
self.device = next(self.model.parameters()).device
|
|
|
|
try:
|
|
if torch.cuda.is_available() and device == "cuda":
|
|
vram_used = torch.cuda.memory_allocated() / 1024**3
|
|
logging.debug(f"GPU memory allocated: {vram_used:.2f} GB")
|
|
except Exception:
|
|
pass
|
|
|
|
def generate(
|
|
self, input_ids, max_new_tokens: int = 512, temperature: float = 0.7, top_p: float = 0.9, do_sample: bool = True, **kwargs
|
|
):
|
|
"""
|
|
Generate text using HuggingFace model.
|
|
|
|
Args:
|
|
input_ids: Token IDs tensor
|
|
max_new_tokens: Maximum number of tokens to generate
|
|
temperature: Sampling temperature
|
|
top_p: Nucleus sampling parameter
|
|
do_sample: Whether to use sampling
|
|
**kwargs: Additional generation parameters (seed will be extracted and used separately)
|
|
|
|
Returns:
|
|
Generated token IDs tensor
|
|
"""
|
|
seed = kwargs.pop("seed", None)
|
|
|
|
if seed is not None:
|
|
torch.manual_seed(seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
if torch.is_tensor(input_ids):
|
|
if isinstance(self.device, dict):
|
|
device = list(self.device.values())[0] if self.device else torch.device("cpu")
|
|
else:
|
|
device = self.device
|
|
input_ids = input_ids.to(device)
|
|
|
|
# Create attention mask if not provided
|
|
attention_mask = kwargs.pop("attention_mask", None)
|
|
if attention_mask is None:
|
|
# Assuming input_ids contains only valid tokens (no padding), create full mask
|
|
attention_mask = torch.ones_like(input_ids)
|
|
|
|
with torch.no_grad():
|
|
generated_ids = self.model.generate(
|
|
input_ids,
|
|
max_new_tokens=max_new_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
do_sample=do_sample,
|
|
attention_mask=attention_mask,
|
|
pad_token_id=self.tokenizer.pad_token_id,
|
|
**kwargs
|
|
)
|
|
|
|
return generated_ids
|
|
|
|
def __del__(self):
|
|
"""Cleanup method for VRAM management."""
|
|
if hasattr(self, "model"):
|
|
del self.model
|
|
if hasattr(self, "tokenizer"):
|
|
del self.tokenizer
|