diff --git a/modules/GPTQ_loader.py b/modules/GPTQ_loader.py index 7045a098..67899547 100644 --- a/modules/GPTQ_loader.py +++ b/modules/GPTQ_loader.py @@ -9,6 +9,7 @@ import modules.shared as shared sys.path.insert(0, str(Path("repositories/GPTQ-for-LLaMa"))) import llama +import llama_inference_offload import opt @@ -24,7 +25,10 @@ def load_quantized(model_name): model_type = shared.args.gptq_model_type.lower() if model_type == 'llama': - load_quant = llama.load_quant + if not shared.args.gptq_pre_layer: + load_quant = llama.load_quant + else: + load_quant = llama_inference_offload.load_quant elif model_type == 'opt': load_quant = opt.load_quant else: @@ -53,24 +57,26 @@ def load_quantized(model_name): print(f"Could not find {pt_model}, exiting...") exit() - model = load_quant(str(path_to_model), str(pt_path), shared.args.gptq_bits) + # Using qwopqwop200's offload + if shared.args.gptq_pre_layer: + model = load_quant(str(path_to_model), str(pt_path), shared.args.gptq_bits, shared.args.gptq_pre_layer) + else: + model = load_quant(str(path_to_model), str(pt_path), shared.args.gptq_bits) - # Multiple GPUs or GPU+CPU - if shared.args.gpu_memory: - memory_map = list(map(lambda x : x.strip(), shared.args.gpu_memory)) - max_cpu_memory = shared.args.cpu_memory.strip() if shared.args.cpu_memory is not None else '99GiB' - max_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_memory['cpu'] = max_cpu_memory + # Using accelerate offload (doesn't work properly) + if shared.args.gpu_memory: + memory_map = list(map(lambda x : x.strip(), shared.args.gpu_memory)) + max_cpu_memory = shared.args.cpu_memory.strip() if shared.args.cpu_memory is not None else '99GiB' + max_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_memory['cpu'] = max_cpu_memory - device_map = accelerate.infer_auto_device_map(model, max_memory=max_memory, no_split_module_classes=["LlamaDecoderLayer"]) - print("Using the following device map for the 4-bit model:", device_map) - # https://huggingface.co/docs/accelerate/package_reference/big_modeling#accelerate.dispatch_model - model = accelerate.dispatch_model(model, device_map=device_map, offload_buffers=True) - - # Single GPU - elif not shared.args.cpu: - model = model.to(torch.device('cuda:0')) + device_map = accelerate.infer_auto_device_map(model, max_memory=max_memory, no_split_module_classes=["LlamaDecoderLayer"]) + print("Using the following device map for the 4-bit model:", device_map) + # https://huggingface.co/docs/accelerate/package_reference/big_modeling#accelerate.dispatch_model + model = accelerate.dispatch_model(model, device_map=device_map, offload_buffers=True) + elif not shared.args.cpu: + model = model.to(torch.device('cuda:0')) return model diff --git a/modules/shared.py b/modules/shared.py index 8cae1079..8d591f4f 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -79,8 +79,9 @@ parser.add_argument('--cai-chat', action='store_true', help='Launch the web UI i parser.add_argument('--cpu', action='store_true', help='Use the CPU to generate text.') parser.add_argument('--load-in-8bit', action='store_true', help='Load the model with 8-bit precision.') parser.add_argument('--load-in-4bit', action='store_true', help='DEPRECATED: use --gptq-bits 4 instead.') -parser.add_argument('--gptq-bits', type=int, default=0, help='Load a pre-quantized model with specified precision. 2, 3, 4 and 8bit are supported. Currently only works with LLaMA and OPT.') -parser.add_argument('--gptq-model-type', type=str, help='Model type of pre-quantized model. Currently only LLaMa and OPT are supported.') +parser.add_argument('--gptq-bits', type=int, default=0, help='GPTQ: Load a pre-quantized model with specified precision. 2, 3, 4 and 8bit are supported. Currently only works with LLaMA and OPT.') +parser.add_argument('--gptq-model-type', type=str, help='GPTQ: Model type of pre-quantized model. Currently only LLaMa and OPT are supported.') +parser.add_argument('--gptq-pre-layer', type=int, default=0, help='GPTQ: The number of layers to preload.') parser.add_argument('--bf16', action='store_true', help='Load the model with bfloat16 precision. Requires NVIDIA Ampere GPU.') parser.add_argument('--auto-devices', action='store_true', help='Automatically split the model across the available GPU(s) and CPU.') parser.add_argument('--disk', action='store_true', help='If the model is too large for your GPU(s) and CPU combined, send the remaining layers to the disk.')