pytorch / pytorch/ignite

Improve idist for gather using nccl and reduce for gloo + gpu

Open
#2,260 1 comment 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
4.8k
Forks
726
Avg merge
5d 21h
Merged PRs (30d)
5

Description

🚀 Feature

Consider the following piece of code

def write_preds_to_file(predictions, filename):
    prediction_tensor = torch.tensor(predictions)
    prediction_tensor = idist.all_gather(prediction_tensor)

    if idist.get_rank() == 0:
        torch.save(prediction_tensor, filename)

The idist.all_gather() is used to collect the tensor from all the processes even if only the rank 0 needs it. The gather() method would be used but the backend nccl does not support it. See here.

The idea here is to implement the gather() method in idist using all_gather() for nccl (and gather() for others backends). Note that reduce() for gloo on GPU could be implemented using all_reduce() in a similar way.

It needs tests + docs

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by locating the idist implementations of gather and reduce and review how they select distributed backends. Add the requested backend-specific behavior for NCCL and Gloo on GPU, then add tests and documentation as noted in the issue; done means both operations work for the described backend combinations.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, documentation, testing-qa
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.