Reduce memory by using all_gather_into_tensor

[muellerzr]

What does this PR do?

torch 1.13 introduced publicly all_gather_into_tensor (before this was _all_gather_base) which is a much more memory efficient version of gather. One thing to note mentioned in the PR here is they did not have this be the base of gather, since it handles uneven inputs automatically. Since Accelerate does this separately and checks, we can safely use this API. Original DeepSpeed PR I discovered showing this

For a general idea of just how much memory can be stored, I ran a small simple test:

import time
import torch
from accelerate import PartialState
from accelerate.utils import gather

def convert_bytes(size):
    "Converts `size` from bytes to the largest possible unit"
    for x in ["bytes", "KB", "MB", "GB", "TB"]:
        if size < 1024.0:
            return f"{round(size, 2)} {x}"
        size /= 1024.0

    return f"{round(size, 2)} PB"

state = PartialState()
tensor = torch.rand((64, 224, 224, 64), device=state.device)

# Using `PartialState`
start_time = time.time()
tensor = gather(tensor)
end_time = time.time()

with state.main_process_first():
    print(f"Process {state.process_index} memory allocated after: {convert_bytes(torch.cuda.max_memory_allocated(state.device))}")
    print(f"Process {state.process_index} time: {end_time - start_time}")

The results can be summarized as such:

Before:
Total CUDA memory allocated: 3.83gb

After:
Total CUDA memory allocated: 2.3gb

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@BenjaminBossan @LysandreJik

[HuggingFaceDocBuilderDev]

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

[BenjaminBossan]

Nice improvement. Before proceeding, I have a few questions, please take a look. Also some minor suggestions, but those are not blockers.

[BenjaminBossan]

Why use torch.zeros instead of torch.empty_like as previously?

[muellerzr]

We're doing something slightly different here with the new API, where this gather works using a different tensor dim than before which is more efficient.

[muellerzr]

Added a comment

[BenjaminBossan]

Not sure, but won't this allocate more memory than previously?

[muellerzr]

It's actually the opposite, hence what we're doing here.

[pacman100]

Thank you @muellerzr for working on this, I have the same comments as Benjamin.

[BenjaminBossan]

In general looks good, thanks for adding support for this feature. I have only one question, but feel free to merge.

[BenjaminBossan]

Not sure, but won't this allocate more memory than previously?