[Frontend] Add FlexibleArgumentParser to support both underscore and dash in names (#5718)

This commit is contained in:
Michael Goin
2024-06-20 19:00:13 -04:00
committed by GitHub
parent 3f3b6b2150
commit 8065a7e220
22 changed files with 72 additions and 45 deletions

View File

@@ -1,4 +1,3 @@
import argparse
import os
import sys
from typing import Optional
@@ -10,6 +9,7 @@ from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.aqlm import (
dequantize_weight, generic_dequantize_gemm, get_int_dtype,
optimized_dequantize_gemm)
from vllm.utils import FlexibleArgumentParser
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
@@ -137,7 +137,7 @@ def dequant_test(k: int, parts: torch.Tensor, nbooks: int, bits: int) -> None:
def main():
parser = argparse.ArgumentParser(description="Benchmark aqlm performance.")
parser = FlexibleArgumentParser(description="Benchmark aqlm performance.")
# Add arguments
parser.add_argument("--nbooks",