text-generation-webui/modules/models.py

336 lines
13 KiB
Python
Raw Normal View History

2023-04-08 02:36:04 +02:00
import gc
2023-02-23 17:28:30 +01:00
import os
import re
2023-02-23 17:28:30 +01:00
import time
from pathlib import Path
import torch
import transformers
2023-03-16 17:34:23 +01:00
from accelerate import infer_auto_device_map, init_empty_weights
2023-06-25 06:44:36 +02:00
from transformers import (
AutoConfig,
AutoModel,
AutoModelForCausalLM,
AutoModelForSeq2SeqLM,
AutoTokenizer,
BitsAndBytesConfig,
LlamaTokenizer
)
2023-02-23 18:41:42 +01:00
import modules.shared as shared
from modules import llama_attn_hijack, sampler_hijack
from modules.logging_colors import logger
from modules.models_settings import infer_loader
2023-02-23 17:28:30 +01:00
transformers.logging.set_verbosity_error()
2023-04-08 02:36:04 +02:00
local_rank = None
2023-02-23 17:28:30 +01:00
if shared.args.deepspeed:
import deepspeed
2023-06-25 06:44:36 +02:00
from transformers.deepspeed import (
HfDeepSpeedConfig,
is_deepspeed_zero3_enabled
)
2023-02-23 18:41:42 +01:00
2023-02-23 17:28:30 +01:00
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"))
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
2023-02-23 17:28:30 +01:00
sampler_hijack.hijack_samplers()
2023-03-13 18:00:38 +01:00
def load_model(model_name, loader=None):
logger.info(f"Loading {model_name}...")
2023-02-23 17:28:30 +01:00
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,
'FlexGen': flexgen_loader,
2023-06-17 01:35:38 +02:00
'RWKV': RWKV_loader,
'ExLlama': ExLlama_loader,
'ExLlama_HF': ExLlama_HF_loader
}
if loader is None:
if shared.args.loader is not None:
loader = shared.args.loader
else:
loader = infer_loader(model_name)
if loader is None:
logger.error('The path to the model does not exist. Exiting.')
return None, None
2023-05-17 00:52:22 +02:00
shared.args.loader = loader
output = load_func_map[loader](model_name)
2023-05-17 00:52:22 +02:00
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)
2023-05-17 00:52:22 +02:00
# 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.\n")
2023-05-17 00:52:22 +02:00
return model, tokenizer
def load_tokenizer(model_name, model):
tokenizer = None
if any(s in model_name.lower() for s in ['gpt-4chan', 'gpt4chan']) and Path(f"{shared.args.model_dir}/gpt-j-6B/").exists():
2023-05-17 00:52:22 +02:00
tokenizer = AutoTokenizer.from_pretrained(Path(f"{shared.args.model_dir}/gpt-j-6B/"))
elif model.__class__.__name__ in ['LlamaForCausalLM', 'LlamaGPTQForCausalLM', 'ExllamaHF']:
2023-05-17 00:52:22 +02:00
# Try to load an universal LLaMA tokenizer
if not any(s in shared.model_name.lower() for s in ['llava', 'oasst']):
2023-05-17 00:52:22 +02:00
for p in [Path(f"{shared.args.model_dir}/llama-tokenizer/"), Path(f"{shared.args.model_dir}/oobabooga_llama-tokenizer/")]:
if p.exists():
logger.info(f"Loading the universal LLaMA tokenizer from {p}...")
2023-05-17 00:52:22 +02:00
tokenizer = LlamaTokenizer.from_pretrained(p, clean_up_tokenization_spaces=True)
return tokenizer
# Otherwise, load it from the model folder and hope that these
# are not outdated tokenizer files.
tokenizer = LlamaTokenizer.from_pretrained(Path(f"{shared.args.model_dir}/{model_name}/"), clean_up_tokenization_spaces=True)
try:
tokenizer.eos_token_id = 2
tokenizer.bos_token_id = 1
tokenizer.pad_token_id = 0
except:
pass
else:
path_to_model = Path(f"{shared.args.model_dir}/{model_name}/")
if path_to_model.exists():
tokenizer = AutoTokenizer.from_pretrained(path_to_model, trust_remote_code=shared.args.trust_remote_code)
2023-05-17 00:52:22 +02:00
return tokenizer
def huggingface_loader(model_name):
path_to_model = Path(f'{shared.args.model_dir}/{model_name}')
if 'chatglm' in model_name.lower():
LoaderClass = AutoModel
else:
config = AutoConfig.from_pretrained(path_to_model, trust_remote_code=shared.args.trust_remote_code)
if config.to_dict().get("is_encoder_decoder", False):
LoaderClass = AutoModelForSeq2SeqLM
shared.is_seq2seq = True
else:
LoaderClass = AutoModelForCausalLM
2023-02-28 03:03:35 +01:00
# Load the model in simple 16-bit mode by default
2023-05-25 06:14:13 +02:00
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]):
2023-05-17 00:52:22 +02:00
model = LoaderClass.from_pretrained(Path(f"{shared.args.model_dir}/{model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.bfloat16 if shared.args.bf16 else torch.float16, trust_remote_code=shared.args.trust_remote_code)
if torch.has_mps:
device = torch.device('mps')
model = model.to(device)
2023-03-18 02:27:26 +01:00
else:
model = model.cuda()
2023-03-18 02:27:26 +01:00
2023-02-23 17:28:30 +01:00
# DeepSpeed ZeRO-3
elif shared.args.deepspeed:
model = LoaderClass.from_pretrained(Path(f"{shared.args.model_dir}/{model_name}"), torch_dtype=torch.bfloat16 if shared.args.bf16 else torch.float16)
2023-02-23 17:28:30 +01:00
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()}")
2023-02-23 17:28:30 +01:00
# Custom
else:
params = {
2023-05-10 03:49:39 +02:00
"low_cpu_mem_usage": True,
2023-05-17 00:52:22 +02:00
"trust_remote_code": shared.args.trust_remote_code
}
2023-05-10 03:49:39 +02:00
2023-03-18 02:56:46 +01:00
if not any((shared.args.cpu, torch.cuda.is_available(), torch.has_mps)):
logger.warning("torch.cuda.is_available() returned False. This means that no GPU has been detected. Falling back to CPU mode.")
2023-02-23 17:28:30 +01:00
shared.args.cpu = True
if shared.args.cpu:
params["torch_dtype"] = torch.float32
2023-02-23 17:28:30 +01:00
else:
params["device_map"] = 'auto'
2023-05-25 06:14:13 +02:00
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.warning("Using the following 4-bit params: " + str(quantization_config_params))
params['quantization_config'] = BitsAndBytesConfig(**quantization_config_params)
elif shared.args.load_in_8bit and any((shared.args.auto_devices, shared.args.gpu_memory)):
params['quantization_config'] = BitsAndBytesConfig(load_in_8bit=True, llm_int8_enable_fp32_cpu_offload=True)
elif shared.args.load_in_8bit:
params['quantization_config'] = BitsAndBytesConfig(load_in_8bit=True)
elif shared.args.bf16:
params["torch_dtype"] = torch.bfloat16
else:
params["torch_dtype"] = torch.float16
2023-02-23 17:28:30 +01:00
2023-05-16 00:38:27 +02:00
params['max_memory'] = get_max_memory_dict()
if shared.args.disk:
params["offload_folder"] = shared.args.disk_cache_dir
checkpoint = Path(f'{shared.args.model_dir}/{model_name}')
if shared.args.load_in_8bit and params.get('max_memory', None) is not None and params['device_map'] == 'auto':
2023-05-17 00:52:22 +02:00
config = AutoConfig.from_pretrained(checkpoint, trust_remote_code=shared.args.trust_remote_code)
with init_empty_weights():
2023-05-17 00:52:22 +02:00
model = LoaderClass.from_config(config, trust_remote_code=shared.args.trust_remote_code)
2023-05-04 02:43:17 +02:00
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
)
model = LoaderClass.from_pretrained(checkpoint, **params)
2023-02-23 17:28:30 +01:00
2023-05-17 00:52:22 +02:00
return model
2023-04-10 04:08:40 +02:00
2023-04-20 02:23:51 +02:00
2023-05-17 00:52:22 +02:00
def flexgen_loader(model_name):
from flexgen.flex_opt import CompressionConfig, ExecutionEnv, OptLM, Policy
2023-04-20 02:23:51 +02:00
2023-05-17 00:52:22 +02:00
# Initialize environment
env = ExecutionEnv.create(shared.args.disk_cache_dir)
# Offloading policy
policy = Policy(1, 1,
shared.args.percent[0], shared.args.percent[1],
shared.args.percent[2], shared.args.percent[3],
shared.args.percent[4], shared.args.percent[5],
overlap=True, sep_layer=True, pin_weight=shared.args.pin_weight,
cpu_cache_compute=False, attn_sparsity=1.0,
compress_weight=shared.args.compress_weight,
comp_weight_config=CompressionConfig(
num_bits=4, group_size=64,
group_dim=0, symmetric=False),
compress_cache=False,
comp_cache_config=CompressionConfig(
num_bits=4, group_size=64,
group_dim=2, symmetric=False))
model = OptLM(f"facebook/{model_name}", env, shared.args.model_dir, policy)
return model
def RWKV_loader(model_name):
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 "cuda")
tokenizer = RWKVTokenizer.from_pretrained(Path(shared.args.model_dir))
return model, tokenizer
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
2023-02-23 17:28:30 +01:00
else:
2023-05-17 00:52:22 +02:00
model_file = list(Path(f'{shared.args.model_dir}/{model_name}').glob('*ggml*.bin'))[0]
2023-02-23 17:28:30 +01:00
logger.info(f"llama.cpp weights detected: {model_file}\n")
2023-05-17 00:52:22 +02:00
model, tokenizer = LlamaCppModel.from_pretrained(model_file)
2023-02-23 17:28:30 +01:00
return model, tokenizer
2023-05-17 00:52:22 +02:00
def GPTQ_loader(model_name):
# Monkey patch
if shared.args.monkey_patch:
2023-05-30 03:40:54 +02:00
logger.warning("Applying the monkey patch for using LoRAs with GPTQ models. It may cause undefined behavior outside its intended scope.")
2023-05-17 00:52:22 +02:00
from modules.monkey_patch_gptq_lora import load_model_llama
model, _ = load_model_llama(model_name)
# No monkey patch
else:
2023-05-17 16:23:13 +02:00
import modules.GPTQ_loader
2023-05-17 00:52:22 +02:00
2023-05-17 16:23:13 +02:00
model = modules.GPTQ_loader.load_quantized(model_name)
2023-05-17 00:52:22 +02:00
return model
2023-05-17 16:12:12 +02:00
def AutoGPTQ_loader(model_name):
2023-05-17 16:23:13 +02:00
import modules.AutoGPTQ_loader
2023-05-17 16:12:12 +02:00
2023-05-17 16:23:13 +02:00
return modules.AutoGPTQ_loader.load_quantized(model_name)
2023-05-17 16:12:12 +02:00
2023-06-17 01:35:38 +02:00
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)
2023-05-16 00:38:27 +02:00
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:
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.")
2023-05-16 00:38:27 +02:00
max_memory = {0: f'{suggestion}GiB', 'cpu': f'{shared.args.cpu_memory or 99}GiB'}
2023-05-17 16:12:12 +02:00
return max_memory if len(max_memory) > 0 else None
2023-05-16 00:38:27 +02:00
2023-04-08 02:36:04 +02:00
def clear_torch_cache():
gc.collect()
if not shared.args.cpu:
torch.cuda.empty_cache()
def unload_model():
shared.model = shared.tokenizer = None
shared.lora_names = []
2023-04-08 02:36:04 +02:00
clear_torch_cache()
def reload_model():
2023-04-08 02:37:41 +02:00
unload_model()
2023-04-08 02:36:04 +02:00
shared.model, shared.tokenizer = load_model(shared.model_name)