Add hvd.grouped_allgather and hvd.grouped_reducescatter

[maxhgerlach]

Checklist before submitting

• Did you read the contributor guide?
• Did you update the docs?
• Did you write any tests to validate this change?
• Did you update the CHANGELOG, if this change affects users?

Description

This adds two new multi-tensor functions to Horovod: hvd.grouped_allgather() and hvd.grouped_reducescatter(). They are very similar to the existing hvd.grouped_allreduce(), introduced with PR #2453, giving users the equivalent extra control over tensor fusion for their Allgather and Reducescatter ops.

To implement this functionality I've added a new attribute output_index to TensorTableEntry. This allows an AllgatherOp or ReducescatterOp that has been enqueued as part of a group to allocate its output tensor properly (in contrast to Allreduce this cannot be done beforehand because the resulting output size is only known after cross-process coordination). I've also added warning log messages in case AllocateOutput ever fails because TensorFlow would crash with just a default "Unknown error." exception message in such a case.

Supported frameworks:

• TensorFlow
• PyTorch
• MXNet

Review process to land

1. All tests and other checks must succeed.
2. At least one member of the technical steering committee must review and approve.
3. If any member of the technical steering committee requests changes, they must be addressed.

[github-actions]

Unit Test Results

  1 011 files   -   36    1 011 suites   - 36   10h 32m 14s ⏱️ - 14m 22s
     814 tests +  27       764 ✔️ +  27       50 💤 ±  0  0 ❌ ±0 
20 527 runs  +268  14 577 ✔️ +250  5 950 💤 +18  0 ❌ ±0 

Results for commit 9bd74dd. ± Comparison against base commit 757883b.

♻️ This comment has been updated with latest results.

[github-actions]

Unit Test Results (with flaky tests)

  1 131 files   -   54    1 131 suites   - 54   11h 2m 1s ⏱️ - 13m 25s
     814 tests +  27       764 ✔️ +  27       50 💤 ±  0  0 ❌ ±0 
23 119 runs  +272  16 039 ✔️ +240  7 080 💤 +32  0 ❌ ±0 

Results for commit 9bd74dd. ± Comparison against base commit 757883b.

♻️ This comment has been updated with latest results.

[maxhgerlach]

Rebased to master to pick up test fixes and #3590.

[maxhgerlach]

I believe this is ready to be reviewed now. For once all the test suites have passed.

For PyTorch and MXNet I have extended the respective OpContext classes so they can refer to all output tensors of a group. But ultimately they only need to allocate memory for one specific output index (to do an Allgather or a Reducescatter). In contrast, the TFOpContext already knew about all outputs of the group op. Maybe the design could be cleaned up a bit here to reduce confusion with how the two outputs of Alltoall are handled (tensor and splits), but that potential issue didn't seem pressing to me so far.

[romerojosh]

Excellent work @maxhgerlach! Very well implemented and I have no comments. Also, thanks for catching adding and cleaning up the documentation.

[maxhgerlach]

Thanks for taking the time to review, @romerojosh! Fortunately, your careful earlier work was easy to transfer, so this was quite straightforward.