wehub-resource-sync
106 lines
3.3 KiB
Python
106 lines
3.3 KiB
Python
import warnings
|
|
from typing import Any
|
|
|
|
import click
|
|
import modal
|
|
from fouroversix.utils import DataType, MatmulBackend, QuantizeBackend, ScaleRule
|
|
|
|
from ..resources import app
|
|
from .coordinators import LocalEvaluationCoordinator, ModalEvaluationCoordinator
|
|
from .utils import EvaluationFramework, PTQMethod
|
|
|
|
|
|
@click.command()
|
|
@click.option(
|
|
"--activation-scale-rule",
|
|
"--a-scale-rule",
|
|
type=ScaleRule,
|
|
default=ScaleRule.mse,
|
|
)
|
|
@click.option("--detach", is_flag=True)
|
|
@click.option("--device", type=str, default="cuda")
|
|
@click.option("--dtype", type=DataType, default=DataType.nvfp4)
|
|
@click.option(
|
|
"--eval-framework",
|
|
"-f",
|
|
type=EvaluationFramework,
|
|
default=EvaluationFramework.lm_eval,
|
|
)
|
|
@click.option("--group-name", type=str, default=None)
|
|
@click.option("--limit", type=int, default=None)
|
|
@click.option("--matmul-backend", type=MatmulBackend, default=None)
|
|
@click.option("--max-length", type=int, default=None)
|
|
@click.option("--modal", is_flag=True)
|
|
@click.option("--modal-gpu", type=str)
|
|
@click.option("--model-name", "-m", type=str, multiple=True, required=True)
|
|
@click.option("--ptq-method", "-p", type=PTQMethod, multiple=True, required=True)
|
|
@click.option("--quantize-backend", type=QuantizeBackend, default=None)
|
|
@click.option("--task", "-t", type=str, multiple=True, default=["wikitext"])
|
|
@click.option("--trust-remote-code", is_flag=True)
|
|
@click.option(
|
|
"--weight-scale-rule",
|
|
"--w-scale-rule",
|
|
type=ScaleRule,
|
|
default=ScaleRule.mse,
|
|
)
|
|
@click.option("--weight-scale-2d", "--w-scale-2d", is_flag=True)
|
|
def cli(
|
|
*,
|
|
detach: bool,
|
|
group_name: str | None,
|
|
modal_gpu: str,
|
|
**kwargs: dict[str, Any],
|
|
) -> None:
|
|
activation_scale_rule = kwargs.get("activation_scale_rule")
|
|
dtype = kwargs.get("dtype")
|
|
weight_scale_rule = kwargs.get("weight_scale_rule")
|
|
|
|
model_names = kwargs.pop("model_name")
|
|
ptq_methods = kwargs.pop("ptq_method")
|
|
tasks = kwargs.pop("task")
|
|
use_modal = kwargs.pop("modal")
|
|
|
|
# Expand shortcuts
|
|
if model_names[0] == "llamaqwen":
|
|
model_names = [
|
|
"meta-llama/Llama-3.2-1B",
|
|
"meta-llama/Llama-3.1-8B",
|
|
"meta-llama/Llama-3.1-70B",
|
|
"Qwen/Qwen3-1.7B",
|
|
"Qwen/Qwen3-8B",
|
|
"Qwen/Qwen3-32B",
|
|
]
|
|
|
|
if isinstance(tasks, tuple):
|
|
tasks = list(tasks)
|
|
|
|
if dtype == DataType.mxfp4 and (
|
|
not activation_scale_rule.is_static() or not weight_scale_rule.is_static()
|
|
):
|
|
msg = (
|
|
"MXFP4 quantization only supports static scale rules. Setting "
|
|
"activation_scale_rule and weight_scale_rule to static_6..."
|
|
)
|
|
warnings.warn(msg, stacklevel=1)
|
|
|
|
kwargs["activation_scale_rule"] = ScaleRule.static_6
|
|
kwargs["weight_scale_rule"] = ScaleRule.static_6
|
|
|
|
if use_modal:
|
|
with modal.enable_output(), app.run(detach=detach):
|
|
coordinator = ModalEvaluationCoordinator(group_name_str=group_name or "")
|
|
coordinator.start.remote(
|
|
model_names,
|
|
ptq_methods,
|
|
tasks,
|
|
modal_gpu=modal_gpu,
|
|
**kwargs,
|
|
)
|
|
else:
|
|
coordinator = LocalEvaluationCoordinator(group_name)
|
|
coordinator.start(model_names, ptq_methods, tasks, **kwargs)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli()
|