From 12861237eb75275493c9ce28d9ffaa0bcbcd3ed7 Mon Sep 17 00:00:00 2001 From: Arnav Garg <106701836+arnavgarg1@users.noreply.github.com> Date: Wed, 27 Sep 2023 21:25:15 +0300 Subject: [PATCH] Place metric functions for BLEU and Rogue scores on correct devices when using multiple GPUs (#3671) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- ludwig/features/text_feature.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/ludwig/features/text_feature.py b/ludwig/features/text_feature.py index 6b1408b62a1..12da5c1409c 100644 --- a/ludwig/features/text_feature.py +++ b/ludwig/features/text_feature.py @@ -299,7 +299,9 @@ def update_metrics( predictions: Dict of tensors returned by predictions(). """ if tokenizer is not None: + # Decode the targets and predictions to compute response-based metrics using the initialized tokenizer. decoded_targets, decoded_predictions = get_decoded_targets_and_predictions(targets, predictions, tokenizer) + for metric_name, metric_fn in self._metric_functions.items(): prediction_key = get_metric_tensor_input(metric_name) try: @@ -308,6 +310,12 @@ def update_metrics( # RESPONSE metrics cannot be computed if decoded texts are not provided. # Decoded texts are only provided using the LLM model type. if decoded_targets is not None and decoded_predictions is not None: + # Move metric function to the device of the predictions. + # For CUDA, it can be computed on any of the GPUs since it uses allgather to collect + # the results from all GPUs and compute the final metric. + # We use 'predictions' as the key since it is always present in the predictions dict. + device = "cuda" if predictions["predictions"].is_cuda else "cpu" + metric_fn = metric_fn.to(device) if metric_name == "bleu": # BLEU takes in targets as a list. metric_fn.update(decoded_predictions, [decoded_targets])