Add a gather_for_metrics capability (#540)

* Add test and full implementation

Author: muellerzr

docs/source/quicktour.mdx

@@ -120,14 +120,13 @@ validation_dataloader = accelerator.prepare(validation_dataloader)
 
 As for your training dataloader, it will mean that (should you run your script on multiple devices) each device will
 only see part of the evaluation data. This means you will need to group your predictions together. This is very easy to
-do with the [`~Accelerator.gather`] method.
+do with the [`~Accelerator.gather_for_metrics`] method.
 
 ```python
 for inputs, targets in validation_dataloader:
     predictions = model(inputs)
     # Gather all predictions and targets
-    all_predictions = accelerator.gather(predictions)
-    all_targets = accelerator.gather(targets)
+    all_predictions, all_targets = accelerator.gather_for_metrics((predictions, targets))
     # Example of use with a *Datasets.Metric*
     metric.add_batch(all_predictions, all_targets)
 ```
@@ -141,11 +140,17 @@ As for the training dataloader, passing your validation dataloader through
 Any instruction using your training dataloader length (for instance if you need the number of total training steps
 to create a learning rate scheduler) should go after the call to [`~Accelerator.prepare`].
 
+As some data at the end of the dataset may be duplicated so the batch can divide equally to all workers, metrics should be 
+calculated through the [`~Accelerator.gather_for_metrics`] method to automatically remove the duplicated data.
+
+If for some reason you don't wish to have this automatically done, [`~Accelerator.gather`] can be used instead to gather 
+the data across all processes and this can manually be done instead.
+
 </Tip>
 
 <Tip warning={true}>
 
-The [`~Accelerator.gather`] method requires the tensors to be all the same size on each process. If
+The [`~Accelerator.gather`] and [`~Accelerator.gather_for_metrics`] methods require the tensors to be all the same size on each process. If
 you have tensors of different sizes on each process (for instance when dynamically padding to the maximum length in
 a batch), you should use the [`~Accelerator.pad_across_processes`] method to pad you tensor to the
 biggest size across processes.

examples/by_feature/checkpointing.py

@@ -238,7 +238,7 @@ def training_function(config, args):
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), eval_dataloader)
             metric.add_batch(
                 predictions=predictions,
                 references=references,

examples/by_feature/cross_validation.py

@@ -207,7 +207,9 @@ def training_function(config, args):
                 with torch.no_grad():
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
-                predictions, references = accelerator.gather((predictions, batch["labels"]))
+                predictions, references = accelerator.gather_for_metrics(
+                    (predictions, batch["labels"]), eval_dataloader
+                )
                 metric.add_batch(
                     predictions=predictions,
                     references=references,
@@ -226,7 +228,7 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), test_dataloader)
             fold_predictions.append(predictions.cpu())
             if i == 0:
                 # We need all of the test predictions

examples/by_feature/fsdp_with_peak_mem_tracking.py

@@ -270,23 +270,16 @@ def collate_fn(examples):
         # context manager to track the peak memory usage during the evaluation
         with TorchTracemalloc() as tracemalloc:
             model.eval()
-            samples_seen = 0
             for step, batch in enumerate(eval_dataloader):
                 # We could avoid this line since we set the accelerator with `device_placement=True`.
                 batch.to(accelerator.device)
                 with torch.no_grad():
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
                 # It is slightly faster to call this once, than multiple times
-                predictions, references = accelerator.gather(
-                    (predictions, batch["labels"])
-                )  # If we are in a multiprocess environment, the last batch has duplicates
-                if accelerator.use_distributed:
-                    if step == len(eval_dataloader) - 1:
-                        predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                        references = references[: len(eval_dataloader.dataset) - samples_seen]
-                    else:
-                        samples_seen += references.shape[0]
+                predictions, references = accelerator.gather_for_metrics(
+                    (predictions, batch["labels"]), eval_dataloader
+                )
                 metric.add_batch(
                     predictions=predictions,
                     references=references,

examples/by_feature/gradient_accumulation.py

@@ -171,7 +171,7 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), eval_dataloader)
             metric.add_batch(
                 predictions=predictions,
                 references=references,

examples/by_feature/memory.py

@@ -183,7 +183,9 @@ def inner_training_loop(batch_size):
                 with torch.no_grad():
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
-                predictions, references = accelerator.gather((predictions, batch["labels"]))
+                predictions, references = accelerator.gather_for_metrics(
+                    (predictions, batch["labels"]), eval_dataloader
+                )
                 metric.add_batch(
                     predictions=predictions,
                     references=references,

examples/by_feature/multi_process_metrics.py

@@ -187,6 +187,8 @@ def training_function(config, args):
                 else:
                     # Otherwise we add the number of samples seen
                     samples_seen += references.shape[0]
+            # All of this can be avoided if you use `Accelerator.gather_for_metrics` instead of `Accelerator.gather`:
+            # accelerator.gather_for_metrics((predictions, references), eval_dataloader)
             metric.add_batch(
                 predictions=predictions,
                 references=references,

examples/by_feature/tracking.py

@@ -197,7 +197,7 @@ def training_function(config, args):
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), eval_dataloader)
             metric.add_batch(
                 predictions=predictions,
                 references=references,

examples/complete_cv_example.py

@@ -233,27 +233,18 @@ def training_function(config, args):
                     accelerator.save_state(output_dir)
         model.eval()
         accurate = 0
-        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch = {k: v.to(accelerator.device) for k, v in batch.items()}
             inputs = (batch["image"] - mean) / std
             with torch.no_grad():
                 outputs = model(inputs)
             predictions = outputs.argmax(dim=-1)
-            predictions, references = accelerator.gather((predictions, batch["label"]))
-            if accelerator.use_distributed:
-                if step == len(eval_dataloader) - 1:
-                    predictions = predictions[: len(eval_dataloader) - samples_seen]
-                    references = references[: len(eval_dataloader) - samples_seen]
-                else:
-                    samples_seen += references.shape[0]
-            else:
-                samples_seen += references.shape[0]
-            accurate_preds = predictions == references
+            predictions, labels = accelerator.gather_for_metrics((predictions, batch["label"]), eval_dataloader)
+            accurate_preds = predictions == labels
             accurate += accurate_preds.long().sum()
 
-        eval_metric = accurate.item() / samples_seen
+        eval_metric = accurate.item() / accelerator.gradient_state.samples_seen
         # Use accelerator.print to print only on the main process.
         accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
         if args.with_tracking:

examples/complete_nlp_example.py

@@ -211,23 +211,14 @@ def collate_fn(examples):
                     accelerator.save_state(output_dir)
 
         model.eval()
-        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather(
-                (predictions, batch["labels"])
-            )  # If we are in a multiprocess environment, the last batch has duplicates
-            if accelerator.use_distributed:
-                if step == len(eval_dataloader) - 1:
-                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                    references = references[: len(eval_dataloader.dataset) - samples_seen]
-                else:
-                    samples_seen += references.shape[0]
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), eval_dataloader)
             metric.add_batch(
                 predictions=predictions,
                 references=references,