Fix (graph/qronos): Normalize contribution to H and G when buffer is disabled (#1440)

Author: JP-Amboage

src/brevitas/graph/qronos.py

@@ -62,19 +62,22 @@ def update_batch(self, module, input, current_layer):
         if not is_quant_enabled:
             # Computing the normalized G matrix
             self.G *= (self.nsamples - batch_size) / self.nsamples
+            inp_processed /= math.sqrt(
+                self.nsamples)  # NOTE: quant_input is normalized before, in the H update
             if self.use_intermediate_buffer:
                 self.B.copy_(inp_processed.bmm(self.quant_input.transpose(2, 1)))
-                self.G += (self.B / self.nsamples)
+                self.G += self.B
             else:
                 self.G += inp_processed.bmm(self.quant_input.transpose(2, 1))
             self.quant_input = None  # NOTE: set back to None now that we've used it
         else:
             # Computing the normalized H matrix
             self.nsamples += batch_size  # NOTE: only increment with quant inputs
             self.H *= (self.nsamples - batch_size) / self.nsamples
+            inp_processed /= math.sqrt(self.nsamples)
             if self.use_intermediate_buffer:
                 self.B.copy_(inp_processed.bmm(inp_processed.transpose(2, 1)))
-                self.H += (self.B / self.nsamples)
+                self.H += self.B
             else:
                 self.H += inp_processed.bmm(inp_processed.transpose(2, 1))
             # store the quantized input for computing the H matrix