Merge pull request #70 from orangetin/inference-py

Update inference to support multi-gpu and offloading

Author: csris

README.md

@@ -186,6 +186,8 @@ Hello human.
 
 Commands are prefixed with a `/`, and the `/quit` command exits.
 
+Please see [the inference README](inference/README.md) for more details about arguments, running on multiple/specific GPUs, and running on consumer hardware.
+
 # Monitoring
 
 By default, the training script simply prints the loss as training proceeds, but it can also output metrics to a file using [loguru](https://github.com/Delgan/loguru) or report them to Weights & Biases.

inference/README.md

@@ -0,0 +1,71 @@
+# OpenChatKit Inference
+This directory contains code for OpenChatKit's inference.
+
+## Arguments
+- `--gpu-id`: Primary GPU device to load inputs onto for inference. Default: `0`
+- `--model`: name/path of the model. Default = `../huggingface_models/GPT-NeoXT-Chat-Base-20B`
+- `--max-tokens`: the maximum number of tokens to generate. Default: `128`
+- `--sample`: indicates whether to sample. Default: `True`
+- `--temperature`: temperature for the LM. Default: `0.6`
+- `--top-k`: top-k for the LM. Default: `40`
+- `--retrieval`: augment queries with context from the retrieval index. Default `False`
+- `-g` `--gpu-vram`: GPU ID and VRAM to allocate to loading the model, separated by a `:` in the format `ID:RAM` where ID is the CUDA ID and RAM is in GiB. `gpu-id` must be present in this list to avoid errors. Accepts multiple values, for example, `-g ID_0:RAM_0 ID_1:RAM_1 ID_N:RAM_N`
+- `-r` `--cpu-ram`: CPU RAM overflow allocation for loading the model. Optional, and only used if the model does not fit onto the GPUs given.
+
+## Hardware requirements for inference
+The GPT-NeoXT-Chat-Base-20B model requires at least 41GB of free VRAM. Used VRAM also goes up by ~100-200 MB per prompt. 
+
+- A **minimum of 80 GB is recommended** 
+
+- A **minimum of 48 GB in VRAM is recommended** for fast responses.
+
+If you'd like to run inference on a GPU with <48 GB VRAM, refer to this section on [running on consumer hardware](#running-on-consumer-hardware).
+
+By default, inference uses only CUDA Device 0.
+
+**NOTE: Inference currently requires at least 1x GPU.**
+
+## Running on multiple GPUs
+Add the argument 
+
+```-g ID0:MAX_VRAM ID1:MAX_VRAM ID2:MAX_VRAM ...``` 
+
+where IDx is the CUDA ID of the device and MAX_VRAM is the amount of VRAM you'd like to allocate to the device.
+
+For example, if you are running this on 4x 48 GB GPUs and want to distribute the model across all devices, add ```-g 0:10 1:12 2:12 3:12 4:12```. In this example, the first device gets loaded to a max of 10 GiB while the others are loaded with a max of 12 GiB.
+
+How it works: The model fills up the max available VRAM on the first device passed and then overflows into the next until the whole model is loaded.
+
+**IMPORTANT: This MAX_VRAM is only for loading the model. It does not account for the additional inputs that are added to the device. It is recommended to set the MAX_VRAM to be at least 1 or 2 GiB less than the max available VRAM on each device, and at least 3GiB less than the max available VRAM on the primary device (set by `gpu-id` default=0).**
+
+**Decrease MAX_VRAM if you run into CUDA OOM. This happens because each input takes up additional space on the device.**
+
+**NOTE: Total MAX_VRAM across all devices must be > size of the model in GB. If not, `bot.py` automatically offloads the rest of the model to RAM and disk. It will use up all available RAM. To allocate a specified amount of RAM: [refer to this section on running on consumer hardware](#running-on-consumer-hardware).**
+
+## Running on specific GPUs
+If you have multiple GPUs but would only like to use a specific device(s), [use the same steps as in this section on running on multiple devices](#running-on-multiple-gpus) and only specify the devices you'd like to use. 
+
+Also, if needed, add the argument `--gpu-id ID` where ID is the CUDA ID of the device you'd like to make the primary device. NOTE: The device specified in `--gpu-id` must be present as one of the ID in the argument `-g` to avoid errors.
+
+- **Example #1**: to run inference on devices 2 and 5 with a max of 25 GiB on each, and make device 5 the primary device, add: `--gpu-id 5 -g 2:25 5:25`. In this example, not adding `--gpu-id 5` will give you an error.
+- **Example #2**: to run inference on devices 0 and 3 with a max of 10GiB on 0 and 40GiB on 3, with device 0 as the primary device, add: `-g 0:10 3:40`. In this example, `--gpu-id` is not required because device 0 is specified in `-g`.
+- **Example #3**: to run inference only on device 1 with a max of 75 GiB, add: `--gpu-id 1 -g 1:75`
+
+
+## Running on consumer hardware
+If you have multiple GPUs, each <48 GB VRAM, [the steps mentioned in this section on running on multiple GPUs](#running-on-multiple-gpus) still apply, unless, any of these apply:
+- Running on just 1x GPU with <48 GB VRAM,
+- <48 GB VRAM combined across multiple GPUs
+- Running into Out-Of-Memory (OOM) issues
+
+In which case, add the flag `-r CPU_RAM` where CPU_RAM is the maximum amount of RAM you'd like to allocate to loading model. Note: This significantly reduces inference speeds. 
+
+The model will load without specifying `-r`, however, it is not recommended because it will allocate all available RAM to the model. To limit how much RAM the model can use, add `-r`.
+
+If the total VRAM + CPU_RAM < the size of the model in GiB, the rest of the model will be offloaded to a folder "offload" at the root of the directory. Note: This significantly reduces inference speeds.
+
+- Example: `-g 0:12 -r 20` will first load up to 12 GiB of the model into the CUDA device 0, then load up to 20 GiB into RAM, and load the rest into the "offload" directory.
+
+How it works: 
+- https://github.com/huggingface/blog/blob/main/accelerate-large-models.md
+- https://www.youtube.com/embed/MWCSGj9jEAo

inference/bot.py

@@ -11,17 +11,47 @@
 import argparse
 import conversation as convo
 import retrieval.wikipedia as wp
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
+from accelerate import infer_auto_device_map, init_empty_weights
+
 
 class ChatModel:
     human_id = "<human>"
     bot_id = "<bot>"
 
-    def __init__(self, model_name, gpu_id):
-        device = torch.device('cuda', gpu_id)
-        self._model = AutoModelForCausalLM.from_pretrained(
-            model_name).half()
-        self._model.to(device)
+    def __init__(self, model_name, gpu_id, max_memory):
+        device = torch.device('cuda', gpu_id)   # TODO: allow sending to cpu
+
+        # recommended default for devices with > 40 GB VRAM
+        # load model onto one device
+        if max_memory is None:
+            self._model = AutoModelForCausalLM.from_pretrained(
+                model_name, torch_dtype=torch.float16, device_map="auto")
+            self._model.to(device)
+        # load the model with the given max_memory config (for devices with insufficient VRAM or multi-gpu)
+        else:
+            config = AutoConfig.from_pretrained(model_name)
+            # load empty weights
+            with init_empty_weights():
+                model_from_conf = AutoModelForCausalLM.from_config(config)
+
+            model_from_conf.tie_weights()
+
+            # create a device_map from max_memory
+            device_map = infer_auto_device_map(
+                model_from_conf,
+                max_memory=max_memory,
+                no_split_module_classes=["GPTNeoXLayer"],
+                dtype="float16"
+            )
+            # load the model with the above device_map
+            self._model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                device_map=device_map,
+                offload_folder="offload",  # optional offload-to-disk overflow directory (auto-created)
+                offload_state_dict=True,
+                torch_dtype=torch.float16
+            )
         self._tokenizer = AutoTokenizer.from_pretrained(model_name)
 
     def do_inference(self, prompt, max_new_tokens, do_sample, temperature, top_k):
@@ -49,7 +79,7 @@ class OpenChatKitShell(cmd.Cmd):
     intro = "Welcome to OpenChatKit shell.   Type /help or /? to list commands.\n"
     prompt = ">>> "
 
-    def __init__(self, gpu_id, model_name_or_path, max_tokens, sample, temperature, top_k, retrieval):
+    def __init__(self, gpu_id, model_name_or_path, max_tokens, sample, temperature, top_k, retrieval, max_memory):
         super().__init__()
         self._gpu_id = int(gpu_id)
         self._model_name_or_path = model_name_or_path
@@ -58,10 +88,11 @@ def __init__(self, gpu_id, model_name_or_path, max_tokens, sample, temperature,
         self._temperature = temperature
         self._top_k = top_k
         self._retrieval = retrieval
+        self._max_memory = max_memory
 
     def preloop(self):
         print(f"Loading {self._model_name_or_path} to cuda:{self._gpu_id}...")
-        self._model = ChatModel(self._model_name_or_path, self._gpu_id)
+        self._model = ChatModel(self._model_name_or_path, self._gpu_id, self._max_memory)
 
         if self._retrieval:
             print(f"Loading retrieval index...")
@@ -139,7 +170,7 @@ def main():
     parser.add_argument(
         '--model',
         default=f"{INFERENCE_DIR}/../huggingface_models/GPT-NeoXT-Chat-Base-20B",
-        help='the ID of the GPU to run on'
+        help='name/path of the model'
     )
     parser.add_argument(
         '--max-tokens',
@@ -168,16 +199,46 @@ def main():
         action='store_true',
         help='augment queries with context from the retrieval index'
     )
+    parser.add_argument(
+        '-g',
+        '--gpu-vram',
+        action='store',
+        help='max VRAM to allocate per GPU',
+        nargs='+',
+        required=False,
+    )
+    parser.add_argument(
+        '-r',
+        '--cpu-ram',
+        default=None,
+        type=int,
+        help='max CPU RAM to allocate',
+        required=False
+    )
     args = parser.parse_args()
 
+    # set max_memory dictionary if given
+    if args.gpu_vram is None:
+        max_memory = None
+    else:
+        max_memory = {}
+        for i in range(len(args.gpu_vram)):
+            # assign CUDA ID as label and XGiB as value
+            max_memory[int(args.gpu_vram[i].split(':')[0])] = f"{args.gpu_vram[i].split(':')[1]}GiB"
+
+        if args.cpu_ram is not None:
+            # add cpu to max-memory if given
+            max_memory['cpu'] = f"{int(args.cpu_ram)}GiB"
+
     OpenChatKitShell(
         args.gpu_id,
         args.model,
         args.max_tokens,
         args.sample,
         args.temperature,
         args.top_k,
-        args.retrieval
+        args.retrieval,
+        max_memory
     ).cmdloop()