feat(global_ptq): wire MDBF QAT into the KL distillation loop

- helpers.detect_quantization_method: recognize MultipathMDBFLinear
  layers as method "mdbf" (priority gptq > dbf > mdbf). Without this a
  MDBF-quantized model was reported as having no quantized layers and
  the distillation loop skipped it entirely.
- core.run_kl_distillation: add the "mdbf" branch (setup/teardown via
  mdbf_adapter, param groups for amplitude/binary params using dbf_lr)
  and expose mdbf_ste_k for the sign STE sharpness.
- core.eval_kl / trainer._GlobalPTQTrainer.compute_loss: support a
  separate teacher_device so the FP16 teacher can be kept on CPU or a
  second GPU while the student trains under DeepSpeed ZeRO-2 with CPU
  optimizer offload. Required to fit 1bpw Llama2-7B/13B QAT on 80GB
  cards.
- global_ptq.py / global_ptq_distributed.py: expose the corresponding
  GlobalPTQConfig(Distributed) fields (mdbf_ste_k, student_device,
  teacher_device) and thread them through to run_kl_distillation and
  the Trainer-based distributed path.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: fujisawa-yoshihiko

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,17 +427,33 @@ 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,
     teacher_model: nn.Module,
     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,17 @@ 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.
 
     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 +662,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 +696,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 +707,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,13 +743,33 @@ 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.")
         if method == "gptq":
             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 +924,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 +988,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,23 +1113,33 @@ 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
 
             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)
 
             # Restore non-EMA params for continued training
             if ema_tracker is not None:
@@ -1103,27 +1174,36 @@ 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":
             write_back_gptq_params(gptq_modules, gptq_optimize_intweight)
         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,14 +1220,18 @@ 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":
         # Final cleanup of differentiable parameters
         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()

global_ptq/onecomp_globalptq/global_ptq/_core/helpers.py

@@ -76,30 +76,34 @@ 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
+    from onecomp.quantizer.mdbf.mdbf_layer import MultipathMDBFLinear
 
     gptq_modules = find_target_modules(model, GPTQLinear)
     dbf_modules = find_target_modules(model, DoubleBinaryLinear)
+    mdbf_modules = find_target_modules(model, MultipathMDBFLinear)
 
-    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, []