wehub-resource-sync
109 lines
3.1 KiB
Python
109 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import click
|
|
import modal
|
|
|
|
from ..resources import Dependency, app, get_image
|
|
|
|
img = get_image(dependencies=[Dependency.transformer_engine, Dependency.fouroversix])
|
|
|
|
with img.imports():
|
|
import torch
|
|
import torch.utils.benchmark as benchmark
|
|
from fouroversix import QuantizationConfig, QuantizeBackend, RoundStyle, ScaleRule
|
|
from fouroversix.quantize.frontend import AVAILABLE_BACKENDS
|
|
|
|
|
|
def run_speedtest(
|
|
*,
|
|
block_scale_2d: bool = False,
|
|
input_shape: str = "1024,1024",
|
|
repeats: int = 100,
|
|
rht: bool = False,
|
|
round_style: str = "nearest",
|
|
scale_rule: str = "mse",
|
|
transpose: bool = False,
|
|
) -> None:
|
|
"""Test speed on a B200 on Modal."""
|
|
|
|
input_shape = tuple(int(dim.strip()) for dim in input_shape.split(","))
|
|
x = torch.randn(input_shape, dtype=torch.bfloat16, device="cuda")
|
|
|
|
print("Testing with config:")
|
|
print(f"- block_scale_2d: {block_scale_2d}")
|
|
print(f"- input_shape: {input_shape}")
|
|
print(f"- rht: {rht}")
|
|
print(f"- round_style: {round_style}")
|
|
print(f"- scale_rule: {scale_rule}")
|
|
print(f"- transpose: {transpose}")
|
|
print()
|
|
|
|
for backend in [
|
|
QuantizeBackend.cuda,
|
|
QuantizeBackend.transformer_engine,
|
|
QuantizeBackend.triton,
|
|
QuantizeBackend.pytorch,
|
|
]:
|
|
config = QuantizationConfig(
|
|
backend=backend,
|
|
block_scale_2d=block_scale_2d,
|
|
rht=rht,
|
|
round_style=RoundStyle(round_style),
|
|
scale_rule=ScaleRule(scale_rule),
|
|
transpose=transpose,
|
|
)
|
|
|
|
backend_cls = AVAILABLE_BACKENDS[backend]
|
|
print(f"{backend.value}: ", end="")
|
|
|
|
if not backend_cls.is_available():
|
|
print("Not available")
|
|
continue
|
|
|
|
if not backend_cls.is_supported(x, config):
|
|
print("Not supported")
|
|
continue
|
|
|
|
config = QuantizationConfig(
|
|
backend=backend,
|
|
rht=rht,
|
|
round_style=RoundStyle(round_style),
|
|
scale_rule=ScaleRule(scale_rule),
|
|
)
|
|
|
|
t = benchmark.Timer(
|
|
setup="from fouroversix import quantize_to_fp4",
|
|
stmt="quantize_to_fp4(x, config)",
|
|
globals={"x": x, "config": config},
|
|
)
|
|
|
|
print(f"{t.timeit(repeats).mean * 1000:.4f}ms")
|
|
|
|
|
|
@app.function(image=img, cpu=4, memory=8 * 1024, gpu="B200")
|
|
def run_speedtest_on_modal(**kwargs: dict[str, Any]) -> None:
|
|
run_speedtest(**kwargs)
|
|
|
|
|
|
@click.command()
|
|
@click.option("--block-scale-2d", is_flag=True)
|
|
@click.option("--input-shape", type=str, default="1024,1024")
|
|
@click.option("--modal", is_flag=True)
|
|
@click.option("--repeats", type=int, default=100)
|
|
@click.option("--rht", is_flag=True)
|
|
@click.option("--round-style", type=RoundStyle, default=RoundStyle.nearest)
|
|
@click.option("--scale-rule", type=ScaleRule, default=ScaleRule.mse)
|
|
@click.option("--transpose", is_flag=True)
|
|
def cli(**kwargs: dict[str, Any]) -> None:
|
|
if kwargs.pop("modal"):
|
|
with modal.enable_output(), app.run():
|
|
run_speedtest_on_modal.remote(**kwargs)
|
|
else:
|
|
run_speedtest(**kwargs)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli()
|