440 lines
17 KiB
Python
440 lines
17 KiB
Python
"""
|
|
Reproducible speed / VRAM / accuracy benchmarks for the toolkit quantization
|
|
backends.
|
|
|
|
Compares bf16 against the custom OstrisLinear backends (convrot8, convrot4 for
|
|
now; add more qtypes to QTYPES as they land).
|
|
|
|
Measures, per qtype:
|
|
- layer inference latency across DiT-representative shapes (vs bf16),
|
|
eager and torch.compile'd
|
|
- layer training latency (forward + backward through the frozen layer),
|
|
eager and torch.compile'd
|
|
- VRAM on a transformer-ish block stack: resident weights, peak during a
|
|
no-grad forward, peak during a train step; forward/train-ckpt peaks also
|
|
under torch.compile
|
|
- accuracy drift vs bf16: output relative error per layer shape and
|
|
accumulated through the block stack
|
|
- weight reconstruction error and one-time quantize (conversion) time
|
|
|
|
Usage:
|
|
python scripts/test_quantizations.py --gpu 1
|
|
python scripts/test_quantizations.py --gpu 1 --qtypes bf16 convrot8
|
|
"""
|
|
|
|
import argparse
|
|
import math
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
# set cuda bus ordering to be pcie
|
|
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
import torch # noqa: E402
|
|
|
|
# (tokens, in_features, out_features) — FLUX/Wan-class projections
|
|
SPEED_SHAPES = [
|
|
(4096, 3072, 3072),
|
|
(4096, 3072, 12288),
|
|
(4096, 12288, 3072),
|
|
(1024, 3072, 12288),
|
|
]
|
|
|
|
# block stack used for the vram/drift tests (mimics a DiT block's linears)
|
|
VRAM_BLOCKS = 8
|
|
VRAM_BLOCK_SHAPES = [(3072, 12288), (12288, 3072), (3072, 3072), (3072, 3072)]
|
|
VRAM_TOKENS = 4096
|
|
|
|
QTYPES = [
|
|
"bf16", "qfloat8", "float8", "convrot8", "convrot4",
|
|
"convrotint7", "convrotint6", "convrotint5", "convrotint4", "convrotint3",
|
|
"convrotint2", "convrotbitnet", "convrotcomfyw4a4",
|
|
]
|
|
|
|
STACK_KEY = f"{VRAM_BLOCKS}-block stack"
|
|
|
|
|
|
def convert(module: torch.nn.Linear, qtype: str) -> torch.nn.Linear:
|
|
"""Quantize a linear with the given qtype. Returns the (possibly replaced)
|
|
module — quanto swaps the module object, the ostris backends convert in place."""
|
|
if qtype == "bf16":
|
|
return module
|
|
from toolkit.util.ostris_quant import convert_linear_to_ostris, get_ostris_quantizer
|
|
|
|
q = get_ostris_quantizer(qtype)
|
|
if q is not None:
|
|
assert convert_linear_to_ostris(module, q), f"conversion refused for {qtype}"
|
|
return module
|
|
|
|
# quanto / torchao qtypes go through the shared toolkit quantize flow; use a
|
|
# holder so quanto's module replacement has a parent to swap into
|
|
from optimum.quanto import freeze
|
|
from toolkit.util.quantize import get_qtype, quantize
|
|
|
|
holder = torch.nn.Sequential(module)
|
|
quantize(holder, weights=get_qtype(qtype))
|
|
freeze(holder)
|
|
return holder[0]
|
|
|
|
|
|
def fp_weight(module: torch.nn.Linear) -> torch.Tensor:
|
|
"""Dequantized weight in float32, whatever the backend."""
|
|
if hasattr(module, "dequantize_weight"):
|
|
return module.dequantize_weight().float()
|
|
w = module.weight
|
|
if hasattr(w, "dequantize"):
|
|
return w.dequantize().float()
|
|
return w.detach().float()
|
|
|
|
|
|
def bench(fn, iters: int, device) -> float:
|
|
for _ in range(max(3, iters // 5)):
|
|
fn()
|
|
torch.cuda.synchronize(device)
|
|
t0 = time.perf_counter()
|
|
for _ in range(iters):
|
|
fn()
|
|
torch.cuda.synchronize(device)
|
|
return (time.perf_counter() - t0) / iters * 1000 # ms
|
|
|
|
|
|
def gb(nbytes: int) -> str:
|
|
return f"{nbytes / 1e9:6.2f} GB"
|
|
|
|
|
|
def make_layer(k: int, n: int, device) -> torch.nn.Linear:
|
|
lin = torch.nn.Linear(k, n, bias=True, dtype=torch.bfloat16, device=device)
|
|
with torch.no_grad():
|
|
lin.weight.mul_(0.02)
|
|
# the train benches model lora-style training: base frozen, grads flow to
|
|
# the input only. without this, bf16/quanto accumulate weight grads that
|
|
# inflate every later vram measurement
|
|
lin.requires_grad_(False)
|
|
return lin
|
|
|
|
|
|
def make_stack(device) -> torch.nn.ModuleList:
|
|
# default nn.Linear init (~1/sqrt(in) std) so block branches contribute at a
|
|
# realistic O(1) scale to the residual stream — scaling weights down further
|
|
# makes accumulated quantization drift look artificially tiny
|
|
torch.manual_seed(0)
|
|
blocks = torch.nn.ModuleList()
|
|
for _ in range(VRAM_BLOCKS):
|
|
blocks.append(torch.nn.ModuleList([
|
|
torch.nn.Linear(k, n, bias=True, dtype=torch.bfloat16, device=device)
|
|
for k, n in VRAM_BLOCK_SHAPES
|
|
]))
|
|
# frozen base (see make_layer)
|
|
blocks.requires_grad_(False)
|
|
return blocks
|
|
|
|
|
|
def block_forward(b, h):
|
|
# pre-norm residual block like a real transformer, so activations stay at a
|
|
# sane scale and quantization drift accumulates realistically across depth
|
|
r = torch.nn.functional.layer_norm(h, h.shape[-1:])
|
|
r = b[0](r) # 3072 -> 12288
|
|
r = torch.nn.functional.gelu(r)
|
|
h = h + b[1](r) # 12288 -> 3072
|
|
r = torch.nn.functional.layer_norm(h, h.shape[-1:])
|
|
return h + b[3](b[2](r)) # 3072 -> 3072 -> 3072
|
|
|
|
|
|
def stack_forward(blocks, x, checkpoint=False):
|
|
h = x
|
|
for b in blocks:
|
|
if checkpoint:
|
|
h = torch.utils.checkpoint.checkpoint(
|
|
block_forward, b, h, use_reentrant=False
|
|
)
|
|
else:
|
|
h = block_forward(b, h)
|
|
return h
|
|
|
|
|
|
def run_speed(qtype: str, device, iters: int, results: dict):
|
|
for m, k, n in SPEED_SHAPES:
|
|
torch.manual_seed(0)
|
|
lin = make_layer(k, n, device)
|
|
lin = convert(lin, qtype)
|
|
x = torch.randn(m, k, device=device, dtype=torch.bfloat16)
|
|
|
|
with torch.no_grad():
|
|
t_inf = bench(lambda: lin(x), iters, device)
|
|
|
|
def train_step():
|
|
xi = x.detach().requires_grad_(True)
|
|
lin(xi).sum().backward()
|
|
|
|
t_train = bench(train_step, max(10, iters // 3), device)
|
|
results[(qtype, "inf", (m, k, n))] = t_inf
|
|
results[(qtype, "train", (m, k, n))] = t_train
|
|
|
|
# compiled variants (compilation happens during bench warmup, so it
|
|
# isn't charged to the timing; a backend that won't compile records
|
|
# nothing and shows as '-')
|
|
lin_c = torch.compile(lin, dynamic=False)
|
|
try:
|
|
with torch.no_grad():
|
|
results[(qtype, "inf_comp", (m, k, n))] = bench(
|
|
lambda: lin_c(x), iters, device
|
|
)
|
|
except Exception as e:
|
|
print(f" [{qtype}] compiled inference failed for {m}x{k}->{n}: {e}")
|
|
|
|
def train_step_c():
|
|
xi = x.detach().requires_grad_(True)
|
|
lin_c(xi).sum().backward()
|
|
|
|
try:
|
|
results[(qtype, "train_comp", (m, k, n))] = bench(
|
|
train_step_c, max(10, iters // 3), device
|
|
)
|
|
except Exception as e:
|
|
print(f" [{qtype}] compiled train failed for {m}x{k}->{n}: {e}")
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def _stack_fwd_peak(blocks, x, device, base) -> int:
|
|
# warm up first so lazy-init allocations (and compilation) are not counted
|
|
# as steady-state peak
|
|
with torch.no_grad():
|
|
stack_forward(blocks, x)
|
|
torch.cuda.synchronize(device)
|
|
torch.cuda.reset_peak_memory_stats(device)
|
|
with torch.no_grad():
|
|
stack_forward(blocks, x)
|
|
torch.cuda.synchronize(device)
|
|
return torch.cuda.max_memory_allocated(device) - base
|
|
|
|
|
|
def _stack_train_peak(blocks, x, device, base, checkpoint) -> int:
|
|
# frozen base; grads flow to the input like lora training
|
|
def train_step():
|
|
xi = x.detach().requires_grad_(True)
|
|
stack_forward(blocks, xi, checkpoint).float().pow(2).mean().backward()
|
|
|
|
train_step()
|
|
torch.cuda.synchronize(device)
|
|
torch.cuda.reset_peak_memory_stats(device)
|
|
train_step()
|
|
torch.cuda.synchronize(device)
|
|
return torch.cuda.max_memory_allocated(device) - base
|
|
|
|
|
|
def run_vram(qtype: str, device, results: dict):
|
|
torch.cuda.empty_cache()
|
|
base = torch.cuda.memory_allocated(device)
|
|
|
|
blocks = make_stack(device)
|
|
for b in blocks:
|
|
for i in range(len(b)):
|
|
b[i] = convert(b[i], qtype)
|
|
torch.cuda.empty_cache()
|
|
results[(qtype, "vram_weights")] = torch.cuda.memory_allocated(device) - base
|
|
|
|
x = torch.randn(VRAM_TOKENS, 3072, device=device, dtype=torch.bfloat16)
|
|
|
|
results[(qtype, "vram_fwd_peak")] = _stack_fwd_peak(blocks, x, device, base)
|
|
# real training checkpoints, but the plain train peak still gets reported
|
|
results[(qtype, "vram_train_peak")] = _stack_train_peak(blocks, x, device, base, False)
|
|
results[(qtype, "vram_train_ckpt_peak")] = _stack_train_peak(blocks, x, device, base, True)
|
|
|
|
# same peaks with every linear compiled (mirrors the trainer's block compile)
|
|
for b in blocks:
|
|
for i in range(len(b)):
|
|
b[i] = torch.compile(b[i], dynamic=False)
|
|
try:
|
|
results[(qtype, "vram_fwd_peak_comp")] = _stack_fwd_peak(blocks, x, device, base)
|
|
results[(qtype, "vram_train_ckpt_peak_comp")] = _stack_train_peak(
|
|
blocks, x, device, base, True
|
|
)
|
|
except Exception as e:
|
|
print(f" [{qtype}] compiled vram measurement failed: {e}")
|
|
|
|
blocks = x = None # release before the allocator accounting of the next run
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def run_drift(qtype: str, device, results: dict):
|
|
"""Output error vs the bf16 reference, per layer shape and through the stack."""
|
|
for m, k, n in SPEED_SHAPES:
|
|
torch.manual_seed(0)
|
|
lin = make_layer(k, n, device)
|
|
x = torch.randn(m, k, device=device, dtype=torch.bfloat16)
|
|
with torch.no_grad():
|
|
y_ref = lin(x).float()
|
|
lin = convert(lin, qtype)
|
|
y_q = lin(x).float()
|
|
results[(qtype, "drift", (m, k, n))] = ((y_q - y_ref).norm() / y_ref.norm()).item()
|
|
torch.cuda.empty_cache()
|
|
|
|
blocks = make_stack(device)
|
|
x = torch.randn(VRAM_TOKENS, 3072, device=device, dtype=torch.bfloat16)
|
|
with torch.no_grad():
|
|
y_ref = stack_forward(blocks, x).float()
|
|
for b in blocks:
|
|
for i in range(len(b)):
|
|
b[i] = convert(b[i], qtype)
|
|
y_q = stack_forward(blocks, x).float()
|
|
results[(qtype, "drift", STACK_KEY)] = ((y_q - y_ref).norm() / y_ref.norm()).item()
|
|
blocks = x = None
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def run_quality_and_quantize_time(qtype: str, device, results: dict):
|
|
torch.manual_seed(0)
|
|
lin = make_layer(3072, 3072, device)
|
|
w0 = lin.weight.detach().float().clone()
|
|
torch.cuda.synchronize(device)
|
|
t0 = time.perf_counter()
|
|
lin = convert(lin, qtype)
|
|
torch.cuda.synchronize(device)
|
|
results[(qtype, "quantize_ms")] = (time.perf_counter() - t0) * 1000
|
|
if qtype == "bf16":
|
|
results[(qtype, "weight_err")] = 0.0
|
|
else:
|
|
wq = fp_weight(lin)
|
|
results[(qtype, "weight_err")] = ((wq - w0).norm() / w0.norm()).item()
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def print_speed_table(title: str, kind: str, qts, results):
|
|
# speedups always reference EAGER bf16, compiled kinds included, so the
|
|
# comp columns answer "what do I gain over plain bf16"
|
|
ref_kind = kind.removesuffix("_comp")
|
|
print(f"\n=== {title} (ms; speedup vs eager bf16) ===")
|
|
print(f"{'M x K -> N':<22}" + "".join(f"{qt:>18}" for qt in qts))
|
|
for shape in SPEED_SHAPES:
|
|
m, k, n = shape
|
|
row = f"{f'{m} x {k} -> {n}':<22}"
|
|
ref = results.get(("bf16", ref_kind, shape))
|
|
for qt in qts:
|
|
t = results.get((qt, kind, shape))
|
|
if t is None:
|
|
row += f"{'-':>18}"
|
|
continue
|
|
is_self_ref = qt == "bf16" and kind == ref_kind
|
|
sp = f" ({ref / t:4.2f}x)" if ref and not is_self_ref else " " * 8
|
|
row += f"{t:8.3f}ms{sp}"
|
|
print(row)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description=__doc__)
|
|
ap.add_argument("--gpu", type=int, default=0, help="cuda device id to run on")
|
|
ap.add_argument("--qtypes", nargs="+", default=QTYPES, help=f"subset of {QTYPES}")
|
|
ap.add_argument("--iters", type=int, default=50, help="timing iterations per case")
|
|
args = ap.parse_args()
|
|
|
|
device = torch.device(f"cuda:{args.gpu}")
|
|
torch.cuda.set_device(device)
|
|
props = torch.cuda.get_device_properties(device)
|
|
print(f"device: cuda:{args.gpu} ({props.name}, sm_{props.major}{props.minor}, "
|
|
f"{props.total_memory / 1e9:.0f} GB)")
|
|
print(f"torch {torch.__version__}\n")
|
|
|
|
# warm the toolkit import chain (module imports + custom-op registration) so it
|
|
# isn't charged to the first qtype's quantize timing
|
|
if any(qt != "bf16" for qt in args.qtypes):
|
|
from toolkit.util.ostris_quant import get_ostris_quantizer
|
|
for qt in args.qtypes:
|
|
if qt != "bf16":
|
|
get_ostris_quantizer(qt)
|
|
|
|
# many distinct module instances share one forward code object; the default
|
|
# per-code cache limit (8) would silently fall back to eager and corrupt the
|
|
# compiled columns
|
|
torch._dynamo.config.cache_size_limit = 4096
|
|
|
|
results = {}
|
|
for qt in args.qtypes:
|
|
print(f"benchmarking {qt} ...")
|
|
torch._dynamo.reset() # drop the previous qtype's compiled artifacts
|
|
run_quality_and_quantize_time(qt, device, results)
|
|
run_drift(qt, device, results)
|
|
run_speed(qt, device, args.iters, results)
|
|
run_vram(qt, device, results)
|
|
|
|
qts = args.qtypes
|
|
print_speed_table("layer latency, inference", "inf", qts, results)
|
|
print_speed_table("layer latency, inference (compiled)", "inf_comp", qts, results)
|
|
print_speed_table("layer latency, train fwd+bwd", "train", qts, results)
|
|
print_speed_table("layer latency, train fwd+bwd (compiled)", "train_comp", qts, results)
|
|
|
|
print(f"\n=== vram on the block stack ({VRAM_BLOCKS} blocks, {VRAM_TOKENS} tokens) ===")
|
|
print(f"{'':<28}" + "".join(f"{qt:>18}" for qt in qts))
|
|
for key, label in (("vram_weights", "weights resident"),
|
|
("vram_fwd_peak", "peak, no-grad fwd"),
|
|
("vram_train_peak", "peak, train step"),
|
|
("vram_train_ckpt_peak", "peak, train step (ckpt)"),
|
|
("vram_fwd_peak_comp", "peak, no-grad fwd (comp)"),
|
|
("vram_train_ckpt_peak_comp", "peak, train ckpt (comp)")):
|
|
row = f"{label:<28}"
|
|
for qt in qts:
|
|
v = results.get((qt, key))
|
|
row += f"{gb(v):>18}" if v is not None else f"{'-':>18}"
|
|
print(row)
|
|
|
|
print("\n=== accuracy drift vs bf16 (output rel err, no-grad) ===")
|
|
print(f"{'':<28}" + "".join(f"{qt:>18}" for qt in qts))
|
|
for shape in SPEED_SHAPES + [STACK_KEY]:
|
|
label = f"{shape[0]} x {shape[1]} -> {shape[2]}" if isinstance(shape, tuple) else shape
|
|
row = f"{label:<28}"
|
|
for qt in qts:
|
|
row += f"{results[(qt, 'drift', shape)]:>18.5f}"
|
|
print(row)
|
|
|
|
print("\n=== quantization ===")
|
|
print(f"{'':<28}" + "".join(f"{qt:>18}" for qt in qts))
|
|
row = f"{'weight rel err':<28}"
|
|
for qt in qts:
|
|
row += f"{results[(qt, 'weight_err')]:>18.5f}"
|
|
print(row)
|
|
row = f"{'quantize time (ms)':<28}"
|
|
for qt in qts:
|
|
row += f"{results[(qt, 'quantize_ms')]:>18.1f}"
|
|
print(row)
|
|
|
|
# ---- clean per-qtype breakdown: speed (geomean over shapes) + accuracy ----
|
|
def geomean_speedup(qt, kind):
|
|
# every speedup references EAGER bf16 (compiled kinds included), so the
|
|
# comp columns answer "what do I gain over plain bf16"
|
|
ref_kind = kind.removesuffix("_comp")
|
|
logs = []
|
|
for shape in SPEED_SHAPES:
|
|
ref = results.get(("bf16", ref_kind, shape))
|
|
t = results.get((qt, kind, shape))
|
|
if ref and t:
|
|
logs.append(math.log(ref / t))
|
|
return math.exp(sum(logs) / len(logs)) if logs else None
|
|
|
|
def fmt_speed(v):
|
|
return f"{v:.2f}x" if v is not None else "-"
|
|
|
|
print("\n=== summary (speed = geomean speedup vs bf16; drift lower is better) ===")
|
|
print(f"{'':<18}{'inference':>12}{'inference comp':>16}{'train':>12}{'train comp':>12}"
|
|
f"{'accuracy drift':>16}{'max vram':>12}{'max vram comp':>15}")
|
|
for qt in qts:
|
|
# real training checkpoints, so the ckpt peak is the meaningful train
|
|
# number; the no-grad fwd peak still matters for sampling
|
|
max_vram = max(results[(qt, "vram_fwd_peak")], results[(qt, "vram_train_ckpt_peak")])
|
|
fwd_c = results.get((qt, "vram_fwd_peak_comp"))
|
|
ckpt_c = results.get((qt, "vram_train_ckpt_peak_comp"))
|
|
max_vram_comp = max(fwd_c, ckpt_c) if fwd_c is not None and ckpt_c is not None else None
|
|
print(f"{qt:<18}"
|
|
f"{fmt_speed(geomean_speedup(qt, 'inf')):>12}"
|
|
f"{fmt_speed(geomean_speedup(qt, 'inf_comp')):>16}"
|
|
f"{fmt_speed(geomean_speedup(qt, 'train')):>12}"
|
|
f"{fmt_speed(geomean_speedup(qt, 'train_comp')):>12}"
|
|
f"{results[(qt, 'drift', STACK_KEY)]:>16.5f}"
|
|
f"{gb(max_vram).strip():>12}"
|
|
f"{(gb(max_vram_comp).strip() if max_vram_comp is not None else '-'):>15}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|