Update tensor_parallel.py

server/text_generation_server/layers/tensor_parallel.py

@@ -66,6 +66,9 @@ class TensorParallelHead(SuperLayer):
             weight = weights.get_tensor(f"{prefix}.weight")
             should_gather = False
 
+        if config.model_type == "baichuan":
+            weight = F.normalize(weight)
+
         return TensorParallelHead(
             get_linear(weight, bias=None),
             process_group=weights.process_group,