wehub-resource-sync
122 lines
3.8 KiB
Python
122 lines
3.8 KiB
Python
"""Command line entrypoint of configuration generation."""
|
|
|
|
from pathlib import Path
|
|
from typing import Union
|
|
|
|
from mlc_llm.interface.gen_config import CONV_TEMPLATES, gen_config
|
|
from mlc_llm.interface.help import HELP
|
|
from mlc_llm.model import MODELS
|
|
from mlc_llm.quantization import QUANTIZATION
|
|
from mlc_llm.support.argparse import ArgumentParser
|
|
from mlc_llm.support.auto_config import detect_config, detect_model_type
|
|
|
|
|
|
def main(argv):
|
|
"""Parse command line argumennts and call `mlc_llm.compiler.gen_config`."""
|
|
parser = ArgumentParser("MLC LLM Configuration Generator")
|
|
|
|
def _parse_output(path: Union[str, Path]) -> Path:
|
|
path = Path(path)
|
|
if not path.is_dir():
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
return path
|
|
|
|
parser.add_argument(
|
|
"config",
|
|
type=detect_config,
|
|
help=HELP["config"] + " (required)",
|
|
)
|
|
parser.add_argument(
|
|
"--quantization",
|
|
type=str,
|
|
required=True,
|
|
choices=list(QUANTIZATION.keys()),
|
|
help=HELP["quantization"] + " (required, choices: %(choices)s)",
|
|
)
|
|
parser.add_argument(
|
|
"--model-type",
|
|
type=str,
|
|
default="auto",
|
|
choices=["auto", *list(MODELS.keys())],
|
|
help=HELP["model_type"] + ' (default: "%(default)s", choices: %(choices)s)',
|
|
)
|
|
parser.add_argument(
|
|
"--conv-template",
|
|
type=str,
|
|
required=True,
|
|
choices=list(CONV_TEMPLATES),
|
|
help=HELP["conv_template"] + " (required, choices: %(choices)s)",
|
|
)
|
|
parser.add_argument(
|
|
"--context-window-size",
|
|
type=int,
|
|
default=None,
|
|
help=HELP["context_window_size"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--sliding-window-size",
|
|
type=int,
|
|
default=None,
|
|
help=HELP["sliding_window_size"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--prefill-chunk-size",
|
|
type=int,
|
|
default=None,
|
|
help=HELP["prefill_chunk_size"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--attention-sink-size",
|
|
type=int,
|
|
default=None,
|
|
help=HELP["attention_sink_size"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--tensor-parallel-shards",
|
|
type=int,
|
|
default=None,
|
|
help=HELP["tensor_parallel_shards"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--pipeline-parallel-stages",
|
|
type=int,
|
|
default=None,
|
|
help=HELP["pipeline_parallel_stages"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--disaggregation",
|
|
type=bool,
|
|
default=None,
|
|
help=HELP["disaggregation"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--max-batch-size",
|
|
type=int,
|
|
default=128,
|
|
help=HELP["max_batch_size"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
"-o",
|
|
type=_parse_output,
|
|
required=True,
|
|
help=HELP["output_gen_mlc_chat_config"] + " (required)",
|
|
)
|
|
parsed = parser.parse_args(argv)
|
|
model = detect_model_type(parsed.model_type, parsed.config)
|
|
gen_config(
|
|
config=parsed.config,
|
|
model=model,
|
|
quantization=QUANTIZATION[parsed.quantization],
|
|
conv_template=parsed.conv_template,
|
|
context_window_size=parsed.context_window_size,
|
|
sliding_window_size=parsed.sliding_window_size,
|
|
prefill_chunk_size=parsed.prefill_chunk_size,
|
|
attention_sink_size=parsed.attention_sink_size,
|
|
tensor_parallel_shards=parsed.tensor_parallel_shards,
|
|
pipeline_parallel_stages=parsed.pipeline_parallel_stages,
|
|
disaggregation=parsed.disaggregation,
|
|
max_batch_size=parsed.max_batch_size,
|
|
output=parsed.output,
|
|
)
|