diff --git a/CHANGELOG.md b/CHANGELOG.md index 39122a78..4c0e5a4d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,24 @@ # Change log +## [v1.1.1+feature/mdbf] 2026-06-18 + +### New Feature: GemLite 1-bit inference for MDBF + +- Added a GemLite-accelerated inference path to `MDBFLinear` / `MultipathMDBFLinear` for the ±1 sign matrices (`A_sign`, `B_sign`), mirroring the DBF layer design; when GemLite is available the sign matmuls are delegated to GemLite 1-bit Triton kernels, and the layer transparently falls back to the dense path otherwise (`onecomp/quantizer/mdbf/mdbf_layer.py`) +- In auto mode (`use_gemlite=None`) GemLite is enabled only for `l == 1`; for `l > 1` the outer rank-`l` amplitude makes the GemLite path slower than dense, so it is skipped unless `use_gemlite=True` forces it +- Freed the redundant GPU packed-sign buffers for any sign matrix served by GemLite, roughly halving the in-memory weight footprint +- Propagated `use_gemlite` through `MDBF.create_inference_layer()` (`onecomp/quantizer/mdbf/_mdbf.py`) + +### Bug Fixes + +- Fixed the Hessian definition in `lowrank_osvd()` (OSVD initialization): the whitening term now applies the full `H^{1/2} = Q diag(sqrt(λ)) Q^T`. The trailing `@ Q^T` was missing, so the previous `W @ Q diag(sqrt(λ))` computed `W @ H^{1/2} @ Q` and did not minimize the intended Hessian-weighted output error (`onecomp/quantizer/mdbf/initialize.py`) +- Fixed `rank_from_bpw()` which effectively hard-coded `scale_bits=0`: `scale_bits` is now an argument defaulting to 16, so the FP16 envelope parameters are counted in the BPW budget, consistent with the paper formula `b = P * [r(n+m) + 16*l*(n+m+2r)] / (nm)`; target and actual BPW now agree (`onecomp/quantizer/mdbf/utils.py`) + +### Tests + +- Added GemLite inference tests to `tests/onecomp/quantizer/mdbf/test_mdbf.py`: `MDBFLinear` / `MultipathMDBFLinear` GemLite output matches the dense path at `l == 1`, `create_inference_layer` GemLite output matches the dequantized weight, dense fallback when GemLite is unavailable, `rank_from_bpw()` consistency with the paper BPW formula (`scale_bits=16`), and a `lowrank_osvd` regression test asserting the OSVD Hessian-weighted error is no worse than plain rank-`r` SVD +- Added `tests/onecomp/quantizer/mdbf/test_osvd_hessian_bug.py`: a numerical proof-of-concept that the previous `H^{1/2}` formulation inflated the Hessian-weighted reconstruction error for non-diagonal Hessians + ## [v1.1.1+feature/mdbf] 2026-05-29 ### Bug Fixes diff --git a/onecomp/quantizer/mdbf/_mdbf.py b/onecomp/quantizer/mdbf/_mdbf.py index 82967a1f..685e3dcc 100644 --- a/onecomp/quantizer/mdbf/_mdbf.py +++ b/onecomp/quantizer/mdbf/_mdbf.py @@ -520,4 +520,5 @@ def create_inference_layer(self, result, linear_module, **kwargs): result=result, bias=bias, device=linear_module.weight.device, + use_gemlite=kwargs.get("use_gemlite"), ) diff --git a/onecomp/quantizer/mdbf/initialize.py b/onecomp/quantizer/mdbf/initialize.py index 294ba5a8..9aedd480 100644 --- a/onecomp/quantizer/mdbf/initialize.py +++ b/onecomp/quantizer/mdbf/initialize.py @@ -220,8 +220,8 @@ def lowrank_osvd( eig_vals = eig_vals.clamp(min=1e-12) sqrt_eig = torch.sqrt(eig_vals) - # W_tilde = W @ H^{1/2} = W @ Q @ diag(sqrt(λ)) - W_tilde = W_fp32 @ eig_vecs @ torch.diag(sqrt_eig) + # W_tilde = W @ H^{1/2} = W @ Q @ diag(sqrt(λ)) @ Q^T + W_tilde = W_fp32 @ eig_vecs @ torch.diag(sqrt_eig) @ eig_vecs.T del H_reg # Rank-r SVD of W_tilde diff --git a/onecomp/quantizer/mdbf/mdbf_layer.py b/onecomp/quantizer/mdbf/mdbf_layer.py index 1c147653..f708cf03 100644 --- a/onecomp/quantizer/mdbf/mdbf_layer.py +++ b/onecomp/quantizer/mdbf/mdbf_layer.py @@ -29,6 +29,14 @@ from .initialize import MDBFParams +# Optional GemLite integration (mirror of dbf/dbf_layer.py) +try: + from onecomp.quantizer.gemlite import create_gemlite_linear, is_gemlite_available + + HAS_GEMLITE_SUPPORT = True +except ImportError: + HAS_GEMLITE_SUPPORT = False + # ============================================================================= # Bit-packing/Unpacking # ============================================================================= @@ -91,9 +99,25 @@ class MDBFLinear(nn.Module): G = S_B * (Q_V_amp @ B_amp^T) Inference: y = x @ W^T = x @ G^T @ F^T + + GemLite acceleration (mirror of dbf/dbf_layer.py): + The two heavy matmuls are against the ±1 sign matrices B_sign (r, m) and + A_sign (n, r). The rank-l amplitudes are separable per scale, so + + y = Σ_k (((x * B_amp[:,k]) @ B_sign^T) * Q_V_amp[:,k] * Q_U_amp[:,k]) ... + + i.e. each sign matmul can be replaced by a 1-bit GemLite kernel and the + amplitudes applied as element-wise scalings around it. GemLite is enabled + per sign matrix (it requires the matmul's in_features to be a multiple of + the group size), falling back to on-the-fly unpack otherwise. """ - def __init__(self, params: MDBFParams): + def __init__( + self, + params: MDBFParams, + device: Optional[torch.device] = None, + use_gemlite: Optional[bool] = None, + ): super().__init__() n, r = params.A_sign.shape @@ -120,6 +144,90 @@ def __init__(self, params: MDBFParams): self.register_buffer("Q_U_amp", params.Q_U_amp.half()) self.register_buffer("Q_V_amp", params.Q_V_amp.half()) + # Optional GemLite kernels for the ±1 sign matmuls. + # Stored in a plain dict (not a submodule) so they stay out of state_dict, + # matching DoubleBinaryLinear._gemlite_layers. + self._gemlite_layers: dict = {} + # CPU stash of packed sign buffers freed from GPU once GemLite serves them + # (see _free_packed_sign). Plain dict -> invisible to .to()/state_dict. + self._packed_cpu: dict = {} + self.use_gemlite = False + self._build_gemlite(params.A_sign, params.B_sign, device, use_gemlite) + + def _build_gemlite( + self, + A_sign: torch.Tensor, + B_sign: torch.Tensor, + device: Optional[torch.device], + use_gemlite: Optional[bool], + ) -> None: + """Build 1-bit GemLite kernels for B_sign (r, m) and A_sign (n, r). + + Note on the multi-scale rank l: + The dense path folds the rank-l amplitude into a single (n, r)/(r, m) + matrix, so its cost is independent of l. The GemLite path keeps the + sign matrices pure ±1 and applies the l amplitude scales outside, so it + issues l separate 1-bit matmuls per sign matrix (cost grows ~O(l)). + Measured on H100 (in=out=4096, r=512): GemLite is ~1.5x faster than + dense at l=1 but ~0.76x (l=2) / ~0.37x (l=4), i.e. a net slowdown. + Therefore, in auto mode (use_gemlite=None) GemLite is only enabled for + l == 1. Pass use_gemlite=True to force it regardless of l. + """ + forced = use_gemlite is True + if use_gemlite is None: + use_gemlite = HAS_GEMLITE_SUPPORT and is_gemlite_available() + if not (use_gemlite and HAS_GEMLITE_SUPPORT): + return + if self.l > 1 and not forced: + # Auto mode: skip GemLite for l>1 (it would be slower than dense). + return + + device_obj = torch.device(device) if device is not None else A_sign.device + # Stage 1 matmul: x @ B_sign^T (in_features = m) + gemlite_B = create_gemlite_linear(B_sign, nbits=1, device=device_obj) + # Stage 2 matmul: u @ A_sign^T (in_features = r) + gemlite_A = create_gemlite_linear(A_sign, nbits=1, device=device_obj) + if gemlite_B is not None: + self._gemlite_layers["B"] = gemlite_B + if gemlite_A is not None: + self._gemlite_layers["A"] = gemlite_A + self.use_gemlite = len(self._gemlite_layers) > 0 + + # Free the redundant GPU packed-sign buffer for any matrix now served by + # GemLite (GemLite keeps its own 1-bit W_q), halving in-memory weight. + for which in self._gemlite_layers: + self._free_packed_sign(which) + + def _free_packed_sign(self, which: str) -> None: + """Move a packed sign buffer off-GPU once GemLite serves that matmul. + + The buffer is un-registered (so ``.to(device)`` won't pull it back to GPU) + and stashed on CPU; ``_save_to_state_dict`` re-emits it, so the saved + state_dict matches the dense layout exactly. + """ + buf_name = f"{which}_sign_packed" + buf = self._buffers.get(buf_name) + if buf is not None: + self._packed_cpu[which] = buf.detach().to("cpu") + del self._buffers[buf_name] + + def _packed_sign(self, which: str, device: torch.device) -> torch.Tensor: + """Return the packed (uint8) sign tensor on *device* (buffer or CPU stash).""" + buf = self._buffers.get(f"{which}_sign_packed") + if buf is None: + buf = self._packed_cpu[which] + return buf.to(device) + + def _save_to_state_dict(self, destination, prefix, keep_vars): + super()._save_to_state_dict(destination, prefix, keep_vars) + # Re-emit any packed-sign buffers freed from GPU in GemLite mode so the + # on-disk format is identical to the dense (non-GemLite) layout. + for which in ("A", "B"): + key = f"{prefix}{which}_sign_packed" + if key not in destination and which in getattr(self, "_packed_cpu", {}): + v = self._packed_cpu[which] + destination[key] = v if keep_vars else v.detach() + def _load_from_state_dict( self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs ): @@ -138,8 +246,9 @@ def _load_from_state_dict( def _get_factor_matrices(self, dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor]: """Compute factor matrices F, G (always on-the-fly unpack)""" - A_sign = unpack_binary(self.A_sign_packed, (self.n, self.r)).to(dtype) - B_sign = unpack_binary(self.B_sign_packed, (self.r, self.m)).to(dtype) + dev = self.A_amp.device + A_sign = unpack_binary(self._packed_sign("A", dev), (self.n, self.r)).to(dtype) + B_sign = unpack_binary(self._packed_sign("B", dev), (self.r, self.m)).to(dtype) amp_A = self.A_amp.to(dtype) @ self.Q_U_amp.to(dtype).T F = A_sign * amp_A @@ -149,7 +258,55 @@ def _get_factor_matrices(self, dtype: torch.dtype) -> Tuple[torch.Tensor, torch. return F, G + def _apply_sign(self, x: torch.Tensor, which: str) -> torch.Tensor: + """Multiply by a ±1 sign matrix, via GemLite when available. + + which="B": x @ B_sign^T (B_sign is (r, m)) + which="A": x @ A_sign^T (A_sign is (n, r)) + """ + gemlite = self._gemlite_layers.get(which) + if gemlite is not None: + return gemlite(x) + if which == "B": + B_sign = unpack_binary(self._packed_sign("B", x.device), (self.r, self.m)).to(x.dtype) + return x @ B_sign.T + A_sign = unpack_binary(self._packed_sign("A", x.device), (self.n, self.r)).to(x.dtype) + return x @ A_sign.T + + def _forward_gemlite(self, x: torch.Tensor) -> torch.Tensor: + """Per-scale forward using GemLite 1-bit kernels for the sign matmuls. + + Exact reformulation of y = x @ G^T @ F^T with + G = B_sign * (Q_V_amp @ B_amp^T), F = A_sign * (A_amp @ Q_U_amp^T): + + u = Σ_k ((x * B_amp[:,k]) @ B_sign^T) * Q_V_amp[:,k] + y = Σ_k ((u * Q_U_amp[:,k]) @ A_sign^T) * A_amp[:,k] + """ + dtype = x.dtype + B_amp = self.B_amp.to(dtype) # (m, l) + Q_V_amp = self.Q_V_amp.to(dtype) # (r, l) + A_amp = self.A_amp.to(dtype) # (n, l) + Q_U_amp = self.Q_U_amp.to(dtype) # (r, l) + + # Stage 1: u = x @ G^T -> (..., r) + u = None + for k in range(self.l): + bk = self._apply_sign(x * B_amp[:, k], "B") + contrib = bk * Q_V_amp[:, k] + u = contrib if u is None else u + contrib + + # Stage 2: y = u @ F^T -> (..., n) + y = None + for k in range(self.l): + ak = self._apply_sign(u * Q_U_amp[:, k], "A") + contrib = ak * A_amp[:, k] + y = contrib if y is None else y + contrib + + return y + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.use_gemlite: + return self._forward_gemlite(x) F, G = self._get_factor_matrices(x.dtype) y = x @ G.T y = y @ F.T @@ -160,6 +317,21 @@ def get_weight(self, dtype: torch.dtype = torch.float32) -> torch.Tensor: F, G = self._get_factor_matrices(dtype) return F @ G + def enable_gemlite( + self, device: Optional[torch.device] = None, force: bool = False + ) -> bool: + """Build GemLite kernels from the packed sign buffers (e.g. after load). + + Returns True if at least one sign matmul is GemLite-accelerated. + """ + if self.use_gemlite and self._gemlite_layers: + return True + dev = self.A_amp.device + A_sign = unpack_binary(self._packed_sign("A", dev), (self.n, self.r)) + B_sign = unpack_binary(self._packed_sign("B", dev), (self.r, self.m)) + self._build_gemlite(A_sign, B_sign, device, True if force else None) + return self.use_gemlite + # ============================================================================= # MultipathMDBFLinear Layer (P-pass) @@ -174,6 +346,7 @@ def __init__( params_list: List[MDBFParams], bias: Optional[torch.Tensor] = None, device=None, + use_gemlite: Optional[bool] = None, ): super().__init__() @@ -184,7 +357,9 @@ def __init__( self.n = params_list[0].A_sign.shape[0] self.m = params_list[0].B_sign.shape[1] - self.paths = nn.ModuleList([MDBFLinear(params) for params in params_list]) + self.paths = nn.ModuleList( + [MDBFLinear(params, device=device, use_gemlite=use_gemlite) for params in params_list] + ) if bias is not None: self.register_buffer("bias", bias.clone().to(torch.float16)) @@ -211,12 +386,26 @@ def get_weight(self, dtype: torch.dtype = torch.float32) -> torch.Tensor: W = W + self.paths[i].get_weight(dtype) return W + def enable_gemlite( + self, device: Optional[torch.device] = None, force: bool = False + ) -> bool: + """Enable GemLite on every path (e.g. after ``from_saved_state``). + + Returns True if any path ended up GemLite-accelerated. + """ + enabled = False + for path in self.paths: + if path.enable_gemlite(device=device, force=force): + enabled = True + return enabled + @classmethod def from_quantization_result( cls, result, bias=None, device=None, + use_gemlite=None, ) -> "MultipathMDBFLinear": """Build MultipathMDBFLinear from MDBFResult. @@ -224,12 +413,13 @@ def from_quantization_result( result: MDBFResult from quantizer. bias: Optional bias tensor (from original Linear). device: Device to place the layer on. + use_gemlite: GemLite acceleration (None=auto, True/False=force). Returns: MultipathMDBFLinear instance. """ params_list = result.get_MDBF_params_list() - return cls(params_list=params_list, bias=bias, device=device) + return cls(params_list=params_list, bias=bias, device=device, use_gemlite=use_gemlite) @classmethod def from_saved_state( @@ -314,6 +504,12 @@ def _t(k): path_layer.l = path_layer.A_amp.shape[1] + # GemLite disabled on load (mirror DoubleBinaryLinear.from_saved_state); + # can be enabled later via enable_gemlite(). + path_layer.use_gemlite = False + path_layer._gemlite_layers = {} + path_layer._packed_cpu = {} + paths.append(path_layer) self.paths = paths diff --git a/onecomp/quantizer/mdbf/utils.py b/onecomp/quantizer/mdbf/utils.py index cd3c8f70..0a51345c 100644 --- a/onecomp/quantizer/mdbf/utils.py +++ b/onecomp/quantizer/mdbf/utils.py @@ -80,6 +80,7 @@ def rank_from_bpw( P: int = 2, min_rank: int = 1, rounding: Literal["floor", "ceil", "round"] = "floor", + scale_bits: int = 16, ) -> int: """ Calculate rank r from target BPW @@ -99,13 +100,16 @@ def rank_from_bpw( - "floor": Round down (ensure b_target is not exceeded) - "ceil": Round up (prioritize approximation accuracy) - "round": Round to nearest (balance) + scale_bits: Bit width of each amplitude (envelope) parameter counted in + the BPW budget. Defaults to 16 (FP16), matching the paper's formula + b = P * [r(n+m) + 16*l*(n+m+2r)] / (nm). Set to 0 to count only the + binary sign matrices. Returns: Calculated rank r """ - # Note: scale_bits=0 is the mode where BPW is calculated only for binary matrices - # To include FP16 scales, change scale_bits to 16 - scale_bits = 0 + # Note: scale_bits=16 counts FP16 envelope parameters in the BPW budget. + # This matches the paper's formula: b = P * [r(n+m) + 16l(n+m+2r)] / (nm) numerator = (b_target * n * m / P) - scale_bits * l * (n + m) denominator = (n + m) + 2 * scale_bits * l diff --git a/tests/onecomp/quantizer/mdbf/test_mdbf.py b/tests/onecomp/quantizer/mdbf/test_mdbf.py index 65ec825a..829901ad 100644 --- a/tests/onecomp/quantizer/mdbf/test_mdbf.py +++ b/tests/onecomp/quantizer/mdbf/test_mdbf.py @@ -14,8 +14,11 @@ sys.path.append(os.path.join(os.path.dirname(__file__), "..")) from onecomp.quantizer.mdbf._mdbf import MDBF, MDBFResult -from onecomp.quantizer.mdbf.initialize import MDBFParams -from onecomp.quantizer.mdbf.utils import reconstruct_weight +from onecomp.quantizer.mdbf.initialize import MDBFParams, lowrank_osvd +from onecomp.quantizer.mdbf import mdbf_layer +from onecomp.quantizer.mdbf.mdbf_layer import MDBFLinear, MultipathMDBFLinear +from onecomp.quantizer.mdbf.utils import bpw_from_rank, rank_from_bpw, reconstruct_weight +from onecomp.quantizer.gemlite import is_gemlite_available from test_module import BaseQuantizeSpec @@ -309,3 +312,229 @@ def apply_quantized_weights(self, module, result, device): ] module.MDBF_params = params_list module.is_quantized = True + + +def _random_sign(shape, device, dtype=torch.float16): + return (torch.randint(0, 2, shape, device=device, dtype=torch.int8) * 2 - 1).to(dtype) + + +def _make_mdbf_params(n, m, r, l, device, dtype=torch.float16): + return MDBFParams( + A_sign=_random_sign((n, r), device, dtype), + B_sign=_random_sign((r, m), device, dtype), + A_amp=torch.randn(n, l, device=device, dtype=dtype), + B_amp=torch.randn(m, l, device=device, dtype=dtype), + Q_U_amp=torch.randn(r, l, device=device, dtype=dtype), + Q_V_amp=torch.randn(r, l, device=device, dtype=dtype), + ) + + +def _assert_gemlite_output_matches_dense(y_dense, y_gemlite): + diff = (y_dense.float() - y_gemlite.float()).abs() + rel = (torch.norm(y_dense.float() - y_gemlite.float()) / torch.norm(y_dense.float())).item() + assert rel < 1e-3, f"GemLite relative output error too large: {rel}" + assert diff.max().item() < 4.0, f"GemLite max abs output error too large: {diff.max().item()}" + + +def _quantize_linear_for_inference_test(in_features, out_features, device, p=1): + layer = torch.nn.Linear(in_features, out_features, bias=False, device=device, dtype=torch.float32) + inp = torch.randn(3, 4, in_features, device=device, dtype=torch.float32) + quantizer = MDBF( + target_bits=1.0, + l=1, + P=p, + svd_mode="svd", + use_admm=False, + use_gradient_refine=False, + ) + hessian, nsamples = quantizer.calculate_hessian(layer, inp) + result = quantizer.quantize_layer(layer, inp, hessian=hessian, nsamples=nsamples) + return layer, inp, quantizer, result + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not is_gemlite_available(), + reason="GemLite unavailable or CUDA not available", +) +def test_mdbflinear_gemlite_matches_dense_forward(): + torch.manual_seed(0) + device = torch.device("cuda") + params = _make_mdbf_params(n=96, m=256, r=128, l=1, device=device) + + dense_layer = MDBFLinear(params, device=device, use_gemlite=False) + gemlite_layer = MDBFLinear(params, device=device, use_gemlite=True) + + assert gemlite_layer.use_gemlite + assert set(gemlite_layer._gemlite_layers) == {"A", "B"} + + x = torch.randn(7, 256, device=device, dtype=torch.float16) + with torch.no_grad(): + y_dense = dense_layer(x) + y_gemlite = gemlite_layer(x) + + assert y_dense.shape == y_gemlite.shape + _assert_gemlite_output_matches_dense(y_dense, y_gemlite) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not is_gemlite_available(), + reason="GemLite unavailable or CUDA not available", +) +def test_multipath_mdbflinear_gemlite_matches_dense_forward(): + torch.manual_seed(1) + device = torch.device("cuda") + params_list = [ + _make_mdbf_params(n=80, m=256, r=128, l=2, device=device), + _make_mdbf_params(n=80, m=256, r=128, l=2, device=device), + ] + bias = torch.randn(80, device=device, dtype=torch.float16) + + dense_layer = MultipathMDBFLinear( + params_list=params_list, + bias=bias, + device=device, + use_gemlite=False, + ) + gemlite_layer = MultipathMDBFLinear( + params_list=params_list, + bias=bias, + device=device, + use_gemlite=True, + ) + + assert all(path.use_gemlite for path in gemlite_layer.paths) + + x = torch.randn(5, 256, device=device, dtype=torch.float16) + with torch.no_grad(): + y_dense = dense_layer(x) + y_gemlite = gemlite_layer(x) + + assert y_dense.shape == y_gemlite.shape + _assert_gemlite_output_matches_dense(y_dense, y_gemlite) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not is_gemlite_available(), + reason="GemLite unavailable or CUDA not available", +) +@pytest.mark.parametrize("p", [1, 2]) +def test_mdbf_create_inference_layer_gemlite_matches_dequantized_forward(p): + torch.manual_seed(7 + p) + device = torch.device("cuda") + layer, inp, quantizer, result = _quantize_linear_for_inference_test( + in_features=256, + out_features=128, + device=device, + p=p, + ) + + dequantized_layer = torch.nn.Linear(256, 128, bias=False, device=device, dtype=torch.float32) + dequantized_layer.weight.data.copy_(result.compute_dequantized_weight().to(device=device, dtype=torch.float32)) + + dense_layer = quantizer.create_inference_layer( + result=result, + linear_module=layer, + use_gemlite=False, + ) + gemlite_layer = quantizer.create_inference_layer( + result=result, + linear_module=layer, + use_gemlite=True, + ) + + assert isinstance(dense_layer, MultipathMDBFLinear) + assert isinstance(gemlite_layer, MultipathMDBFLinear) + assert all(path.use_gemlite for path in gemlite_layer.paths) + + with torch.no_grad(): + y_dequantized = dequantized_layer(inp).float() + y_dense = dense_layer(inp.to(torch.float16)).float() + y_gemlite = gemlite_layer(inp.to(torch.float16)).float() + + _assert_gemlite_output_matches_dense(y_dense, y_gemlite) + _assert_gemlite_output_matches_dense(y_dequantized, y_gemlite) + + +def test_mdbflinear_falls_back_to_dense_when_gemlite_unavailable(monkeypatch): + """When GemLite is not installed, forcing use_gemlite=True must fall back to + the dense path instead of raising (Test plan: dense fallback).""" + # Simulate an environment without GemLite support. + monkeypatch.setattr(mdbf_layer, "HAS_GEMLITE_SUPPORT", False) + + device = torch.device("cpu") + torch.manual_seed(0) + params = _make_mdbf_params(n=16, m=32, r=8, l=1, device=device, dtype=torch.float32) + + forced_layer = MDBFLinear(params, device=device, use_gemlite=True) + dense_layer = MDBFLinear(params, device=device, use_gemlite=False) + + # GemLite is unavailable, so even use_gemlite=True must not enable a kernel. + assert forced_layer.use_gemlite is False + assert forced_layer._gemlite_layers == {} + + x = torch.randn(4, 32, device=device, dtype=torch.float32) + with torch.no_grad(): + y_forced = forced_layer(x) + y_dense = dense_layer(x) + + # The fallback must take the exact same dense forward path. + assert y_forced.shape == (4, 16) + assert torch.equal(y_forced, y_dense) + + +def test_rank_from_bpw_matches_paper_formula(): + """rank_from_bpw() with scale_bits=16 must be consistent with the paper's BPW + formula b = P * [r(n+m) + 16*l*(n+m+2r)] / (nm) (Test plan: rank_from_bpw).""" + n, m, l, P = 512, 2048, 8, 1 + + for r_true in (16, 64, 128, 256): + # bpw_from_rank uses scale_bits=16, the paper default. + b = bpw_from_rank(n, m, r_true, l=l, P=P) + # floor rounding must never exceed the requested budget. + r_floor = rank_from_bpw(n, m, b, l=l, P=P, scale_bits=16, rounding="floor") + assert r_floor <= r_true + assert bpw_from_rank(n, m, r_floor, l=l, P=P) <= b + 1e-9 + # round-trip: rounding to nearest recovers the exact rank. + r_round = rank_from_bpw(n, m, b, l=l, P=P, scale_bits=16, rounding="round") + assert r_round == r_true + + # scale_bits=0 drops the envelope cost, so more rank fits in the same budget. + b = bpw_from_rank(n, m, 64, l=l, P=P) + r_with_scale = rank_from_bpw(n, m, b, l=l, P=P, scale_bits=16, rounding="round") + r_without_scale = rank_from_bpw(n, m, b, l=l, P=P, scale_bits=0, rounding="round") + assert r_without_scale > r_with_scale + + +def test_lowrank_osvd_beats_plain_svd_in_hessian_error(): + """OSVD (H^{1/2}=Q diag(sqrt(λ)) Q^T whitening) must achieve a Hessian-weighted + output error no larger than plain rank-r SVD for a non-diagonal H + (Test plan: OSVD in activation-aware mode).""" + torch.manual_seed(0) + n, m, r = 32, 24, 4 + W = torch.randn(n, m, dtype=torch.float64) + + # Non-diagonal SPD Hessian (the H-weighting only matters when Q != I). + A = torch.randn(m, m, dtype=torch.float64) + H = A @ A.T + 0.1 * torch.eye(m, dtype=torch.float64) + + def hessian_error(W_hat): + E = W - W_hat + return torch.trace(E @ H @ E.T).item() + + # OSVD reconstruction: W_hat = U' @ V'^T. + U_prime, V_prime = lowrank_osvd(W, H, r, ridge=0.0) + err_osvd = hessian_error(U_prime @ V_prime.T) + + # Plain rank-r SVD (ignores H). + U_s, S_s, Vh_s = torch.linalg.svd(W, full_matrices=False) + W_svd = (U_s[:, :r] * S_s[:r]) @ Vh_s[:r, :] + err_svd = hessian_error(W_svd) + + # OSVD is tailored to the H-weighted objective, so it must not be worse. + assert err_osvd <= err_svd + 1e-6 * abs(err_svd), ( + f"OSVD H-weighted error ({err_osvd:.6e}) should be <= plain SVD ({err_svd:.6e})" + ) + # For a genuinely non-diagonal H the two solutions must differ. + assert abs(err_osvd - err_svd) > 1e-8 * abs(err_svd), ( + "OSVD and plain SVD errors are identical: the H-weighting is not being exercised" + ) diff --git a/tests/onecomp/quantizer/mdbf/test_osvd_hessian_bug.py b/tests/onecomp/quantizer/mdbf/test_osvd_hessian_bug.py new file mode 100644 index 00000000..5047b406 --- /dev/null +++ b/tests/onecomp/quantizer/mdbf/test_osvd_hessian_bug.py @@ -0,0 +1,236 @@ +""" +Regression test for the H^{1/2} whitening in lowrank_osvd (OSVD). + +OSVD minimizes the Hessian-weighted output error + min_{U,V} tr((W - U V^T) H (W - U V^T)^T) +by whitening with H^{1/2} before the SVD and un-whitening with H^{-1/2} +afterwards. For H = Q Λ Q^T the matrix square root is + H^{1/2} = Q Λ^{1/2} Q^T, +so the whitening step must apply the full Q Λ^{1/2} Q^T; in particular the +trailing Q^T is required. + +A formulation that drops the trailing Q^T (whitening with W @ Q Λ^{1/2} +instead of W @ Q Λ^{1/2} Q^T) computes W @ H^{1/2} @ Q. Its right singular +vectors are then Q-rotated relative to those of W @ H^{1/2}, + V_r_no_qt = Q^T @ V_r, +and applying the full H^{-1/2} = Q Λ^{-1/2} Q^T afterwards leaves a spurious +Q^T @ Q^T = (Q @ Q)^{-1} != I factor. Its reconstruction is therefore not the +H-weighted best rank-r approximation of W. + +These tests verify that the full-Q^T formulation is optimal for the H-weighted +objective and that dropping Q^T is strictly worse for a non-diagonal H (the two +coincide only when Q = I, i.e. H is diagonal). +""" + +import pytest +import torch + + +def _hessian_error(W: torch.Tensor, W_hat: torch.Tensor, H: torch.Tensor) -> float: + """tr((W - W_hat) H (W - W_hat)^T)""" + E = W - W_hat + return torch.trace(E @ H @ E.T).item() + + +def _osvd_full(W: torch.Tensor, H: torch.Tensor, r: int) -> torch.Tensor: + """ + OSVD with the full whitening H^{1/2} = Q diag(sqrt(λ)) Q^T. + + W_tilde = W @ H^{1/2} -> SVD -> V' = H^{-1/2} @ V_r @ diag(sqrt(Σ)) + """ + eig_vals, eig_vecs = torch.linalg.eigh(H) + eig_vals = eig_vals.clamp(min=1e-12) + sqrt_eig = torch.sqrt(eig_vals) + inv_sqrt_eig = 1.0 / sqrt_eig + + # Full H^{1/2}: Q diag(sqrt(λ)) Q^T + W_tilde = W @ eig_vecs @ torch.diag(sqrt_eig) @ eig_vecs.T + + U_w, S_w, Vh_w = torch.linalg.svd(W_tilde, full_matrices=False) + r = min(r, S_w.numel()) + U_r, S_r, V_r = U_w[:, :r], S_w[:r], Vh_w[:r, :].T + sqrt_S = torch.sqrt(S_r.clamp(min=1e-12)) + + U_prime = U_r * sqrt_S[None, :] + # V' = H^{-1/2} @ V_r @ diag(sqrt(Σ)) (H^{-1/2} = Q diag(1/sqrt(λ)) Q^T) + V_prime = eig_vecs @ torch.diag(inv_sqrt_eig) @ eig_vecs.T @ V_r @ torch.diag(sqrt_S) + return U_prime @ V_prime.T + + +def _osvd_no_qt(W: torch.Tensor, H: torch.Tensor, r: int) -> torch.Tensor: + """ + Variant that drops the trailing Q^T from the whitening step. + + W_tilde = W @ Q @ diag(sqrt(λ)) (= W @ H^{1/2} @ Q, no trailing Q^T) + V' = Q @ diag(1/sqrt(λ)) @ Q^T @ V_r @ diag(sqrt(Σ)) (full H^{-1/2}) + + Because W_tilde is W @ H^{1/2} @ Q, its right singular vectors V_r equal + Q^T @ V_r_full, so applying the full H^{-1/2} in the second step leaves a + spurious Q^T @ Q^T = (Q^2)^{-1} != I factor and does not minimize the + H-weighted objective. + """ + eig_vals, eig_vecs = torch.linalg.eigh(H) + eig_vals = eig_vals.clamp(min=1e-12) + sqrt_eig = torch.sqrt(eig_vals) + inv_sqrt_eig = 1.0 / sqrt_eig + + # Drop the trailing Q^T -> this whitens as W @ H^{1/2} @ Q + W_tilde = W @ eig_vecs @ torch.diag(sqrt_eig) + + U_w, S_w, Vh_w = torch.linalg.svd(W_tilde, full_matrices=False) + r = min(r, S_w.numel()) + U_r, S_r, V_r = U_w[:, :r], S_w[:r], Vh_w[:r, :].T + sqrt_S = torch.sqrt(S_r.clamp(min=1e-12)) + + U_prime = U_r * sqrt_S[None, :] + # Full H^{-1/2}, but V_r is already Q-rotated, so the rotations do not cancel + V_prime = eig_vecs @ torch.diag(inv_sqrt_eig) @ eig_vecs.T @ V_r @ torch.diag(sqrt_S) + return U_prime @ V_prime.T + + +def _make_pd_matrix(m: int, seed: int, noise: float = 0.1) -> torch.Tensor: + """Generate a non-diagonal positive-definite matrix (a diagonal H makes Q = I)""" + torch.manual_seed(seed) + A = torch.randn(m, m) + return A @ A.T + noise * torch.eye(m) + + +class TestOsvdHessianWhitening: + """ + OSVD H^{1/2} whitening: numerical verification. + """ + + @pytest.mark.parametrize("n,m,r,seed", [ + (16, 12, 3, 0), + (32, 24, 4, 1), + (8, 6, 2, 42), + ]) + def test_no_qt_has_larger_hessian_error(self, n, m, r, seed): + """ + The no-Q^T variant yields a larger H-weighted output error than the full + formulation. + + The full OSVD is the optimal solution of the minimization problem, so + full error <= no-Q^T error holds. The two are equal only when H is + diagonal (i.e. Q = I). + """ + torch.manual_seed(seed) + W = torch.randn(n, m, dtype=torch.float64) + H = _make_pd_matrix(m, seed=seed + 100).to(torch.float64) + + W_hat_full = _osvd_full(W, H, r) + W_hat_no_qt = _osvd_no_qt(W, H, r) + + err_full = _hessian_error(W, W_hat_full, H) + err_no_qt = _hessian_error(W, W_hat_no_qt, H) + + print( + f"\n[n={n},m={m},r={r},seed={seed}]" + f" full={err_full:.6e} no_qt={err_no_qt:.6e}" + f" degradation={100*(err_no_qt-err_full)/max(abs(err_full),1e-12):.2f}%" + ) + + # The full formulation is optimal, so it is always <= the no-Q^T variant + assert err_full <= err_no_qt + 1e-6 * abs(err_full), ( + f"full ({err_full:.6e}) should be <= no_qt ({err_no_qt:.6e})" + ) + # For non-diagonal H the two must genuinely differ (else Q^T does not matter here) + assert abs(err_no_qt - err_full) > 1e-8 * abs(err_full), ( + "Both errors are equal: the trailing Q^T is not being exercised " + "(check whether H is effectively diagonal)" + ) + + def test_reconstruction_differs(self): + """ + The full and no-Q^T reconstructions are different matrices. + """ + torch.manual_seed(7) + n, m, r = 10, 8, 3 + W = torch.randn(n, m, dtype=torch.float64) + H = _make_pd_matrix(m, seed=77).to(torch.float64) + + W_hat_full = _osvd_full(W, H, r) + W_hat_no_qt = _osvd_no_qt(W, H, r) + + max_diff = (W_hat_full - W_hat_no_qt).abs().max().item() + print(f"\nmax |W_hat_full - W_hat_no_qt| = {max_diff:.6e}") + + assert max_diff > 1e-6, ( + f"Expected a visible difference between full and no-Q^T W_hat, got {max_diff:.2e}" + ) + + def test_diagonal_H_matches(self): + """ + With a diagonal H whose eigenvalues are ascending-sorted (i.e. Q is the + identity matrix) the two formulations coincide. + + torch.linalg.eigh returns eigenvalues in ascending order, so eig_vecs = I + holds exactly when H is diagonal with already-ascending diagonal entries; + then the trailing Q^T is the identity and both formulations agree. + + This confirms that the trailing Q^T only matters through the off-diagonal + (or permutation) component of Q. + """ + torch.manual_seed(5) + n, m, r = 10, 8, 3 + W = torch.randn(n, m, dtype=torch.float64) + # Diagonal H: sorting eigenvalues ascending makes the Q returned by eigh the identity + diag_vals = torch.sort(torch.rand(m, dtype=torch.float64) + 0.5).values + H = torch.diag(diag_vals) + + # Pre-check that eigh returns the identity matrix + eig_vals, eig_vecs = torch.linalg.eigh(H) + assert torch.allclose(eig_vecs.abs(), torch.eye(m, dtype=torch.float64), atol=1e-10), \ + "For sorted diagonal H, eig_vecs should be identity" + + W_hat_full = _osvd_full(W, H, r) + W_hat_no_qt = _osvd_no_qt(W, H, r) + + max_diff = (W_hat_full - W_hat_no_qt).abs().max().item() + err_diff = abs( + _hessian_error(W, W_hat_full, H) - _hessian_error(W, W_hat_no_qt, H) + ) + print(f"\nDiagonal H (sorted): max diff = {max_diff:.2e}, error diff = {err_diff:.2e}") + + # When Q = I the two agree (the trailing Q^T becomes the identity) + assert max_diff < 1e-5, ( + "With sorted diagonal H (Q=I), full and no-Q^T should agree " + f"(got max diff {max_diff:.2e})" + ) + + def test_no_qt_inconsistency_analytically(self): + """ + Analytic confirmation of the inconsistency: + W_tilde_no_qt = W @ Q @ diag(sqrt(λ)) equals W @ H^{1/2} @ Q. + Therefore V_r_no_qt = Q^T @ V_r_full. + Multiplying by H^{-1/2} = Q diag(1/sqrt(λ)) Q^T in the second step gives + Q diag(1/sqrt(λ)) Q^T @ Q^T @ V_r_full + = Q diag(1/sqrt(λ)) @ (Q @ Q)^{-1} @ V_r_full (Q^T @ Q^T != I), + leaving a spurious (Q^2)^{-1} factor. + """ + torch.manual_seed(99) + m = 6 + H = _make_pd_matrix(m, seed=99).to(torch.float64) + + eig_vals, Q = torch.linalg.eigh(H) + eig_vals = eig_vals.clamp(min=1e-12) + sqrt_eig = torch.sqrt(eig_vals) + + # Confirm H^{1/2} @ Q = Q @ diag(sqrt(λ)) + H_half = Q @ torch.diag(sqrt_eig) @ Q.T + assert torch.allclose(H_half @ Q, Q @ torch.diag(sqrt_eig), atol=1e-10), \ + "H^{1/2} @ Q should equal Q @ diag(sqrt_eig)" + + # W_tilde_no_qt = W @ Q @ diag(sqrt(λ)) = W @ H^{1/2} @ Q != W @ H^{1/2} + torch.manual_seed(13) + W = torch.randn(8, m, dtype=torch.float64) + W_tilde_no_qt = W @ Q @ torch.diag(sqrt_eig) + W_tilde_full = W @ H_half + + assert not torch.allclose(W_tilde_no_qt, W_tilde_full, atol=1e-8), \ + "W_tilde_no_qt and W_tilde_full should differ for non-diagonal H" + + # Confirm Q^T @ Q^T != I (the source of the spurious factor) + QtQt = Q.T @ Q.T + assert not torch.allclose(QtQt, torch.eye(m, dtype=torch.float64), atol=1e-6), \ + "Q^T @ Q^T should not be identity (confirming the extra rotation is spurious)"