Add convrot4 and convrot8 quantizations to the ui. A lot of improvements to the convrot quantization methods.
This commit is contained in:
Jaret Burkett
parent
933ca1c517
commit
095d6e7418
|
|
@ -2109,6 +2109,12 @@ class BaseSDTrainProcess(BaseTrainProcess):
|
|||
if user_set_cache_limit:
|
||||
torch._dynamo.config.cache_size_limit = cache_size_limit
|
||||
torch._dynamo.config.suppress_errors = False
|
||||
# torch 2.9 inductor bug: the new memory-coalescing tiling analysis
|
||||
# crashes on some dynamic-shape index expressions (sympy PowByNatural
|
||||
# "assert p >= 0", seen with Qwen Image). The analysis doesn't apply
|
||||
# to dynamic shapes anyway, so turn it off.
|
||||
if hasattr(torch._inductor.config.triton, 'coalesce_tiling_analysis'):
|
||||
torch._inductor.config.triton.coalesce_tiling_analysis = False
|
||||
|
||||
compile_mode = getattr(self.model_config, 'compile_mode', 'default')
|
||||
compile_dynamic = getattr(self.model_config, 'compile_dynamic', True)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,339 @@
|
|||
"""
|
||||
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)
|
||||
- layer training latency (forward + backward through the frozen layer)
|
||||
- VRAM on a transformer-ish block stack: resident weights, peak during a
|
||||
no-grad forward, peak during a train step
|
||||
- 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
|
||||
|
||||
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", "orbit4", "orbitvq4", "convrot8", "convrot4"]
|
||||
|
||||
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)
|
||||
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
|
||||
]))
|
||||
return blocks
|
||||
|
||||
|
||||
def stack_forward(blocks, x):
|
||||
# pre-norm residual blocks like a real transformer, so activations stay at a
|
||||
# sane scale and quantization drift accumulates realistically across depth
|
||||
h = x
|
||||
for b in blocks:
|
||||
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:])
|
||||
h = h + b[3](b[2](r)) # 3072 -> 3072 -> 3072
|
||||
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
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
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)
|
||||
|
||||
# no-grad forward peak (sampling); warm up first so lazy-init allocations 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)
|
||||
results[(qtype, "vram_fwd_peak")] = torch.cuda.max_memory_allocated(device) - base
|
||||
|
||||
# train step peak (frozen base; grads flow to the input like lora training)
|
||||
def train_step():
|
||||
xi = x.detach().requires_grad_(True)
|
||||
stack_forward(blocks, xi).float().pow(2).mean().backward()
|
||||
|
||||
train_step()
|
||||
torch.cuda.synchronize(device)
|
||||
torch.cuda.reset_peak_memory_stats(device)
|
||||
train_step()
|
||||
torch.cuda.synchronize(device)
|
||||
results[(qtype, "vram_train_peak")] = torch.cuda.max_memory_allocated(device) - base
|
||||
|
||||
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):
|
||||
print(f"\n=== {title} (ms; speedup vs 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", kind, shape))
|
||||
for qt in qts:
|
||||
t = results.get((qt, kind, shape))
|
||||
if t is None:
|
||||
row += f"{'-':>18}"
|
||||
continue
|
||||
sp = f" ({ref / t:4.2f}x)" if ref and qt != "bf16" 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)
|
||||
|
||||
results = {}
|
||||
for qt in args.qtypes:
|
||||
print(f"benchmarking {qt} ...")
|
||||
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, train fwd+bwd", "train", 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")):
|
||||
row = f"{label:<28}"
|
||||
for qt in qts:
|
||||
row += f"{gb(results[(qt, key)]):>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):
|
||||
logs = []
|
||||
for shape in SPEED_SHAPES:
|
||||
ref = results.get(("bf16", 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 float("nan")
|
||||
|
||||
print("\n=== summary (speed = geomean speedup vs bf16; drift lower is better) ===")
|
||||
print(f"{'':<12}{'inference':>18}{'train':>18}{'accuracy drift':>20}{'max vram':>16}")
|
||||
for qt in qts:
|
||||
max_vram = max(results[(qt, "vram_fwd_peak")], results[(qt, "vram_train_peak")])
|
||||
print(f"{qt:<12}"
|
||||
f"{geomean_speedup(qt, 'inf'):>17.2f}x"
|
||||
f"{geomean_speedup(qt, 'train'):>17.2f}x"
|
||||
f"{results[(qt, 'drift', STACK_KEY)]:>20.5f}"
|
||||
f"{gb(max_vram):>16}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -66,11 +66,9 @@ F4_MAX = 6.0
|
|||
F8_E4M3_MAX = 448.0
|
||||
BLOCK = 16 # nvfp4 scale block
|
||||
|
||||
_E2M1_VALS = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
|
||||
_E2M1_EDGES = [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0]
|
||||
|
||||
_hadamard_cache = {}
|
||||
_vals_cache = {}
|
||||
_edges_cache = {}
|
||||
_skip_warned = set()
|
||||
|
||||
|
|
@ -162,15 +160,22 @@ def dequantize_nvfp4(
|
|||
K: int,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
# single-pass triton path when available: the torch chain below is ~7 full-size
|
||||
# elementwise passes with fp32 intermediates, which made every convrot4 training
|
||||
# backward pay a dequant cost comparable to the gradient matmul itself
|
||||
if _triton_available() and packed.is_cuda and dtype in (torch.bfloat16, torch.float16, torch.float32):
|
||||
return _fp4_dequant_op(
|
||||
packed, scales.view(torch.uint8), pts.reshape(1).view(torch.uint8),
|
||||
str(dtype).split(".")[-1],
|
||||
)
|
||||
codes = torch.stack([packed & 15, packed >> 4], dim=-1).view(rows, K)
|
||||
vals = _cached(
|
||||
_vals_cache,
|
||||
str(packed.device),
|
||||
lambda: torch.tensor(_E2M1_VALS, device=packed.device),
|
||||
)
|
||||
mag = torch.index_select(vals, 0, (codes & 7).flatten().to(torch.int32)).view(
|
||||
rows, K
|
||||
# the lookup table is built inline (NOT module-cached): this function runs inside
|
||||
# custom-op backwards, which torch.compile traces with fake tensors where a
|
||||
# pre-existing real tensor is illegal; an in-trace constructed constant is fine
|
||||
vals = torch.tensor(
|
||||
[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], device=packed.device
|
||||
)
|
||||
mag = torch.index_select(vals, 0, (codes & 7).flatten().to(torch.int32)).view(rows, K)
|
||||
v = mag * torch.where((codes & 8) > 0, -1.0, 1.0)
|
||||
v = v.view(rows, K // BLOCK, BLOCK) * (scales.float() * pts).unsqueeze(-1)
|
||||
return v.view(rows, K).to(dtype)
|
||||
|
|
@ -229,7 +234,9 @@ def _get_kernel():
|
|||
x = tl.load(x_ptr + pid_m * K + offs, mask=mask, other=0.0).to(tl.float32)
|
||||
xb = tl.reshape(x, (BLOCK_K // 16, 16))
|
||||
amax = tl.max(tl.abs(xb), axis=1)
|
||||
scale8 = (amax / (6.0 * pts)).to(tl.float8e4nv)
|
||||
# clamp to the e4m3 max so an oversized block can't overflow the fp8
|
||||
# conversion (matters now that activations use a fixed pts=1)
|
||||
scale8 = tl.minimum(amax / (6.0 * pts), 448.0).to(tl.float8e4nv)
|
||||
denom = scale8.to(tl.float32) * pts
|
||||
denom = tl.where(denom > 0, denom, 1.0)
|
||||
# note: triton fp32 division on this backend is ~1ulp off ieee (even with
|
||||
|
|
@ -298,8 +305,11 @@ def _nvfp4_act_quant_op(x: torch.Tensor) -> list[torch.Tensor]:
|
|||
rows, K = x.shape
|
||||
rows_pad = -(-rows // 16) * 16
|
||||
x = x.contiguous()
|
||||
pts = x.float().abs().amax() / (F4_MAX * F8_E4M3_MAX)
|
||||
pts = torch.where(pts > 0, pts, torch.ones_like(pts))
|
||||
# activations use a FIXED per-tensor scale of 1: their block scales fit the
|
||||
# e4m3 range natively (unlike tiny weight magnitudes, which keep dynamic pts),
|
||||
# measured quality-neutral on realistic distributions — and it removes a full
|
||||
# activation read (global amax) plus a device sync from every forward
|
||||
pts = torch.ones((), device=x.device)
|
||||
packed = torch.empty(rows_pad, K // 2, device=x.device, dtype=torch.uint8)
|
||||
if rows_pad != rows:
|
||||
packed[rows:].zero_()
|
||||
|
|
@ -333,11 +343,12 @@ def quantize_nvfp4_fused(x: torch.Tensor, blocked_scales: bool = False):
|
|||
rows, K = x.shape
|
||||
if not (_triton_available() and x.is_cuda and K % 16 == 0):
|
||||
if blocked_scales:
|
||||
# match the custom op: rows padded to a multiple of 16 for _scaled_mm
|
||||
# match the custom op: rows padded to a multiple of 16 for _scaled_mm,
|
||||
# fixed pts=1 for activations (see _nvfp4_act_quant_op)
|
||||
rows_pad = -(-rows // 16) * 16
|
||||
if rows_pad != rows:
|
||||
x = F.pad(x, (0, 0, 0, rows_pad - rows))
|
||||
packed, scales, pts = quantize_nvfp4(x)
|
||||
packed, scales, pts = quantize_nvfp4(x, pts=torch.ones((), device=x.device))
|
||||
return packed, to_blocked(scales), pts
|
||||
return quantize_nvfp4(x)
|
||||
if blocked_scales:
|
||||
|
|
@ -353,6 +364,77 @@ def quantize_nvfp4_fused(x: torch.Tensor, blocked_scales: bool = False):
|
|||
return packed, scales, pts
|
||||
|
||||
|
||||
|
||||
# ---------------- fp4 dequant kernel (backward hot path) ----------------
|
||||
|
||||
_dequant_kernel = None
|
||||
|
||||
|
||||
def _get_dequant_kernel():
|
||||
global _dequant_kernel
|
||||
if _dequant_kernel is not None:
|
||||
return _dequant_kernel
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
@triton.jit
|
||||
def nvfp4_dequant_kernel(
|
||||
q_ptr, s_ptr, pts_ptr, out_ptr, K,
|
||||
BLOCK_B: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0)
|
||||
pid_k = tl.program_id(1)
|
||||
pts = tl.load(pts_ptr)
|
||||
offs_b = pid_k * BLOCK_B + tl.arange(0, BLOCK_B)
|
||||
bmask = offs_b < K // 2
|
||||
byte = tl.load(q_ptr + row * (K // 2) + offs_b, mask=bmask, other=0)
|
||||
codes = tl.interleave(byte & 15, byte >> 4) # (2*BLOCK_B,), column order
|
||||
m = (codes & 7).to(tl.float32)
|
||||
# arithmetic e2m1 decode ([0, .5, 1, 1.5, 2, 3, 4, 6]), exact
|
||||
mag = tl.where(m < 2, m * 0.5, tl.exp2(tl.floor(m / 2) - 1) * (1 + (m % 2) * 0.5))
|
||||
v = tl.where((codes & 8) > 0, -mag, mag)
|
||||
n_s: tl.constexpr = (2 * BLOCK_B) // 16
|
||||
offs_s = pid_k * n_s + tl.arange(0, n_s)
|
||||
s = tl.load(s_ptr + row * (K // 16) + offs_s, mask=offs_s < K // 16, other=0.0)
|
||||
vb = tl.reshape(v, (n_s, 16)) * (s.to(tl.float32) * pts)[:, None]
|
||||
out = tl.reshape(vb, (2 * BLOCK_B,))
|
||||
offs_v = pid_k * (2 * BLOCK_B) + tl.arange(0, 2 * BLOCK_B)
|
||||
tl.store(out_ptr + row * K + offs_v, out.to(out_ptr.dtype.element_ty), mask=offs_v < K)
|
||||
|
||||
_dequant_kernel = nvfp4_dequant_kernel
|
||||
return _dequant_kernel
|
||||
|
||||
|
||||
# custom op so the kernel stays opaque where it matters most: inside the fp4
|
||||
# training op's registered backward, which torch.compile traces with fake tensors
|
||||
@torch.library.custom_op("ostris::convrot_fp4_dequant", mutates_args=())
|
||||
def _fp4_dequant_op(
|
||||
packed: torch.Tensor,
|
||||
scales_u8: torch.Tensor,
|
||||
pts_u8: torch.Tensor,
|
||||
out_dtype: str,
|
||||
) -> torch.Tensor:
|
||||
rows, half = packed.shape
|
||||
out = torch.empty(rows, half * 2, device=packed.device, dtype=getattr(torch, out_dtype))
|
||||
kernel = _get_dequant_kernel()
|
||||
block_b = 1024
|
||||
grid = (rows, -(-half // block_b))
|
||||
kernel[grid](
|
||||
packed.contiguous(),
|
||||
scales_u8.view(torch.float8_e4m3fn),
|
||||
pts_u8.view(torch.float32),
|
||||
out, half * 2,
|
||||
BLOCK_B=block_b, num_warps=4,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
@_fp4_dequant_op.register_fake
|
||||
def _fp4_dequant_fake(packed, scales_u8, pts_u8, out_dtype):
|
||||
rows, half = packed.shape
|
||||
return torch.empty(rows, half * 2, device=packed.device, dtype=getattr(torch, out_dtype))
|
||||
|
||||
|
||||
# ---------------- backend ----------------
|
||||
|
||||
|
||||
|
|
@ -379,6 +461,67 @@ def _fp4_gemm_supported(device) -> bool:
|
|||
return supported
|
||||
|
||||
|
||||
# the fp4 training-path linear: forward VALUE is the real fp4 tensor-core gemm
|
||||
# (bit-identical to the inference path), gradient is the straight-through estimate
|
||||
# d y / d x_rot ~= dequant(W'). the backward re-dequantizes the weight from the fp4
|
||||
# codes instead of saving a bf16 copy per layer (F.linear in the old fake-quant path
|
||||
# retained a dequantized weight for every layer of the graph — a full bf16 model of
|
||||
# extra train-step vram) and x is not saved at all.
|
||||
@torch.library.custom_op("ostris::convrot_fp4_linear_ste", mutates_args=())
|
||||
def _fp4_linear_ste_op(
|
||||
x2d: torch.Tensor,
|
||||
qdata: torch.Tensor,
|
||||
scales_u8: torch.Tensor,
|
||||
scales_blocked_u8: torch.Tensor,
|
||||
pts_u8: torch.Tensor,
|
||||
bias: Optional[torch.Tensor],
|
||||
out_dtype: str,
|
||||
) -> torch.Tensor:
|
||||
m = x2d.shape[0]
|
||||
aq, a_scales_blocked, a_pts = quantize_nvfp4_fused(x2d, blocked_scales=True)
|
||||
out = torch._scaled_mm(
|
||||
aq.view(torch.float4_e2m1fn_x2),
|
||||
qdata.view(torch.float4_e2m1fn_x2).t(),
|
||||
a_scales_blocked.view(torch.float8_e4m3fn),
|
||||
scales_blocked_u8.view(torch.float8_e4m3fn),
|
||||
out_dtype=getattr(torch, out_dtype),
|
||||
)
|
||||
if out.shape[0] != m:
|
||||
out = out[:m]
|
||||
s = (a_pts * pts_u8.view(torch.float32).reshape(())).to(out.dtype)
|
||||
if bias is not None:
|
||||
return torch.addcmul(bias, out, s)
|
||||
return out * s
|
||||
|
||||
|
||||
@_fp4_linear_ste_op.register_fake
|
||||
def _fp4_linear_ste_fake(x2d, qdata, scales_u8, scales_blocked_u8, pts_u8, bias, out_dtype):
|
||||
return torch.empty(
|
||||
x2d.shape[0], qdata.shape[0], device=x2d.device, dtype=getattr(torch, out_dtype)
|
||||
)
|
||||
|
||||
|
||||
def _fp4_linear_ste_setup(ctx, inputs, output):
|
||||
x2d, qdata, scales_u8, scales_blocked_u8, pts_u8, bias, out_dtype = inputs
|
||||
ctx.save_for_backward(qdata, scales_u8, pts_u8)
|
||||
|
||||
|
||||
def _fp4_linear_ste_backward(ctx, grad):
|
||||
qdata, scales_u8, pts_u8 = ctx.saved_tensors
|
||||
out_f, in_half = qdata.shape
|
||||
w = dequantize_nvfp4(
|
||||
qdata, scales_u8.view(torch.float8_e4m3fn),
|
||||
pts_u8.view(torch.float32).reshape(()),
|
||||
out_f, in_half * 2, grad.dtype,
|
||||
)
|
||||
return grad @ w, None, None, None, None, None, None
|
||||
|
||||
|
||||
_fp4_linear_ste_op.register_autograd(
|
||||
_fp4_linear_ste_backward, setup_context=_fp4_linear_ste_setup
|
||||
)
|
||||
|
||||
|
||||
class ConvRotQuantizer(OstrisQuantizer):
|
||||
"""ConvRot W4A4 backend. One instance per qtype, shareable across modules."""
|
||||
|
||||
|
|
@ -451,16 +594,24 @@ class ConvRotQuantizer(OstrisQuantizer):
|
|||
def forward(self, module, x: torch.Tensor) -> torch.Tensor:
|
||||
rot = module.cr_rot_size
|
||||
in_f, out_f = module.in_features, module.out_features
|
||||
x_rot = rotate(x, rot)
|
||||
x2d = x_rot.reshape(-1, in_f)
|
||||
m = x2d.shape[0]
|
||||
m = x.numel() // in_f
|
||||
|
||||
if x.requires_grad:
|
||||
# training path: straight-through fake-quant of the activation so
|
||||
# adapters see the same W4A4 numerics as deployment, then a
|
||||
# differentiable bf16 matmul against the dequantized rotated weight.
|
||||
# gated on requires_grad alone (not is_grad_enabled) so both passes of
|
||||
# gradient checkpointing take the same branch and recompute identically
|
||||
# training path, gated on requires_grad alone (not is_grad_enabled) so
|
||||
# both passes of gradient checkpointing take the same branch
|
||||
if _fp4_gemm_supported(x.device):
|
||||
# fp4 tensor-core forward (bit-identical to the inference path)
|
||||
# with a straight-through analytic backward
|
||||
x2d = rotate(x, rot).reshape(-1, in_f)
|
||||
out = _fp4_linear_ste_op(
|
||||
x2d, module.cr_qdata, module.cr_scales,
|
||||
module.cr_scales_blocked, module.cr_pts, module.bias,
|
||||
str(x.dtype).split(".")[-1],
|
||||
)
|
||||
return out.reshape(*x.shape[:-1], out_f)
|
||||
# no fp4 hardware: straight-through fake-quant of the activation and a
|
||||
# differentiable bf16 matmul against the dequantized rotated weight
|
||||
x2d = rotate(x, rot).reshape(-1, in_f)
|
||||
with torch.no_grad():
|
||||
aq, a_scales, a_pts = quantize_nvfp4(x2d.detach())
|
||||
x_dq = dequantize_nvfp4(aq, a_scales, a_pts, m, in_f, x.dtype)
|
||||
|
|
@ -471,8 +622,14 @@ class ConvRotQuantizer(OstrisQuantizer):
|
|||
|
||||
if _fp4_gemm_supported(x.device):
|
||||
# row padding for _scaled_mm happens inside the act-quant op (compile
|
||||
# safety); slice the mm output back to m rows (a contiguous prefix)
|
||||
aq, a_scales_blocked, a_pts = quantize_nvfp4_fused(x2d, blocked_scales=True)
|
||||
# safety); slice the mm output back to m rows (a contiguous prefix).
|
||||
# NOTE: no fused rotate+quant here (unlike convrot8): the e2m1 packing
|
||||
# needs tl.reshape/tl.split of dot-derived tensors, which this triton
|
||||
# backend miscompiles when the dot sits in/after a loop (element order
|
||||
# scrambles). the int8 kernel avoids those primitives and fuses safely.
|
||||
aq, a_scales_blocked, a_pts = quantize_nvfp4_fused(
|
||||
rotate(x, rot).reshape(-1, in_f), blocked_scales=True
|
||||
)
|
||||
out = torch._scaled_mm(
|
||||
aq.view(torch.float4_e2m1fn_x2),
|
||||
module.cr_qdata.view(torch.float4_e2m1fn_x2).t(),
|
||||
|
|
@ -491,7 +648,7 @@ class ConvRotQuantizer(OstrisQuantizer):
|
|||
|
||||
# no fp4 hardware: dequantized matmul (correct, no speedup)
|
||||
w = self._dequantize_rotated(module, x.dtype)
|
||||
out = F.linear(x2d, w, module.bias)
|
||||
out = F.linear(rotate(x, rot).reshape(-1, in_f), w, module.bias)
|
||||
return out.reshape(*x.shape[:-1], out_f)
|
||||
|
||||
|
||||
|
|
@ -803,22 +960,26 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
|
|||
def forward(self, module, x: torch.Tensor) -> torch.Tensor:
|
||||
rot = module.cr8_rot_size
|
||||
in_f, out_f = module.in_features, module.out_features
|
||||
x_rot = rotate(x, rot)
|
||||
x2d = x_rot.reshape(-1, in_f)
|
||||
m = x2d.shape[0]
|
||||
m = x.numel() // in_f
|
||||
|
||||
if x.requires_grad:
|
||||
# training: gated on requires_grad alone so both gradient-checkpoint
|
||||
# passes take the same branch
|
||||
if _int8_gemm_supported(x.device):
|
||||
# int8 tensor-core forward (bit-identical to the inference path)
|
||||
# with a straight-through analytic backward
|
||||
# with a straight-through analytic backward. the rotation stays a
|
||||
# cublas matmul: an in-kernel tl.dot rotation was tried and measured
|
||||
# SLOWER on every shape (small-tile dots at low tensor-core
|
||||
# utilization, computed twice for the amax and quant passes, cost
|
||||
# more than the activation round-trips they saved)
|
||||
out = _int8_linear_ste_op(
|
||||
x2d, module.cr8_qdata, module.cr8_scales, module.bias,
|
||||
rotate(x, rot).reshape(-1, in_f),
|
||||
module.cr8_qdata, module.cr8_scales, module.bias,
|
||||
str(x.dtype).split(".")[-1],
|
||||
)
|
||||
return out.reshape(*x.shape[:-1], out_f)
|
||||
# no int8 hardware: straight-through fake-quant + bf16 matmul
|
||||
x2d = rotate(x, rot).reshape(-1, in_f)
|
||||
with torch.no_grad():
|
||||
aq, a_s = quantize_int8_rows(x2d.detach())
|
||||
x_dq = (aq.float() * a_s.unsqueeze(1)).to(x.dtype)
|
||||
|
|
@ -830,11 +991,12 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
|
|||
if _int8_gemm_supported(x.device):
|
||||
# row padding for _int_mm happens inside the act-quant op (compile
|
||||
# safety); slice the mm output back to m rows (a contiguous prefix)
|
||||
aq, a_s = _int8_act_quant_padded(x2d)
|
||||
aq, a_s = _int8_act_quant_padded(rotate(x, rot).reshape(-1, in_f))
|
||||
i32 = torch._int_mm(aq, module.cr8_qdata.t())
|
||||
out = _int8_epilogue(i32[:m], a_s[:m], self._scales(module), module.bias, x.dtype)
|
||||
return out.reshape(*x.shape[:-1], out_f)
|
||||
|
||||
w = self._dequantize_rotated(module, x.dtype)
|
||||
out = F.linear(x2d, w, module.bias)
|
||||
out = F.linear(rotate(x, rot).reshape(-1, in_f), w, module.bias)
|
||||
return out.reshape(*x.shape[:-1], out_f)
|
||||
|
||||
|
|
|
|||
|
|
@ -28,6 +28,10 @@ class OstrisQuantizer:
|
|||
other methods. One backend instance may be shared by many modules.
|
||||
"""
|
||||
|
||||
# the qtype string this instance was resolved from (stamped by
|
||||
# get_ostris_quantizer); pre-quantized saves need it to restore the backend
|
||||
qtype: Optional[str] = None
|
||||
|
||||
def can_quantize(self, module: torch.nn.Linear) -> bool:
|
||||
"""Whether this backend can quantize the given linear (e.g. shape constraints)."""
|
||||
return True
|
||||
|
|
@ -105,13 +109,16 @@ def get_ostris_quantizer(qtype: str) -> Optional[OstrisQuantizer]:
|
|||
from toolkit.util.orbit_vq_quant import ORBIT_VQ_QTYPES, OrbitVQQuantizer
|
||||
from toolkit.util.convrot_quant import CONVROT_QTYPES, get_convrot_quantizer
|
||||
|
||||
quantizer = None
|
||||
if qtype in ORBIT_QTYPES:
|
||||
return OrbitQuantizer(ORBIT_QTYPES[qtype])
|
||||
if qtype in ORBIT_VQ_QTYPES:
|
||||
return OrbitVQQuantizer(**ORBIT_VQ_QTYPES[qtype])
|
||||
if qtype in CONVROT_QTYPES:
|
||||
return get_convrot_quantizer(qtype)
|
||||
return None
|
||||
quantizer = OrbitQuantizer(ORBIT_QTYPES[qtype])
|
||||
elif qtype in ORBIT_VQ_QTYPES:
|
||||
quantizer = OrbitVQQuantizer(**ORBIT_VQ_QTYPES[qtype])
|
||||
elif qtype in CONVROT_QTYPES:
|
||||
quantizer = get_convrot_quantizer(qtype)
|
||||
if quantizer is not None:
|
||||
quantizer.qtype = qtype
|
||||
return quantizer
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
|
|
|
|||
|
|
@ -1236,7 +1236,10 @@ export const groupedModelOptions: GroupedSelectOption[] = modelArchs.reduce((acc
|
|||
|
||||
export const quantizationOptions: SelectOption[] = [
|
||||
{ value: '', label: '- NONE -' },
|
||||
{ value: 'qfloat8', label: 'float8 (default)' },
|
||||
{ value: 'qfloat8', label: 'qfloat8 (default)' },
|
||||
{ value: 'float8', label: 'float8' },
|
||||
{ value: 'convrot8', label: '8bit convrot' },
|
||||
{ value: 'convrot4', label: '4bit convrot' },
|
||||
{ value: 'uint7', label: '7 bit' },
|
||||
{ value: 'uint6', label: '6 bit' },
|
||||
{ value: 'uint5', label: '5 bit' },
|
||||
|
|
|
|||
|
|
@ -1 +1 @@
|
|||
VERSION = "0.10.22"
|
||||
VERSION = "0.10.23"
|
||||
|
|
|
|||
Loading…
Reference in New Issue