|
| 1 | +import argparse |
| 2 | +import gc |
| 3 | +import os |
| 4 | + |
| 5 | +import deepspeed |
| 6 | +import torch |
| 7 | + |
| 8 | +import utils |
| 9 | +from ds_inference import DSInferenceModel |
| 10 | +from ds_zero import DSZeROModel |
| 11 | +from hf_accelerate import HFAccelerateModel |
| 12 | +from utils import ( |
| 13 | + BENCHMARK, |
| 14 | + DS_INFERENCE, |
| 15 | + DS_ZERO, |
| 16 | + HF_ACCELERATE, |
| 17 | + GenerateRequest, |
| 18 | + Model, |
| 19 | + get_argument_parser, |
| 20 | + get_dummy_batch, |
| 21 | + parse_generate_kwargs, |
| 22 | + print_rank_n, |
| 23 | + run_and_log_time |
| 24 | +) |
| 25 | + |
| 26 | + |
| 27 | +def benchmark_generation(model: Model, |
| 28 | + request: GenerateRequest, |
| 29 | + cycles: int = 5): |
| 30 | + total_new_tokens_generated = 0 |
| 31 | + for _ in range(cycles): |
| 32 | + response = model.generate(request) |
| 33 | + total_new_tokens_generated += sum( |
| 34 | + new_tokens for new_tokens in response.num_generated_tokens) |
| 35 | + return total_new_tokens_generated |
| 36 | + |
| 37 | + |
| 38 | +def get_benchmark_results(benchmark_time: float, |
| 39 | + initialization_time: float, |
| 40 | + total_new_tokens_generated: int, |
| 41 | + batch_size: int, |
| 42 | + cycles: int) -> str: |
| 43 | + throughput = total_new_tokens_generated / benchmark_time |
| 44 | + latency = benchmark_time / cycles |
| 45 | + return f""" |
| 46 | +*** Performance stats: |
| 47 | +Throughput (including tokenization) = {throughput:.2f} tokens/sec |
| 48 | +Throughput (including tokenization) = {1000 / throughput:.2f} msecs/token |
| 49 | +Model loading time = {initialization_time:.2f} secs |
| 50 | +Total tokens generated = {total_new_tokens_generated} with batch size = {batch_size} |
| 51 | +Latency = {latency:.2f} secs |
| 52 | +Model loading time + generation time per batch = {initialization_time + latency:.2f} secs |
| 53 | +""" |
| 54 | + |
| 55 | + |
| 56 | +def benchmark_end_to_end(args: argparse.Namespace, |
| 57 | + model_class: Model, |
| 58 | + zero_activated: bool = False) -> None: |
| 59 | + model, initialization_time = run_and_log_time( |
| 60 | + (model_class, {"args": args}) |
| 61 | + ) |
| 62 | + |
| 63 | + request = parse_generate_kwargs( |
| 64 | + get_dummy_batch(args.batch_size), |
| 65 | + args.generate_kwargs |
| 66 | + ) |
| 67 | + |
| 68 | + print_rank_n(f"generate_kwargs = {args.generate_kwargs}") |
| 69 | + print_rank_n(f"batch_size = {args.batch_size}") |
| 70 | + |
| 71 | + # warmup is a must if measuring speed as it's when all the optimizations are performed |
| 72 | + # e.g. on 8x80 a100 the first pass of 100 tokens takes 23sec, and the next one is 4secs |
| 73 | + response = model.generate(request) |
| 74 | + |
| 75 | + for i, (o, _) in zip(request.text, zip(response.text, response.num_generated_tokens)): |
| 76 | + print_rank_n(f"{'-' * 60}\nin = {i}\nout = {o}\n") |
| 77 | + |
| 78 | + if (args.benchmark_cycles > 0): |
| 79 | + print_rank_n(f"*** Running benchmark") |
| 80 | + |
| 81 | + torch.cuda.empty_cache() |
| 82 | + gc.collect() |
| 83 | + |
| 84 | + # warm up |
| 85 | + model.generate(request) |
| 86 | + torch.cuda.synchronize() |
| 87 | + |
| 88 | + # benchmark |
| 89 | + total_new_tokens_generated, benchmark_time = run_and_log_time( |
| 90 | + ( |
| 91 | + benchmark_generation, |
| 92 | + { |
| 93 | + "model": model, |
| 94 | + "request": request, |
| 95 | + "cycles": args.benchmark_cycles |
| 96 | + } |
| 97 | + ) |
| 98 | + ) |
| 99 | + |
| 100 | + # with ZeRO every GPU is generating batch_size * sequence_length tokens |
| 101 | + if (zero_activated): |
| 102 | + world_size = int(os.getenv('WORLD_SIZE', '1')) |
| 103 | + total_new_tokens_generated *= world_size |
| 104 | + |
| 105 | + print_rank_n( |
| 106 | + get_benchmark_results( |
| 107 | + benchmark_time, |
| 108 | + initialization_time, |
| 109 | + total_new_tokens_generated, |
| 110 | + args.batch_size, |
| 111 | + args.benchmark_cycles |
| 112 | + ) |
| 113 | + ) |
| 114 | + |
| 115 | + |
| 116 | +def get_args() -> argparse.Namespace: |
| 117 | + parser = get_argument_parser() |
| 118 | + |
| 119 | + group = parser.add_argument_group(title="launch config") |
| 120 | + group.add_argument("--benchmark_cycles", type=int, |
| 121 | + default=0, help="additionally run benchmark") |
| 122 | + group.add_argument("--local_rank", required=False, |
| 123 | + type=int, help="used by dist launchers") |
| 124 | + group.add_argument("--batch_size", default=1, type=int, help="batch size") |
| 125 | + group.add_argument("--cpu_offload", action="store_true", |
| 126 | + help="whether to activate CPU offload for DS ZeRO") |
| 127 | + |
| 128 | + args = utils.get_args(parser, BENCHMARK) |
| 129 | + |
| 130 | + launched_with_deepspeed = args.deployment_framework in [ |
| 131 | + DS_INFERENCE, DS_ZERO] |
| 132 | + |
| 133 | + if (not launched_with_deepspeed): |
| 134 | + assert args.local_rank == None, "local_rank must be None if not launched with DeepSpeed" |
| 135 | + |
| 136 | + if (args.cpu_offload): |
| 137 | + assert args.deployment_framework == DS_ZERO, "cpu_offload only works with DS_ZeRO" |
| 138 | + |
| 139 | + return args |
| 140 | + |
| 141 | + |
| 142 | +def main() -> None: |
| 143 | + args = get_args() |
| 144 | + |
| 145 | + if (args.deployment_framework == HF_ACCELERATE): |
| 146 | + benchmark_end_to_end(args, HFAccelerateModel) |
| 147 | + elif (args.deployment_framework == DS_INFERENCE): |
| 148 | + deepspeed.init_distributed("nccl") |
| 149 | + benchmark_end_to_end(args, DSInferenceModel) |
| 150 | + elif (args.deployment_framework == DS_ZERO): |
| 151 | + deepspeed.init_distributed("nccl") |
| 152 | + benchmark_end_to_end(args, DSZeROModel, zero_activated=True) |
| 153 | + else: |
| 154 | + raise ValueError( |
| 155 | + f"Unknown deployment framework {args.deployment_framework}") |
| 156 | + |
| 157 | + |
| 158 | +if (__name__ == "__main__"): |
| 159 | + main() |
0 commit comments