diff --git a/global_ptq/onecomp_globalptq/global_ptq/_core/core.py b/global_ptq/onecomp_globalptq/global_ptq/_core/core.py index e5a5e3f7..c27883c3 100644 --- a/global_ptq/onecomp_globalptq/global_ptq/_core/core.py +++ b/global_ptq/onecomp_globalptq/global_ptq/_core/core.py @@ -54,6 +54,15 @@ write_back_dbf_binary, write_back_dbf_scaling, ) +from .mdbf_adapter import ( + load_mdbf_state, + restore_mdbf_original, + save_mdbf_state, + setup_mdbf_differentiable, + setup_mdbf_forwards_only, + write_back_mdbf_amp, + write_back_mdbf_binary, +) logger = getLogger(__name__) @@ -418,6 +427,20 @@ def cosine_warmup_lr_lambda( # --------------------------------------------------------------------------- +@torch.no_grad() +def _teacher_logits( + teacher_model: nn.Module, + input_ids: torch.Tensor, + teacher_dev: torch.device, + student_dev: torch.device, +) -> torch.Tensor: + """Run teacher forward; move logits to *student_dev* if devices differ.""" + if teacher_dev == student_dev: + return get_logits(teacher_model(input_ids)) + logits_t = get_logits(teacher_model(input_ids.to(teacher_dev))) + return logits_t.to(student_dev) + + @torch.no_grad() def eval_kl( model: nn.Module, @@ -425,10 +448,12 @@ def eval_kl( dataloader: List[Dict[str, torch.Tensor]], dev: torch.device, temperature: float = 1.0, + teacher_dev: Optional[torch.device] = None, ) -> float: """Mean KL divergence over *dataloader* batches.""" was_training = model.training model.eval() + teacher_dev = teacher_dev or dev total, n = 0.0, 0 for batch in dataloader: input_ids = batch["input_ids"].to(dev) @@ -437,7 +462,7 @@ def eval_kl( attention_mask = attention_mask.to(dev) logits_s = get_logits(model(input_ids)) - logits_t = get_logits(teacher_model(input_ids)) + logits_t = _teacher_logits(teacher_model, input_ids, teacher_dev, dev) total += compute_kl_loss( logits_t, logits_s, temperature, attention_mask=attention_mask, ).item() @@ -587,6 +612,7 @@ def run_kl_distillation( gptq_intweight_lr: float = 1e-4, optimize_binary: bool = False, ste_k: float = 100.0, + mdbf_ste_k: float = 2.0, calibration_dataset=None, num_calibration_samples: int = 128, max_length: int = 2048, @@ -616,12 +642,23 @@ def run_kl_distillation( early_stopping_patience: int = 0, use_mixed_precision: bool = False, grad_accum_steps: int = 1, + student_device: Optional[str] = None, + teacher_device: Optional[str] = None, ) -> Dict: - """Run KL-distillation global PTQ on a GPTQ or DBF quantized model. + """Run KL-distillation global PTQ on a GPTQ, DBF or MDBF quantized model. + + The quantization method is auto-detected from the layer types present in + *quantized_model* (see :func:`detect_quantization_method`). For MDBF the + per-path amplitude factors are trained with ``dbf_lr``; with + ``optimize_binary=True`` the +/-1 sign matrices are trained as well through + a smooth sign STE whose sharpness is controlled by ``mdbf_ste_k``. The model is modified **in-place**. Returns a results dict. """ - dev = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dev = torch.device( + student_device or ("cuda" if torch.cuda.is_available() else "cpu") + ) + teacher_dev = torch.device(teacher_device) if teacher_device else dev # ------------------------------------------------------------------ # 1. Detect method @@ -631,7 +668,7 @@ def run_kl_distillation( logger.warning("No quantized layers detected — skipping global PTQ.") return {"global_executed": False, "reason": "not_quantized"} - if method not in ("gptq", "dbf"): + if method not in ("gptq", "dbf", "mdbf"): logger.info("Method '%s' detected — not supported.", method) return {"global_executed": False, "reason": f"unsupported_method_{method}"} @@ -665,7 +702,8 @@ def run_kl_distillation( teacher_model.eval() for p in teacher_model.parameters(): p.requires_grad = False - teacher_model.to(dev) + if teacher_dev.type != "cpu": + teacher_model.to(teacher_dev) # ------------------------------------------------------------------ # 4. Move student to GPU and set up differentiable parameters @@ -675,6 +713,7 @@ def run_kl_distillation( gptq_modules: list = [] dbf_modules: list = [] + mdbf_modules: list = [] original_forwards: Dict[str, object] = {} param_groups: list = [] binary_params: list = [] @@ -710,6 +749,24 @@ def run_kl_distillation( f", {len(binary_params)} binary" if binary_params else "", ) + elif method == "mdbf": + mdbf_modules = detected_modules + original_forwards, scaling_params, binary_params = setup_mdbf_differentiable( + mdbf_modules, optimize_binary, ste_k=mdbf_ste_k, + ) + logger.info("MDBF binary STE sharpness mdbf_ste_k=%.4g", mdbf_ste_k) + all_mdbf_params = list(scaling_params) + if binary_params: + all_mdbf_params += binary_params + param_groups = [{"params": all_mdbf_params, "lr": dbf_lr}] + + logger.info( + "Trainable: %d amp params%s across %d MDBF modules", + len(scaling_params), + f", {len(binary_params)} binary" if binary_params else "", + len(mdbf_modules), + ) + total_trainable = sum(len(pg["params"]) for pg in param_groups) if total_trainable == 0: logger.warning("No trainable parameters — skipping.") @@ -717,6 +774,8 @@ def run_kl_distillation( restore_gptq_original(gptq_modules, original_forwards) elif method == "dbf": restore_dbf_original(dbf_modules, original_forwards) + elif method == "mdbf": + restore_mdbf_original(mdbf_modules, original_forwards) quantized_model.cpu() del teacher_model gc.collect() @@ -871,17 +930,24 @@ def run_kl_distillation( if method == "gptq": initial_state = save_gptq_state(gptq_modules) restore_gptq_original(gptq_modules, original_forwards) - else: + elif method == "dbf": initial_state = save_dbf_state(dbf_modules) restore_dbf_original(dbf_modules, original_forwards) + else: # mdbf + initial_state = save_mdbf_state(mdbf_modules) + restore_mdbf_original(mdbf_modules, original_forwards) - initial_kl = eval_kl(quantized_model, teacher_model, dataloader, dev, temperature) + initial_kl = eval_kl( + quantized_model, teacher_model, dataloader, dev, temperature, teacher_dev, + ) logger.info("Initial KL = %.6f", initial_kl) if method == "gptq": setup_gptq_forwards_only(gptq_modules, original_forwards, gptq_optimize_intweight) elif method == "dbf": setup_dbf_forwards_only(dbf_modules, original_forwards) + elif method == "mdbf": + setup_mdbf_forwards_only(mdbf_modules, original_forwards) # ------------------------------------------------------------------ # 7. Training loop @@ -928,8 +994,9 @@ def _forward_and_loss() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: with amp_ctx: logits_s = get_logits(quantized_model(input_ids)) - with torch.no_grad(): - logits_t = get_logits(teacher_model(input_ids)) + logits_t = _teacher_logits( + teacher_model, input_ids, teacher_dev, dev, + ) kl = compute_kl_loss( logits_t, logits_s, temperature, attention_mask=attention_mask, @@ -1052,16 +1119,24 @@ def _forward_and_loss() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: elif method == "dbf": write_back_dbf_binary(dbf_modules) restore_dbf_original(dbf_modules, original_forwards) + elif method == "mdbf": + write_back_mdbf_binary(mdbf_modules) + write_back_mdbf_amp(mdbf_modules) + restore_mdbf_original(mdbf_modules, original_forwards) - current_kl = eval_kl(quantized_model, teacher_model, dataloader, dev, temperature) + current_kl = eval_kl( + quantized_model, teacher_model, dataloader, dev, temperature, teacher_dev, + ) if current_kl < best_kl: best_kl = current_kl patience_counter = 0 if method == "gptq": best_state = save_gptq_state(gptq_modules) - else: + elif method == "dbf": best_state = save_dbf_state(dbf_modules) + else: # mdbf + best_state = save_mdbf_state(mdbf_modules) else: patience_counter += 1 @@ -1069,6 +1144,8 @@ def _forward_and_loss() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: setup_gptq_forwards_only(gptq_modules, original_forwards, gptq_optimize_intweight) elif method == "dbf": setup_dbf_forwards_only(dbf_modules, original_forwards) + elif method == "mdbf": + setup_mdbf_forwards_only(mdbf_modules, original_forwards) # Restore non-EMA params for continued training if ema_tracker is not None: @@ -1103,15 +1180,19 @@ def _forward_and_loss() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if best_state is not None and best_kl < initial_kl: if method == "gptq": load_gptq_state(gptq_modules, best_state) - else: + elif method == "dbf": load_dbf_state(dbf_modules, best_state) + else: # mdbf + load_mdbf_state(mdbf_modules, best_state) logger.info("Loaded best state (KL=%.6f)", best_kl) elif best_kl >= initial_kl: logger.info("No improvement — rolling back to initial state.") if method == "gptq": load_gptq_state(gptq_modules, initial_state) - else: + elif method == "dbf": load_dbf_state(dbf_modules, initial_state) + else: # mdbf + load_mdbf_state(mdbf_modules, initial_state) best_kl = initial_kl else: if method == "gptq": @@ -1119,11 +1200,16 @@ def _forward_and_loss() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: elif method == "dbf": write_back_dbf_binary(dbf_modules) write_back_dbf_scaling(dbf_modules) + elif method == "mdbf": + write_back_mdbf_binary(mdbf_modules) + write_back_mdbf_amp(mdbf_modules) if method == "gptq": restore_gptq_original(gptq_modules, original_forwards, cleanup=False) elif method == "dbf": restore_dbf_original(dbf_modules, original_forwards, cleanup=False) + elif method == "mdbf": + restore_mdbf_original(mdbf_modules, original_forwards, cleanup=False) # Cleanup hooks if use_inter_loss: @@ -1140,7 +1226,9 @@ def _forward_and_loss() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: # Final evaluation quantized_model.eval() - final_kl = eval_kl(quantized_model, teacher_model, dataloader, dev, temperature) + final_kl = eval_kl( + quantized_model, teacher_model, dataloader, dev, temperature, teacher_dev, + ) # Cleanup if method == "gptq": @@ -1148,6 +1236,8 @@ def _forward_and_loss() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: restore_gptq_original(gptq_modules, original_forwards, cleanup=True) elif method == "dbf": restore_dbf_original(dbf_modules, original_forwards, cleanup=True) + elif method == "mdbf": + restore_mdbf_original(mdbf_modules, original_forwards, cleanup=True) del teacher_model gc.collect() diff --git a/global_ptq/onecomp_globalptq/global_ptq/_core/helpers.py b/global_ptq/onecomp_globalptq/global_ptq/_core/helpers.py index af32b360..5a5b2bc4 100644 --- a/global_ptq/onecomp_globalptq/global_ptq/_core/helpers.py +++ b/global_ptq/onecomp_globalptq/global_ptq/_core/helpers.py @@ -76,30 +76,43 @@ def detect_quantization_method( """Auto-detect the quantization method applied to *model*. Returns: - (method, modules) where *method* is ``"gptq"``, ``"dbf"``, or - ``None``, and *modules* is the list of ``(name, module)`` pairs - for the detected quantized layers. + (method, modules) where *method* is ``"gptq"``, ``"dbf"``, + ``"mdbf"``, or ``None``, and *modules* is the list of + ``(name, module)`` pairs for the detected quantized layers. - When both GPTQ and DBF layers are present (mixed quantization), - a warning is emitted and only GPTQ layers are returned. + Priority: GPTQ > DBF > MDBF. When multiple types coexist a warning is + emitted and only the highest-priority layers are returned. """ from onecomp.quantizer.gptq.gptq_layer import GPTQLinear from onecomp.quantizer.dbf.dbf_layer import DoubleBinaryLinear + try: + from onecomp.quantizer.mdbf.mdbf_layer import MultipathMDBFLinear + except ImportError: + # MDBF quantizer is optional; without it only GPTQ/DBF are detectable. + MultipathMDBFLinear = None + gptq_modules = find_target_modules(model, GPTQLinear) dbf_modules = find_target_modules(model, DoubleBinaryLinear) + mdbf_modules = ( + find_target_modules(model, MultipathMDBFLinear) + if MultipathMDBFLinear is not None + else [] + ) - if gptq_modules and dbf_modules: + if gptq_modules and (dbf_modules or mdbf_modules): logger.warning( - "Mixed GPTQ + DBF model detected (gptq=%d, dbf=%d). " + "Mixed GPTQ + DBF/MDBF model detected (gptq=%d, dbf=%d, mdbf=%d). " "Global PTQ currently optimises GPTQ layers only; " - "DBF layers will be skipped.", - len(gptq_modules), len(dbf_modules), + "other layers will be skipped.", + len(gptq_modules), len(dbf_modules), len(mdbf_modules), ) if gptq_modules: return "gptq", gptq_modules if dbf_modules: return "dbf", dbf_modules + if mdbf_modules: + return "mdbf", mdbf_modules return None, [] diff --git a/global_ptq/onecomp_globalptq/global_ptq/_core/mdbf_adapter.py b/global_ptq/onecomp_globalptq/global_ptq/_core/mdbf_adapter.py new file mode 100644 index 00000000..4dda578f --- /dev/null +++ b/global_ptq/onecomp_globalptq/global_ptq/_core/mdbf_adapter.py @@ -0,0 +1,376 @@ +"""MDBF differentiable parameter management for global PTQ. + +Makes MultipathMDBFLinear layers trainable by exposing amplitude parameters +(A_amp, B_amp, Q_U_amp, Q_V_amp) per path as optimisable tensors, and +optionally enabling differentiable binary-sign optimisation via smooth sign STE. + +Architecture recap (per MDBFLinear path): + F = A_sign * (A_amp @ Q_U_amp^T) shape: (n, r) + G = B_sign * (Q_V_amp @ B_amp^T) shape: (r, m) + y = x @ G^T @ F^T + +Trainable (continuous): + A_amp, B_amp, Q_U_amp, Q_V_amp — amplitude/scale factors per path. +Trainable (discrete, opt-in): + A_sign, B_sign — ±1 binary factor matrices per path, via sign STE. + +Copyright 2025-2026 Fujitsu Ltd. + +Authors: Keiji Kimura + +""" + +from types import MethodType +from typing import Dict, List, Tuple + +import torch +import torch.nn as nn + +from .helpers import smooth_sign_ste + +_AMP_ATTRS = ("A_amp", "B_amp", "Q_U_amp", "Q_V_amp") +_BINARY_SIGN_NAMES = ("A", "B") + +# Sharpness for sign-STE: same rationale as DBF adapter (values near ±1, +# tanh saturation avoided by using k=2 instead of the GPTQ default k=100). +_BINARY_STE_K = 2.0 + + +# --------------------------------------------------------------------------- +# Finding MDBF modules +# --------------------------------------------------------------------------- + + +def find_mdbf_modules(model: nn.Module) -> List[Tuple[str, nn.Module]]: + """Return all ``MultipathMDBFLinear`` modules as ``(name, module)`` pairs.""" + from onecomp.quantizer.mdbf.mdbf_layer import MultipathMDBFLinear + + from .helpers import find_target_modules + + return find_target_modules(model, MultipathMDBFLinear) + + +# --------------------------------------------------------------------------- +# Differentiable forward (per MDBFLinear path) +# --------------------------------------------------------------------------- + + +def _make_mdbf_differentiable_forward(): + """Build a differentiable ``forward`` for a ``MultipathMDBFLinear``. + + Each path's computation graph is reconstructed from the (possibly + optimisable) amplitude parameters and, when ``_opt_A_sign_{p}`` / + ``_opt_B_sign_{p}`` exist, through :func:`smooth_sign_ste` so that + gradients flow to the float sign-weight tensors as well. + """ + from onecomp.quantizer.mdbf.mdbf_layer import unpack_binary + + def differentiable_forward(self, x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + k = getattr(self, "_binary_ste_k", _BINARY_STE_K) + + y = None + for path in self.paths: + # ---- amplitude parameters (always trainable) ---- + A_amp = getattr(path, "_opt_A_amp", path.A_amp).to(dtype) + B_amp = getattr(path, "_opt_B_amp", path.B_amp).to(dtype) + Q_U_amp = getattr(path, "_opt_Q_U_amp", path.Q_U_amp).to(dtype) + Q_V_amp = getattr(path, "_opt_Q_V_amp", path.Q_V_amp).to(dtype) + + # ---- sign matrices (STE or packed buffer) ---- + if hasattr(path, "_opt_A_sign"): + A_sign = smooth_sign_ste(path._opt_A_sign, k=k).to(dtype) + else: + A_sign = unpack_binary(path._packed_sign("A", x.device), (path.n, path.r)).to( + dtype + ) + + if hasattr(path, "_opt_B_sign"): + B_sign = smooth_sign_ste(path._opt_B_sign, k=k).to(dtype) + else: + B_sign = unpack_binary(path._packed_sign("B", x.device), (path.r, path.m)).to( + dtype + ) + + # F = A_sign * (A_amp @ Q_U_amp^T) shape: (n, r) + F = A_sign * (A_amp @ Q_U_amp.T) + # G = B_sign * (Q_V_amp @ B_amp^T) shape: (r, m) + G = B_sign * (Q_V_amp @ B_amp.T) + + # y += x @ G^T @ F^T + path_out = x @ G.T @ F.T + + y = path_out if y is None else y + path_out + + if self.bias is not None: + y = y + self.bias.to(dtype) + return y + + return differentiable_forward + + +# --------------------------------------------------------------------------- +# Parameter setup +# --------------------------------------------------------------------------- + + +def setup_mdbf_differentiable( + mdbf_modules: List[Tuple[str, nn.Module]], + optimize_binary: bool = False, + ste_k: float = _BINARY_STE_K, +) -> Tuple[Dict[str, object], List[torch.Tensor], List[torch.Tensor]]: + """Make MDBF amplitude (and optionally sign) parameters trainable. + + For each ``MultipathMDBFLinear`` module, the original ``forward`` is + replaced with a differentiable version. Amplitude parameters + (A_amp, B_amp, Q_U_amp, Q_V_amp) of every path are promoted to + float32 ``nn.Parameter`` objects stored as ``_opt_*`` attributes on the + individual ``MDBFLinear`` path modules. + + Args: + mdbf_modules: List of ``(name, module)`` pairs from + :func:`find_mdbf_modules`. + optimize_binary: When True, also expose unpacked ±1 sign matrices + as float tensors with ``requires_grad=True`` so that gradients + flow through :func:`smooth_sign_ste`. + ste_k: Sharpness for binary sign STE (``tanh(k*x)`` backward). + Default is :data:`_BINARY_STE_K` (2.0). + + Returns: + ``(original_forwards, amp_params, binary_params)`` + + *original_forwards* maps module name → original forward (for restore). + *amp_params* is a flat list of float32 ``nn.Parameter`` objects. + *binary_params* is a flat list of float tensors (empty when + *optimize_binary* is ``False``). + """ + from onecomp.quantizer.mdbf.mdbf_layer import unpack_binary + + original_forwards: Dict[str, object] = {} + amp_params: List[torch.Tensor] = [] + binary_params: List[torch.Tensor] = [] + + for name, mod in mdbf_modules: + for path in mod.paths: + # ---- continuous amplitude parameters ---- + for attr in _AMP_ATTRS: + buf = getattr(path, attr) + fp32 = buf.data.detach().clone().float() + new_param = nn.Parameter(fp32, requires_grad=True) + setattr(path, f"_opt_{attr}", new_param) + amp_params.append(new_param) + + # ---- discrete sign parameters (optional) ---- + if optimize_binary: + for sign in _BINARY_SIGN_NAMES: + shape = (path.n, path.r) if sign == "A" else (path.r, path.m) + packed_key = f"{sign}_sign_packed" + packed = path._buffers.get(packed_key) + if packed is None: + # GemLite mode: stashed on CPU + packed = path._packed_cpu.get(sign) + if packed is None: + continue + # GemLite removes its redundant packed buffers from the + # module and stashes them on CPU. The differentiable sign + # parameter must nevertheless live with the layer; keeping + # it on CPU makes the reconstructed CUDA forward fail with + # a device mismatch. + target_device = path.A_amp.device + unpacked = ( + unpack_binary(packed.to(target_device), shape) + .float() + .detach() + .clone() + ) + new_param = nn.Parameter(unpacked, requires_grad=True) + setattr(path, f"_opt_{sign}_sign", new_param) + binary_params.append(new_param) + + original_forwards[name] = mod.forward + mod._binary_ste_k = ste_k + mod.forward = MethodType(_make_mdbf_differentiable_forward(), mod) + + return original_forwards, amp_params, binary_params + + +# --------------------------------------------------------------------------- +# Forward restore / re-install +# --------------------------------------------------------------------------- + + +def restore_mdbf_original( + mdbf_modules: List[Tuple[str, nn.Module]], + original_forwards: Dict[str, object], + cleanup: bool = False, +) -> None: + """Restore every module's original ``forward`` method.""" + for name, mod in mdbf_modules: + if name in original_forwards: + mod.__dict__.pop("forward", None) + if not hasattr(mod, "forward") or mod.forward != original_forwards[name]: + mod.forward = original_forwards[name] + + if cleanup: + if hasattr(mod, "_binary_ste_k"): + delattr(mod, "_binary_ste_k") + for path in mod.paths: + for attr in _AMP_ATTRS: + opt_attr = f"_opt_{attr}" + if hasattr(path, opt_attr): + delattr(path, opt_attr) + for sign in _BINARY_SIGN_NAMES: + opt_attr = f"_opt_{sign}_sign" + if hasattr(path, opt_attr): + delattr(path, opt_attr) + + +def setup_mdbf_forwards_only( + mdbf_modules: List[Tuple[str, nn.Module]], + original_forwards: Dict[str, object], +) -> None: + """Re-install differentiable forwards for continued training after eval.""" + for name, mod in mdbf_modules: + if name not in original_forwards: + original_forwards[name] = mod.forward + mod.forward = MethodType(_make_mdbf_differentiable_forward(), mod) + + +# --------------------------------------------------------------------------- +# Write-back +# --------------------------------------------------------------------------- + + +def _refresh_gemlite_sign_kernels(path: nn.Module) -> None: + """Repack GemLite kernels after their source sign matrices changed. + + MDBF's GemLite kernels keep a private packed copy of each sign matrix. + Updating ``*_sign_packed`` (or the CPU stash) alone therefore leaves the + inference kernel stale. Rebuild only paths which were already using + GemLite. Repacking failures are propagated so finalization cannot + silently leave the path with missing or stale inference kernels. + """ + gemlite_layers = getattr(path, "_gemlite_layers", None) + if not getattr(path, "use_gemlite", False) and not gemlite_layers: + return + + enable_gemlite = getattr(path, "enable_gemlite", None) + if enable_gemlite is None: + return + + # ``enable_gemlite`` returns early while old kernels are registered, so + # discard them before asking MDBFLinear to repack from the updated signs. + path._gemlite_layers = {} + path.use_gemlite = False + enable_gemlite(device=path.A_amp.device, force=True) + + +def write_back_mdbf_binary(mdbf_modules: List[Tuple[str, nn.Module]]) -> None: + """Write optimised float sign tensors back to packed uint8 buffers.""" + from onecomp.quantizer.mdbf.mdbf_layer import pack_binary + + with torch.no_grad(): + for _name, mod in mdbf_modules: + for path in mod.paths: + signs_changed = False + for sign in _BINARY_SIGN_NAMES: + opt_attr = f"_opt_{sign}_sign" + if not hasattr(path, opt_attr): + continue + w = getattr(path, opt_attr) + q = w.sign() + q[q == 0] = 1 + shape = (path.n, path.r) if sign == "A" else (path.r, path.m) + packed, _ = pack_binary(q.to(torch.int8).reshape(shape)) + buf_key = f"{sign}_sign_packed" + if buf_key in path._buffers: + path._buffers[buf_key].copy_(packed.to(path._buffers[buf_key].device)) + elif sign in path._packed_cpu: + path._packed_cpu[sign].copy_(packed.cpu()) + signs_changed = True + if signs_changed: + _refresh_gemlite_sign_kernels(path) + + +def write_back_mdbf_amp(mdbf_modules: List[Tuple[str, nn.Module]]) -> None: + """Copy float32 optimised amp params back to fp16 buffers for inference.""" + with torch.no_grad(): + for _name, mod in mdbf_modules: + for path in mod.paths: + for attr in _AMP_ATTRS: + opt_attr = f"_opt_{attr}" + if not hasattr(path, opt_attr): + continue + opt_param = getattr(path, opt_attr) + buf = getattr(path, attr) + buf.copy_(opt_param.data.half()) + + +def finalize_mdbf_differentiable( + mdbf_modules: List[Tuple[str, nn.Module]], + original_forwards: Dict[str, object], + write_back: bool, +) -> None: + """Finalize MDBF QAT without overwriting a state restored for rollback.""" + if write_back: + write_back_mdbf_binary(mdbf_modules) + write_back_mdbf_amp(mdbf_modules) + restore_mdbf_original( + mdbf_modules, original_forwards, cleanup=True, + ) + + +# --------------------------------------------------------------------------- +# State save / load (for rollback) +# --------------------------------------------------------------------------- + + +def save_mdbf_state(mdbf_modules: List[Tuple[str, nn.Module]]) -> Dict: + """Snapshot amplitude buffers and packed sign buffers.""" + state: Dict[str, dict] = {} + for name, mod in mdbf_modules: + paths_state = {} + for p, path in enumerate(mod.paths): + d: dict = {} + for attr in _AMP_ATTRS: + d[attr] = getattr(path, attr).data.clone() + for sign in _BINARY_SIGN_NAMES: + buf_key = f"{sign}_sign_packed" + if buf_key in path._buffers: + d[buf_key] = path._buffers[buf_key].clone() + elif sign in path._packed_cpu: + d[buf_key] = path._packed_cpu[sign].clone() + paths_state[p] = d + state[name] = paths_state + return state + + +def load_mdbf_state( + mdbf_modules: List[Tuple[str, nn.Module]], + state: Dict, +) -> None: + """Restore a previously saved snapshot.""" + with torch.no_grad(): + for name, mod in mdbf_modules: + if name not in state: + continue + paths_state = state[name] + for p, path in enumerate(mod.paths): + if p not in paths_state: + continue + d = paths_state[p] + signs_changed = False + for attr in _AMP_ATTRS: + if attr in d: + getattr(path, attr).copy_(d[attr]) + for sign in _BINARY_SIGN_NAMES: + buf_key = f"{sign}_sign_packed" + if buf_key not in d: + continue + if buf_key in path._buffers: + path._buffers[buf_key].copy_(d[buf_key]) + elif sign in path._packed_cpu: + path._packed_cpu[sign].copy_(d[buf_key].cpu()) + signs_changed = True + if signs_changed: + _refresh_gemlite_sign_kernels(path) diff --git a/global_ptq/onecomp_globalptq/global_ptq/_core/trainer.py b/global_ptq/onecomp_globalptq/global_ptq/_core/trainer.py index 6a712535..33b54868 100644 --- a/global_ptq/onecomp_globalptq/global_ptq/_core/trainer.py +++ b/global_ptq/onecomp_globalptq/global_ptq/_core/trainer.py @@ -31,6 +31,12 @@ setup_dbf_forwards_only, write_back_dbf_binary, ) +from .mdbf_adapter import ( + restore_mdbf_original, + setup_mdbf_forwards_only, + write_back_mdbf_binary, + write_back_mdbf_amp, +) logger = getLogger(__name__) @@ -91,9 +97,11 @@ def __init__( self, *, teacher_model: nn.Module, + teacher_device=None, method: str, gptq_modules: list, dbf_modules: list, + mdbf_modules: list = None, original_forwards: dict, optimize_intweight: bool, optimize_binary: bool, @@ -105,9 +113,11 @@ def __init__( ): super().__init__(**kwargs) self.teacher_model = teacher_model + self.teacher_device = teacher_device self.method = method self.gptq_modules = gptq_modules self.dbf_modules = dbf_modules + self.mdbf_modules = mdbf_modules or [] self.original_forwards = original_forwards self.optimize_intweight = optimize_intweight self.optimize_binary = optimize_binary @@ -157,9 +167,19 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs): loss = torch.tensor(0.0, device=logits_s.device) if self.w_distill > 0 and self.teacher_model is not None: with torch.no_grad(): - # Teacher also gets all available inputs - teacher_outputs = self.teacher_model(**inputs) - logits_t = get_logits(teacher_outputs) + if ( + self.teacher_device is not None + and self.teacher_device != logits_s.device + ): + teacher_inputs = { + k: v.to(self.teacher_device) + if isinstance(v, torch.Tensor) else v + for k, v in inputs.items() + } + teacher_outputs = self.teacher_model(**teacher_inputs) + else: + teacher_outputs = self.teacher_model(**inputs) + logits_t = get_logits(teacher_outputs).to(logits_s.device) loss = loss + self.w_distill * compute_kl_loss( logits_t, logits_s, self.temperature, @@ -191,6 +211,10 @@ def evaluate(self, eval_dataset=None, ignore_keys=None, elif self.method == "dbf": write_back_dbf_binary(self.dbf_modules) restore_dbf_original(self.dbf_modules, self.original_forwards) + elif self.method == "mdbf": + write_back_mdbf_binary(self.mdbf_modules) + write_back_mdbf_amp(self.mdbf_modules) + restore_mdbf_original(self.mdbf_modules, self.original_forwards) result = super().evaluate(eval_dataset, ignore_keys, metric_key_prefix) @@ -203,5 +227,9 @@ def evaluate(self, eval_dataset=None, ignore_keys=None, setup_dbf_forwards_only( self.dbf_modules, self.original_forwards, ) + elif self.method == "mdbf": + setup_mdbf_forwards_only( + self.mdbf_modules, self.original_forwards, + ) return result diff --git a/global_ptq/onecomp_globalptq/global_ptq/global_ptq.py b/global_ptq/onecomp_globalptq/global_ptq/global_ptq.py index 3f875b29..f301d334 100644 --- a/global_ptq/onecomp_globalptq/global_ptq/global_ptq.py +++ b/global_ptq/onecomp_globalptq/global_ptq/global_ptq.py @@ -76,8 +76,10 @@ class GlobalPTQ(PostQuantizationProcess): ste_k (float): Smoothness parameter for GPTQ integer-weight Smooth STE rounding. Only used when ``gptq_optimize_intweight=True``. - DBF binary STE uses a fixed internal sharpness (k=2). Default is 100.0. + mdbf_ste_k (float): + Sharpness for MDBF binary sign STE (``tanh(k*x)`` backward). + Default is 2.0. calibration_dataset (list or None): List of text strings to use as calibration data. If ``None`` (default), the AllenAI C4 dataset is @@ -159,7 +161,6 @@ class GlobalPTQ(PostQuantizationProcess): optimiser update. Default is 1 (no accumulation). Incompatible with ``use_sam=True``; when both are set, this value is silently forced to 1. - Examples: >>> from onecomp import Runner, ModelConfig, GPTQ >>> from onecomp_globalptq import GlobalPTQ @@ -184,6 +185,7 @@ class GlobalPTQ(PostQuantizationProcess): dbf_lr: float = 5e-5 optimize_binary: bool = False ste_k: float = 100.0 + mdbf_ste_k: float = 2.0 calibration_dataset: Optional[List[str]] = None num_calibration_samples: int = 128 max_length: int = 2048 @@ -236,6 +238,10 @@ class GlobalPTQ(PostQuantizationProcess): # --- Gradient Accumulation --- grad_accum_steps: int = 1 + # --- Device placement (multi-GPU / CPU teacher) --- + student_device: Optional[str] = None + teacher_device: Optional[str] = None + def __post_init__(self): super().__post_init__() if self.epochs < 1: @@ -251,11 +257,20 @@ def __post_init__(self): f"Available: {list(_VALID_CALIBRATION_STRATEGIES)}" ) - def run( + # OneComp < 1.3.1 made ``run`` abstract. Newer versions provide a + # validated/audited public ``run`` wrapper and make ``_run`` abstract. + # Define the legacy entry point only when the installed base requires it, + # so current OneComp can retain its wrapper instead of being overridden. + if "_run" not in getattr(PostQuantizationProcess, "__abstractmethods__", set()): + + def run(self, quantized_model: nn.Module, model_config: ModelConfig) -> None: + return self._run(quantized_model, model_config) + + def _run( self, quantized_model: nn.Module, model_config: ModelConfig, - ) -> None: + ) -> dict: """Execute global PTQ on the quantized model. Modifies *quantized_model* in-place. The model is returned on @@ -286,6 +301,7 @@ def run( gptq_intweight_lr=self.gptq_intweight_lr, optimize_binary=self.optimize_binary, ste_k=self.ste_k, + mdbf_ste_k=self.mdbf_ste_k, calibration_dataset=self.calibration_dataset, num_calibration_samples=self.num_calibration_samples, max_length=self.max_length, @@ -315,6 +331,8 @@ def run( early_stopping_patience=self.early_stopping_patience, use_mixed_precision=self.use_mixed_precision, grad_accum_steps=self.grad_accum_steps, + student_device=self.student_device, + teacher_device=self.teacher_device, ) except Exception: diff --git a/global_ptq/onecomp_globalptq/global_ptq/global_ptq_distributed.py b/global_ptq/onecomp_globalptq/global_ptq/global_ptq_distributed.py index 7ca9dc54..2835d9b8 100644 --- a/global_ptq/onecomp_globalptq/global_ptq/global_ptq_distributed.py +++ b/global_ptq/onecomp_globalptq/global_ptq/global_ptq_distributed.py @@ -68,8 +68,10 @@ class GlobalPTQDistributed(PostQuantizationProcess): ste_k (float): Smoothness parameter for GPTQ integer-weight Smooth STE rounding. Only used when ``gptq_optimize_intweight=True``. - DBF binary STE uses a fixed internal sharpness (k=2). Default is 100.0. + mdbf_ste_k (float): + Sharpness for MDBF binary sign STE (``tanh(k*x)`` backward). + Default is 2.0. dbf_lr (float): Learning rate for DBF scaling parameters. Default is 5e-5. @@ -169,9 +171,10 @@ class GlobalPTQDistributed(PostQuantizationProcess): gptq_intweight_lr: float = 1e-4 ste_k: float = 100.0 - # --- DBF --- + # --- DBF / MDBF --- dbf_lr: float = 5e-5 optimize_binary: bool = False + mdbf_ste_k: float = 2.0 # --- Calibration --- calibration_dataset: Optional[List[str]] = None @@ -192,6 +195,7 @@ class GlobalPTQDistributed(PostQuantizationProcess): # --- Distributed --- deepspeed_config: Optional[str] = None + teacher_device: Optional[str] = None # --- Output / Logging / Checkpointing --- output_dir: Optional[str] = None @@ -238,7 +242,16 @@ def __post_init__(self): self.save_steps, self.save_strategy, ) - def run( + # OneComp < 1.3.1 made ``run`` abstract. Newer versions provide a + # validated/audited public ``run`` wrapper and make ``_run`` abstract. + # Define the legacy entry point only when the installed base requires it, + # so current OneComp can retain its wrapper instead of being overridden. + if "_run" not in getattr(PostQuantizationProcess, "__abstractmethods__", set()): + + def run(self, quantized_model: nn.Module, model_config: ModelConfig) -> None: + return self._run(quantized_model, model_config) + + def _run( self, quantized_model: nn.Module, model_config: ModelConfig, @@ -274,6 +287,12 @@ def run( write_back_dbf_scaling, restore_dbf_original, ) + from ._core.mdbf_adapter import ( + finalize_mdbf_differentiable, + load_mdbf_state, + save_mdbf_state, + setup_mdbf_differentiable, + ) from onecomp import CalibrationConfig from onecomp.calibration import prepare_calibration_dataset from transformers import TrainingArguments, default_data_collator @@ -292,7 +311,7 @@ def run( if method is None: logger.warning("No quantized layers detected — skipping.") return - if method not in ("gptq", "dbf"): + if method not in ("gptq", "dbf", "mdbf"): logger.info("Method '%s' not supported — skipping.", method) return @@ -333,6 +352,7 @@ def run( gptq_modules = [] dbf_modules = [] + mdbf_modules = [] original_forwards = {} param_groups = [] @@ -375,6 +395,31 @@ def run( f", {len(binary_params)} binary" if binary_params else "", ) + elif method == "mdbf": + mdbf_modules = detected_modules + original_forwards, amp_params, binary_params = ( + setup_mdbf_differentiable( + mdbf_modules, + self.optimize_binary, + ste_k=self.mdbf_ste_k, + ) + ) + logger.info("MDBF binary STE sharpness mdbf_ste_k=%.4g", self.mdbf_ste_k) + all_mdbf_params = list(amp_params) + if binary_params: + all_mdbf_params += binary_params + param_groups = [{ + "params": all_mdbf_params, + "lr": self.dbf_lr, + "weight_decay": 0.0, + }] + logger.info( + "Trainable: %d amp params%s across %d MDBF modules", + len(amp_params), + f", {len(binary_params)} binary" if binary_params else "", + len(mdbf_modules), + ) + # DeepSpeed ZeRO requires contiguous tensors for all-reduce. for pg in param_groups: for p in pg["params"]: @@ -388,6 +433,8 @@ def run( restore_gptq_original(gptq_modules, original_forwards, cleanup=True) elif method == "dbf": restore_dbf_original(dbf_modules, original_forwards) + elif method == "mdbf": + restore_mdbf_original(mdbf_modules, original_forwards, cleanup=True) quantized_model.cpu() return @@ -403,13 +450,26 @@ def run( # 5. Teacher model # ------------------------------------------------------------------ need_teacher = self.w_distill > 0 + teacher_dev = dev if need_teacher: - logger.info("Loading FP16 teacher model...") + world_size = int(os.environ.get("WORLD_SIZE", "1")) + resolved_teacher = self.teacher_device + if resolved_teacher is None and ( + self.deepspeed_config or world_size > 1 + ): + resolved_teacher = "cpu" + if resolved_teacher is not None: + teacher_dev = torch.device(resolved_teacher) + logger.info( + "Loading FP16 teacher model (teacher_device=%s)...", + teacher_dev, + ) teacher_model = model_config.load_model(device_map="cpu") teacher_model.eval() for p in teacher_model.parameters(): p.requires_grad = False - teacher_model.to(dev) + if teacher_dev.type != "cpu": + teacher_model.to(teacher_dev) else: logger.info( "w_distill=0 — skipping teacher model load (pure QAT mode)." @@ -417,7 +477,7 @@ def run( # ------------------------------------------------------------------ # 6. TrainingArguments # ------------------------------------------------------------------ - lr = self.gptq_lr if method == "gptq" else self.dbf_lr + lr = self.gptq_lr if method == "gptq" else self.dbf_lr # dbf_lr used for both dbf and mdbf resolved_output_dir = ( self.output_dir if self.output_dir is not None @@ -459,8 +519,10 @@ def run( # Save initial state for rollback if training degrades quality if method == "gptq": _initial_state = save_gptq_state(gptq_modules) - else: + elif method == "dbf": _initial_state = save_dbf_state(dbf_modules) + else: # mdbf + _initial_state = save_mdbf_state(mdbf_modules) # ------------------------------------------------------------------ # 7. Train @@ -468,9 +530,11 @@ def run( trainer = _GlobalPTQTrainer( model=quantized_model, teacher_model=teacher_model, + teacher_device=teacher_dev if need_teacher else None, method=method, gptq_modules=gptq_modules, dbf_modules=dbf_modules, + mdbf_modules=mdbf_modules, original_forwards=original_forwards, optimize_intweight=self.gptq_optimize_intweight, optimize_binary=self.optimize_binary, @@ -523,8 +587,10 @@ def run( ) if method == "gptq": load_gptq_state(gptq_modules, _initial_state) - else: + elif method == "dbf": load_dbf_state(dbf_modules, _initial_state) + else: # mdbf + load_mdbf_state(mdbf_modules, _initial_state) else: logger.info( "Best eval_loss: %.6f (initial=%.6f) at step %d.", @@ -539,6 +605,11 @@ def run( write_back_dbf_binary(dbf_modules) write_back_dbf_scaling(dbf_modules) restore_dbf_original(dbf_modules, original_forwards) + elif method == "mdbf": + finalize_mdbf_differentiable( + mdbf_modules, original_forwards, + write_back=not rollback_happened, + ) finally: if original_use_cache is not None: quantized_model.config.use_cache = original_use_cache diff --git a/global_ptq/pyproject.toml b/global_ptq/pyproject.toml index 3a3846ca..2e79c166 100644 --- a/global_ptq/pyproject.toml +++ b/global_ptq/pyproject.toml @@ -13,7 +13,7 @@ authors = [ ] requires-python = ">=3.12, <3.14" dependencies = [ - "onecomp >= 1.1.0", + "onecomp >= 1.3.1", "torch", "transformers >= 5.3.0", ] diff --git a/global_ptq/tests/test_global_ptq.py b/global_ptq/tests/test_global_ptq.py index faea2682..2ce032d3 100644 --- a/global_ptq/tests/test_global_ptq.py +++ b/global_ptq/tests/test_global_ptq.py @@ -450,6 +450,323 @@ def test_state_save_load_roundtrip(self): # =========================================================================== +def _make_synthetic_mdbf_linear( + in_dim=16, + out_dim=16, + rank=8, + l=2, + P=1, + device="cpu", + use_gemlite=False, + bias=None, +): + from onecomp.quantizer.mdbf.initialize import MDBFParams + from onecomp.quantizer.mdbf.mdbf_layer import MultipathMDBFLinear + + params_list = [] + for _ in range(P): + A_sign = torch.sign(torch.randn(out_dim, rank)) + A_sign[A_sign == 0] = 1 + B_sign = torch.sign(torch.randn(rank, in_dim)) + B_sign[B_sign == 0] = 1 + params_list.append( + MDBFParams( + A_sign=A_sign, + B_sign=B_sign, + A_amp=torch.randn(out_dim, l).abs() + 0.01, + B_amp=torch.randn(in_dim, l).abs() + 0.01, + Q_U_amp=torch.randn(rank, l).abs() + 0.01, + Q_V_amp=torch.randn(rank, l).abs() + 0.01, + ) + ) + return MultipathMDBFLinear( + params_list, + bias=bias, + device=device, + use_gemlite=use_gemlite, + ) + + +class _TinyMDBFModel(nn.Module): + def __init__(self, in_dim=16, out_dim=16, rank=8, l=2, P=1, device="cpu"): + super().__init__() + self.layer1 = _make_synthetic_mdbf_linear(in_dim, out_dim, rank, l, P, device) + self.layer2 = _make_synthetic_mdbf_linear(in_dim, out_dim, rank, l, P, device) + + def forward(self, x): + return self.layer2(self.layer1(x)) + + +class TestMdbfAdapter: + def test_find_modules(self): + from onecomp_globalptq.global_ptq._core.mdbf_adapter import find_mdbf_modules + assert len(find_mdbf_modules(_TinyMDBFModel())) == 2 + assert find_mdbf_modules(nn.Linear(10, 10)) == [] + + def test_amp_params_require_grad(self): + from onecomp_globalptq.global_ptq._core.mdbf_adapter import ( + find_mdbf_modules, setup_mdbf_differentiable, + ) + model = _TinyMDBFModel() + modules = find_mdbf_modules(model) + _, amp, binary = setup_mdbf_differentiable(modules, optimize_binary=False) + # 2 layers x 1 path x 4 amplitude tensors + assert len(amp) == 8 + assert len(binary) == 0 + assert all(p.requires_grad for p in amp) + + def test_binary_params_with_optimize_binary(self): + from onecomp_globalptq.global_ptq._core.mdbf_adapter import ( + find_mdbf_modules, setup_mdbf_differentiable, + ) + model = _TinyMDBFModel() + modules = find_mdbf_modules(model) + _, amp, binary = setup_mdbf_differentiable(modules, optimize_binary=True) + assert len(amp) == 8 + # 2 layers x 1 path x (A_sign, B_sign) + assert len(binary) == 4 + assert all(bp.requires_grad for bp in binary) + + def test_amp_params_are_float32(self): + from onecomp_globalptq.global_ptq._core.mdbf_adapter import ( + find_mdbf_modules, setup_mdbf_differentiable, + ) + modules = find_mdbf_modules(_TinyMDBFModel()) + _, amp, _ = setup_mdbf_differentiable(modules) + assert all(p.dtype == torch.float32 for p in amp) + + def test_multipath_param_counts(self): + from onecomp_globalptq.global_ptq._core.mdbf_adapter import ( + find_mdbf_modules, setup_mdbf_differentiable, + ) + modules = find_mdbf_modules(_TinyMDBFModel(P=2)) + _, amp, binary = setup_mdbf_differentiable(modules, optimize_binary=True) + assert len(amp) == 16 + assert len(binary) == 8 + + @pytest.mark.parametrize("optimize_binary", [False, True]) + def test_differentiable_forward_matches_inference(self, optimize_binary): + from onecomp_globalptq.global_ptq._core.mdbf_adapter import ( + find_mdbf_modules, setup_mdbf_differentiable, + ) + torch.manual_seed(0) + model = _make_synthetic_mdbf_linear( + in_dim=16, out_dim=12, rank=8, l=3, P=2, + bias=torch.randn(12), + ) + x = torch.randn(4, 16) + expected = model(x) + modules = find_mdbf_modules(model) + setup_mdbf_differentiable( + modules, optimize_binary=optimize_binary, + ) + actual = model(x) + assert torch.allclose(actual, expected, rtol=1e-4, atol=1e-4) + + def test_all_parameters_receive_finite_nonzero_gradients(self): + from onecomp_globalptq.global_ptq._core.mdbf_adapter import ( + find_mdbf_modules, setup_mdbf_differentiable, + ) + torch.manual_seed(1) + model = _TinyMDBFModel(P=2, l=3) + modules = find_mdbf_modules(model) + _, amp, binary = setup_mdbf_differentiable( + modules, optimize_binary=True, + ) + model(torch.randn(3, 16)).square().mean().backward() + for param in amp + binary: + assert param.grad is not None + assert torch.isfinite(param.grad).all() + assert torch.count_nonzero(param.grad) > 0 + + def test_amp_gradient_flows(self): + from onecomp_globalptq.global_ptq._core.mdbf_adapter import ( + find_mdbf_modules, setup_mdbf_differentiable, + ) + model = _TinyMDBFModel() + modules = find_mdbf_modules(model) + _, amp, _ = setup_mdbf_differentiable(modules) + model(torch.randn(2, 16)).sum().backward() + assert any(p.grad is not None and p.grad.abs().sum() > 0 for p in amp) + + def test_binary_weight_receives_gradient(self): + from onecomp_globalptq.global_ptq._core.mdbf_adapter import ( + find_mdbf_modules, restore_mdbf_original, setup_mdbf_differentiable, + ) + model = _TinyMDBFModel() + modules = find_mdbf_modules(model) + orig, _, binary = setup_mdbf_differentiable(modules, optimize_binary=True) + model(torch.randn(2, 16)).sum().backward() + restore_mdbf_original(modules, orig) + assert any(bp.grad is not None and bp.grad.abs().sum() > 0 for bp in binary) + + def test_write_back_amp_restores_fp16_buffers(self): + from onecomp_globalptq.global_ptq._core.mdbf_adapter import ( + find_mdbf_modules, setup_mdbf_differentiable, write_back_mdbf_amp, + ) + modules = find_mdbf_modules(_TinyMDBFModel()) + setup_mdbf_differentiable(modules) + expected = [] + for _, mod in modules: + for path in mod.paths: + for attr in ("A_amp", "B_amp", "Q_U_amp", "Q_V_amp"): + opt = getattr(path, f"_opt_{attr}") + opt.data.add_(0.125) + expected.append(opt.data.half().clone()) + write_back_mdbf_amp(modules) + index = 0 + for _, mod in modules: + for path in mod.paths: + for attr in ("A_amp", "B_amp", "Q_U_amp", "Q_V_amp"): + actual = getattr(path, attr) + assert actual.dtype == torch.float16 + assert torch.equal(actual, expected[index]) + index += 1 + + def test_write_back_binary_keeps_signs_packed(self): + from onecomp.quantizer.mdbf.mdbf_layer import unpack_binary + from onecomp_globalptq.global_ptq._core.mdbf_adapter import ( + find_mdbf_modules, setup_mdbf_differentiable, write_back_mdbf_binary, + ) + modules = find_mdbf_modules(_TinyMDBFModel()) + setup_mdbf_differentiable(modules, optimize_binary=True) + expected = [] + for _, mod in modules: + for path in mod.paths: + for sign in ("A", "B"): + opt = getattr(path, f"_opt_{sign}_sign") + opt.data.neg_() + expected.append(opt.data.sign().clone()) + write_back_mdbf_binary(modules) + index = 0 + for _, mod in modules: + for path in mod.paths: + for sign, shape in ( + ("A", (path.n, path.r)), + ("B", (path.r, path.m)), + ): + packed = path._buffers[f"{sign}_sign_packed"] + assert packed.dtype == torch.uint8 + actual = unpack_binary(packed, shape).float() + assert torch.equal(actual, expected[index]) + index += 1 + + def test_cleanup_removes_shadow_parameters(self): + """After teardown no ``_opt_*`` shadows must remain in the state dict.""" + from onecomp_globalptq.global_ptq._core.mdbf_adapter import ( + find_mdbf_modules, restore_mdbf_original, setup_mdbf_differentiable, + ) + model = _TinyMDBFModel() + modules = find_mdbf_modules(model) + orig, _, _ = setup_mdbf_differentiable(modules, optimize_binary=True) + assert any(k.startswith("layer1.paths.0._opt_") for k in model.state_dict()) + restore_mdbf_original(modules, orig, cleanup=True) + assert not [k for k in model.state_dict() if "_opt_" in k] + assert not [n for n, _ in model.named_parameters() if "_opt_" in n] + + def test_state_save_load_roundtrip(self): + """The same input and every persisted MDBF buffer survive a roundtrip.""" + from onecomp_globalptq.global_ptq._core.mdbf_adapter import ( + find_mdbf_modules, load_mdbf_state, save_mdbf_state, + ) + torch.manual_seed(2) + model = _TinyMDBFModel() + modules = find_mdbf_modules(model) + state = save_mdbf_state(modules) + x = torch.randn(2, 16) + before = model(x).clone() + for _, mod in modules: + for path in mod.paths: + for attr in ("A_amp", "B_amp", "Q_U_amp", "Q_V_amp"): + getattr(path, attr).zero_() + for sign in ("A", "B"): + key = f"{sign}_sign_packed" + if key in path._buffers: + path._buffers[key].zero_() + elif sign in path._packed_cpu: + path._packed_cpu[sign].zero_() + load_mdbf_state(modules, state) + after = model(x) + assert torch.equal(after, before) + restored = save_mdbf_state(modules) + assert restored.keys() == state.keys() + for name, paths in state.items(): + for path_index, buffers in paths.items(): + for key, expected in buffers.items(): + assert torch.equal( + restored[name][path_index][key], expected, + ) + + def test_rollback_finalize_does_not_overwrite_restored_state(self): + from onecomp_globalptq.global_ptq._core.mdbf_adapter import ( + finalize_mdbf_differentiable, + find_mdbf_modules, + load_mdbf_state, + save_mdbf_state, + setup_mdbf_differentiable, + ) + model = _TinyMDBFModel() + modules = find_mdbf_modules(model) + initial = save_mdbf_state(modules) + original_forwards, _, _ = setup_mdbf_differentiable( + modules, optimize_binary=True, + ) + for _, mod in modules: + for path in mod.paths: + path._opt_A_amp.data.add_(10) + path._opt_A_sign.data.neg_() + load_mdbf_state(modules, initial) + finalize_mdbf_differentiable( + modules, original_forwards, write_back=False, + ) + restored = save_mdbf_state(modules) + for name, paths in initial.items(): + for path_index, buffers in paths.items(): + for key, expected in buffers.items(): + assert torch.equal( + restored[name][path_index][key], expected, + ) + + @pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA not available", + ) + def test_gemlite_binary_shadows_and_repacked_kernels(self): + from onecomp_globalptq.global_ptq._core.mdbf_adapter import ( + find_mdbf_modules, + restore_mdbf_original, + setup_mdbf_differentiable, + write_back_mdbf_binary, + ) + layer = _make_synthetic_mdbf_linear( + in_dim=256, out_dim=256, rank=128, l=1, P=1, + device="cuda:0", use_gemlite=True, + ) + path = layer.paths[0] + if not path.use_gemlite: + pytest.skip("GemLite is unavailable for this CUDA environment") + + x = torch.randn(2, 256, device="cuda:0", dtype=torch.float16) + before = layer(x) + old_kernels = dict(path._gemlite_layers) + modules = find_mdbf_modules(layer) + original_forwards, _, binary = setup_mdbf_differentiable( + modules, optimize_binary=True, + ) + assert binary + assert all(param.device == path.A_amp.device for param in binary) + path._opt_A_sign.data.neg_() + write_back_mdbf_binary(modules) + restore_mdbf_original(modules, original_forwards, cleanup=True) + + assert path.use_gemlite + assert all( + path._gemlite_layers[name] is not kernel + for name, kernel in old_kernels.items() + ) + after = layer(x) + assert not torch.equal(after, before) + + class TestDetectQuantizationMethod: def test_detects_gptq(self): from onecomp_globalptq.global_ptq._core.helpers import detect_quantization_method @@ -463,6 +780,25 @@ def test_detects_dbf(self): assert method == "dbf" assert len(modules) == 2 + def test_detects_mdbf(self): + from onecomp_globalptq.global_ptq._core.helpers import detect_quantization_method + method, modules = detect_quantization_method(_TinyMDBFModel()) + assert method == "mdbf" + assert len(modules) == 2 + + def test_gptq_takes_priority_over_mdbf(self): + from onecomp_globalptq.global_ptq._core.helpers import detect_quantization_method + + class _Mixed(nn.Module): + def __init__(self): + super().__init__() + self.gptq = _TinyGPTQModel() + self.mdbf = _TinyMDBFModel() + + method, modules = detect_quantization_method(_Mixed()) + assert method == "gptq" + assert len(modules) == 2 + def test_plain_model_returns_none(self): from onecomp_globalptq.global_ptq._core.helpers import detect_quantization_method method, modules = detect_quantization_method(nn.Linear(10, 10)) diff --git a/global_ptq/uv.lock b/global_ptq/uv.lock index 3c6b55fd..23c7fbdf 100644 --- a/global_ptq/uv.lock +++ b/global_ptq/uv.lock @@ -2200,8 +2200,8 @@ wheels = [ [[package]] name = "onecomp" -version = "1.1.1" -source = { git = "https://github.com/FujitsuResearch/OneCompression.git?branch=main#da41b49e3ea578570ab5963f73efc90c911f3202" } +version = "1.3.1" +source = { git = "https://github.com/FujitsuResearch/OneCompression.git?branch=main#53c0be25fc4bd79eb51e29ec267c32d30adbe71e" } dependencies = [ { name = "accelerate" }, { name = "datasets" },