x_max -> x_mean and w_max -> w_mean name changes and some comments (#378

)

awq/quantize/quantizer.py

@@ -244,17 +244,23 @@ def _search_best_scale(
  # Put x on the right device
  inp = inp.to(next(module2inspect.parameters()).device)
 
- # [STEP 1]: Compute maximum of weight
+ # [STEP 1]: Compute per-channel mean of normalised weights
+ # All layer weights are concatted together
  weight = torch.cat([_m.weight for _m in layers], dim=0)
  org_shape = weight.shape
+ # The weights are reshaped to be organised by quantization group
  weight = weight.view(-1, self.group_size)
+ # Calculates the relative magnitude of the weights within each of the quantization groups, 
+ # and rescales each group individually so that each group has weights on a 0-1 scale.
  w_scale = weight.abs() / weight.abs().amax(dim=1, keepdim=True)
+ # Resizes the rescaled weight matrix back up to its original dimensions
  w_scale = w_scale.view(org_shape)
- w_max = w_scale.mean(0)
+ # Gets the average rescaled magnitude for each output channel
+ w_mean = w_scale.mean(0)
  clear_memory(weight)
 
- # [STEP 2]: Compute maximum of x
- x_max = inp.abs().view(-1, inp.shape[-1]).mean(0)
+ # [STEP 2]: Compute per-channel mean of the input activation
+ x_mean = inp.abs().view(-1, inp.shape[-1]).mean(0)
 
  # [STEP 3]: Compute output of module
  with torch.no_grad():
@@ -266,7 +272,7 @@ def _search_best_scale(
 
  # [STEP 4]: Compute loss
  best_scales = self._compute_best_scale(
- inp, w_max, x_max, module2inspect, layers, fp16_output, module_kwargs
+ inp, w_mean, x_mean, module2inspect, layers, fp16_output, module_kwargs
  )
 
  return (
@@ -278,8 +284,8 @@ def _search_best_scale(
  def _compute_best_scale(
  self,
  x,
- w_max,
- x_max,
+ w_mean,
+ x_mean,
  module2inspect,
  linears2scale: List[nn.Linear],
  fp16_output,
@@ -303,18 +309,18 @@ def _compute_best_scale(
  org_sd = {k: v.cpu() for k, v in module2inspect.state_dict().items()}
 
  device = x.device
- x_max = x_max.view(-1).to(device)
- w_max = w_max.view(-1).to(device)
+ x_mean = x_mean.view(-1).to(device)
+ w_mean = w_mean.view(-1).to(device)
 
  for ratio in range(n_grid):
  # create new scales
  ratio = ratio / n_grid
 
  # NOTE: s^-1 * x is fused here, according to paper
  if self.duo_scaling:
- scales = (x_max.pow(ratio) / w_max.pow(1 - ratio)).clamp(min=1e-4)
+ scales = (x_mean.pow(ratio) / w_mean.pow(1 - ratio)).clamp(min=1e-4)
  else:
- scales = x_max.pow(ratio).clamp(min=1e-4).view(-1)
+ scales = x_mean.pow(ratio).clamp(min=1e-4).view(-1)
  scales = scales / (scales.max() * scales.min()).sqrt()
  scales_view = scales.view(1, -1).to(device)