从命令行解析参数
()
| 57 | |
| 58 | |
| 59 | def parse_args(): |
| 60 | """ |
| 61 | 从命令行解析参数 |
| 62 | """ |
| 63 | parser = argparse.ArgumentParser("Cache transfer manager") |
| 64 | parser.add_argument( |
| 65 | "--splitwise_role", |
| 66 | type=str, |
| 67 | default="mixed", |
| 68 | help="splitwise role, can be decode, prefill or mixed", |
| 69 | ) |
| 70 | parser.add_argument("--rank", type=int, default=0, help="local tp rank") |
| 71 | parser.add_argument("--device_id", type=int, default=0, help="device id") |
| 72 | parser.add_argument("--max_model_len", type=int, default=32768, help="max model length") |
| 73 | parser.add_argument("--num_layers", type=int, default=1, help="model num layers") |
| 74 | parser.add_argument("--mp_num", type=int, default=1, help="number of model parallel") |
| 75 | parser.add_argument( |
| 76 | "--cache_dtype", |
| 77 | type=str, |
| 78 | default="bfloat16", |
| 79 | choices=["uint8", "bfloat16", "block_wise_fp8"], |
| 80 | help="cache dtype", |
| 81 | ) |
| 82 | parser.add_argument( |
| 83 | "--default_dtype", |
| 84 | type=str, |
| 85 | default="bfloat16", |
| 86 | choices=["float16", "bfloat16", "uint8"], |
| 87 | help="paddle default dtype, swap_cache_batch only support float16、bfloat16 and uint8 now", |
| 88 | ) |
| 89 | parser.add_argument("--key_cache_shape", type=str, default="", help="key cache shape") |
| 90 | parser.add_argument("--value_cache_shape", type=str, default="", help="value cache shape") |
| 91 | parser.add_argument("--cache_queue_port", type=int, default=9923, help="cache queue port") |
| 92 | parser.add_argument("--enable_splitwise", type=int, default=0, help="enable splitwise ") |
| 93 | parser.add_argument("--pod_ip", type=str, default="0.0.0.0", help="pod ip") |
| 94 | parser.add_argument( |
| 95 | "--engine_worker_queue_port", |
| 96 | type=int, |
| 97 | default=9923, |
| 98 | help="engine worker queue port", |
| 99 | ) |
| 100 | parser.add_argument("--num_cpu_blocks", type=int, default=4, help="cpu cache block number") |
| 101 | parser.add_argument( |
| 102 | "--protocol", |
| 103 | type=str, |
| 104 | default="ipc", |
| 105 | help="cache transfer protocol, only support ipc now", |
| 106 | ) |
| 107 | parser.add_argument("--local_data_parallel_id", type=int, default=0) |
| 108 | parser.add_argument("--rdma_port", type=str, default="", help="rmda port") |
| 109 | parser.add_argument( |
| 110 | "--speculative_config", |
| 111 | type=json.loads, |
| 112 | default="{}", |
| 113 | help="speculative config", |
| 114 | ) |
| 115 | parser.add_argument("--create_cache_tensor", action="store_true") |
| 116 | parser.add_argument( |
no test coverage detected