From 095d6e7418a3e36d956c7069c23541f143c48ba9 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sat, 11 Jul 2026 13:49:41 -0600 Subject: [PATCH] Add convrot4 and convrot8 quantizations to the ui. A lot of improvements to the convrot quantization methods. --- jobs/process/BaseSDTrainProcess.py | 6 + scripts/test_quantizations.py | 339 +++++++++++++++++++++++++++++ toolkit/util/convrot_quant.py | 226 ++++++++++++++++--- toolkit/util/ostris_quant.py | 19 +- ui/src/app/jobs/new/options.ts | 5 +- version.py | 2 +- 6 files changed, 557 insertions(+), 40 deletions(-) create mode 100644 scripts/test_quantizations.py diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index c0f36d8e..127a810d 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -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) diff --git a/scripts/test_quantizations.py b/scripts/test_quantizations.py new file mode 100644 index 00000000..b708e639 --- /dev/null +++ b/scripts/test_quantizations.py @@ -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() diff --git a/toolkit/util/convrot_quant.py b/toolkit/util/convrot_quant.py index af605e78..3ab4069e 100644 --- a/toolkit/util/convrot_quant.py +++ b/toolkit/util/convrot_quant.py @@ -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) + diff --git a/toolkit/util/ostris_quant.py b/toolkit/util/ostris_quant.py index 0ef213e3..f2794e6f 100644 --- a/toolkit/util/ostris_quant.py +++ b/toolkit/util/ostris_quant.py @@ -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() diff --git a/ui/src/app/jobs/new/options.ts b/ui/src/app/jobs/new/options.ts index 7e816285..7a4f0261 100644 --- a/ui/src/app/jobs/new/options.ts +++ b/ui/src/app/jobs/new/options.ts @@ -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' }, diff --git a/version.py b/version.py index 4d445c57..7338a1bc 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.10.22" +VERSION = "0.10.23"