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

modelopt/torch/quantization/ggml/gptq.py

@@ -19,6 +19,7 @@
 
 import torch
 
+from ..extensions import get_cuda_ext_ggml
 from ..utils.calib_utils import GPTQHelper, register_gptq_helper
 from .common import floored_second_moments, pin_packed_weight
 from .registry import GGML_FORMAT_REGISTRY
@@ -78,11 +79,20 @@ 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):
+    def updates_on_any_device(self):
+        """Whether the CUDA encoders are available.
+
+        The update reads only the weight, its Hessian and the format's codec, but the PyTorch
+        fallback's scratch grows with ``block_chunk_size``, past the room the update is checked for.
+        """
+        return get_cuda_ext_ggml() is not None
+
+    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

@@ -21,7 +21,6 @@
 import warnings
 from collections import Counter, defaultdict
 from collections.abc import Callable, Mapping, Sequence
-from concurrent.futures import ThreadPoolExecutor
 from contextlib import nullcontext
 from functools import partial
 from typing import Any, TypeAlias
@@ -79,7 +78,9 @@
     _IMATRIX_ENCODER_REGISTRY,
     GPTQHelper,
     bfloat16_hessian_gemm,
+    run_on_each_gpu,
     track_calib_forward_memory,
+    update_expert_weights,
 )
 
 __all__ = [
@@ -2648,15 +2649,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
+            work = [
+                (expert_handles.pop(id(quantizer)), weight)
+                for weight, quantizer in module.iter_weights_for_calibration()
+                if id(quantizer) in expert_handles
+            ]
+            updated = len(work)
+            unreached = sum(handle.n_samples == 0 for handle, _ in work)
+            update_expert_weights(work, block_size, perc_damp)
+            # Its weight slices would keep this module's onloaded weights alive past the window.
+            del work
         if updated:
             print_rank_0(
                 f"[{name}] GPTQ updated {updated} expert weights, {unreached} of them never "
@@ -2794,25 +2797,11 @@ def _encode_on_each_gpu(by_device, encode):
         if device not in gpus:
             for item in items:
                 encode(*item)
-    if not gpus:
-        return
-    # The first encode runs here, so lazy setup that is not thread-safe, such as loading a format's
-    # CUDA extension, happens once before the workers start.
-    first, *by_device[gpus[0]] = by_device[gpus[0]]
-    encode(*first)
-    inference = torch.is_inference_mode_enabled()
-
-    def worker(device):
-        # The device, grad mode and inference mode are per thread: take the caller's.
-        torch.cuda.set_device(device)
-        with torch.inference_mode() if inference else torch.no_grad():
-            for item in by_device[device]:
-                encode(*item)
-        torch.cuda.synchronize(device)
-
-    with ThreadPoolExecutor(len(gpus)) as pool:
-        for future in [pool.submit(worker, gpu) for gpu in gpus]:
-            future.result()
+    if gpus:
+        pending = {device: iter(by_device[device]) for device in gpus}
+        run_on_each_gpu(
+            gpus, lambda device: next(pending[device], None), lambda item, _: encode(*item)
+        )
 
 
 def _run_weight_scale_calibration(model, forward_loop, scale_algorithm):

modelopt/torch/quantization/utils/calib_utils.py

@@ -39,6 +39,8 @@
 """GPTQ helper and Hessian utilities for calibration."""
 
 import functools
+import queue
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 
 import torch
@@ -106,13 +108,103 @@ def _hessian_offload_device(nbytes: int) -> torch.device:
     rescale of the whole matrix. Only GPUs this process already holds memory on are candidates,
     and only outside distributed runs, so no other job's or rank's GPU is touched.
     """
-    if not (dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1):
-        for gpu in sorted(get_process_gpus(), key=get_available_gpu_memory, reverse=True):
-            if _gpu_has_room(gpu, nbytes):
-                return torch.device("cuda", gpu)
+    for gpu in sorted(_own_gpus(), key=get_available_gpu_memory, reverse=True):
+        if _gpu_has_room(gpu, nbytes):
+            return torch.device("cuda", gpu)
     return torch.device("cpu")
 
 
+def _own_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 get_process_gpus()
+
+
+def _update_room(helper, weight, block_size) -> int:
+    """Bytes a GPU needs to run ``helper``'s update of ``weight``, as measured with the CUDA codecs.
+
+    The Hessian, the inverse and their temporaries; the float32 copy of the weight with a
+    dtype-conversion temporary (a projection much wider than its input makes the weight the larger
+    part); the block's errors and up to four temporaries of a group no wider than the block, each
+    as large as the weight when the input is no wider than a block; and the codec's tables and
+    scratch, about 35 MiB on a GPU's first update.
+    """
+    columns = weight.shape[-1]
+    rows = weight.numel() // columns
+    block = min(block_size, columns)
+    return 2**26 + 4 * (8 * helper.hessian.numel() + rows * (2 * columns + 5 * block))
+
+
+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 w.is_cuda and h.updates_on_any_device()]
+    # A worker holds one expert at a time.
+    room = max((_update_room(h, w, block_size) for h, w in movable), default=0)
+    gpus = [torch.device("cuda", g) for g in _own_gpus() if get_available_gpu_memory(g) > room]
+    if len(gpus) < 2 or len(movable) < 2:
+        movable = []
+
+    def update(item, device):
+        helper, weight = item
+        helper.update_weights(block_size, perc_damp, weight, device=device)
+        helper.free()
+
+    moved = {id(helper) for helper, _ in movable}
+    for item in work:
+        if id(item[0]) not in moved:
+            update(item, None)
+    if not movable:
+        return
+
+    pending: queue.SimpleQueue = queue.SimpleQueue()
+    for item in movable:
+        pending.put(item)
+
+    def next_item(device):
+        try:
+            return pending.get_nowait()
+        except queue.Empty:
+            return None
+
+    run_on_each_gpu(gpus, next_item, update)
+    for device in {weight.device for _, weight in movable}:
+        torch.cuda.synchronize(device)
+
+
+def run_on_each_gpu(gpus, next_item, run):
+    """Call ``run(item, device)`` on the items ``next_item(device)`` hands out, one thread per GPU.
+
+    ``next_item`` returns the next item for that GPU, or ``None`` when there is none left; each GPU's
+    items can be its own or come from a queue all of them share. The first item runs in the caller's
+    thread with ``device=None``, on the device its data is on, so lazy setup that is not thread-safe,
+    such as loading a format's CUDA extension, happens once before the workers start.
+    """
+    first = next_item(gpus[0])
+    if first is not None:
+        run(first, None)
+    inference = torch.is_inference_mode_enabled()
+
+    def worker(device):
+        # The device, grad mode and inference mode are per thread: take the caller's.
+        torch.cuda.set_device(device)
+        with torch.inference_mode() if inference else torch.no_grad():
+            while (item := next_item(device)) is not None:
+                run(item, device)
+        torch.cuda.synchronize(device)
+
+    with ThreadPoolExecutor(len(gpus)) as pool:
+        for future in [pool.submit(worker, device) for device in gpus]:
+            future.result()
+
+
 def update_hessian(input, hessian, n_samples, input_dtype=torch.float32):
     """Update hessian matrix with new input samples using incremental formula.
 
@@ -359,20 +451,31 @@ def free(self):
         self.weight = None
         self.h_inv = None
 
-    def update_weights(self, block_size, perc_damp, weight=None):
+    def updates_on_any_device(self) -> bool:
+        """Whether :meth:`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, and fits the room :func:`update_expert_weights` checks for.
+        """
+        return False
+
+    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 whose :meth:`updates_on_any_device` is true 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/_test_utils/torch/quantization/models.py

@@ -174,15 +174,13 @@ def get_input(self, batch_size=2):
 class ToyFusedExperts(nn.Module):
     """HF transformers 5 fused-experts layout: one ``F.linear`` pair per routed expert."""
 
-    def __init__(self):
+    def __init__(self, hidden=TOY_MOE_HIDDEN, intermediate=TOY_MOE_INTERMEDIATE):
         super().__init__()
         self.num_experts = TOY_MOE_EXPERTS
         self.gate_up_proj = nn.Parameter(
-            torch.randn(TOY_MOE_EXPERTS, 2 * TOY_MOE_INTERMEDIATE, TOY_MOE_HIDDEN) * 0.02
-        )
-        self.down_proj = nn.Parameter(
-            torch.randn(TOY_MOE_EXPERTS, TOY_MOE_HIDDEN, TOY_MOE_INTERMEDIATE) * 0.02
+            torch.randn(TOY_MOE_EXPERTS, 2 * intermediate, hidden) * 0.02
         )
+        self.down_proj = nn.Parameter(torch.randn(TOY_MOE_EXPERTS, hidden, intermediate) * 0.02)
 
     def forward(self, hidden_states, expert_of_token):
         output = torch.zeros_like(hidden_states)
@@ -198,20 +196,20 @@ def forward(self, hidden_states, expert_of_token):
 class ToyMoEBlock(nn.Module):
     """Routes token ``t`` to expert ``t % 3``, so the last expert receives no token."""
 
-    def __init__(self):
+    def __init__(self, hidden=TOY_MOE_HIDDEN, intermediate=TOY_MOE_INTERMEDIATE):
         super().__init__()
-        self.experts = ToyFusedExperts()
-        self.shared_expert = nn.Linear(TOY_MOE_HIDDEN, TOY_MOE_HIDDEN, bias=False)
+        self.experts = ToyFusedExperts(hidden, intermediate)
+        self.shared_expert = nn.Linear(hidden, hidden, bias=False)
 
     def forward(self, hidden_states):
         routed = self.experts(hidden_states, torch.arange(len(hidden_states)) % 3)
         return routed + self.shared_expert(hidden_states)
 
 
 class ToyMoE(nn.Module):
-    def __init__(self):
+    def __init__(self, hidden=TOY_MOE_HIDDEN, intermediate=TOY_MOE_INTERMEDIATE):
         super().__init__()
-        self.mlp = ToyMoEBlock()
+        self.mlp = ToyMoEBlock(hidden, intermediate)
 
     def forward(self, hidden_states):
         return self.mlp(hidden_states)