[Distributed] metric calculation supports tp logits

Author: SylarTiaNII

paddlenlp/trainer/trainer.py

@@ -2780,6 +2780,12 @@ def evaluation_loop(
 
         # Metrics!
         if self.compute_metrics is not None and all_preds is not None and all_labels is not None:
+            if self.args.tensor_parallel_degree > 1 and all_preds.shape != all_labels.shape:
+                hcg = fleet.get_hybrid_communicate_group()
+                model_parallel_group = hcg.get_model_parallel_group()
+                gathered_predictions = []
+                dist.all_gather(gathered_predictions, all_preds, group=model_parallel_group)
+                all_preds = paddle.concat(gathered_predictions, axis=0)
             metrics = self.compute_metrics(EvalPrediction(predictions=all_preds, label_ids=all_labels))
         else:
             metrics = {}