()
| 81 | |
| 82 | |
| 83 | def parse_args(): |
| 84 | parser = argparse.ArgumentParser(description="bitsandbytes inference benchmark tool") |
| 85 | |
| 86 | parser.add_argument("model_id", type=str, help="The model checkpoint to use.") |
| 87 | |
| 88 | parser.add_argument( |
| 89 | "--configs", |
| 90 | nargs="+", |
| 91 | choices=["bf16", "fp16", "nf4", "nf4-dq", "int8", "int8-decomp"], |
| 92 | default=["nf4", "int8", "int8-decomp"], |
| 93 | ) |
| 94 | parser.add_argument("--bf16", dest="configs", action="append_const", const="bf16") |
| 95 | parser.add_argument("--fp16", dest="configs", action="append_const", const="fp16") |
| 96 | parser.add_argument("--nf4", dest="configs", action="append_const", const="nf4") |
| 97 | parser.add_argument("--nf4-dq", dest="configs", action="append_const", const="nf4-dq") |
| 98 | parser.add_argument("--int8", dest="configs", action="append_const", const="int8") |
| 99 | parser.add_argument("--int8-decomp", dest="configs", action="append_const", const="int8-decomp") |
| 100 | |
| 101 | parser.add_argument("--batches", nargs="+", type=int, default=[1, 8, 16, 32]) |
| 102 | parser.add_argument("--input-length", type=int, default=64) |
| 103 | |
| 104 | parser.add_argument("--out-dir", type=str, default="reports") |
| 105 | |
| 106 | parser.add_argument("--iterations", type=int, default=10, help="Number of iterations for each benchmark run") |
| 107 | parser.add_argument( |
| 108 | "--warmup-runs", type=int, default=10, help="Number of warmup runs to discard before measurement" |
| 109 | ) |
| 110 | parser.add_argument( |
| 111 | "--output-length", |
| 112 | type=int, |
| 113 | default=64, |
| 114 | help="If set, `max_new_tokens` and `min_new_tokens` will be set to this value.", |
| 115 | ) |
| 116 | |
| 117 | return parser.parse_args() |
| 118 | |
| 119 | |
| 120 | def run_benchmark(args, config, batch_size): |
no outgoing calls
no test coverage detected