Add generation server scripts using HF accelerate and DS-inference

inference/benchmark.py

@@ -0,0 +1,189 @@
+import argparse
+import gc
+import os
+import time
+from typing import Any, List, Tuple, Union
+
+import deepspeed
+import torch
+
+import constants
+import utils
+from ds_inference import DSInferenceModel
+from ds_zero import DSZeROModel
+from hf_accelerate import HFAccelerateModel
+from utils import (
+    Execute,
+    GenerateRequest,
+    Model,
+    get_argument_parser,
+    get_dummy_batch,
+    parse_generate_kwargs,
+    print_rank_n
+)
+
+
+def run_and_log_time(execs: Union[List[Execute], Execute]) -> Tuple[Union[List[Any], Any], float]:
+    """
+    runs a list of Execute objects and returns a list of outputs and the time taken
+    """
+    start_time = time.time()
+
+    if (type(execs) == list):
+        results = []
+        for e in execs:
+            results.append(e())
+    else:
+        results = execs()
+
+    time_elapsed = time.time() - start_time
+    return results, time_elapsed
+
+
+def benchmark_generation(model: Model,
+                         request: GenerateRequest,
+                         cycles: int = 5):
+    total_new_tokens_generated = 0
+    for _ in range(cycles):
+        response = model.generate(request)
+        total_new_tokens_generated += sum(
+            new_tokens for new_tokens in response.num_generated_tokens)
+    return total_new_tokens_generated
+
+
+def get_benchmark_results(benchmark_time: float,
+                          initialization_time: float,
+                          total_new_tokens_generated: int,
+                          batch_size: int,
+                          cycles: int) -> str:
+    throughput = total_new_tokens_generated / benchmark_time
+    latency = benchmark_time / cycles
+    return f"""
+*** Performance stats:
+Throughput (including tokenization) = {throughput:.2f} tokens/sec
+Throughput (including tokenization) = {1000 / throughput:.2f} msecs/token
+Model loading time = {initialization_time:.2f} secs
+Total tokens generated = {total_new_tokens_generated} with batch size = {batch_size}
+Latency = {latency:.2f} secs
+Model loading time + generation time per batch = {initialization_time + latency:.2f} secs
+"""
+
+
+def benchmark_end_to_end(args: argparse.Namespace,
+                         model_class: Model,
+                         zero_activated: bool = False) -> None:
+    model, initialization_time = run_and_log_time(
+        Execute(model_class, {"args": args})
+    )
+
+    request = parse_generate_kwargs(
+        get_dummy_batch(args.batch_size),
+        args.generate_kwargs
+    )
+
+    print_rank_n(f"generate_kwargs = {request}")
+    print_rank_n(f"batch_size = {args.batch_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
+    response = model.generate(request)
+
+    for i, (o, _) in zip(request.text, zip(response.text, response.num_generated_tokens)):
+        print_rank_n(f"{'-' * 60}\nin = {i}\nout = {o}\n")
+
+    if (args.benchmark_cycles > 0):
+        print_rank_n(f"*** Running benchmark")
+
+        torch.cuda.empty_cache()
+        gc.collect()
+
+        # warm up
+        model.generate(request)
+        torch.cuda.synchronize()
+
+        # benchmark
+        total_new_tokens_generated, benchmark_time = run_and_log_time(
+            Execute(
+                benchmark_generation,
+                {
+                    "model": model,
+                    "request": request,
+                    "cycles": args.benchmark_cycles
+                }
+            )
+        )
+
+        # with ZeRO every GPU is generating batch_size * sequence_length tokens
+        if (zero_activated):
+            world_size = int(os.getenv('WORLD_SIZE', '1'))
+            total_new_tokens_generated *= world_size
+
+        print_rank_n(
+            get_benchmark_results(
+                benchmark_time,
+                initialization_time,
+                total_new_tokens_generated,
+                args.batch_size,
+                args.benchmark_cycles
+            )
+        )
+
+
+def get_args() -> argparse.Namespace:
+    parser = get_argument_parser()
+
+    group = parser.add_argument_group(title="launch config")
+    group.add_argument(
+        "--deployment_framework",
+        type=str,
+        choices=[
+            constants.HF_ACCELERATE,
+            constants.DS_INFERENCE,
+            constants.DS_ZERO
+        ],
+        default=constants.HF_ACCELERATE
+    )
+    group.add_argument("--benchmark_cycles", type=int,
+                       default=0, help="additionally run benchmark")
+    group.add_argument("--local_rank", required=False,
+                       type=int, help="used by dist launchers")
+    group.add_argument("--batch_size", default=1, type=int, help="batch size")
+    group.add_argument("--save_mp_checkpoint_path", required=False,
+                       type=str, help="MP checkpoints path for DS inference")
+    group.add_argument("--cpu_offload", action="store_true",
+                       help="whether to activate CPU offload for DS ZeRO")
+
+    args = utils.get_args(parser)
+
+    launched_with_deepspeed = args.deployment_framework in [
+        constants.DS_INFERENCE, constants.DS_ZERO]
+
+    if (not launched_with_deepspeed):
+        assert args.local_rank == None, "local_rank must be None if not launched with DeepSpeed"
+
+    if (args.save_mp_checkpoint_path):
+        assert args.deployment_framework == constants.DS_INFERENCE, "save_mp_checkpoint_path only works with DS inference"
+
+    if (args.cpu_offload):
+        assert args.deployment_framework == constants.DS_ZERO, "cpu_offload only works with DS_ZeRO"
+
+    return args
+
+
+def main() -> None:
+    args = get_args()
+
+    if (args.deployment_framework == constants.HF_ACCELERATE):
+        benchmark_end_to_end(args, HFAccelerateModel)
+    elif (args.deployment_framework == constants.DS_INFERENCE):
+        deepspeed.init_distributed("nccl")
+        benchmark_end_to_end(args, DSInferenceModel)
+    elif (args.deployment_framework == constants.DS_ZERO):
+        benchmark_end_to_end(args, DSZeROModel, zero_activated=True)
+    else:
+        raise ValueError(
+            f"Unknown deployment framework {args.deployment_framework}")
+
+
+if (__name__ == "__main__"):
+    main()

inference/cache_ds_checkpoints.py

@@ -0,0 +1,27 @@
+import argparse
+
+import utils
+from ds_inference import cache_ds_checkpoints
+from utils import get_argument_parser
+
+
+def get_args() -> argparse.Namespace:
+    parser = get_argument_parser()
+
+    group = parser.add_argument_group(title="launch config")
+    group.add_argument("--local_rank", required=False,
+                       type=int, help="used by dist launchers")
+    group.add_argument("--save_mp_checkpoint_path", required=True,
+                       type=str, help="MP checkpoints path for DS inference")
+
+    args = utils.get_args(parser)
+
+    return args
+
+
+def main() -> None:
+    cache_ds_checkpoints(get_args())
+
+
+if (__name__ == "__main__"):
+    main()

inference/cli.py

@@ -0,0 +1,71 @@
+import argparse
+import json
+
+import constants
+import utils
+from ds_inference import DSInferenceGRPCServer
+from hf_accelerate import HFAccelerateModel
+from utils import get_argument_parser, parse_generate_kwargs, print_rank_n
+
+
+def get_args() -> argparse.Namespace:
+    parser = get_argument_parser()
+
+    group = parser.add_argument_group(title="launch config")
+    group.add_argument(
+        "--deployment_framework",
+        type=str,
+        choices=[
+            constants.HF_ACCELERATE,
+            constants.DS_INFERENCE
+        ],
+        default=constants.HF_ACCELERATE
+    )
+    group.add_argument("--save_mp_checkpoint_path", required=False,
+                       type=str, help="MP checkpoints path for DS inference")
+    group.add_argument("--shutdown_command", required=False,
+                       type=str, default="__shutdown__", help="This string will exit the script")
+
+    args = utils.get_args(parser)
+
+    if (args.save_mp_checkpoint_path):
+        assert args.deployment_framework == constants.DS_INFERENCE, "save_mp_checkpoint_path only works with DS inference"
+
+    return args
+
+
+def main() -> None:
+    args = get_args()
+
+    if (args.deployment_framework == constants.HF_ACCELERATE):
+        model = HFAccelerateModel(args)
+    elif (args.deployment_framework == constants.DS_INFERENCE):
+        model = DSInferenceGRPCServer(args)
+    else:
+        raise ValueError(
+            f"Unknown deployment framework {args.deployment_framework}")
+
+    generate_kwargs = args.generate_kwargs
+
+    while (True):
+        # currently only 1 process is running so its
+        # fine but might need to run_rank_n for this
+        # if running a deployment_framework with
+        # multiple processes
+        input_text = input("Input text: ")
+
+        if (input_text == args.shutdown_command):
+            model.shutdown()
+
+        if (input("change generate_kwargs? [y/n] ") == "y"):
+            generate_kwargs = json.loads(input("Generate kwargs: "))
+
+        request = parse_generate_kwargs(input_text, generate_kwargs)
+        response = model.generate(request)
+
+        print_rank_n("Output text:", response.text)
+        print_rank_n("Generated tokens:", response.num_generated_tokens)
+
+
+if (__name__ == "__main__"):
+    main()

inference/constants.py

@@ -0,0 +1,3 @@
+HF_ACCELERATE = "hf_accelerate"
+DS_INFERENCE = "ds_inference"
+DS_ZERO = "ds_zero"

inference/ds_inference/__init__.py

@@ -0,0 +1,3 @@
+from .cache import cache_ds_checkpoints
+from .grpc_server import DSInferenceGRPCServer
+from .model import DSInferenceModel

inference/ds_inference/cache.py

@@ -0,0 +1,67 @@
+import argparse
+import os
+import shutil
+
+import deepspeed
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM
+
+from utils import print_rank_n, run_rank_n
+
+from .model import write_checkponts_json
+
+
+def cache_ds_checkpoints(args: argparse.Namespace) -> None:
+    print_rank_n("Loading model...")
+    world_size = int(os.getenv("WORLD_SIZE", "1"))
+
+    # Load model
+    with deepspeed.OnDevice(dtype=args.dtype, device="meta"):
+        model = AutoModelForCausalLM.from_config(
+            AutoConfig.from_pretrained(args.model_name),
+            torch_dtype=torch.bfloat16
+        )
+    model = model.eval()
+
+    # Write checkpoints.json
+    tmp_directory = "tmp"
+    run_rank_n(
+        os.makedirs,
+        {
+            "name": tmp_directory,
+            "exist_ok": True
+        }
+    )
+    checkpoints_json = os.path.join(tmp_directory, "checkpoints.json")
+    run_rank_n(
+        write_checkponts_json,
+        {
+            "checkpoints_json": checkpoints_json,
+            "model_name": args.model_name
+        },
+        barrier=True
+    )
+
+    run_rank_n(
+        os.makedirs,
+        {
+            "name": args.save_mp_checkpoint_path,
+            "exist_ok": True
+        },
+        barrier=True
+    )
+
+    if (args.dtype == torch.float16):
+        model = deepspeed.init_inference(
+            model,
+            mp_size=world_size,
+            dtype=args.dtype,
+            checkpoint=checkpoints_json,
+            replace_with_kernel_inject=True,
+            save_mp_checkpoint_path=args.save_mp_checkpoint_path
+        )
+    elif (args.dtype == torch.bfloat16):
+        raise NotImplementedError("bfloat16 is not yet supported")
+
+    run_rank_n(shutil.rmtree, {"path": tmp_directory})
+    print_rank_n("Model loaded")