[data][train] Fix deadlocks caused by streaming_split

[raulchen]

Why are these changes needed?

Fix a deadlock issue for training jobs. The issue happens in the following situation:

• The output blocks of streaming_split are assigned to multiple splits (output_split_idx).
• When one split has finished reading all blocks, it won't stop the iteration until all the other splits have all finished, because of this.
• This is usually fine. But when the unfinished splits are waiting for the finished splits (e.g., there is a gradient synchronization), there will be a dead lock due to circular dependencies.

This PR makes the finished splits can finish iteration immediately without waiting for others.

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[stephanie-wang]

Hmm sorry but I don't quite understand the deadlock situation in the PR description and the proposed fix. Doesn't SplitCoordinator explicitly require all the consumers to read at the same time? Is the deadlock situation in the PR description somehow different?

[raulchen]

Update:
The issue happens when there is a gradient sync per batch. E.g.,

for batch in it.iter_batches():
    all_reduce()

We suspect it's because streaming_split returns different number of rows for each train worker and makes the all_reduce misaligned.
We'll merge this PR first, as it doesn't make sense to have this dependency in the first place.

[c21]

late LGTM

[genesis-jamin]

Hi @raulchen , thanks for pushing this fix -- this actually fixed a NCCL timeout error that we were seeing when doing multi-node distributed training. The behavior there was that sometimes randomly at the start of a train epoch, we would hit a NCCL timeout error because all of the ranks except one were trying to allreduce the gradients.

I'm also confused by the deadlock explanation though. Have you / the team thought more about how exactly this would have created a deadlock with gradient synchronization? We iterate over our data using the iter_torch_batches function, and I thought this would call into streaming_split with equal=True (so each worker gets an equal number of rows).

If it's helpful, we only started seeing this issue when we scaled up the model size (probably because gradient synchronization took longer).