MLP: Fix edge case when hidden_size == interm_size (happens in some TP setups)

Author: turboderp

exllamav3/modules/mlp.py

@@ -546,9 +546,9 @@ def load_local(self, device: torch.Device, load_slice: int, **kwargs):
         if self.num_slices == 1 and self.multi_gu[0] is not None and self.downs[0].inner.bc is not None:
             mgu = self.multi_gu[0]
             self.bsz1_pa_args = [
-                (device, (2, 1, self.hidden_size), self.interm_dtype),
-                (device, (2, 1, mgu.out_features), self.interm_dtype),
-                (device, (1, 1, 1, mgu.out_features), torch.half)
+                (device, (2, 1, self.hidden_size), self.interm_dtype, "gu"),
+                (device, (2, 1, mgu.out_features), self.interm_dtype, "a1"),
+                (device, (1, 1, 1, mgu.out_features), torch.half, "a2")
             ]
             self.bc = ext.BC_GatedMLP(
                 *(g_tensor_cache.get(*arg) for arg in self.bsz1_pa_args),