Add -gptq-preload for 4-bit offloading (#460)

This works in a 4GB card now:

```
python server.py --model llama-7b-hf --gptq-bits 4 --gptq-pre-layer 20
```

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)
-
-    # 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 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)
 
-        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)
+        # 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
 
-    # 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

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.')