Run fused-MoE expert GPTQ updates on all of this process's GPUs

modelopt/torch/quantization/ggml/gptq.py

@@ -78,11 +78,15 @@ class GGMLGPTQHelper(GPTQHelper):
     those exact codes rather than encoding the GPTQ'd weight again, which would not return them.
     """
 
-    def update_weights(self, block_size, perc_damp, weight=None):
+    # The update reads only the weight, its Hessian and the format's codec.
+    UPDATES_ON_ANY_DEVICE = True
+
+    def update_weights(self, block_size, perc_damp, weight=None, device=None):
         """Run the GPTQ update, then pin the chosen payload to the weight quantizer."""
-        super().update_weights(block_size, perc_damp, weight)
+        super().update_weights(block_size, perc_damp, weight, device)
         quantizer = self.weight_quantizer
-        pin_packed_weight(quantizer, quantizer.num_bits, self._packed)
+        weight = self.module.weight if weight is None else weight
+        pin_packed_weight(quantizer, quantizer.num_bits, self._packed.to(weight.device))
         self._packed = None
 
     def _blockwise_update(self, block_size):

modelopt/torch/quantization/model_calib.py

@@ -69,7 +69,12 @@
     persistent_materialization,
     promote_static_block_weight_quantizers,
 )
-from .utils.calib_utils import _GPTQ_HELPER_REGISTRY, _IMATRIX_ENCODER_REGISTRY, GPTQHelper
+from .utils.calib_utils import (
+    _GPTQ_HELPER_REGISTRY,
+    _IMATRIX_ENCODER_REGISTRY,
+    GPTQHelper,
+    update_expert_weights,
+)
 
 __all__ = [
     "CalibratorFactory",
@@ -2466,18 +2471,17 @@ def _expert_hook(_input_quantizer, _args, output):
             handle.update_weights(block_size, perc_damp)
         handle.free()
     for name, module in fused_experts:
-        updated = unreached = 0
         with enable_weight_access_and_writeback(module, model, names):
-            for weight, quantizer in module.iter_weights_for_calibration():
-                handle = expert_handles.pop(id(quantizer), None)
-                if handle is not None:
-                    unreached += handle.n_samples == 0
-                    handle.update_weights(block_size, perc_damp, weight)
-                    handle.free()
-                    updated += 1
-        if updated:
+            work = [
+                (expert_handles.pop(id(quantizer)), weight)
+                for weight, quantizer in module.iter_weights_for_calibration()
+                if id(quantizer) in expert_handles
+            ]
+            unreached = sum(handle.n_samples == 0 for handle, _ in work)
+            update_expert_weights(work, block_size, perc_damp)
+        if work:
             print_rank_0(
-                f"[{name}] GPTQ updated {updated} expert weights, {unreached} of them never "
+                f"[{name}] GPTQ updated {len(work)} expert weights, {unreached} of them never "
                 "reached by calibration"
             )
     del gptq_handles, expert_handles

modelopt/torch/quantization/utils/calib_utils.py

@@ -38,6 +38,9 @@
 
 """GPTQ helper and Hessian utilities for calibration."""
 
+import queue
+from concurrent.futures import ThreadPoolExecutor
+
 import torch
 import torch.distributed as dist
 
@@ -62,16 +65,72 @@ def _hessian_offload_device(nbytes: int) -> torch.device:
     and only outside distributed runs, so no other job's or rank's GPU is touched. A candidate
     must also have the Hessian's ``nbytes`` free.
     """
-    if not (dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1):
-        gpus = [i for i in range(torch.cuda.device_count()) if torch.cuda.memory_reserved(i)]
-        for gpu in sorted(gpus, key=get_used_gpu_mem_fraction):
-            if _gpu_is_full(gpu):
-                break
-            if torch.cuda.mem_get_info(gpu)[0] > nbytes:
-                return torch.device("cuda", gpu)
+    for gpu in sorted(_process_gpus(), key=get_used_gpu_mem_fraction):
+        if _gpu_is_full(gpu):
+            break
+        if torch.cuda.mem_get_info(gpu)[0] > nbytes:
+            return torch.device("cuda", gpu)
     return torch.device("cpu")
 
 
+def _process_gpus() -> list[int]:
+    """GPUs this process already holds memory on; none in a multi-rank run, where they may be another rank's."""
+    if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
+        return []
+    return [i for i in range(torch.cuda.device_count()) if torch.cuda.memory_reserved(i)]
+
+
+def update_expert_weights(work, block_size, perc_damp):
+    """Run the GPTQ updates of ``(helper, weight)`` pairs, spread over this process's GPUs.
+
+    Expert updates are independent, and in a layerwise run of a model sharded over several GPUs
+    only the current layer's GPU works. Helpers whose update can run on any device run in one
+    worker thread per GPU this process uses that has room for an update, each writing its result
+    back to its weight; the others, or all of them when fewer than two GPUs qualify, run in turn on
+    their weight's device.
+    """
+    movable = [(h, w) for h, w in work if h.UPDATES_ON_ANY_DEVICE and w.is_cuda]
+    # A worker holds one expert at a time: its weight, Hessian, inverse and their temporaries.
+    room = 8 * max((4 * h.hessian.numel() for h, _ in movable), default=0)
+    gpus = [gpu for gpu in _process_gpus() if torch.cuda.mem_get_info(gpu)[0] > room]
+    if len(gpus) < 2 or len(movable) < 2:
+        movable = []
+    moved = {id(helper) for helper, _ in movable}
+    for helper, weight in work:
+        if id(helper) not in moved:
+            helper.update_weights(block_size, perc_damp, weight)
+            helper.free()
+    if not movable:
+        return
+
+    # The first update runs here, so lazy setup such as loading a format's CUDA extension, which is
+    # not thread-safe, happens once before the workers start.
+    helper, weight = movable[0]
+    helper.update_weights(block_size, perc_damp, weight)
+    helper.free()
+    pending: queue.SimpleQueue = queue.SimpleQueue()
+    for item in movable[1:]:
+        pending.put(item)
+
+    def worker(gpu):
+        torch.cuda.set_device(gpu)  # per thread
+        device = torch.device("cuda", gpu)
+        while True:
+            try:
+                helper, weight = pending.get_nowait()
+            except queue.Empty:
+                break
+            helper.update_weights(block_size, perc_damp, weight, device=device)
+            helper.free()
+        torch.cuda.synchronize(device)
+
+    with ThreadPoolExecutor(len(gpus)) as pool:
+        for future in [pool.submit(worker, gpu) for gpu in gpus]:
+            future.result()
+    for device in {weight.device for _, weight in movable}:
+        torch.cuda.synchronize(device)
+
+
 def update_hessian(input, hessian, n_samples):
     """Update hessian matrix with new input samples using incremental formula.
 
@@ -156,6 +215,9 @@ class GPTQHelper:
     """
 
     CACHE_NAME = "_forward_no_gptq_hessian"
+    # Whether update_weights() may run on a device other than its weight's: true only for helpers
+    # whose update reads no quantizer state, which lives on the weight's device.
+    UPDATES_ON_ANY_DEVICE = False
 
     def __init__(
         self,
@@ -253,20 +315,22 @@ def free(self):
         self.weight = None
         self.h_inv = None
 
-    def update_weights(self, block_size, perc_damp, weight=None):
+    def update_weights(self, block_size, perc_damp, weight=None, device=None):
         """Run GPTQ blockwise weight update on this helper's weight.
 
         Populates ``self.weight`` and ``self.h_inv``, runs the blockwise update, logs MSE, and
         overwrites ``weight`` (default ``module.weight``) in place, so an expert's slice of a
-        fused 3-D parameter is written back into that parameter.
+        fused 3-D parameter is written back into that parameter. The update runs on ``device``,
+        by default the weight's; only a helper with ``UPDATES_ON_ANY_DEVICE`` may pass another.
         """
         weight = self.module.weight if weight is None else weight
-        hessian = self.hessian.to(weight.device)
-        self.weight = weight.data.float().clone()
+        device = weight.device if device is None else device
+        hessian = self.hessian.to(device)
+        self.weight = weight.data.to(device, torch.float32, copy=True)
         self._prepare_hessian_inverse(hessian, perc_damp)
         self._blockwise_update(block_size)
         if self.log_mse:
-            self._print_mse_error(hessian, weight)
+            self._print_mse_error(hessian, weight.data.to(device))
         weight.data.copy_(self.weight.reshape(weight.shape))
 
     # ------------------------------------------------------------------

tests/gpu/torch/quantization/test_gptq.py

@@ -18,12 +18,16 @@
 
 import pytest
 import torch
+from _test_utils.torch.quantization.models import TOY_MOE_HIDDEN, ToyFusedExperts, ToyMoE
 from _test_utils.torch.transformers_models import get_tiny_llama
 from conftest import requires_triton
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.unified_export_hf import _export_quantized_weight
+from modelopt.torch.quantization.ggml.common import pinned_packed_weight
+from modelopt.torch.quantization.ggml.gptq import GGMLGPTQHelper
 from modelopt.torch.quantization.model_calib import gptq
+from modelopt.torch.quantization.nn import QuantModuleRegistry
 from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor
 from modelopt.torch.quantization.utils import calib_utils
 from modelopt.torch.quantization.utils.calib_utils import (
@@ -57,6 +61,55 @@ def used_fraction(device):
     torch.testing.assert_close(helper.hessian.to("cuda:0"), reference)
 
 
+def test_expert_updates_on_another_gpu_match_one_gpu(need_2_gpus, monkeypatch):
+    """GGML expert updates spread over two GPUs give the one-GPU weights and codes, pinned at home."""
+    pytest.importorskip("transformers")  # the fused-experts quant wrapper is the HF plugin's
+    torch.manual_seed(RAND_SEED)
+    reference = ToyMoE().cuda()
+    inputs = torch.randn(96, TOY_MOE_HIDDEN, device="cuda")
+    shard = torch.zeros(1, device="cuda:1")  # noqa: F841  # the model also lives on cuda:1
+    config = {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "*weight_quantizer",
+                "cfg": {"num_bits": "iq2_xxs", "block_sizes": {-1: 256}, "backend": "ggml"},
+            },
+        ],
+        "algorithm": {"method": "gptq", "block_size": 256, "perc_damp": 0.3},
+    }
+    devices = []
+    update_weights = GGMLGPTQHelper.update_weights
+
+    def recording_update_weights(self, *args, device=None, **kwargs):
+        devices.append(device)
+        return update_weights(self, *args, device=device, **kwargs)
+
+    monkeypatch.setattr(GGMLGPTQHelper, "update_weights", recording_update_weights)
+    experts = []
+    try:
+        for gpus in ([0], [0, 1]):
+            monkeypatch.setattr(calib_utils, "_process_gpus", lambda gpus=gpus: gpus)
+            model = copy.deepcopy(reference)
+            mtq.quantize(model, config, forward_loop=lambda m: m(inputs))
+            experts.append(model.mlp.experts)
+    finally:
+        QuantModuleRegistry.unregister(ToyFusedExperts)
+
+    assert torch.device("cuda", 1) in devices
+    one_gpu, two_gpus = experts
+    for name in ("gate_up_proj", "down_proj"):
+        assert torch.equal(getattr(one_gpu, name), getattr(two_gpus, name))
+        for q1, q2 in zip(
+            getattr(one_gpu, f"{name}_weight_quantizers"),
+            getattr(two_gpus, f"{name}_weight_quantizers"),
+        ):
+            assert pinned_packed_weight(q2, "iq2_xxs").device == getattr(two_gpus, name).device
+            assert torch.equal(
+                pinned_packed_weight(q1, "iq2_xxs"), pinned_packed_weight(q2, "iq2_xxs")
+            )
+
+
 @pytest.mark.parametrize("blocker", ["idle_gpu_full", "idle_gpu_too_small", "multi_rank"])
 def test_hessian_of_a_full_gpu_falls_back_to_the_cpu(need_2_gpus, monkeypatch, blocker):
     """With no other GPU that can take it, or on any rank of a multi-rank run, it goes to the CPU."""