()
| 56 | |
| 57 | |
| 58 | def parse_args(): |
| 59 | parser = argparse.ArgumentParser() |
| 60 | parser.add_argument("batch_size", type=int) |
| 61 | parser.add_argument("seqlen", type=int) |
| 62 | parser.add_argument("head_num", type=int) |
| 63 | parser.add_argument("head_size", type=int) |
| 64 | parser.add_argument("--n_layers", default=1, type=int, help="number of transformer layers") |
| 65 | parser.add_argument("--avg_seqlen", default=0, type=int, help="average seqlen, 0 for equal to seqlen") |
| 66 | parser.add_argument("--dtype", choices=["fp16", "fp32"], default="fp16") |
| 67 | parser.add_argument("--export_data", action="store_true", help="whether to export test data. " |
| 68 | "if true, only run pytorch to generate data, don't run ByteTransformer") |
| 69 | parser.add_argument("--export_data_path", default="./", type=str, help="path to export test data") |
| 70 | parser.add_argument("--lib_path", default="./lib/libths_bytetransformer.so", |
| 71 | type=str, help="lib path of torch op ext") |
| 72 | parser.add_argument("--iters", default=100, type=int, help="perf iterations") |
| 73 | return parser.parse_args() |
| 74 | |
| 75 | |
| 76 | if __name__ == "__main__": |
no outgoing calls
no test coverage detected