tune_models.py 1.89 KB
Newer Older
Alan Turner's avatar
Alan Turner committed
1
import os, subprocess, argparse, time
Alan Turner's avatar
Alan Turner committed
2
3
import tune_ck as tc

Alan Turner's avatar
Alan Turner committed
4

Alan Turner's avatar
Alan Turner committed
5
def parse_args():
Alan Turner's avatar
Alan Turner committed
6
7
    parser = argparse.ArgumentParser(
        description="Tune CK GEMMs for one or more ONNX models")
Alan Turner's avatar
Alan Turner committed
8
    parser.add_argument('--models',
Alan Turner's avatar
Alan Turner committed
9
10
11
                        '-m',
                        nargs='+',
                        help='ONNX models to be tuned',
Alan Turner's avatar
Alan Turner committed
12
13
                        required=True)
    parser.add_argument('--batch_sizes',
Alan Turner's avatar
Alan Turner committed
14
15
16
                        '-b',
                        nargs='+',
                        help='Batch sizes to tune',
Alan Turner's avatar
Alan Turner committed
17
                        required=True)
Alan Turner's avatar
Alan Turner committed
18
19
20
21
    parser.add_argument('--sequence_length',
                        '-s',
                        type=int,
                        default=384,
Alan Turner's avatar
Alan Turner committed
22
                        help='Sequence length for transformer models')
Alan Turner's avatar
Alan Turner committed
23
24
25
    parser.add_argument('-n',
                        type=int,
                        default=16,
Alan Turner's avatar
Alan Turner committed
26
27
28
29
                        help='Number of instances to tune')
    args = parser.parse_args()
    return args

Alan Turner's avatar
Alan Turner committed
30

Alan Turner's avatar
Alan Turner committed
31
32
33
34
35
36
def tune_models(models, batch_sizes, seq_len, n):
    time_stamp = time.strftime("%Y_%m_%d_%H_%M")
    log_file = "ck_tuning_{}.log".format(time_stamp)
    json_file = "ck_tuning_{}.json".format(time_stamp)
    for model in models:
        for batch in batch_sizes:
Alan Turner's avatar
Alan Turner committed
37
38
39
40
41
42
43
            out = subprocess.run(
                'MIGRAPHX_LOG_CK_GEMM=1 ../build/bin/driver run {} -g --fill1 input_ids --input-dim @input_ids {} {}  | grep \'ck_gemm.*: \[{{\' | sort -u >> {}'
                .format(model, batch, seq_len, log_file),
                capture_output=True,
                check=True,
                shell=True)

Alan Turner's avatar
Alan Turner committed
44
    tc.tune(log_file, n, json_file)
Alan Turner's avatar
Alan Turner committed
45
46
    tuning_path = os.path.abspath(json_file)
    os.environ["MIGRAPHX_CK_TUNING"] = tuning_path
Alan Turner's avatar
Alan Turner committed
47
48
    print("\nTuning results have been saved to:\n{}\n".format(json_file))

Alan Turner's avatar
Alan Turner committed
49

Alan Turner's avatar
Alan Turner committed
50
51
52
def run(args):
    tune_models(args.models, args.batch_sizes, args.sequence_length, args.n)

Alan Turner's avatar
Alan Turner committed
53
54

run(parse_args())