cheery-picked final PR changes

Author: realAsma

modelopt/torch/quantization/algorithms.py

@@ -28,15 +28,14 @@
 import regex as re
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from tqdm import tqdm
 
 from modelopt.torch.opt.conversion import ModeloptStateManager
 from modelopt.torch.opt.hparam import CustomHPType, Hparam, HPType
 from modelopt.torch.opt.searcher import LPS, BaseSearcher, SearchConfig, SearchStateDict
 from modelopt.torch.opt.utils import get_hparam, named_hparams
 from modelopt.torch.utils import create_param_grad_clear_hook, print_rank_0, report_memory
-from modelopt.torch.utils.distributed import DistributedProcessGroup, is_master
+from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState, is_master
 
 from . import config as mtq_config
 from . import model_calib
@@ -944,19 +943,72 @@ def run_search_with_stats(self, max_weight_size, verbose=False):
         return best_recipes, is_satisfied
 
 
+# TODO: does torch compile improves speed?
 @torch.compile
-def _get_kl_div_loss(logits_unquant: torch.Tensor, logits_quant: torch.Tensor) -> torch.Tensor:
-    # TODO: Support TensorParallel
-    prob_unquant = F.softmax(logits_unquant, dim=-1)
-    log_prob_quant = F.log_softmax(logits_quant, dim=-1)
-    return F.kl_div(log_prob_quant, prob_unquant, reduction="sum", log_target=False)
+def _get_softmax_dist(
+    logits: torch.Tensor, tp_group, return_log_prob: bool = False
+) -> torch.Tensor:
+    # TODO: test this
+    dtype = logits.dtype
+    max_logits = torch.amax(logits, dim=-1, keepdim=True)
+    torch.distributed.all_reduce(max_logits, op=torch.distributed.ReduceOp.MAX, group=tp_group)
+    logits = (logits - max_logits).float()
+    sum_exp_logits = torch.exp(torch.logsumexp(logits, dim=-1, keepdim=True))
+    torch.distributed.all_reduce(sum_exp_logits, op=torch.distributed.ReduceOp.SUM, group=tp_group)
+    logits = logits - torch.log(sum_exp_logits)
+    if return_log_prob:
+        return logits.to(dtype)
+    else:
+        return torch.exp(logits).to(dtype)
+
+
+@torch.compile
+def _get_softmax(logits: torch.Tensor, return_log_prob: bool = False) -> torch.Tensor:
+    # TODO: do we need to do log_softmax in float32?
+    # log_softmax is supposed to be numerically stable implementation
+    log_prob = torch.log_softmax(logits.float(), dim=-1)
+    if return_log_prob:
+        return log_prob
+    else:
+        return torch.exp(log_prob)
+
+
+@torch.compile
+def _get_p_log_q(p: torch.Tensor, log_q: torch.Tensor) -> torch.Tensor:
+    return torch.sum(p * log_q).float()
+
+
+def _get_prob_from_logits(
+    logits: torch.Tensor, return_log_prob: bool = False, lm_head: nn.Module = None
+) -> torch.Tensor:
+    parallel_state: ParallelState | None = (
+        getattr(lm_head, "parallel_state", None) if lm_head is not None else None
+    )
+    if parallel_state is not None and parallel_state.tensor_parallel_group.is_initialized():
+        return _get_softmax_dist(
+            logits, parallel_state.tensor_parallel_group.group, return_log_prob
+        )
+    return _get_softmax(logits, return_log_prob)
+
+
+def _get_kl_div_loss(
+    prob_unquant: torch.Tensor, logits_quant: torch.Tensor, lm_head: nn.Module = None
+) -> torch.Tensor:
+    log_prob_quant = _get_prob_from_logits(logits_quant, return_log_prob=True, lm_head=lm_head)
+    # We dont need to calculate the full kl div loss here, just get p*log_q
+    return _get_p_log_q(prob_unquant, log_prob_quant)
+
+
+def _get_lm_head(model: nn.Module) -> nn.Module:
+    for name, module in model.named_modules():
+        if name.endswith(("lm_head", "output_layer")):  # HF transformers models or Megatron models
+            return module
+    return None
 
 
 class AutoQuantizeKLDivSearcher(_AutoQuantizeBaseSearcher):
     """A searcher for AutoQuantize algorithm that uses KL-Divergence loss based score estimation."""
 
-    score_module_rules: list[str | Callable] = [lambda name: ""]
-
     @property
     def default_search_config(self):
         """Get the default config for the searcher."""
@@ -973,9 +1025,10 @@ def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig:
         config = config or {}
         for ignored_key in ["score_func", "loss_func", "forward_backward_step"]:
             if ignored_key in config:
-                warnings.warn(
-                    f"`{ignored_key}` is ignored for KL-Divergence loss based `auto_quantize`."
-                )
+                if config[ignored_key] is not None:
+                    warnings.warn(
+                        f"`{ignored_key}` is ignored for KL-Divergence loss based `auto_quantize`."
+                    )
                 config.pop(ignored_key)
         config = super().sanitize_search_config(config)
         assert config["forward_step"] is not None, (
@@ -984,21 +1037,12 @@ def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig:
         )
         return config
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def estimate_sensitivity_scores(self):
         """Estimate the sensitivity scores for the model.
 
         Higher score means more sensitive to quantization.
         """
-        # Check if tensor parallelism is being used
-        for name, module in self.model.named_modules():
-            if hasattr(module, "parallel_state"):
-                if hasattr(module.parallel_state, "tensor_parallel_group"):
-                    if module.parallel_state.tensor_parallel_group.is_initialized():
-                        warnings.warn(
-                            "Tensor Parallel is not supported for KL-Divergence based auto_quantize. "
-                        )
-                        break
 
         def set_to_unquantized():
             for name, hparam in named_hparams(self.model, unique=True):
@@ -1016,17 +1060,27 @@ def set_to_unquantized():
         ):
             set_to_unquantized()
             logits_unquant = self.config["forward_step"](self.model, data)
+            prob_unquant = _get_prob_from_logits(
+                logits_unquant,
+                return_log_prob=False,
+                lm_head=_get_lm_head(self.model),
+            )
 
-            for name, hparam in named_hparams(self.model, configurable=True):
+            for name, hparam in tqdm(
+                list(named_hparams(self.model, configurable=True)), desc="Evaluating hparams"
+            ):
                 if not isinstance(hparam, QuantRecipeHparam):
                     continue
                 for recipe in hparam.choices:
                     if recipe == QuantRecipe(quant_cfg=None):
                         continue
                     hparam.active = recipe
                     logits_quant = self.config["forward_step"](self.model, data)
-                    score = _get_kl_div_loss(logits_unquant, logits_quant)
-                    hparam._importance_dict[recipe][hparam.score_modules[0]] = score
+                    score = _get_kl_div_loss(prob_unquant, logits_quant, _get_lm_head(self.model))
+                    if hparam._importance_dict[recipe][hparam.score_modules[0]] is None:
+                        hparam._importance_dict[recipe][hparam.score_modules[0]] = score
+                    else:
+                        hparam._importance_dict[recipe][hparam.score_modules[0]] += score
                 hparam.active = QuantRecipe(quant_cfg=None)
 
     def run_search_with_stats(self, max_weight_size, verbose=False):