[GPTQ][ddp] enabling DDP for GPTQ (vllm-project#2333)

After the changes in
vllm-project/compressed-tensors#572
vllm-project/compressed-tensors#534
vllm-project#2340 we're ready to
start rolling out DDP implementations of various modifiers

# API:
The Api we've landed on attempts to maintain the normal flow with
minimal changes necessary to enable DDP:

1) the user will call torchrun --nproc_per_node==<num_threads> script.py
to start the script
2) the user will initialize the distributed context, (they can use the
helper init_dist to do this)
3) the user will load the model using the new context manager, setting
the device map as outlined
[here](vllm-project/compressed-tensors#572).
(For most users this will be "auto_offload")
4) (optional) the user can partition the dataset at load time using
get_rank_partition or just load as normal and oneshot will partition the
data later (will load 1 copy of dataset into cpu memory for each rank
which may be onerous)
```python
from compressed_tensors.offload import load_offloaded_model, init_dist
init_dist()
with load_offloaded_model(): 
    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto_offload")
...
ds = load_dataset(
    DATASET_ID, split=get_rank_partition(DATASET_SPLIT, NUM_CALIBRATION_SAMPLES)
```

# Implementation
Adding the DDP process to GPTQ has relatively straightforward though
optimizing it for speed was a bit trickier. There are 4 steps

1) assigning each module to a rank which it will be compressed by
2) for each module assigned to a rank, having all hessian information
sent by other ranks to the assigned rank
3) each rank compresses the modules that it was assigned
4) broadcast the final quantized values to all ranks

Step 1 required the largest optimization, without any load balancing, we
ran into situations where 1 rank could be doing twice as much work as
another. Thus we implemented basic load balancing and time estimation
that seems to be working well in practice. The other major optimization
was using asynchronous ops for thread to thread communication. Before
these optimizations, 2 thread GPTQ was as fast as 1 thread GPTQ for
llama3-8B, afterward it results in a 27% speedup despite being a
relatively small model.

| model_id | world_size | max_time | max_memory | save_time |
flex_extract | eval_time |

|----------|-------------|----------|------------|-----------|--------------|-----------|
| Meta-Llama-3-8B-Instruct | 1 | 745.03 | 5.82 | 19.57 | 0.7066 | 95.28
|
| Meta-Llama-3-8B-Instruct | 2 | 372.20 | 5.57 | 49.10 | 0.7089 | 95.24
|
| Meta-Llama-3-8B-Instruct | 4 | 264.07 | 5.82 | 52.50 | 0.7180 | 96.74
|
| Qwen3-30B-A3B | 1 | 14207.53 | 6.56 | 748.23 | 0.8704 | 209.93 |
| Qwen3-30B-A3B | 2 | 7018.25 | 6.36 | 696.65 | 0.8810 | 205.89 |
| Qwen3-30B-A3B | 4 | 3694.46 | 6.36 | 723.05 | 0.8832 | 217.62 |

# GPTQ Changes

while validating numerical accuracy of the DDP technique, we noticed
that accuracy improved significantly for each thread added. After some
debugging we realized this was because the existing [hessian
calculation](https://github.com/vllm-project/llm-compressor/pull/2333/changes#diff-18d1319f01629ca65cc54f955dc6177f6dd025f057013932b2ed29842854f3ecL61-L65)
was causing an accumulation of floating point errors. By rewriting the
hessian calculation to sum the intermediate hessians and only divide by
num_samples at the end, we improved the GSM8K evaluation from (.67, .66)
to (.71, .71). You can repro these results
[here](https://github.com/vllm-project/llm-compressor/pull/2333/changes#diff-d31ce0453051853c17ba2a5225b3d1bfab548e095bab0967d6acfd1b3ce1b35d)

---------

Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

Author: HDCharles

examples/quantization_w4a16/llama3_ddp_example.py

@@ -0,0 +1,112 @@
+#############################################################################
+# This script is adapted from ./llama3_example.py and adds DDP functionality.
+# run this with `torchrun --nproc_per_node=2 llama3_ddp_example.py`
+# or change nproc_per_node to your desired configuration
+# to adapt other examples to use DDP, see the 2 altered sections below
+#############################################################################
+
+import time
+
+import torch
+from compressed_tensors.offload import dispatch_model, init_dist, load_offloaded_model
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from llmcompressor import oneshot
+from llmcompressor.datasets.utils import get_rank_partition
+from llmcompressor.modifiers.quantization import GPTQModifier
+
+model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
+
+###### DDP MODEL LOAD CHANGE #####
+init_dist()
+with load_offloaded_model():
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id, dtype="auto", device_map="auto_offload"
+    )
+##################################
+
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+DATASET_ID = "HuggingFaceH4/ultrachat_200k"
+DATASET_SPLIT = "train_sft"
+NUM_CALIBRATION_SAMPLES = 512
+MAX_SEQUENCE_LENGTH = 2048
+
+###### DDP DATA LOAD CHANGE #####
+ds = load_dataset(
+    DATASET_ID, split=get_rank_partition(DATASET_SPLIT, NUM_CALIBRATION_SAMPLES)
+)
+##########################
+
+ds = ds.shuffle(seed=42)
+
+
+def preprocess(example):
+    return {
+        "text": tokenizer.apply_chat_template(
+            example["messages"],
+            tokenize=False,
+        )
+    }
+
+
+ds = ds.map(preprocess)
+
+
+def tokenize(sample):
+    return tokenizer(
+        sample["text"],
+        padding=False,
+        max_length=MAX_SEQUENCE_LENGTH,
+        truncation=True,
+        add_special_tokens=False,
+    )
+
+
+ds = ds.map(tokenize, remove_columns=ds.column_names)
+
+recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"])
+
+
+torch.cuda.reset_peak_memory_stats()
+start_time = time.time()
+
+
+# Apply algorithms.
+oneshot(
+    model=model,
+    dataset=ds,
+    recipe=recipe,
+    max_seq_length=MAX_SEQUENCE_LENGTH,
+    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
+)
+
+elapsed_time = time.time() - start_time
+peak_memory_gb = torch.cuda.max_memory_allocated() / (1024**3)
+print("Quantization Complete")
+print(f"Time: {elapsed_time / 60:.2f} minutes ({elapsed_time:.2f} seconds)")
+print(f"Peak GPU Memory: {peak_memory_gb:.2f} GB")
+
+
+# Confirm generations of the quantized model look sane.
+print("\n\n")
+print("========== SAMPLE GENERATION ==============")
+dispatch_model(model)
+sample = tokenizer("Hello my name is", return_tensors="pt")
+sample = {key: value.to(model.device) for key, value in sample.items()}
+output = model.generate(**sample, max_new_tokens=100)
+print(tokenizer.decode(output[0]))
+print("==========================================\n\n")
+
+print("Saving...")
+# Save to disk compressed.
+SAVE_DIR = (
+    model_id.rstrip("/").split("/")[-1]
+    + "-W4A16-G128-DDP"
+    + str(torch.distributed.get_world_size())
+)
+model.save_pretrained(SAVE_DIR, save_compressed=True)
+tokenizer.save_pretrained(SAVE_DIR)
+
+torch.distributed.destroy_process_group()

src/llmcompressor/modifiers/quantization/gptq/base.py

@@ -2,6 +2,7 @@
 from typing import Dict, List, Optional, Tuple, Union
 
 import torch
+from compressed_tensors.offload.dist_utils import is_distributed
 from compressed_tensors.quantization import (
     QuantizationConfig,
     QuantizationScheme,
@@ -17,6 +18,7 @@
 )
 from loguru import logger
 from pydantic import PrivateAttr
+from torch import distributed as dist
 
 from llmcompressor.core import Event, EventType, State
 from llmcompressor.modifiers import Modifier
@@ -29,10 +31,13 @@
 from llmcompressor.modifiers.quantization.quantization import QuantizationMixin
 from llmcompressor.modifiers.utils import update_fused_layer_weight_global_scales
 from llmcompressor.sentinel import Sentinel
+from llmcompressor.utils import greedy_bin_packing, wait_for_comms
 from llmcompressor.utils.metric_logging import CompressionLogger
 
 __all__ = ["GPTQModifier"]
 
+_GPTQ_Q_PARAMS = ["weight", "weight_scale", "weight_zero_point", "weight_g_idx"]
+
 
 class GPTQModifier(Modifier, QuantizationMixin):
     """
@@ -123,7 +128,9 @@ class GPTQModifier(Modifier, QuantizationMixin):
     # private variables
     _module_names: Dict[torch.nn.Module, str] = PrivateAttr(default_factory=dict)
     _hessians: Dict[torch.nn.Module, torch.Tensor] = PrivateAttr(default_factory=dict)
-    _num_samples: Dict[torch.nn.Module, int] = PrivateAttr(default_factory=dict)
+    _num_samples: Dict[torch.nn.Module, torch.Tensor] = PrivateAttr(
+        default_factory=dict
+    )
 
     def resolve_quantization_config(self) -> QuantizationConfig:
         config = super().resolve_quantization_config()
@@ -248,7 +255,9 @@ def calibrate_module(
                 "cpu" if self.offload_hessians else get_execution_device(module)
             )
             self._hessians[module] = make_empty_hessian(module, device=init_device)
-            self._num_samples[module] = 0
+            self._num_samples[module] = torch.zeros(
+                tuple(), device=get_execution_device(module)
+            )
 
         # Accumulate hessian with input with optional offloading
         with self._maybe_onload_hessian(module):
@@ -263,7 +272,32 @@ def compress_modules(self):
         """
         Quantize modules which have been calibrated
         """
-        for module in list(self._num_samples.keys()):
+        ### Not Distributed
+        if not is_distributed():
+            self.compress_module_list(list(self._num_samples.keys()))
+            return
+
+        ### Distributed
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+
+        # Assign modules to ranks
+        module_list, rank_to_modules, module_to_rank = greedy_bin_packing(
+            list(self._hessians.keys()),
+            world_size,
+            item_weight_fn=lambda mod: self._hessians[mod].shape[0],
+        )
+
+        # send hessians to assigned ranks
+        self._reduce_hessian_to_target_rank(module_list, module_to_rank)
+
+        self.compress_module_list(rank_to_modules[rank])
+
+        # broadcast compressed modules to each rank
+        self._broadcast_quantized_params(module_list, module_to_rank)
+
+    def compress_module_list(self, module_list):
+        for module in module_list:
             name = self._module_names[module]
             num_samples = self._num_samples[module]
             quant_args = getattr_chain(module, "quantization_scheme.weights")
@@ -275,23 +309,59 @@ def compress_modules(self):
                 self._maybe_onload_hessian(module),
                 CompressionLogger(module) as comp_logger,
             ):
-                loss, quantized_weight, scale, zero_point, g_idx = quantize_weight(
+                loss, q_param_dict = quantize_weight(
                     module=module,
                     quant_args=quant_args,
-                    hessians_dict=self._hessians,
+                    hessian=self._hessians.pop(module) / self._num_samples.pop(module),
                     blocksize=self.block_size,
                     percdamp=self.dampening_frac,
                 )
                 comp_logger.set_loss(loss)
 
-            update_offload_parameter(module, "weight", quantized_weight)
-            update_offload_parameter(module, "weight_scale", scale)
-            update_offload_parameter(module, "weight_zero_point", zero_point)
-            if g_idx is not None:
-                update_offload_parameter(module, "weight_g_idx", g_idx)
-
-            # self._hessians[module] already deleted by quantize_weight
-            del self._num_samples[module]
+            for attr, val in q_param_dict.items():
+                update_offload_parameter(module, attr, val)
+
+    def _reduce_hessian_to_target_rank(self, module_list, module_to_rank):
+        rank = dist.get_rank()
+        pending_comms = []
+        for module in module_list:
+            target_rank = module_to_rank[module]
+            with self._maybe_onload_hessian(module):
+                pending_comms.append(
+                    dist.reduce(
+                        self._hessians[module],
+                        op=dist.ReduceOp.SUM,
+                        dst=target_rank,
+                        async_op=True,
+                    )
+                )
+                pending_comms.append(
+                    dist.reduce(
+                        self._num_samples[module],
+                        op=dist.ReduceOp.SUM,
+                        dst=target_rank,
+                        async_op=True,
+                    )
+                )
+                if rank != target_rank:
+                    self._hessians.pop(module, None)
+                    self._num_samples.pop(module, None)
+        wait_for_comms(pending_comms)
+
+    def _broadcast_quantized_params(self, module_list, module_to_rank):
+        pending_comms = []
+        for module in module_list:
+            src_rank = module_to_rank[module]
+
+            # Get parameters from module
+            for attr in _GPTQ_Q_PARAMS:
+                if getattr(module, attr, None) is not None:
+                    pending_comms.append(
+                        dist.broadcast(
+                            getattr(module, attr), src=src_rank, async_op=True
+                        )
+                    )
+        wait_for_comms(pending_comms)
 
     def on_end(self, state: State, event: Event, **kwargs):
         """

src/llmcompressor/modifiers/quantization/gptq/gptq_quantize.py

@@ -34,8 +34,8 @@ def accumulate_hessian(
     inp: torch.Tensor,
     module: torch.nn.Module,
     H: torch.Tensor | None,
-    num_samples: int,
-) -> tuple[torch.Tensor, int]:
+    num_samples: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
     inp = inp.to(device=H.device)
     if len(inp.shape) == 2:
         inp = inp.unsqueeze(0)
@@ -58,11 +58,10 @@ def accumulate_hessian(
             inp = inp.permute([1, 0, 2])
             inp = inp.flatten(1)
 
-    H *= num_samples / (num_samples + num_added)
     num_samples += num_added
 
     inp = inp.to(dtype=GPTQ_PRECISION)
-    inp = math.sqrt(2 / num_samples) * inp
+    inp = math.sqrt(2) * inp
     H += inp.matmul(inp.t())
 
     return H, num_samples
@@ -71,7 +70,7 @@ def accumulate_hessian(
 def quantize_weight(
     module: torch.nn.Module,
     quant_args: QuantizationArgs,
-    hessians_dict: dict[torch.nn.Module, torch.Tensor],
+    hessian: torch.Tensor,
     blocksize: int = 128,
     percdamp: float = 0.01,
 ) -> tuple[float, torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor]:
@@ -91,8 +90,7 @@ def quantize_weight(
     final_shape = module.weight.shape
     final_dtype = module.weight.dtype
     W = module.weight.clone()
-    H = hessians_dict[module]  # unfortunately python does not have a `move` keyword
-    del hessians_dict[module]  # so we have to delete the original reference manually
+    H = hessian
 
     # create observer for calculating quantization parameters
     observer = Observer.load_from_registry(
@@ -279,13 +277,14 @@ def quantize_weight(
     W = W.reshape(final_shape).to(final_dtype)
 
     loss = torch.sum(losses).item()
-    return (
-        loss,
-        W,
-        scale.to(dtype=final_dtype),
-        zero_point.to(dtype=quant_args.zp_dtype),
-        g_idx,
-    )
+    q_param_dict = {
+        "weight": W,
+        "weight_scale": scale.to(dtype=final_dtype),
+        "weight_zero_point": zero_point.to(dtype=quant_args.zp_dtype),
+    }
+    if g_idx is not None:
+        q_param_dict["weight_g_idx"] = g_idx
+    return (loss, q_param_dict)
 
 
 def _apply_activation_ordering(

src/llmcompressor/transformers/compression/compressed_tensors_utils.py

@@ -9,6 +9,7 @@
     SparsityCompressionConfig,
 )
 from compressed_tensors.config import CompressionFormat
+from compressed_tensors.offload import is_rank0
 from loguru import logger
 from transformers import PreTrainedModel
 
@@ -88,22 +89,23 @@ def save_pretrained_wrapper(
             if compressor is not None:
                 compressor.compress_model(model)
 
-            # save (compressed) model structure
-            original_save_pretrained.__get__(model, model_class)(
-                save_directory,
-                safe_serialization=safe_serialization,
-                **kwargs,
-            )
+            if is_rank0():
+                # save (compressed) model structure
+                original_save_pretrained.__get__(model, model_class)(
+                    save_directory,
+                    safe_serialization=safe_serialization,
+                    **kwargs,
+                )
 
-            # update config to reflect compression
-            if compressor is not None:
-                compressor.update_config(save_directory)
+                # update config to reflect compression
+                if compressor is not None:
+                    compressor.update_config(save_directory)
 
-            # update existing recipe
-            update_and_save_recipe(model.name_or_path, save_directory)
+                # update existing recipe
+                update_and_save_recipe(model.name_or_path, save_directory)
 
-            # copy python files from cache dir to save_path if any
-            copy_python_files_from_model_cache(model, save_directory)
+                # copy python files from cache dir to save_path if any
+                copy_python_files_from_model_cache(model, save_directory)
 
         save_pretrained_wrapper._overridden = True
         return save_pretrained_wrapper

src/llmcompressor/utils/__init__.py

@@ -7,3 +7,4 @@
 from .transformers import *
 from .dev import *
 from .helpers import *
+from .dist import *