Add a gather_for_metrics capability

[muellerzr]

Introduce a gather_for_metrics function

What does this add?

This PR adds a new function to Accelerator called gather_for_metrics, which assists with calculating the right metric in distributed setups

Who is it for?

Users of accelerate that want to ensure that their reported metrics are fully accurate

Why is it needed?

To assist with making sure all the batches have the right batch size, Accelerate will pad the length on the last batch to be duplicates of the last sample. These need to be dropped when calculating the final metrics on the last batch, and currently it looks something like:

if accelerator.use_distributed:
    # Then see if we're on the last batch of our eval dataloader
    if step == len(eval_dataloader) - 1:
        # Last batch needs to be truncated on distributed systems as it contains additional samples
        predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
        references = references[: len(eval_dataloader.dataset) - samples_seen]
    else:
        # Otherwise we add the number of samples seen
        samples_seen += references.shape[0]

This PR adds a new utility called accelerate.gather_for_metrics which will handle this check for us, entirely thanks to the GradientState capability.

Note: this method currently doesn't work for TPU's, as it needs to be eval_dataloader._loader.dataset. This PR fixes this as well

What parts of the API does this impact?

User-facing:

• A new Accelerator.gather_for_metrics function was added

Internal structure:

• Preprocessed dataloaders now have a new total_dataset_length attribute
• GradientState now keeps track of the number of samples seen

Basic Usage Example(s):

When calculating metrics, users can now do the following to properly calculate their metrics:

input, target = next(iter(dataloader))
with torch.no_grad():
    logits = ddp_model(ddp_input)
    logits, target = accelerator.gather_for_metrics((logits, ddp_target), dataloader)
    accuracy_multi = accuracy(logits.argmax(dim=-1), target)

When would I use it, and when wouldn't I?

Since this works on distributed and non-distributed systems, always if the evaluation dataset has been prepared by Accelerator. Users should just add this to any script that calculates metrics.

TODO:

• Update the other examples to use this new API. multiprocess_metrics will serve as a lower-level example

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[sgugger]

Very clever! It's really great, just make sure to update the docs and the examples :-)

[sgugger]

Nice doc!

[plamb-viso]

is this functionality in 0.11.0?

[muellerzr]

It is not, we suggest not using it for now and doing the check manually as shown in the metric example script as some bugs were discovered: #575

[plamb-viso]

Cool, thank you, excited for this change when it happens