add group size as a parameter

layers.py

@@ -36,26 +36,27 @@ class QuantizedLinear(nn.Module):
         self,
         in_features: int,
         out_features: int,
+        group_size: int = 128,
         dtype: torch.dtype = torch.uint8,
     ):
-        # TODO: Take group_size as an input instead of hardcoding it here.
         super().__init__()
         self.in_features = in_features
         self.out_features = out_features
+        self.group_size = group_size
         self.weight = nn.ParameterDict(
             {
                 "packed": nn.Parameter(
                     torch.empty(
-                        out_features * in_features // (128 * 2), 128, dtype=dtype
+                        out_features * in_features // (group_size * 2), group_size, dtype=dtype
                     ),
                     requires_grad=False,
                 ),
                 "scale": nn.Parameter(
-                    torch.empty(out_features * in_features // 128, 1),
+                    torch.empty(out_features * in_features // group_size, 1),
                     requires_grad=False,
                 ),
                 "zero_point": nn.Parameter(
-                    torch.empty(out_features * in_features // 128, 1),
+                    torch.empty(out_features * in_features // group_size, 1),
                     requires_grad=False,
                 ),
             }
@@ -86,7 +87,7 @@ class QuantizedLinear(nn.Module):
         )
 
         del self.weight, self.bias
-        quantize_(self, int4_weight_only(group_size=128))
+        quantize_(self, int4_weight_only(group_size=self.group_size))
         self.unpacked = True
         torch.cuda.empty_cache()
 

moondream.py

@@ -80,31 +80,33 @@ class MoondreamModel(nn.Module):
         self.region = nn.ModuleDict(
             {
                 "coord_encoder": QuantizedLinear(
-                    config.region.coord_feat_dim, config.region.dim, dtype=dtype
+                    config.region.coord_feat_dim, config.region.dim, group_size=config.text.group_size, dtype=dtype
                 ),
                 "coord_decoder": nn.ModuleDict(
                     {
                         "fc1": QuantizedLinear(
-                            config.region.dim, config.region.inner_dim, dtype=dtype
+                            config.region.dim, config.region.inner_dim, group_size=config.text.group_size, dtype=dtype
                         ),
                         "fc2": QuantizedLinear(
                             config.region.inner_dim,
                             config.region.coord_out_dim,
+                            group_size=config.text.group_size,
                             dtype=dtype,
                         ),
                     }
                 ),
                 "size_encoder": QuantizedLinear(
-                    config.region.size_feat_dim, config.region.dim, dtype=dtype
+                    config.region.size_feat_dim, config.region.dim, group_size=config.text.group_size, dtype=dtype
                 ),
                 "size_decoder": nn.ModuleDict(
                     {
                         "fc1": QuantizedLinear(
-                            config.region.dim, config.region.inner_dim, dtype=dtype
+                            config.region.dim, config.region.inner_dim, group_size=config.text.group_size, dtype=dtype
                         ),
                         "fc2": QuantizedLinear(
                             config.region.inner_dim,
                             config.region.size_out_dim,
+                            group_size=config.text.group_size,
                             dtype=dtype,
                         ),
                     }

text.py

@@ -164,19 +164,19 @@ def build_text_model(config: TextConfig, dtype: torch.dtype = torch.float16) ->
                             "ln": nn.LayerNorm(config.dim, dtype=dtype),
                             "attn": nn.ModuleDict(
                                 {
-                                    "qkv": QuantizedLinear(config.dim, qkv_dim, dtype=dtype),
+                                    "qkv": QuantizedLinear(config.dim, qkv_dim, group_size=config.text.group_size, dtype=dtype),
                                     "proj": QuantizedLinear(
-                                        config.dim, config.dim, dtype=dtype
+                                        config.dim, config.dim, group_size=config.text.group_size, dtype=dtype
                                     ),
                                 }
                             ),
                             "mlp": nn.ModuleDict(
                                 {
                                     "fc1": QuantizedLinear(
-                                        config.dim, config.ff_dim, dtype=dtype
+                                        config.dim, config.ff_dim, group_size=config.text.group_size, dtype=dtype
                                     ),
                                     "fc2": QuantizedLinear(
-                                        config.ff_dim, config.dim, dtype=dtype
+                                        config.ff_dim, config.dim, group_size=config.text.group_size, dtype=dtype
                                     ),
                                 }
                             ),