OneCompression/onecomp/cli.py at d3d3e926e4ccbf6e42c3eacb87f65def6c54cef3 · FujitsuResearch/OneCompression · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
"""

Copyright 2025-2026 Fujitsu Ltd.

Author: Keiji Kimura

"""

import argparse

from .__version__ import __version__


def main():
    parser = argparse.ArgumentParser(
        prog="onecomp",
        description="OneComp: One-liner LLM quantization (AutoBit + QEP)",
    )
    parser.add_argument(
        "model_id",
        help="Hugging Face model ID or local path",
    )
    parser.add_argument(
        "--wbits",
        type=float,
        default=None,
        help="target bitwidth (default: auto-estimated from VRAM)",
    )
    parser.add_argument(
        "--total-vram-gb",
        type=float,
        default=None,
        help="VRAM budget in GB for bitwidth estimation (default: auto-detect)",
    )
    parser.add_argument(
        "--groupsize",
        type=int,
        default=128,
        help="GPTQ group size (default: 128, -1 to disable)",
    )
    parser.add_argument(
        "--device",
        default="cuda:0",
        help="device to place the model on (default: cuda:0)",
    )
    parser.add_argument(
        "--no-qep",
        action="store_true",
        help="disable QEP (enabled by default)",
    )
    parser.add_argument(
        "--no-eval",
        action="store_true",
        help="skip perplexity and accuracy evaluation",
    )
    parser.add_argument(
        "--eval-original",
        action="store_true",
        help="also evaluate the original (unquantized) model",
    )
    parser.add_argument(
        "--save-dir",
        default="auto",
        help='save directory (default: auto-generated, "none" to skip)',
    )
    parser.add_argument(
        "--check-env",
        action="store_true",
        help=(
            "Print an environment and memory report before quantization. "
            "Exits with code 1 if OOM risk is 'danger'."
        ),
    )
    parser.add_argument(
        "--version",
        action="version",
        version=f"%(prog)s {__version__}",
    )

    args = parser.parse_args()

    save_dir = None if args.save_dir.lower() == "none" else args.save_dir

    # Lazy import to keep --help fast
    from .runner import Runner  # pylint: disable=import-outside-toplevel

    if args.check_env:
        import sys  # pylint: disable=import-outside-toplevel
        from .utils.vram_estimator import (  # pylint: disable=import-outside-toplevel
            check_environment,
            print_env_report,
        )

        env_result = check_environment(
            args.model_id,
            total_vram_gb=args.total_vram_gb,
            group_size=args.groupsize,
            save_dir=save_dir if isinstance(save_dir, str) and save_dir != "auto" else None,
        )
        print_env_report(env_result, total_vram_gb_override=args.total_vram_gb)
        if env_result.risk == "danger":
            sys.exit(1)

    Runner.auto_run(
        model_id=args.model_id,
        wbits=args.wbits,
        total_vram_gb=args.total_vram_gb,
        groupsize=args.groupsize,
        device=args.device,
        qep=not args.no_qep,
        evaluate=not args.no_eval,
        eval_original_model=args.eval_original,
        save_dir=save_dir,
    )