Feat (mx): automatic group_dim in layerwise quant (#1012)

notebooks/minifloat_mx_tutorial.ipynb

@@ -163,8 +163,8 @@
  " pass\n",
  "\n",
  "class MXFloat8Act(MXFloat8e4m3Act, Fp8e4m3Mixin):\n",
- " # It is necessary to specify the group dimension for the activation quantization\n",
- " group_dim = 1\n",
+ " # In layerwise quantization, groupdim is automatically determined\n",
+ " pass\n",
  "\n",
  "\n",
  "class MXModel(nn.Module):\n",
@@ -221,9 +221,8 @@
  " group_size = 8\n",
  "\n",
  "class MXFloat8ActNoPadding(MXFloat8e4m3Act, Fp8e4m3Mixin):\n",
- " # It is necessary to specify the group dimension for the activation quantization\n",
+ " # In layerwise quantization, groupdim is automatically determined\n",
  " group_size = 8\n",
- " group_dim = 1\n",
  "\n",
  "\n",
  "class MXModelNoPadding(nn.Module):\n",
@@ -277,8 +276,8 @@
  " pass\n",
  "\n",
  "class MXFloat8Act(MXFloat8e4m3Act, Fp8e4m3Mixin):\n",
- " # It is necessary to specify the group dimension for the activation quantization\n",
- " group_dim = 1\n",
+ " # In layerwise quantization, groupdim is automatically determined\n",
+ " pass\n",
  "\n",
  "\n",
  "class MXModel(nn.Module):\n",
@@ -314,12 +313,11 @@
  "class MXInt8Weight(MXInt8Weight):\n",
  " # The group dimension for the weights it is automatically identified based on the layer type\n",
  " # If a new layer type is used, it can be manually specified\n",
- " bit_width = 8\n",
+ " pass\n",
  "\n",
  "class MXInt8Act(MXInt8Act):\n",
- " # It is necessary to specify the group dimension for the activation quantization\n",
- " group_dim = 1\n",
- " bit_width = 8\n",
+ " # In layerwise quantization, groupdim is automatically determined\n",
+ " pass\n",
  "\n",
  "class MXModel(nn.Module):\n",
  " def __init__(self):\n",

src/brevitas/quant/solver/act.py

@@ -1,6 +1,8 @@
 # Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
 # SPDX-License-Identifier: BSD-3-Clause
 
+from warnings import warn
+
 import torch
 from torch import nn
 from torch import Tensor
@@ -111,6 +113,33 @@ def scaling_shape(scaling_per_output):
  elif scaling_per_output == ScalingPerOutputType.TENSOR:
  return SCALAR_SHAPE
 
+ @value
+ def group_dim(module=None, group_size=None):
+ # Avoid circular import
+ from brevitas.nn.quant_layer import QuantWeightBiasInputOutputLayer
+
+ if group_size is not None and module is not None:
+ if isinstance(module, QuantWeightBiasInputOutputLayer):
+ if isinstance(module, nn.Linear):
+ return -1
+ elif isinstance(module,
+ (nn.Conv1d,
+ nn.Conv2d,
+ nn.Conv3d,
+ nn.ConvTranspose1d,
+ nn.ConvTranspose2d,
+ nn.ConvTranspose3d)):
+ warn(
+ "Group dim is being selected assuming batched input. Using unbatched input will fail and requires manually specification of group_dim"
+ )
+ # We are assuming batched input
+ return 1
+ else:
+ raise RuntimeError("Cannot determine automatically group_dim. Please specify")
+ else:
+ raise RuntimeError(
+ f"Cannot determine automatically group_dim for {type(module)}. Please specify")
+
 
 class SolveActScalingPerOutputChannelShape(ExtendedInjector):