__main__.py

# Copyright (c) 2024 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import json
import os
import re
import sys

from auto_round.auto_scheme import AutoScheme
from auto_round.compressors import BaseCompressor
from auto_round.eval.eval_cli import EvalArgumentParser, eval, eval_task_by_task
from auto_round.eval.evaluation import run_model_evaluation
from auto_round.schemes import PRESET_SCHEMES
from auto_round.utils import (
    clear_memory,
    get_device_and_parallelism,
    get_model_dtype,
    parse_layer_config_arg,
)

RECIPES = {
    "default": {"batch_size": 8, "iters": 200, "seqlen": 2048, "nsamples": 128, "lr": None},
    "best": {"batch_size": 8, "iters": 1000, "seqlen": 2048, "nsamples": 512, "lr": None},
    "light": {"batch_size": 8, "iters": 50, "seqlen": 2048, "nsamples": 128, "lr": 5e-3},
    "fast": {"batch_size": 4, "iters": 200, "seqlen": 512, "nsamples": 128, "lr": None},
}


class BasicArgumentParser(argparse.ArgumentParser):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.add_argument(
            "model",
            default=None,
            nargs="?",
            help="Path to the pre-trained model or model identifier from huggingface.co/models. "
            "Examples: 'facebook/opt-125m', 'bert-base-uncased', or local path like '/path/to/model'",
        )
        basic = self.add_argument_group("Basic Arguments")
        basic.add_argument(
            "--model_name",
            "--model",
            "--model_name_or_path",
            default="facebook/opt-125m",
            help="Path to the pre-trained model or model identifier from huggingface.co/models. "
            "Examples: 'facebook/opt-125m', 'bert-base-uncased', or local path like '/path/to/model'",
        )
        basic.add_argument("--model_dtype", default=None, help="model dtype used to load the pre-trained model")
        basic.add_argument(
            "--platform",
            default="hf",
            help="Platform to load the pre-trained model. Options: [hf, model_scope]."
            " hf stands for huggingface and model_scope stands for model scope.",
        )
        basic.add_argument(
            "--scheme",
            default="W4A16",
            type=str,
            # choices=["W4A16", "W2A16", "W3A16", "W8A16", "MXFP4", "MXFP8", "NVFP4", "FPW8A16", "FP8_STATIC"],
            help="Quantization scheme to use. "
            "W4A16: 4-bit weights with 16-bit activations (default). "
            "Other options include W2A16, W3A16, W8A16 for different bit widths, "
            "and MXFP4/MXFP8/NVFP4 for different data type.",
        )
        basic.add_argument(
            "--batch_size",
            "--train_bs",
            "--bs",
            default=None,
            type=int,
            help="The batch size for tuning/calibration."
            "Larger batch sizes may improve stability but require more memory.",
        )
        basic.add_argument(
            "--avg_bits", "--target_bits", default=None, type=float, help="for auto scheme, number of avg weight bits"
        )
        basic.add_argument(
            "--options", default=None, type=str, help="for auto scheme, options for auto scheme, e.g. 'W4A16,W8A16'"
        )

        basic.add_argument(
            "--iters",
            "--iter",
            default=None,
            type=int,
            help="Number of iterations to tune each block. "
            "More iterations may lead to better quantization quality but take longer.",
        )
        basic.add_argument(
            "--seqlen",
            "--seq_len",
            default=None,
            type=int,
            help="Sequence length of the calibration samples"
            "Longer sequences capture more context but use more memory.",
        )
        basic.add_argument(
            "--nsamples",
            "--nsample",
            default=None,
            type=int,
            help="Number of calibration samples to use for quantization.",
        )
        basic.add_argument(
            "--device_map",
            "--device",
            "--devices",
            default="0",
            type=str,
            help="The device to be used for tuning. "
            "Currently, device settings support CPU, GPU, and HPU."
            "The default is set to cuda:0,"
            "allowing for automatic detection and switch to HPU or CPU."
            "set --device 0,1,2 to use multiple cards.",
        )
        basic.add_argument(
            "--dataset",
            default="NeelNanda/pile-10k",
            type=str,
            help="Calibration dataset for quantization. "
            "Should be a dataset from huggingface datasets or local path. ",
        )
        basic.add_argument("--seed", default=42, type=int, help="Random seed for reproducibility.")
        basic.add_argument("--adam", action="store_true", help="Use Adam optimizer instead of SignSGD.")
        basic.add_argument(
            "--low_gpu_mem_usage",
            action="store_true",
            help="Enable memory-efficient mode by offloading intermediate features to CPU. "
            "Useful when working with large models that don't fit in GPU memory.",
        )
        basic.add_argument(
            "--low_cpu_mem_usage",
            action="store_true",
            help=(
                "Deprecated: low CPU memory mode is enabled by default. "
                "This flag is kept only for backward compatibility and has no effect "
                "beyond explicitly re-enabling the default behavior."
            ),
        )
        basic.add_argument(
            "--disable_low_cpu_mem_usage",
            action="store_true",
            help=("Disable low CPU memory mode. " "Use this flag to turn off the default low CPU memory behavior."),
        )
        basic.add_argument(
            "--format",
            "--formats",
            default="auto_round",
            type=str,
            help="Output format for the quantized model."
            "'auto_round' is the recommended format"
            "use command `auto_round list format` to show all supported formats with support scheme.",
        )
        basic.add_argument(
            "--output_dir",
            default="./tmp_autoround",
            type=str,
            help="Directory to save the quantized model and related files",
        )
        basic.add_argument(
            "--not_use_best_mse",
            action="store_true",
            help="Disable using the iteration with best MSE loss during tuning.",
        )
        basic.add_argument(
            "--enable_torch_compile", action="store_true", help="Enable PyTorch compilation for faster execution. "
        )
        basic.add_argument(
            "--disable_trust_remote_code",
            action="store_true",
            help="Disable trusting remote code when loading models. "
            "Use for security if you don't trust the model source.",
        )

        tuning = self.add_argument_group("Tuning Arguments")
        tuning.add_argument(
            "--ignore_scale_zp_bits",
            action="store_true",
            help="for auto scheme whether ignore scale zp bits calculation ",
        )
        tuning.add_argument(
            "--lr",
            default=None,
            type=float,
            help="Learning rate for tuning. " "If None, automatically sets to 1.0/iters. ",
        )
        tuning.add_argument(
            "--minmax_lr",
            default=None,
            type=float,
            help="Learning rate specifically for min-max tuning. " "If None, uses the same value as --lr. ",
        )
        tuning.add_argument(
            "--momentum",
            default=0,
            type=float,
            help="Momentum factor for the optimizer. Default is 0 (no momentum).",
        )
        tuning.add_argument(
            "--gradient_accumulate_steps",
            default=1,
            type=int,
            help="Number of steps to accumulate gradients before updating weights. "
            "Effectively increases batch size without requiring more GPU memory. "
            "Useful for large models with limited memory.",
        )
        tuning.add_argument(
            "--nblocks",
            default=1,
            type=int,
            help="Number of blocks to tune simultaneously. "
            "Higher values may speed up tuning but require more memory. "
            "Recommended to keep at 1 for stability with large models.",
        )
        tuning.add_argument(
            "--scale_dtype",
            default=None,
            choices=["fp16", "float16", "bf16", "bfloat16", "fp32", "float32"],
            help="Data type for quantization scales. "
            "fp16/bf16: lower memory, fp32: higher precision. "
            "Choose based on your hardware support and accuracy requirements.",
        )
        tuning.add_argument(
            "--disable_amp",
            action="store_true",
            help="Disable Automatic Mixed Precision (AMP). "
            "AMP speeds up training but may affect numerical stability in some cases.",
        )
        tuning.add_argument(
            "--disable_minmax_tuning",
            action="store_true",
            help="Disable weight min-max range tuning. "
            "Not recommended as it may significantly reduce quantization accuracy.",
        )
        tuning.add_argument(
            "--enable_norm_bias_tuning", action="store_true", help="Enable normalization layer bias tuning. "
        )
        tuning.add_argument(
            "--disable_quanted_input",
            action="store_true",
            help="Use original (non-quantized) inputs for each block instead of"
            " quantized outputs from previous blocks. ",
        )
        tuning.add_argument(
            "--to_quant_block_names",
            default=None,
            type=str,
            help="Specific blocks to quantize, separated by commas. "
            "Example: 'block1,block2,block3'. "
            "If None, all blocks will be quantized.",
        )
        tuning.add_argument(
            "--enable_alg_ext",
            action="store_true",
            help="Enable experimental algorithms that may provide better quantization results. "
            "These are newer methods that might improve accuracy but are less tested.",
        )
        tuning.add_argument(
            "--disable_deterministic_algorithms",
            action="store_true",
            help="deprecated, disable torch deterministic algorithms.",
        )
        tuning.add_argument(
            "--enable_deterministic_algorithms",
            action="store_true",
            help="Enable PyTorch deterministic algorithms for reproducible results. ",
        )
        group_opt_rtn = tuning.add_mutually_exclusive_group()
        group_opt_rtn.add_argument(
            "--disable_opt_rtn",
            action="store_const",
            const=True,
            dest="disable_opt_rtn",
            default=None,
            help="Disable optimization for RTN (Round-To-Nearest) mode when iters=0. "
            "RTN is fast but less accurate; keeping optimization enabled is recommended.",
        )

        group_opt_rtn.add_argument(
            "--enable_opt_rtn",
            action="store_const",
            const=False,
            dest="disable_opt_rtn",
            help="Enable optimization for RTN mode when iters=0.",
        )

        scheme = self.add_argument_group("Scheme Arguments")
        scheme.add_argument("--bits", default=None, type=int, help="Number of bits for weight quantization. ")
        scheme.add_argument(
            "--group_size",
            default=None,
            type=lambda s: int(s) if s.lstrip("-").isdigit() else tuple([int(x.strip()) for x in s.split(",")]),
            help="Group size for weight quantization.",
        )
        scheme.add_argument("--asym", action="store_true", help="Use asymmetric quantization instead of symmetric.")
        scheme.add_argument(
            "--act_asym", action="store_true", help="Use asymmetric quantization for activation instead of symmetric."
        )
        scheme.add_argument(
            "--data_type",
            "--dtype",
            default=None,
            help="Data type for quantization. Options: 'int' for integer, 'mx_fp' for mixed floating-point, etc.",
        )
        scheme.add_argument(
            "--act_bits",
            default=None,
            type=int,
            help="Number of bits for activation quantization. "
            "Activation quantization significantly impacts performance and accuracy.",
        )
        scheme.add_argument(
            "--act_group_size",
            default=None,
            type=int,
            help="Group size for activation quantization. " "Similar to weight group size but for activations.",
        )
        scheme.add_argument(
            "--act_data_type", "--act_dtype", default=None, type=str, help="Data type for activation quantization. "
        )
        scheme.add_argument(
            "--disable_act_dynamic", action="store_true", help="Use static instead of dynamic activation quantization. "
        )
        scheme.add_argument(
            "--layer_config",
            default=None,
            type=str,
            help="Per-layer quantization config for missing tensors (e.g., MTP layers) as a JSON string. "
            "Keys are name prefixes, values are config dicts with optional bits/group_size/sym. "
            'Example: "{mtp:{bits:8,data_type:int},mtp.fc:{bits:16}}". '
            "These settings are saved to extra_config and override the global quantization config.",
        )
        scheme.add_argument(
            "--shared_layers",
            type=str,
            nargs="+",
            action="append",
            default=None,
            help="[mix-precision] ensure that listed layers are using same data type for quantization",
        )
        scheme.add_argument(
            "--quant_lm_head",
            action="store_true",
            help="Quantize the lm_head. " "Usually kept in higher precision for better output quality.",
        )
        scheme.add_argument(
            "--ignore_layers",
            "--fp_layers",
            default="",
            type=str,
            help="List of layer names to keep in original precision (not quantized). "
            "Useful for preserving critical layers. Separate multiple names with commas.",
        )
        scheme.add_argument(
            "--static_kv_dtype",
            default=None,
            type=str,
            choices=["fp8", "float8_e4m3fn"],
            help="Data type for static quantize key and value. ",
        )

        scheme.add_argument(
            "--static_attention_dtype",
            default=None,
            type=str,
            choices=["fp8", "float8_e4m3fn"],
            help="Data type for static quantize attention. ",
        )
        scheme.add_argument(
            "--rotation_type",
            default=None,
            type=str,
            choices=["hadamard", "random_hadamard", "quarot_hadamard"],
            help="Research feature: applies a rotation (e.g., Hadamard) to reduce activation/weight outliers",
        )
        gguf = self.add_argument_group("Double Quant Arguments")
        gguf.add_argument(
            "--super_group_size", default=None, type=int, help="Super group size for double quantization."
        )
        gguf.add_argument(
            "--super_bits",
            default=None,
            type=int,
            help="Number of bits for scale and zero-point quantization in double quantization. ",
        )

        ## ======================= eval =======================
        eval_args = self.add_argument_group("eval arguments")
        eval_args.add_argument(
            "--tasks",
            "--task",
            nargs="?",
            const="lambada_openai,hellaswag,winogrande,piqa,mmlu,wikitext,truthfulqa_mc1,"
            "openbookqa,boolq,arc_easy,arc_challenge",
            default=None,
            help="LM-Evaluation-Harness tasks to run. "
            "Specify specific tasks like 'mmlu,wikitext' for custom evaluation.",
        )
        eval_args.add_argument("--eval_bs", default=None, type=int, help="Batch size for evaluation.")
        eval_args.add_argument(
            "--limit",
            type=float,
            default=None,
            metavar="N|0<N<1",
            help="Limit the number of examples per task. "
            "Integer: exact number of examples (e.g., 1000). "
            "Float between 0-1: fraction of total examples.",
        )
        eval_args.add_argument(
            "--eval_task_by_task", action="store_true", help="Evaluate tasks sequentially instead of batching. "
        )
        eval_args.add_argument(
            "--eval_backend",
            default="hf",
            type=str,
            choices=["hf", "vllm"],
            help="Backend to use for model evaluation. Use hf backend for evaluation by default.",
        )
        eval_args.add_argument(
            "--vllm_args",
            default=None,
            type=str,
            help="(for vllm) Custom vllm arguments in format: '--arg1=value1,--arg2=value2'. "
            "Example: '--tensor_parallel_size=2,--gpu_memory_utilization=0.9'",
        )
        eval_args.add_argument(
            "--eval_model_dtype",
            default=None,
            type=str,
            help="Torch data type for model loading during evaluation. "
            "Options: 'float16', 'bfloat16', 'float32'. "
            "Should match your hardware capabilities for best performance.",
        )
        eval_args.add_argument("--add_bos_token", action="store_true", help="add BOS token")

        ## ======================= MLLM =======================
        mllm_args = self.add_argument_group("Multimodal Large Language Model(MLLM) arguments")
        mllm_args.add_argument(
            "--mllm",
            action="store_true",
            help="[Deprecated] AutoRound now automatically detects and uses MLLM mode when needed.",
        )
        mllm_args.add_argument(
            "--quant_nontext_module",
            action="store_true",
            help="Quantize non-text modules (vision/audio/video components). "
            "Enables full multimodal model quantization but may affect visual quality.",
        )
        mllm_args.add_argument(
            "--extra_data_dir",
            default=None,
            type=str,
            help="Directory containing multimodal data (images/audio/videos). "
            "Can be a single directory or specify types: "
            "'image=/path/to/images,video=/path/to/videos,audio=/path/to/audio'. "
            "If not found locally, will attempt to download standard datasets.",
        )
        mllm_args.add_argument(
            "--template",
            default=None,
            type=str,
            help="Custom template for building training datasets. "
            "Useful for specialized multimodal tasks or custom data formats.",
        )

        ## ======================= diffusion model eval =======================
        diffusion_args = self.add_argument_group("diffusion model arguments")
        diffusion_args.add_argument(
            "--prompt_file",
            default=None,
            type=str,
            help="File containing prompts for evaluation, one per line. "
            "Use this for batch evaluation with multiple prompts.",
        )
        diffusion_args.add_argument(
            "--prompt",
            default=None,
            type=str,
            help="Single prompt for quick testing. " "Overrides prompt_file if both are specified.",
        )
        diffusion_args.add_argument(
            "--metrics",
            "--metric",
            default="clip",
            help="Evaluation metrics for generated images. "
            "'clip': CLIP score measuring text-image alignment. "
            "'clip-iqa': CLIP-based image quality assessment. "
            "'imagereward': Learned metric for image quality.",
        )
        diffusion_args.add_argument(
            "--image_save_dir",
            default="./tmp_image_save",
            type=str,
            help="Directory to save generated images during evaluation. " "Useful for visual inspection of results.",
        )
        diffusion_args.add_argument(
            "--guidance_scale",
            default=7.5,
            type=float,
            help="Classifier-free guidance scale for diffusion models. "
            "Higher values (7-20) make the model follow the prompt more closely. "
            "Lower values give more creative/random results.",
        )

        diffusion_args.add_argument(
            "--num_inference_steps",
            default=50,
            type=int,
            help="Number of denoising steps in the diffusion process. "
            "More steps (50-100) usually give better quality but take longer. "
            "Fewer steps (10-30) are faster but lower quality.",
        )

        diffusion_args.add_argument(
            "--generator_seed",
            default=None,
            type=int,
            help="Random seed for image generation reproducibility. "
            "Using the same seed produces identical results across runs.",
        )


def list_item():
    args = argparse.ArgumentParser()
    args.add_argument("item", type=str, help="item to list, e.g., format")
    args = args.parse_args()
    if args.item == "format" or args.item == "formats":
        from auto_round.formats import OutputFormat

        print("AutoRound supported output formats and quantization scheme:")
        print(OutputFormat.get_support_matrix())


def start(recipe="default"):
    recipe = RECIPES[recipe]
    parser = BasicArgumentParser()
    args = parser.parse_args()
    for k, v in recipe.items():
        if getattr(args, k) is None:
            setattr(args, k, v)
    tune(args)


def tune(args):
    assert args.model or args.model_name, "[model] or --model MODEL_NAME should be set."
    if args.model is None:
        args.model = args.model_name
    if args.eval_bs is None:
        args.eval_bs = "auto"
    from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
    from transformers.utils.versions import require_version

    if args.tasks is not None:
        require_version(
            "lm_eval>=0.4.2",
            "lm-eval is required for evaluation, please install it with `pip install 'lm-eval>=0.4.2'`",
        )

    from auto_round.utils import detect_device, get_library_version, logger

    if args.low_cpu_mem_usage:
        logger.warning(
            "`low_cpu_mem_usage` is deprecated and is now enabled by default. "
            "To disable it, use `--disable_low_cpu_mem_usage`."
        )

    if args.format is None:
        args.format = "auto_round"

    formats = args.format.lower().replace(" ", "").split(",")
    from auto_round.utils import SUPPORTED_FORMATS

    for format in formats:
        if format not in SUPPORTED_FORMATS:
            raise ValueError(f"{format} is not supported, we only support {SUPPORTED_FORMATS}")

    if "auto_gptq" in args.format and args.asym is True:
        logger.warning(
            "the auto_gptq kernel has issues with asymmetric quantization. "
            "It is recommended to use sym quantization or --format='auto_round'"
        )

    if "marlin" in args.format and args.asym is True:
        raise RuntimeError("marlin backend only supports sym quantization, please remove --asym")

    device_str, use_auto_mapping = get_device_and_parallelism(args.device_map)

    import torch

    if args.enable_torch_compile:
        logger.info(
            "`torch.compile` is enabled to reduce tuning costs. "
            "If it causes issues, you can disable it by removing `--enable_torch_compile` argument."
        )

    model_name = args.model
    if model_name[-1] == "/":
        model_name = model_name[:-1]
    logger.info(f"start to quantize {model_name}")

    from auto_round import AutoRound

    if "bloom" in model_name:
        args.low_gpu_mem_usage = False

    if args.quant_lm_head:
        for format in formats:
            if "auto_round" not in format and "fake" not in format:
                auto_round_formats = [s for s in SUPPORTED_FORMATS if s.startswith("auto_round")]
                raise ValueError(
                    f"{format} is not supported for lm-head quantization, please change to {auto_round_formats}"
                )

    enable_torch_compile = True if "--enable_torch_compile" in sys.argv else False
    sym = None  # the default value should be None now
    if args.asym:  # if the scheme is asym, how to set it to sym is an issue
        sym = False
    act_dynamic = None
    if args.disable_act_dynamic:
        act_dynamic = False
    scheme = args.scheme.upper()
    if scheme not in PRESET_SCHEMES:
        raise ValueError(f"{scheme} is not supported. only {PRESET_SCHEMES.keys()} are supported ")
    if args.disable_deterministic_algorithms:
        logger.warning(
            "default not use deterministic_algorithms. disable_deterministic_algorithms is deprecated,"
            " please use enable_deterministic_algorithms instead. "
        )

    from auto_round.compressors import (
        DiffusionExtraConfig,
        ExtraConfig,
        MLLMExtraConfig,
        SchemeExtraConfig,
        TuningExtraConfig,
    )

    extra_config = ExtraConfig()
    tuning_config = TuningExtraConfig(
        amp=not args.disable_amp,
        disable_opt_rtn=args.disable_opt_rtn,
        enable_alg_ext=args.enable_alg_ext,
        enable_minmax_tuning=not args.disable_minmax_tuning,
        enable_norm_bias_tuning=args.enable_norm_bias_tuning,
        enable_quanted_input=not args.disable_quanted_input,
        enable_deterministic_algorithms=args.enable_deterministic_algorithms,
        lr=args.lr,
        minmax_lr=args.minmax_lr,
        nblocks=args.nblocks,
        to_quant_block_names=args.to_quant_block_names,
        scale_dtype=args.scale_dtype,
    )
    scheme_config = SchemeExtraConfig(
        bits=args.bits,
        group_size=args.group_size,
        sym=sym,
        data_type=args.data_type,
        act_bits=args.act_bits,
        act_group_size=args.act_group_size,
        act_data_type=args.act_data_type,
        act_dynamic=act_dynamic,
        act_sym=None if not args.asym else False,
        super_bits=args.super_bits,
        super_group_size=args.super_group_size,
        quant_lm_head=args.quant_lm_head,
        ignore_layers=args.ignore_layers,
        static_kv_dtype=args.static_kv_dtype,
        static_attention_dtype=args.static_attention_dtype,
    )
    mllm_config = MLLMExtraConfig(
        quant_nontext_module=args.quant_nontext_module, extra_data_dir=args.extra_data_dir, template=args.template
    )
    diffusion_config = DiffusionExtraConfig(
        guidance_scale=args.guidance_scale,
        num_inference_steps=args.num_inference_steps,
        generator_seed=args.generator_seed,
    )
    extra_config.tuning_config = tuning_config
    extra_config.scheme_config = scheme_config
    extra_config.mllm_config = mllm_config
    extra_config.diffusion_config = diffusion_config

    layer_config = {}
    if args.layer_config:
        layer_config = parse_layer_config_arg(args.layer_config)
        args.layer_config = layer_config

    low_cpu_mem_usage = True
    if args.disable_low_cpu_mem_usage:
        low_cpu_mem_usage = False

    if args.avg_bits is not None:
        if args.options is None:
            raise ValueError("please set --options for auto scheme")
        if enable_torch_compile:
            logger.warning(
                "`enable_torch_compile=True` with AutoScheme may cause compile errors "
                "on some models. If so, try removing `--enable_torch_compile`."
            )
        scheme = AutoScheme(
            options=args.options,
            avg_bits=args.avg_bits,
            shared_layers=args.shared_layers,
            ignore_scale_zp_bits=args.ignore_scale_zp_bits,
            low_gpu_mem_usage=True,  # force it to be True as it uses much smaller vram but similar time cost
            low_cpu_mem_usage=low_cpu_mem_usage,
        )
    rot_config = None
    if args.rotation_type:
        from auto_round.experimental.transform.rotation_config import RotationConfig

        rot_config = RotationConfig(hadamard_type=args.rotation_type)

    autoround: BaseCompressor = AutoRound(
        model=model_name,
        platform=args.platform,
        scheme=scheme,
        dataset=args.dataset,
        iters=args.iters,
        seqlen=args.seqlen,
        nsamples=args.nsamples,
        batch_size=args.batch_size,
        gradient_accumulate_steps=args.gradient_accumulate_steps,
        low_gpu_mem_usage=args.low_gpu_mem_usage,
        low_cpu_mem_usage=low_cpu_mem_usage,
        device_map=args.device_map,
        enable_torch_compile=enable_torch_compile,
        seed=args.seed,
        not_use_best_mse=args.not_use_best_mse,
        enable_adam=args.adam,
        extra_config=extra_config,
        layer_config=layer_config,
        model_dtype=args.model_dtype,
        momentum=args.momentum,
        trust_remote_code=not args.disable_trust_remote_code,
        rotation_config=rot_config,
    )

    model_name = args.model.rstrip("/")

    if model_name.split("/")[-1].strip(".") == "" and "gguf" not in args.format:
        if autoround.group_size <= 0:
            if "fp" in autoround.act_data_type:
                suffix = f"afp{autoround.act_bits}"
            else:
                suffix = f"a{autoround.act_bits}"
        else:
            suffix = f"g{autoround.group_size}"
        export_dir = os.path.join(args.output_dir, f"w{autoround.bits}{suffix}")
    elif model_name.split("/")[-1].strip(".") == "" and "gguf" in args.format:
        export_dir = args.output_dir
    elif model_name.split("./")[-1].strip("./") != "" and "gguf" in args.format:
        export_dir = os.path.join(args.output_dir, model_name.split("/")[-1] + "-gguf")
    else:
        if isinstance(autoround.group_size, tuple):
            assert len(autoround.group_size) == 2, f"Only support 2D group_size, but get {autoround.group_size}"
            suffix = f"g{autoround.group_size[0]}x{autoround.group_size[1]}"
        else:
            if autoround.group_size <= 0:
                if "fp" in autoround.act_data_type:
                    suffix = f"afp{autoround.act_bits}"
                else:
                    suffix = f"a{autoround.act_bits}"
            else:
                suffix = f"g{autoround.group_size}"
        prefix = (
            autoround.data_type.lower().replace("_", "")
            if "int" not in autoround.data_type or "mx" in autoround.data_type
            else ""
        )
        export_dir = os.path.join(
            args.output_dir,
            model_name.split("/")[-1] + (f"-{prefix}" if prefix else "") + f"-w{autoround.bits}{suffix}",
        )

    # ======================= Quantize and save model =======================
    model, folders = autoround.quantize_and_save(export_dir, format=args.format)  # pylint: disable=E1101
    tokenizer = autoround.tokenizer  # pylint: disable=E1101

    model.eval()
    clear_memory()

    # ======================= Model evaluation =======================
    run_model_evaluation(model, tokenizer, autoround, folders, formats, device_str, args)


def setup_eval_parser():
    parser = EvalArgumentParser()
    args = parser.parse_args()
    return args


def run_eval():
    from auto_round.logger import logger
    from auto_round.utils import is_gguf_model, is_mllm_model

    args = setup_eval_parser()
    assert args.model or args.model_name, "[model] or --model MODEL_NAME should be set."

    if args.model is None:
        args.model = args.model_name
    if "llama" in args.model.lower() and not args.add_bos_token:
        logger.warning("set add_bos_token=True for llama model.")
        args.add_bos_token = True
    if not is_gguf_model(args.model) and is_mllm_model(args.model):
        args.mllm = True

    if args.eval_task_by_task:
        eval_task_by_task(
            model=args.model,
            device=args.device_map,
            limit=args.limit,
            tasks=args.tasks,
            batch_size=args.eval_bs,
            trust_remote_code=not args.disable_trust_remote_code,
            eval_model_dtype=args.eval_model_dtype,
            add_bos_token=args.add_bos_token,
        )
    else:
        eval(args)


def run():
    if "list" in sys.argv or "--list" in sys.argv:
        if "list" in sys.argv:
            sys.argv.remove("list")
        if "--list" in sys.argv:
            sys.argv.remove("--list")
        list_item()
        exit()
    if "--eval" in sys.argv or "eval" in sys.argv:
        if "--eval" in sys.argv:
            sys.argv.remove("--eval")
        if "eval" in sys.argv:
            sys.argv.remove("eval")
        run_eval()
    else:
        start()


def run_best():
    start("best")


def run_light():
    start("light")


def run_fast():
    start("fast")


def run_mllm():
    run()


if __name__ == "__main__":
    run()