Use torch.compile to speed up GPTQ algo

Signed-off-by: aladerran <aladerran@gmail.com>

Author: aladerran

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

@@ -3,6 +3,7 @@
 from typing import Dict, Optional, Tuple, Union
 
 import torch
+import torch._dynamo.config
 import transformers
 from compressed_tensors.quantization import (
     ActivationOrdering,
@@ -16,6 +17,8 @@
 from llmcompressor.observers.base import Observer
 from llmcompressor.pytorch.utils.helpers import tensor_sparsity
 
+torch._dynamo.config.capture_scalar_outputs = True
+
 GPTQ_PRECISION = torch.float32
 
 __all__ = ["make_empty_hessian", "accumulate_hessian", "quantize_weight"]
@@ -68,6 +71,7 @@ def accumulate_hessian(
     return H, num_samples
 
 
+@torch.compile
 def quantize_weight(
     module: torch.nn.Module,
     quant_args: QuantizationArgs,