mirror of
https://github.com/oobabooga/text-generation-webui.git
synced 2024-10-01 01:26:03 -04:00
427 lines
15 KiB
Python
427 lines
15 KiB
Python
import gc
|
|
import os
|
|
import re
|
|
import time
|
|
import traceback
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import transformers
|
|
from accelerate import (
|
|
infer_auto_device_map,
|
|
init_empty_weights,
|
|
is_ccl_available,
|
|
is_xpu_available
|
|
)
|
|
from transformers import (
|
|
AutoConfig,
|
|
AutoModel,
|
|
AutoModelForCausalLM,
|
|
AutoModelForSeq2SeqLM,
|
|
AutoTokenizer,
|
|
BitsAndBytesConfig,
|
|
GPTQConfig
|
|
)
|
|
|
|
import modules.shared as shared
|
|
from modules import RoPE, llama_attn_hijack, sampler_hijack
|
|
from modules.logging_colors import logger
|
|
from modules.models_settings import get_model_metadata
|
|
|
|
transformers.logging.set_verbosity_error()
|
|
|
|
local_rank = None
|
|
if shared.args.deepspeed:
|
|
import deepspeed
|
|
from transformers.deepspeed import (
|
|
HfDeepSpeedConfig,
|
|
is_deepspeed_zero3_enabled
|
|
)
|
|
|
|
from modules.deepspeed_parameters import generate_ds_config
|
|
|
|
# Distributed setup
|
|
local_rank = shared.args.local_rank if shared.args.local_rank is not None else int(os.getenv("LOCAL_RANK", "0"))
|
|
world_size = int(os.getenv("WORLD_SIZE", "1"))
|
|
if is_xpu_available() and is_ccl_available():
|
|
torch.xpu.set_device(local_rank)
|
|
deepspeed.init_distributed(backend="ccl")
|
|
else:
|
|
torch.cuda.set_device(local_rank)
|
|
deepspeed.init_distributed()
|
|
ds_config = generate_ds_config(shared.args.bf16, 1 * world_size, shared.args.nvme_offload_dir)
|
|
dschf = HfDeepSpeedConfig(ds_config) # Keep this object alive for the Transformers integration
|
|
|
|
sampler_hijack.hijack_samplers()
|
|
|
|
|
|
def load_model(model_name, loader=None):
|
|
logger.info(f"Loading {model_name}...")
|
|
t0 = time.time()
|
|
|
|
shared.is_seq2seq = False
|
|
load_func_map = {
|
|
'Transformers': huggingface_loader,
|
|
'AutoGPTQ': AutoGPTQ_loader,
|
|
'GPTQ-for-LLaMa': GPTQ_loader,
|
|
'llama.cpp': llamacpp_loader,
|
|
'llamacpp_HF': llamacpp_HF_loader,
|
|
'RWKV': RWKV_loader,
|
|
'ExLlama': ExLlama_loader,
|
|
'ExLlama_HF': ExLlama_HF_loader,
|
|
'ExLlamav2': ExLlamav2_loader,
|
|
'ExLlamav2_HF': ExLlamav2_HF_loader,
|
|
'ctransformers': ctransformers_loader,
|
|
'AutoAWQ': AutoAWQ_loader,
|
|
}
|
|
|
|
if loader is None:
|
|
if shared.args.loader is not None:
|
|
loader = shared.args.loader
|
|
else:
|
|
loader = get_model_metadata(model_name)['loader']
|
|
if loader is None:
|
|
logger.error('The path to the model does not exist. Exiting.')
|
|
return None, None
|
|
|
|
shared.args.loader = loader
|
|
output = load_func_map[loader](model_name)
|
|
if type(output) is tuple:
|
|
model, tokenizer = output
|
|
else:
|
|
model = output
|
|
if model is None:
|
|
return None, None
|
|
else:
|
|
tokenizer = load_tokenizer(model_name, model)
|
|
|
|
# Hijack attention with xformers
|
|
if any((shared.args.xformers, shared.args.sdp_attention)):
|
|
llama_attn_hijack.hijack_llama_attention()
|
|
|
|
logger.info(f"Loaded the model in {(time.time()-t0):.2f} seconds.")
|
|
return model, tokenizer
|
|
|
|
|
|
def load_tokenizer(model_name, model):
|
|
tokenizer = None
|
|
path_to_model = Path(f"{shared.args.model_dir}/{model_name}/")
|
|
if any(s in model_name.lower() for s in ['gpt-4chan', 'gpt4chan']) and Path(f"{shared.args.model_dir}/gpt-j-6B/").exists():
|
|
tokenizer = AutoTokenizer.from_pretrained(Path(f"{shared.args.model_dir}/gpt-j-6B/"))
|
|
elif path_to_model.exists():
|
|
if shared.args.use_fast:
|
|
logger.info('Loading the tokenizer with use_fast=True.')
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
path_to_model,
|
|
trust_remote_code=shared.args.trust_remote_code,
|
|
use_fast=shared.args.use_fast
|
|
)
|
|
|
|
return tokenizer
|
|
|
|
|
|
def huggingface_loader(model_name):
|
|
|
|
path_to_model = Path(f'{shared.args.model_dir}/{model_name}')
|
|
params = {
|
|
'low_cpu_mem_usage': True,
|
|
'trust_remote_code': shared.args.trust_remote_code,
|
|
'torch_dtype': torch.bfloat16 if shared.args.bf16 else torch.float16
|
|
}
|
|
config = AutoConfig.from_pretrained(path_to_model, trust_remote_code=params['trust_remote_code'])
|
|
|
|
if 'chatglm' in model_name.lower():
|
|
LoaderClass = AutoModel
|
|
else:
|
|
if config.to_dict().get('is_encoder_decoder', False):
|
|
LoaderClass = AutoModelForSeq2SeqLM
|
|
shared.is_seq2seq = True
|
|
else:
|
|
LoaderClass = AutoModelForCausalLM
|
|
|
|
# Load the model in simple 16-bit mode by default
|
|
if not any([shared.args.cpu, shared.args.load_in_8bit, shared.args.load_in_4bit, shared.args.auto_devices, shared.args.disk, shared.args.deepspeed, shared.args.gpu_memory is not None, shared.args.cpu_memory is not None, shared.args.compress_pos_emb > 1, shared.args.alpha_value > 1, shared.args.disable_exllama]):
|
|
model = LoaderClass.from_pretrained(path_to_model, **params)
|
|
if torch.backends.mps.is_available():
|
|
device = torch.device('mps')
|
|
model = model.to(device)
|
|
elif is_xpu_available():
|
|
device = torch.device("xpu")
|
|
model = model.to(device)
|
|
else:
|
|
model = model.cuda()
|
|
|
|
# DeepSpeed ZeRO-3
|
|
elif shared.args.deepspeed:
|
|
model = LoaderClass.from_pretrained(path_to_model, torch_dtype=params['torch_dtype'])
|
|
model = deepspeed.initialize(model=model, config_params=ds_config, model_parameters=None, optimizer=None, lr_scheduler=None)[0]
|
|
model.module.eval() # Inference
|
|
logger.info(f'DeepSpeed ZeRO-3 is enabled: {is_deepspeed_zero3_enabled()}')
|
|
|
|
# Load with quantization and/or offloading
|
|
else:
|
|
|
|
if not any((shared.args.cpu, torch.cuda.is_available(), is_xpu_available(), torch.backends.mps.is_available())):
|
|
logger.warning('torch.cuda.is_available() and is_xpu_available() returned False. This means that no GPU has been detected. Falling back to CPU mode.')
|
|
|
|
shared.args.cpu = True
|
|
|
|
if shared.args.cpu:
|
|
params['torch_dtype'] = torch.float32
|
|
else:
|
|
params['device_map'] = 'auto'
|
|
params['max_memory'] = get_max_memory_dict()
|
|
if shared.args.load_in_4bit:
|
|
# See https://github.com/huggingface/transformers/pull/23479/files
|
|
# and https://huggingface.co/blog/4bit-transformers-bitsandbytes
|
|
quantization_config_params = {
|
|
'load_in_4bit': True,
|
|
'bnb_4bit_compute_dtype': eval("torch.{}".format(shared.args.compute_dtype)) if shared.args.compute_dtype in ["bfloat16", "float16", "float32"] else None,
|
|
'bnb_4bit_quant_type': shared.args.quant_type,
|
|
'bnb_4bit_use_double_quant': shared.args.use_double_quant,
|
|
}
|
|
|
|
logger.info('Using the following 4-bit params: ' + str(quantization_config_params))
|
|
params['quantization_config'] = BitsAndBytesConfig(**quantization_config_params)
|
|
|
|
elif shared.args.load_in_8bit:
|
|
if any((shared.args.auto_devices, shared.args.gpu_memory)):
|
|
params['quantization_config'] = BitsAndBytesConfig(load_in_8bit=True, llm_int8_enable_fp32_cpu_offload=True)
|
|
else:
|
|
params['quantization_config'] = BitsAndBytesConfig(load_in_8bit=True)
|
|
|
|
if params['max_memory'] is not None:
|
|
with init_empty_weights():
|
|
model = LoaderClass.from_config(config, trust_remote_code=params['trust_remote_code'])
|
|
|
|
model.tie_weights()
|
|
params['device_map'] = infer_auto_device_map(
|
|
model,
|
|
dtype=torch.int8,
|
|
max_memory=params['max_memory'],
|
|
no_split_module_classes=model._no_split_modules
|
|
)
|
|
|
|
if shared.args.disk:
|
|
params['offload_folder'] = shared.args.disk_cache_dir
|
|
|
|
if shared.args.disable_exllama:
|
|
try:
|
|
gptq_config = GPTQConfig(bits=config.quantization_config.get('bits', 4), disable_exllama=True)
|
|
params['quantization_config'] = gptq_config
|
|
logger.info('Loading with ExLlama kernel disabled.')
|
|
except:
|
|
exc = traceback.format_exc()
|
|
logger.error('Failed to disable exllama. Does the config.json for this model contain the necessary quantization info?')
|
|
print(exc)
|
|
|
|
if shared.args.compress_pos_emb > 1:
|
|
params['rope_scaling'] = {'type': 'linear', 'factor': shared.args.compress_pos_emb}
|
|
elif shared.args.alpha_value > 1:
|
|
params['rope_scaling'] = {'type': 'dynamic', 'factor': RoPE.get_alpha_value(shared.args.alpha_value, shared.args.rope_freq_base)}
|
|
|
|
model = LoaderClass.from_pretrained(path_to_model, **params)
|
|
|
|
return model
|
|
|
|
|
|
def llamacpp_loader(model_name):
|
|
from modules.llamacpp_model import LlamaCppModel
|
|
|
|
path = Path(f'{shared.args.model_dir}/{model_name}')
|
|
if path.is_file():
|
|
model_file = path
|
|
else:
|
|
model_file = list(Path(f'{shared.args.model_dir}/{model_name}').glob('*.gguf'))[0]
|
|
|
|
logger.info(f"llama.cpp weights detected: {model_file}")
|
|
model, tokenizer = LlamaCppModel.from_pretrained(model_file)
|
|
return model, tokenizer
|
|
|
|
|
|
def llamacpp_HF_loader(model_name):
|
|
from modules.llamacpp_hf import LlamacppHF
|
|
|
|
for fname in [model_name, "oobabooga_llama-tokenizer", "llama-tokenizer"]:
|
|
path = Path(f'{shared.args.model_dir}/{fname}')
|
|
if all((path / file).exists() for file in ['tokenizer_config.json', 'special_tokens_map.json', 'tokenizer.model']):
|
|
logger.info(f'Using tokenizer from: {path}')
|
|
break
|
|
else:
|
|
logger.error("Could not load the model because a tokenizer in transformers format was not found. Please download oobabooga/llama-tokenizer.")
|
|
return None, None
|
|
|
|
if shared.args.use_fast:
|
|
logger.info('Loading the tokenizer with use_fast=True.')
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
path,
|
|
trust_remote_code=shared.args.trust_remote_code,
|
|
use_fast=shared.args.use_fast
|
|
)
|
|
|
|
model = LlamacppHF.from_pretrained(model_name)
|
|
return model, tokenizer
|
|
|
|
|
|
def ctransformers_loader(model_name):
|
|
from modules.ctransformers_model import CtransformersModel
|
|
|
|
path = Path(f'{shared.args.model_dir}/{model_name}')
|
|
ctrans = CtransformersModel()
|
|
if ctrans.model_type_is_auto():
|
|
model_file = path
|
|
else:
|
|
if path.is_file():
|
|
model_file = path
|
|
else:
|
|
entries = Path(f'{shared.args.model_dir}/{model_name}')
|
|
gguf = list(entries.glob('*.gguf'))
|
|
bin = list(entries.glob('*.bin'))
|
|
if len(gguf) > 0:
|
|
model_file = gguf[0]
|
|
elif len(bin) > 0:
|
|
model_file = bin[0]
|
|
else:
|
|
logger.error("Could not find a model for ctransformers.")
|
|
return None, None
|
|
|
|
logger.info(f'ctransformers weights detected: {model_file}')
|
|
model, tokenizer = ctrans.from_pretrained(model_file)
|
|
return model, tokenizer
|
|
|
|
|
|
def AutoAWQ_loader(model_name):
|
|
from awq import AutoAWQForCausalLM
|
|
|
|
model_dir = Path(f'{shared.args.model_dir}/{model_name}')
|
|
|
|
model = AutoAWQForCausalLM.from_quantized(
|
|
quant_path=model_dir,
|
|
max_new_tokens=shared.args.max_seq_len,
|
|
trust_remote_code=shared.args.trust_remote_code,
|
|
fuse_layers=not shared.args.no_inject_fused_attention,
|
|
max_memory=get_max_memory_dict(),
|
|
batch_size=1,
|
|
safetensors=any(model_dir.glob('*.safetensors')),
|
|
)
|
|
|
|
return model
|
|
|
|
|
|
def GPTQ_loader(model_name):
|
|
|
|
# Monkey patch
|
|
if shared.args.monkey_patch:
|
|
logger.warning("Applying the monkey patch for using LoRAs with GPTQ models. It may cause undefined behavior outside its intended scope.")
|
|
from modules.monkey_patch_gptq_lora import load_model_llama
|
|
|
|
model, _ = load_model_llama(model_name)
|
|
|
|
# No monkey patch
|
|
else:
|
|
import modules.GPTQ_loader
|
|
|
|
model = modules.GPTQ_loader.load_quantized(model_name)
|
|
|
|
return model
|
|
|
|
|
|
def AutoGPTQ_loader(model_name):
|
|
import modules.AutoGPTQ_loader
|
|
|
|
return modules.AutoGPTQ_loader.load_quantized(model_name)
|
|
|
|
|
|
def ExLlama_loader(model_name):
|
|
from modules.exllama import ExllamaModel
|
|
|
|
model, tokenizer = ExllamaModel.from_pretrained(model_name)
|
|
return model, tokenizer
|
|
|
|
|
|
def ExLlama_HF_loader(model_name):
|
|
from modules.exllama_hf import ExllamaHF
|
|
|
|
return ExllamaHF.from_pretrained(model_name)
|
|
|
|
|
|
def ExLlamav2_loader(model_name):
|
|
from modules.exllamav2 import Exllamav2Model
|
|
|
|
model, tokenizer = Exllamav2Model.from_pretrained(model_name)
|
|
return model, tokenizer
|
|
|
|
|
|
def ExLlamav2_HF_loader(model_name):
|
|
from modules.exllamav2_hf import Exllamav2HF
|
|
|
|
return Exllamav2HF.from_pretrained(model_name)
|
|
|
|
|
|
def RWKV_loader(model_name):
|
|
'''
|
|
This loader is not currently maintained as RWKV can now be loaded
|
|
through the transformers library.
|
|
'''
|
|
from modules.RWKV import RWKVModel, RWKVTokenizer
|
|
|
|
model = RWKVModel.from_pretrained(
|
|
Path(f'{shared.args.model_dir}/{model_name}'),
|
|
dtype="fp32" if shared.args.cpu else "bf16" if shared.args.bf16 else "fp16",
|
|
device="cpu" if shared.args.cpu else "xpu" if is_xpu_available() else "cuda"
|
|
)
|
|
|
|
tokenizer = RWKVTokenizer.from_pretrained(Path(shared.args.model_dir))
|
|
return model, tokenizer
|
|
|
|
|
|
def get_max_memory_dict():
|
|
max_memory = {}
|
|
if shared.args.gpu_memory:
|
|
memory_map = list(map(lambda x: x.strip(), shared.args.gpu_memory))
|
|
for i in range(len(memory_map)):
|
|
max_memory[i] = f'{memory_map[i]}GiB' if not re.match('.*ib$', memory_map[i].lower()) else memory_map[i]
|
|
|
|
max_cpu_memory = shared.args.cpu_memory.strip() if shared.args.cpu_memory is not None else '99GiB'
|
|
max_memory['cpu'] = f'{max_cpu_memory}GiB' if not re.match('.*ib$', max_cpu_memory.lower()) else max_cpu_memory
|
|
|
|
# If --auto-devices is provided standalone, try to get a reasonable value
|
|
# for the maximum memory of device :0
|
|
elif shared.args.auto_devices:
|
|
if is_xpu_available():
|
|
total_mem = (torch.xpu.get_device_properties(0).total_memory / (1024 * 1024))
|
|
else:
|
|
total_mem = (torch.cuda.get_device_properties(0).total_memory / (1024 * 1024))
|
|
suggestion = round((total_mem - 1000) / 1000) * 1000
|
|
if total_mem - suggestion < 800:
|
|
suggestion -= 1000
|
|
|
|
suggestion = int(round(suggestion / 1000))
|
|
logger.warning(f"Auto-assiging --gpu-memory {suggestion} for your GPU to try to prevent out-of-memory errors. You can manually set other values.")
|
|
max_memory = {0: f'{suggestion}GiB', 'cpu': f'{shared.args.cpu_memory or 99}GiB'}
|
|
|
|
return max_memory if len(max_memory) > 0 else None
|
|
|
|
|
|
def clear_torch_cache():
|
|
gc.collect()
|
|
if not shared.args.cpu:
|
|
if is_xpu_available():
|
|
torch.xpu.empty_cache()
|
|
else:
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def unload_model():
|
|
shared.model = shared.tokenizer = None
|
|
shared.lora_names = []
|
|
shared.model_dirty_from_training = False
|
|
clear_torch_cache()
|
|
|
|
|
|
def reload_model():
|
|
unload_model()
|
|
shared.model, shared.tokenizer = load_model(shared.model_name)
|