From c69695d7b7d5f93f5bcc82a8bffe28b1797a40fb Mon Sep 17 00:00:00 2001 From: Mayank Mishra Date: Mon, 8 Aug 2022 22:56:37 +0530 Subject: [PATCH 1/3] server --- scripts/inference/README.md | 59 +++++- .../inference/bloom-accelerate-inference.py | 93 +++----- scripts/inference/bloom-accelerate-server.py | 199 ++++++++++++++++++ scripts/inference/bloom-ds-inference.py | 130 +++--------- scripts/inference/bloom-ds-server.py | 199 ++++++++++++++++++ scripts/inference/bloom-ds-zero-inference.py | 56 ++--- scripts/inference/cache-ds-model.py | 80 +++++++ 7 files changed, 627 insertions(+), 189 deletions(-) create mode 100644 scripts/inference/bloom-accelerate-server.py create mode 100644 scripts/inference/bloom-ds-server.py create mode 100644 scripts/inference/cache-ds-model.py diff --git a/scripts/inference/README.md b/scripts/inference/README.md index 44e98f9fb..fe54518bc 100644 --- a/scripts/inference/README.md +++ b/scripts/inference/README.md @@ -1,5 +1,60 @@ # Inference scripts for BLOOM +To run a server using HuggingFace (requires [accelerate](https://github.com/huggingface/accelerate) to be installed): +``` +python scripts/inference/bloom-accelerate-server.py --model_name bigscience/bloom --dtype bf16 --log_file data.log --host $ADDRESS --port $PORT +``` + +To run a server using deepspeed (requires [DeepSpeed MII](https://github.com/microsoft/DeepSpeed-mii) to be installed): +``` +export DS_CACHE= + +deepspeed --num_gpus 8 scripts/inference/cache-ds-model.py --model_name bigscience/bloom --dtype fp16 + +python scripts/inference/bloom-ds-server.py --model_name bigscience/bloom --dtype fp16 --log_file data.log --host $ADDRESS --port $PORT +``` + +Usage: +Currently, the script supports 3 method: +1. The main generate method +``` +curl -H "Content-Type: application/json" -X POST -d '{ "input_text": "India is a country of", "top_k": "5", "top_p": "0.9", "temperature": "0.7", "min_length": "1", "max_new_tokens": "40" }' http://$ADDRESS:$PORT/generate/ +``` +returns +``` +{"output_text":" many languages and cultures. The country is a melting pot of different cultures and languages. The country is a home to more than 1.2 billion people. The country is a home to more than 22","query_id":8,"total_time_taken":"19.358 s"} +``` +2. Method that returns the model description +``` +curl -H "Content-Type: application/json" -X GET http://$ADDRESS:$PORT/about/ +``` +returns +``` +Please don't send any personal information to this endpoint. We are logging your data. + +Usage: +A request object should look like: +{ + input_text: "Hello, I'm a model", + "top_k": 5, + "top_p": 0.9, + "temperature": 0.7, + "min_length": 1, + "max_new_tokens": 40, +} + +Default values (use if not provided in request object): +top_k = 50 +top_p = 1 +temperature = 1 +min_length = 1 +max_new_tokens = 40 +``` +3. Method to check GPU usage +``` +curl -H "Content-Type: application/json" -X GET http://$ADDRESS:$PORT/gpu/ +``` +returns the nvidia-smi output ## BLOOM Inference solutions Here are some stats on JeanZay's 8x80GB A100 node w/ 512GB of CPU memory: @@ -14,7 +69,7 @@ Throughput in msecs: | project \ bs | 1 | 8 | 16 | 32 | 64 | 128 | | :----------- | :---- | :---- | :---- | :---- | :---- | :--- | -| accelerate | 230.38 | 31.78 | 17.84 | 10.89 | oom | omm | +| accelerate | 230.38 | 31.78 | 17.84 | 10.89 | oom | oom | | ds-inference | 40.57 | 5.23 | | | 2.77 | 0.66 | | ds-zero | 283 | 34.88 | oom | oom | oom | oom | @@ -192,4 +247,4 @@ $ python scripts/inference/bloom-accelerate-inference.py --name bigscience/bloom [...] -``` +``` \ No newline at end of file diff --git a/scripts/inference/bloom-accelerate-inference.py b/scripts/inference/bloom-accelerate-inference.py index 415b2f765..6b91f9e86 100644 --- a/scripts/inference/bloom-accelerate-inference.py +++ b/scripts/inference/bloom-accelerate-inference.py @@ -4,7 +4,9 @@ import gc import torch import math -from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM +from transformers import AutoTokenizer, AutoModelForCausalLM +from utils import generate_, get_max_memory_per_gpu_dict + def get_args(): parser = argparse.ArgumentParser() @@ -18,46 +20,6 @@ def get_args(): return parser.parse_args() -def get_max_memory_per_gpu_dict(dtype, model_name): - """ try to generate the memory map based on what we know about the model and the available hardware """ - - # figure out the memory map - the minimum per gpu required to load the model - n_gpus = torch.cuda.device_count() - - if model_name == "bigscience/bloom" and n_gpus == 8 and torch.cuda.get_device_properties(0).total_memory > 79*2**30: - # hand crafted optimized memory map for 8x80 setup over BLOOM - # this works with bs=40 - return {0: '0GIB', 1: '51GIB', 2: '51GIB', 3: '51GIB', 4: '51GIB', 5: '51GIB', 6: '51GIB', 7: '51GIB'} - - try: - # model_params calculation, as we don't have a model yet to do: - #model_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values()) - - config = AutoConfig.from_pretrained(model_name) - h = config.n_embed - l = config.n_layer - v = config.vocab_size - # from https://github.com/bigscience-workshop/bigscience/tree/6917a3b5fefcf439d3485ca184b4d9f6ab605150/math#model-sizing - model_params = l*(12*h**2 + 13*h) + v*h + 4*h - except: - print(f"The model {model_name} has a broken config file. Please notify the owner") - raise - - bytes = torch.finfo(dtype).bits / 8 - param_memory_total_in_bytes = model_params * bytes - # add 5% since weight sizes aren't the same and some GPU may need more memory - param_memory_per_gpu_in_bytes = int(param_memory_total_in_bytes / n_gpus * 1.05) - print(f"Estimating {param_memory_per_gpu_in_bytes/2**30:0.2f}GB per gpu for weights") - - # check the real available memory - # load cuda kernels first and only measure the real free memory after loading (shorter by ~2GB) - torch.ones(1).cuda() - max_memory_per_gpu_in_bytes = torch.cuda.mem_get_info(0)[0] - if max_memory_per_gpu_in_bytes < param_memory_per_gpu_in_bytes: - raise ValueError(f"Unable to generate the memory map automatically as the needed estimated memory per gpu ({param_memory_per_gpu_in_bytes/2**30:0.2f}GB) is bigger than the available per gpu memory ({max_memory_per_gpu_in_bytes/2**30:0.2f}GB)") - - return {i: param_memory_per_gpu_in_bytes for i in range(torch.cuda.device_count())} - t_start = time.time() num_tokens = 100 @@ -122,30 +84,25 @@ def get_max_memory_per_gpu_dict(dtype, model_name): if rank == 0: print(f"Generate args {generate_kwargs}") inputs = input_sentences[:args.batch_size] -def generate(): - """ returns a list of zipped inputs, outputs and number of new tokens """ - - input_tokens = tokenizer.batch_encode_plus(inputs, return_tensors="pt", padding=True) - for t in input_tokens: - if torch.is_tensor(input_tokens[t]): - input_tokens[t] = input_tokens[t].to("cuda:0") - - outputs = model.generate(**input_tokens, **generate_kwargs) - - input_tokens_lengths = [x.shape[0] for x in input_tokens.input_ids] - output_tokens_lengths = [x.shape[0] for x in outputs] - - total_new_tokens = [o-i for i,o in zip(input_tokens_lengths, output_tokens_lengths)] - outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True) - - return zip(inputs, outputs, total_new_tokens) # warmup is a must if measuring speed as it's when all the optimizations are performed # e.g. on 8x80 a100 the first pass of 100 tokens takes 23sec, and the next one is 4secs -_ = generate() +_ = generate_( + inputs, + model, + tokenizer, + generate_kwargs, + "cuda:0" +) t_generate_start = time.time() -generated = generate() +generated = generate_( + inputs, + model, + tokenizer, + generate_kwargs, + "cuda:0" +) t_generate_span = time.time() - t_generate_start if rank == 0: for i,o,_ in generated: @@ -164,7 +121,13 @@ def generate(): # warm up for i in range(1): - _ = generate() + _ = generate_( + inputs, + model, + tokenizer, + generate_kwargs, + "cuda:0" + ) torch.cuda.synchronize() # benchmark @@ -172,7 +135,13 @@ def generate(): cycles = 5 total_new_tokens_generated = 0 for i in range(cycles): - generated = generate() + generated = generate_( + inputs, + model, + tokenizer, + generate_kwargs, + "cuda:0" + ) total_new_tokens_generated += sum(new_tokens for _,_,new_tokens in generated) torch.cuda.synchronize() if rank == 0: diff --git a/scripts/inference/bloom-accelerate-server.py b/scripts/inference/bloom-accelerate-server.py new file mode 100644 index 000000000..09e74f63a --- /dev/null +++ b/scripts/inference/bloom-accelerate-server.py @@ -0,0 +1,199 @@ +import argparse +import logging +import sys +import time +from typing import List, Union + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +import utils +from flask import Flask, request +from utils import MaxTokensError, get_max_memory_per_gpu_dict, get_stack_trace, get_torch_dtype, parse_input +from waitress import serve + + +def get_args(): + parser = argparse.ArgumentParser(description="Text generation server") + + group = parser.add_argument_group(title="model") + group.add_argument("--model_name", type=str, + required=True, help="model to use") + group.add_argument("--dtype", type=str, required=True, + choices=["bf16", "fp16"], help="dtype for model") + + group = parser.add_argument_group(title="launch config") + group.add_argument("--log_file", type=str, help="log data") + group.add_argument("--host", type=str, required=True, help="host address") + group.add_argument("--port", type=int, required=True, help="port number") + + group = parser.add_argument_group(title="limitting values") + group.add_argument("--allowed_max_new_tokens", type=int, + default=100, help="max allowed tokens") + + group = parser.add_argument_group(title="default values") + group.add_argument("--top_k", type=int, default=50, help="default top_k") + group.add_argument("--top_p", type=float, default=1, help="default top_p") + group.add_argument("--temperature", type=float, + default=1, help="default temperature") + group.add_argument("--min_length", type=int, default=1, help="min length") + group.add_argument("--max_new_tokens", type=int, + default=40, help="max new tokens") + + args = parser.parse_args() + + args.dtype = get_torch_dtype(args.dtype) + + return args + + +class Model: + def __init__(self, args: argparse.Namespace) -> None: + print("Loading model...") + + self.tokenizer = AutoTokenizer.from_pretrained(args.model_name) + + self.model = AutoModelForCausalLM.from_pretrained( + args.model_name, + device_map="auto", + max_memory=get_max_memory_per_gpu_dict( + args.dtype, args.model_name), + torch_dtype=args.dtype + ) + + self.model.eval() + self.input_device = "cuda:0" + + # optimize model by generating once (optimization happens on the first run) + self.generate("Hi, I'm a model", 5, 0.9, 0.7, 1, 40) + + print("Model loaded") + + def generate(self, + text: Union[str, List[str]], + top_k: int, + top_p: float, + temperature: float, + min_length: int, + max_new_tokens: int) -> Union[str, List[str]]: + return_format = type(text) + if (return_format == str): + text = [text] + + x = self.tokenizer(text, return_tensors="pt", padding=True) + + input_ids = x["input_ids"].to(self.input_device) + attention_mask = x["attention_mask"].to(self.input_device) + + with torch.no_grad(): + output = self.model.generate( + input_ids=input_ids, + attention_mask=attention_mask, + top_k=top_k, + top_p=top_p, + temperature=temperature, + min_length=min_length, + max_new_tokens=max_new_tokens + ) + + output_text = self.tokenizer.batch_decode( + output, skip_special_tokens=True) + + if (return_format == str): + return output_text[0] + return output_text + + +#################################################################################### +args = get_args() +app = Flask(__name__) + +# Setup logging +logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d : %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + filename=args.log_file +) +logger = logging.getLogger(__name__) + +model = Model(args) +query_id = 0 +#################################################################################### + + +@app.route("/gpu_status/", methods=["GET"]) +def gpu_status() -> str: + return utils.gpu_status() + + +@app.route("/about/", methods=["GET"]) +def about() -> str: + return utils.about(args.log_file) + + +@app.route("/generate/", methods=["POST"]) +def generate() -> dict: + # needs to be global since it is updated + global query_id + + try: + start_time = time.time() + json_obj = request.get_json() + + (input_text, + top_k, + top_p, + temperature, + min_length, + max_new_tokens) = parse_input(json_obj, args) + + if (max_new_tokens > args.allowed_max_new_tokens): + raise MaxTokensError(max_new_tokens, args.allowed_max_new_tokens) + + output_text = model.generate( + input_text, + top_k, + top_p, + temperature, + min_length, + max_new_tokens + ) + json_obj["query_id"] = query_id + + total_time_taken = time.time() - start_time + output = { + "output_text": output_text, + "total_time_taken": "{:.3f} s".format(total_time_taken), + "query_id": query_id + } + if (args.log_file): + logger.info(json_obj) + logger.info(output) + except Exception: + e_type, e_message, e_stack_trace = sys.exc_info() + output = { + "error": { + "error": str(e_type.__name__), + "message": str(e_message), + "stack_trace": get_stack_trace(e_stack_trace) + }, + "time_taken": "{} s".format(str(time.time() - start_time)), + "query_id": query_id + } + if (args.log_file): + logger.info(json_obj) + logger.error(output) + del output["error"]["stack_trace"] + + query_id += 1 + + return output + + +def main(): + serve(app, host=args.host, port=args.port) + + +if (__name__ == "__main__"): + main() diff --git a/scripts/inference/bloom-ds-inference.py b/scripts/inference/bloom-ds-inference.py index 3be81a27c..e038f2141 100644 --- a/scripts/inference/bloom-ds-inference.py +++ b/scripts/inference/bloom-ds-inference.py @@ -17,19 +17,16 @@ from argparse import ArgumentParser from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig -from transformers.deepspeed import HfDeepSpeedConfig from transformers.models.bloom.modeling_bloom import BloomBlock as BloomBlock import deepspeed import gc -import glob -import io -import json import math import os -import sys import time import torch import torch.distributed as dist +from utils import generate_, write_checkponts_json + t_start = time.time() @@ -52,62 +49,6 @@ ### Model loading and instantiating on GPU (via ZeRO) -def get_checkpoint_files(pretrained_model_name_or_path): - # XXX: I just hacked this one together to automatically handle the fetching of the model file or - # shards into cache and returning the cached entries - note that I removed most arguments - - from transformers.utils import WEIGHTS_NAME, WEIGHTS_INDEX_NAME, cached_path, hf_bucket_url, is_offline_mode - from transformers.utils.hub import EntryNotFoundError - from transformers.modeling_utils import get_checkpoint_shard_files - - cache_dir = None - is_sharded = False - - # XXX: preparation for revision branches if needed - revision = None - #revision = "sharded" - - # this supports nodes with no network (so you need to pre-cache the model and the tokenizer with - # python -c "from transformers import AutoModel; AutoModel.from_pretrained('bigscience/bloom')" - if is_offline_mode(): - print("Offline mode: forcing local_files_only=True") - local_files_only = True - else: - local_files_only = False - - filename = WEIGHTS_NAME - archive_file = hf_bucket_url(pretrained_model_name_or_path, filename=filename, revision=revision) - - try: - resolved_archive_file = cached_path(archive_file, cache_dir=cache_dir, local_files_only=local_files_only,) - return [resolved_archive_file] - - except (EntryNotFoundError, FileNotFoundError): - if filename == WEIGHTS_NAME: - # Maybe the checkpoint is sharded, we try to grab the index name in this case. - archive_file = hf_bucket_url( - pretrained_model_name_or_path, - filename=WEIGHTS_INDEX_NAME, - revision=revision, - ) - resolved_archive_file = cached_path( - archive_file, - cache_dir=cache_dir, - local_files_only=local_files_only, - ) - is_sharded = True - - if is_sharded: - # resolved_archive_file becomes a list of files that point to the different checkpoint shards in this case. - resolved_archive_file, sharded_metadata = get_checkpoint_shard_files( - pretrained_model_name_or_path, - resolved_archive_file, - cache_dir=cache_dir, - revision=revision - ) - - return resolved_archive_file - model_name = args.name #print(get_checkpoint_files(model_name)) @@ -157,25 +98,9 @@ def get_checkpoint_files(pretrained_model_name_or_path): ### Deepspeed-Inference Loading checkpoints_json = "checkpoints.json" -def write_checkponts_json(): - - with io.open(checkpoints_json, 'w', encoding='utf-8') as f: - - #checkpoint_dir = "/gpfsscratch/rech/six/commun/uan68tv-model-conversion/bloom" - #checkpoint_files = glob.glob(f"{checkpoint_dir}/*bin") - checkpoint_files = get_checkpoint_files(model_name) - - #print("Checkpoint files:", checkpoint_files) - - data = { - "type": "BLOOM-176B", - "checkpoints": checkpoint_files, - "version": 1.0 - } - json.dump(data, f) if rank == 0: - write_checkponts_json() + write_checkponts_json(model_name, checkpoints_json) dist.barrier() if args.benchmark: @@ -233,31 +158,26 @@ def write_checkponts_json(): if rank == 0: print(f"Generate args {generate_kwargs}") inputs = input_sentences[:args.batch_size] -def generate(): - """ returns a list of zipped inputs, outputs and number of new tokens """ - - input_tokens = tokenizer.batch_encode_plus(inputs, return_tensors="pt", padding=True) - for t in input_tokens: - if torch.is_tensor(input_tokens[t]): - input_tokens[t] = input_tokens[t].to(torch.cuda.current_device()) - - outputs = model.generate(**input_tokens, **generate_kwargs) - - input_tokens_lengths = [x.shape[0] for x in input_tokens.input_ids] - output_tokens_lengths = [x.shape[0] for x in outputs] - - total_new_tokens = [o-i for i,o in zip(input_tokens_lengths, output_tokens_lengths)] - outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True) - - return zip(inputs, outputs, total_new_tokens) # warmup is a must if measuring speed as it's when all the optimizations are performed # e.g. on 8x80 a100 the first pass of 100 tokens takes 23sec, and the next one is 4secs -_ = generate() +_ = generate_( + inputs, + model, + tokenizer, + generate_kwargs, + torch.cuda.current_device() +) t_generate_start = time.time() -generated = generate() +generated = generate_( + inputs, + model, + tokenizer, + generate_kwargs, + torch.cuda.current_device() +) t_generate_span = time.time() - t_generate_start if rank == 0: for i,o,_ in generated: @@ -277,7 +197,13 @@ def generate(): # warm up for i in range(1): - _ = generate() + _ = generate_( + inputs, + model, + tokenizer, + generate_kwargs, + torch.cuda.current_device() + ) torch.cuda.synchronize() # benchmark @@ -285,7 +211,13 @@ def generate(): cycles = 5 total_new_tokens_generated = 0 for i in range(cycles): - generated = generate() + generated = generate_( + inputs, + model, + tokenizer, + generate_kwargs, + torch.cuda.current_device() + ) total_new_tokens_generated += sum(new_tokens for _,_,new_tokens in generated) torch.cuda.synchronize() if rank == 0: diff --git a/scripts/inference/bloom-ds-server.py b/scripts/inference/bloom-ds-server.py new file mode 100644 index 000000000..b23530c8a --- /dev/null +++ b/scripts/inference/bloom-ds-server.py @@ -0,0 +1,199 @@ +import argparse +import logging +import os +import sys +import time +from typing import List, Union + +from transformers import AutoTokenizer + +import mii +import utils +from flask import Flask, request +from utils import MaxTokensError, get_stack_trace, parse_input +from waitress import serve + + +def get_args(): + parser = argparse.ArgumentParser(description="Text generation server") + + group = parser.add_argument_group(title="model") + group.add_argument("--model_name", type=str, + required=True, help="model to use") + group.add_argument("--dtype", type=str, required=True, + choices=["bf16", "fp16"], help="dtype for model") + + group = parser.add_argument_group(title="launch config") + group.add_argument("--log_file", type=str, help="log data") + group.add_argument("--host", type=str, required=True, help="host address") + group.add_argument("--port", type=int, required=True, help="port number") + group.add_argument("--local_rank", required=False, + type=int, help="used by dist launchers") + + group = parser.add_argument_group(title="limitting values") + group.add_argument("--allowed_max_new_tokens", type=int, + default=100, help="max allowed tokens") + + group = parser.add_argument_group(title="default values") + group.add_argument("--top_k", type=int, default=50, help="default top_k") + group.add_argument("--top_p", type=float, default=1, help="default top_p") + group.add_argument("--temperature", type=float, + default=1, help="default temperature") + group.add_argument("--min_length", type=int, default=1, help="min length") + group.add_argument("--max_new_tokens", type=int, + default=40, help="max new tokens") + + args = parser.parse_args() + args.deployment_name = args.model_name + "_deployment" + + return args + + +class Model: + def __init__(self, args: argparse.Namespace) -> None: + if (args.dtype == "fp16"): + mii.deploy( + task="text-generation", + model=args.model_name, + deployment_name=args.deployment_name, + mii_config={ + "dtype": args.dtype, + "tensor_parallel": 8, + "port_number": 50950, + "checkpoint_dict": { + "checkpoints": ["BLOOM-176B-non-tp.pt"] * 8 + [f'BLOOM-176B-tp_0{i}.pt' for i in range(8)], + "parallelization": "tp", + "version": 1.0, + "type": "BLOOM" + } + }, + model_path=os.getenv("DS_CACHE") + ) + else: + raise NotImplementedError("This is not yet supported") + + self.tokenizer = AutoTokenizer.from_pretrained(args.model_name) + self.model = mii.mii_query_handle(args.deployment_name) + + def generate(self, + text: Union[str, List[str]], + top_k: int, + top_p: float, + temperature: float, + min_length: int, + max_new_tokens: int) -> Union[str, List[str]]: + return_format = type(text) + if (return_format == str): + text = [text] + + output_text = self.model.query( + { + "query": text + }, + top_k=top_k, + top_p=top_p, + temperature=temperature, + min_length=min_length, + max_new_tokens=max_new_tokens + ).response + + output_text = [_ for _ in output_text] + + if (return_format == str): + return output_text[0] + return output_text + + +#################################################################################### +args = get_args() +app = Flask(__name__) + +# Setup logging +logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d : %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + filename=args.log_file +) +logger = logging.getLogger(__name__) + +model = Model(args) +query_id = 0 +#################################################################################### + + +@app.route("/gpu_status/", methods=["GET"]) +def gpu_status() -> str: + return utils.gpu_status() + + +@app.route("/about/", methods=["GET"]) +def about() -> str: + return utils.about(args.log_file) + + +@app.route("/generate/", methods=["POST"]) +def generate() -> dict: + # needs to be global since it is updated + global query_id + + try: + start_time = time.time() + json_obj = request.get_json() + + (input_text, + top_k, + top_p, + temperature, + min_length, + max_new_tokens) = parse_input(json_obj, args) + + if (max_new_tokens > args.allowed_max_new_tokens): + raise MaxTokensError(max_new_tokens, args.allowed_max_new_tokens) + + output_text = model.generate( + input_text, + top_k, + top_p, + temperature, + min_length, + max_new_tokens + ) + json_obj["query_id"] = query_id + + total_time_taken = time.time() - start_time + output = { + "output_text": output_text, + "total_time_taken": "{:.3f} s".format(total_time_taken), + "query_id": query_id + } + if (args.log_file): + logger.info(json_obj) + logger.info(output) + except Exception: + e_type, e_message, e_stack_trace = sys.exc_info() + output = { + "error": { + "error": str(e_type.__name__), + "message": str(e_message), + "stack_trace": get_stack_trace(e_stack_trace) + }, + "time_taken": "{} s".format(str(time.time() - start_time)), + "query_id": query_id + } + if (args.log_file): + logger.info(json_obj) + logger.error(output) + del output["error"]["stack_trace"] + + query_id += 1 + + return output + + +def main(): + serve(app, host=args.host, port=args.port) + + +if (__name__ == "__main__"): + main() diff --git a/scripts/inference/bloom-ds-zero-inference.py b/scripts/inference/bloom-ds-zero-inference.py index 043b4967f..4631ba549 100644 --- a/scripts/inference/bloom-ds-zero-inference.py +++ b/scripts/inference/bloom-ds-zero-inference.py @@ -21,15 +21,13 @@ from transformers.models.bloom.modeling_bloom import BloomBlock as BloomBlock import deepspeed import gc -import glob -import io -import json import math import os -import sys import time import torch import torch.distributed as dist +from utils import generate_ + t_start = time.time() @@ -140,32 +138,27 @@ if rank == 0: print(f"Generate args {generate_kwargs}") inputs = input_sentences[:args.batch_size] -def generate(): - """ returns a list of zipped inputs, outputs and number of new tokens """ - - input_tokens = tokenizer.batch_encode_plus(inputs, return_tensors="pt", padding=True) - for t in input_tokens: - if torch.is_tensor(input_tokens[t]): - input_tokens[t] = input_tokens[t].to(torch.cuda.current_device()) - - outputs = model.generate(**input_tokens, **generate_kwargs) - - input_tokens_lengths = [x.shape[0] for x in input_tokens.input_ids] - output_tokens_lengths = [x.shape[0] for x in outputs] - - total_new_tokens = [o-i for i,o in zip(input_tokens_lengths, output_tokens_lengths)] - outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True) - - return zip(inputs, outputs, total_new_tokens) # XXX: this is currently doing world_size streams on world_size gpus, so we can feed it different inputs on each! and hence the time can be divided by world_size # warmup is a must if measuring speed as it's when all the optimizations are performed # e.g. on 8x80 a100 the first pass of 100 tokens takes 23sec, and the next one is 4secs -_ = generate() +_ = generate_( + inputs, + model, + tokenizer, + generate_kwargs, + torch.cuda.current_device() +) t_generate_start = time.time() -pairs = generate() +pairs = generate_( + inputs, + model, + tokenizer, + generate_kwargs, + torch.cuda.current_device() +) t_generate_span = time.time() - t_generate_start if rank == 0: for i,o,_ in pairs: @@ -185,7 +178,13 @@ def generate(): # warm up for i in range(1): - _ = generate() + _ = generate_( + inputs, + model, + tokenizer, + generate_kwargs, + torch.cuda.current_device() + ) torch.cuda.synchronize() # benchmark @@ -193,7 +192,13 @@ def generate(): cycles = 5 total_new_tokens_generated = 0 for i in range(cycles): - generated = generate() + generated = generate_( + inputs, + model, + tokenizer, + generate_kwargs, + torch.cuda.current_device() + ) total_new_tokens_generated += sum(new_tokens for _,_,new_tokens in generated) torch.cuda.synchronize() @@ -208,4 +213,3 @@ def generate(): Tokenize and generate {total_new_tokens_generated} (bs={args.batch_size}) tokens: {t_generate_span:.3f} secs Start to finish: {t_ready - t_start + t_generate_span:.3f} secs """) - diff --git a/scripts/inference/cache-ds-model.py b/scripts/inference/cache-ds-model.py new file mode 100644 index 000000000..964658bf2 --- /dev/null +++ b/scripts/inference/cache-ds-model.py @@ -0,0 +1,80 @@ +import argparse +import os + +import deepspeed +import torch.distributed as dist +from transformers import AutoConfig, AutoModelForCausalLM + +import mii +from utils import get_torch_dtype, print_rank_n, run_rank_n, write_checkponts_json + + +def ParseArgs(): + parser = argparse.ArgumentParser(description="Text generation server") + + group = parser.add_argument_group(title="model") + group.add_argument("--model_name", type=str, + required=True, help="model to use") + group.add_argument("--dtype", type=str, required=True, + choices=["bf16", "fp16"], help="dtype for model") + + group = parser.add_argument_group(title="launch config") + group.add_argument("--local_rank", required=False, + type=int, help="used by dist launchers") + + args = parser.parse_args() + + args.dtype = get_torch_dtype(args.dtype) + + return args + + +def main(): + deepspeed.init_distributed("nccl") + args = ParseArgs() + + checkpoints_json = "checkpoints.json" + + if (os.path.isdir(os.getenv("DS_CACHE"))): + print_rank_n("Found cached model at {}".format(os.getenv("DS_CACHE"))) + exit() + + # ensure all processes didn't find cache + dist.barrier() + + print_rank_n("Caching model at {}".format(os.getenv("DS_CACHE"))) + + run_rank_n( + write_checkponts_json, + { + "model_name": args.model_name, + "checkpoints_json": checkpoints_json + } + ) + + run_rank_n( + os.makedirs, + { + "name": os.getenv("DS_CACHE") + } + ) + + with deepspeed.OnDevice(dtype=args.dtype, device="meta"): + model = AutoModelForCausalLM.from_config( + AutoConfig.from_pretrained(args.model_name), + torch_dtype=args.dtype + ) + model.eval() + + deepspeed.init_inference( + model, + mp_size=int(os.getenv("WORLD_SIZE", "1")), + dtype=args.dtype, + checkpoint=checkpoints_json, + replace_with_kernel_inject=True, + save_mp_checkpoint_path=os.getenv("DS_CACHE") + ) + + +if (__name__ == "__main__"): + main() From ff9b8b007fb5dd2eb9d2ed51114384691292066e Mon Sep 17 00:00:00 2001 From: Mayank Mishra Date: Mon, 8 Aug 2022 23:12:05 +0530 Subject: [PATCH 2/3] add utils --- scripts/inference/utils.py | 273 +++++++++++++++++++++++++++++++++++++ 1 file changed, 273 insertions(+) create mode 100644 scripts/inference/utils.py diff --git a/scripts/inference/utils.py b/scripts/inference/utils.py new file mode 100644 index 000000000..66fc0515d --- /dev/null +++ b/scripts/inference/utils.py @@ -0,0 +1,273 @@ +import io +import json +import subprocess +import traceback +from argparse import Namespace +from typing import Any, List, Tuple + +import torch +import torch.distributed as dist +from transformers import AutoConfig + + +def gpu_status(): + try: + info = subprocess.check_output(["nvidia-smi"]) + info = info.decode("utf8") + except Exception as e: + info = "Executing nvidia-smi failed: " + str(e) + return info + + +def about(log_file: str) -> str: + if (log_file): + description = "Please don't send any personal information to this endpoint. We are logging your data.\n\n" + else: + description = "" + description += '''Usage: +A request object should look like: +{ + input_text: "Hello, I'm a model", + "top_k": 5, + "top_p": 0.9, + "temperature": 0.7, + "min_length": 1, + "max_new_tokens": 40 +} +Default values (use if not provided in request object): +top_k = 50 +top_p = 1 +temperature = 1 +min_length = 1 +max_new_tokens = 40 +''' + return description + + +# TODO remove when bloom-inference is merged into main +def get_checkpoint_files(pretrained_model_name_or_path): + # XXX: I just hacked this one together to automatically handle the fetching of the model file or + # shards into cache and returning the cached entries - note that I removed most arguments + + from transformers.modeling_utils import get_checkpoint_shard_files + from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME, cached_path, hf_bucket_url, is_offline_mode + from transformers.utils.hub import EntryNotFoundError + + cache_dir = None + is_sharded = False + + # XXX: preparation for revision branches if needed + revision = None + #revision = "sharded" + + # this supports nodes with no network (so you need to pre-cache the model and the tokenizer with + # python -c "from transformers import AutoModel; AutoModel.from_pretrained('bigscience/bloom')" + if is_offline_mode(): + print("Offline mode: forcing local_files_only=True") + local_files_only = True + else: + local_files_only = False + + filename = WEIGHTS_NAME + archive_file = hf_bucket_url( + pretrained_model_name_or_path, filename=filename, revision=revision) + + try: + resolved_archive_file = cached_path( + archive_file, cache_dir=cache_dir, local_files_only=local_files_only,) + return [resolved_archive_file] + + except (EntryNotFoundError, FileNotFoundError): + if filename == WEIGHTS_NAME: + # Maybe the checkpoint is sharded, we try to grab the index name in this case. + archive_file = hf_bucket_url( + pretrained_model_name_or_path, + filename=WEIGHTS_INDEX_NAME, + revision=revision, + ) + resolved_archive_file = cached_path( + archive_file, + cache_dir=cache_dir, + local_files_only=local_files_only, + ) + is_sharded = True + + if is_sharded: + # resolved_archive_file becomes a list of files that point to the different checkpoint shards in this case. + resolved_archive_file, sharded_metadata = get_checkpoint_shard_files( + pretrained_model_name_or_path, + resolved_archive_file, + cache_dir=cache_dir, + revision=revision + ) + + return resolved_archive_file + + +def get_stack_trace(e_stack_trace): + trace_back = traceback.extract_tb(e_stack_trace) + + # Format stacktrace + stack_trace = [] + for trace in trace_back: + stack_trace.append("File : {}, Line : {}, Func.Name : {}, Message : {}".format( + trace[0], trace[1], trace[2], trace[3])) + + return stack_trace + + +class MaxTokensError(Exception): + def __init__(self, max_new_tokens: int, allowed_max_new_tokens: int) -> None: + super().__init__("max_new_tokens = {} > {} is not supported.".format( + max_new_tokens, allowed_max_new_tokens)) + + +def run_rank_n(func: callable, + kwargs: dict, + barrier: bool = False, + rank: int = 0) -> Any: + if (dist.is_initialized()): + if (dist.get_rank() == rank): + output = func(**kwargs) + if (barrier): + dist.barrier() + return output + else: + if (barrier): + dist.barrier() + else: + return func(**kwargs) + + +def print_rank_n(*values, rank: int = 0) -> None: + if (dist.is_initialized()): + if (dist.get_rank() == rank): + print(*values) + else: + print(*values) + + +def get_max_memory_per_gpu_dict(dtype, model_name): + """ try to generate the memory map based on what we know about the model and the available hardware """ + + # figure out the memory map - the minimum per gpu required to load the model + n_gpus = torch.cuda.device_count() + + if model_name == "bigscience/bloom" and n_gpus == 8 and torch.cuda.get_device_properties(0).total_memory > 79*2**30: + # hand crafted optimized memory map for 8x80 setup over BLOOM + # this works with bs=40 + return {0: '0GIB', 1: '51GIB', 2: '51GIB', 3: '51GIB', 4: '51GIB', 5: '51GIB', 6: '51GIB', 7: '51GIB'} + + try: + # model_params calculation, as we don't have a model yet to do: + #model_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values()) + + config = AutoConfig.from_pretrained(model_name) + h = config.n_embed + l = config.n_layer + v = config.vocab_size + # from https://github.com/bigscience-workshop/bigscience/tree/6917a3b5fefcf439d3485ca184b4d9f6ab605150/math#model-sizing + model_params = l*(12*h**2 + 13*h) + v*h + 4*h + except: + print( + f"The model {model_name} has a broken config file. Please notify the owner") + raise + + bytes = torch.finfo(dtype).bits / 8 + param_memory_total_in_bytes = model_params * bytes + # add 5% since weight sizes aren't the same and some GPU may need more memory + param_memory_per_gpu_in_bytes = int( + param_memory_total_in_bytes / n_gpus * 1.05) + print( + f"Estimating {param_memory_per_gpu_in_bytes/2**30:0.2f}GB per gpu for weights") + + # check the real available memory + # load cuda kernels first and only measure the real free memory after loading (shorter by ~2GB) + torch.ones(1).cuda() + max_memory_per_gpu_in_bytes = torch.cuda.mem_get_info(0)[0] + if max_memory_per_gpu_in_bytes < param_memory_per_gpu_in_bytes: + raise ValueError( + f"Unable to generate the memory map automatically as the needed estimated memory per gpu ({param_memory_per_gpu_in_bytes/2**30:0.2f}GB) is bigger than the available per gpu memory ({max_memory_per_gpu_in_bytes/2**30:0.2f}GB)") + + return {i: param_memory_per_gpu_in_bytes for i in range(torch.cuda.device_count())} + + +def write_checkponts_json(model_name: str, checkpoints_json: str) -> None: + with io.open(checkpoints_json, 'w', encoding='utf-8') as f: + #checkpoint_dir = "/gpfsscratch/rech/six/commun/uan68tv-model-conversion/bloom" + #checkpoint_files = glob.glob(f"{checkpoint_dir}/*bin") + checkpoint_files = get_checkpoint_files(model_name) + + #print("Checkpoint files:", checkpoint_files) + + data = { + "type": "BLOOM-176B", + "checkpoints": checkpoint_files, + "version": 1.0 + } + json.dump(data, f) + + +def generate_(inputs, + model, + tokenizer, + generate_kwargs, + input_device): + """ returns a list of zipped inputs, outputs and number of new tokens """ + + input_tokens = tokenizer.batch_encode_plus( + inputs, return_tensors="pt", padding=True) + for t in input_tokens: + if torch.is_tensor(input_tokens[t]): + input_tokens[t] = input_tokens[t].to(input_device) + + outputs = model.generate(**input_tokens, **generate_kwargs) + + input_tokens_lengths = [x.shape[0] for x in input_tokens.input_ids] + output_tokens_lengths = [x.shape[0] for x in outputs] + + total_new_tokens = [o-i for i, + o in zip(input_tokens_lengths, output_tokens_lengths)] + outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True) + + return zip(inputs, outputs, total_new_tokens) + + +def get_torch_dtype(dtype_str: str): + if (dtype_str == "bf16"): + return torch.bfloat16 + elif (dtype_str == "fp16"): + return torch.float16 + + +def get_str_dtype(dtype_str: str): + if (dtype_str == torch.bfloat16): + return "bf16" + elif (dtype_str == torch.float16): + return "fp16" + + +def parse_input(json_obj: dict, + args: Namespace) -> Tuple[List[str], + int, + float, + float, + int, + int, + str]: + input_text = json_obj["input_text"] + top_k = int(json_obj.get("top_k", args.top_k)) + top_p = float(json_obj.get("top_p", args.top_p)) + temperature = float(json_obj.get("temperature", args.temperature)) + min_length = int(json_obj.get("min_length", args.min_length)) + max_new_tokens = int(json_obj.get( + "max_new_tokens", args.max_new_tokens)) + + return ( + input_text, + top_k, + top_p, + temperature, + min_length, + max_new_tokens + ) From 9cd56e91112547b727f4bdcb945dec3dfc51bc63 Mon Sep 17 00:00:00 2001 From: Mayank Mishra Date: Thu, 11 Aug 2022 00:03:39 +0530 Subject: [PATCH 3/3] fix comment --- scripts/inference/bloom-ds-inference.py | 3 --- scripts/inference/utils.py | 2 +- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/scripts/inference/bloom-ds-inference.py b/scripts/inference/bloom-ds-inference.py index e2a0ce5c3..1ad1b4794 100644 --- a/scripts/inference/bloom-ds-inference.py +++ b/scripts/inference/bloom-ds-inference.py @@ -46,9 +46,6 @@ deepspeed.init_distributed('nccl') rank = dist.get_rank() - -### Model loading and instantiating on GPUs - model_name = args.name #print(get_checkpoint_files(model_name)) diff --git a/scripts/inference/utils.py b/scripts/inference/utils.py index 66fc0515d..5de1a5171 100644 --- a/scripts/inference/utils.py +++ b/scripts/inference/utils.py @@ -44,7 +44,7 @@ def about(log_file: str) -> str: return description -# TODO remove when bloom-inference is merged into main +### Model loading and instantiating on GPUs def get_checkpoint_files(pretrained_model_name_or_path): # XXX: I just hacked this one together to automatically handle the fetching of the model file or # shards into cache and returning the cached entries - note that I removed most arguments