Normalize hidden state after adding control vectors to preserve L2 norm

exl2_wrapper.py

@@ -78,7 +78,9 @@ class ExLlamaV2ModuleWrapper:
     def wrapped_forward(self, *args, **kwargs):
         x = self.module.forward(*args, **kwargs)
         try:
+            prev_norm = torch.norm(x, p=2)
             x += self.control_vector[self.module.layer_idx].clone().to(x.device)
+            x *= prev_norm / torch.norm(x, p=2)
         except IndexError:
             pass
         return x