[core][compiled graphs] Adjust overlap algorithm

[ruisearch42]

Why are these changes needed?

Adjust gpu communication overlap algorithm to be less conservative, see details in #48953

Related issue number

Closes #48953

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]

Sorry, I didn't really understand the upstream/downstream NCCL actors change. Can you either clarify with some examples/comments or see if the suggested change could work instead?

[stephanie-wang]

I found this logic hard to understand. Reading the issue, would it work if we instead modified this to is_nccl_write_to_peer(self, upstream_actor), and make sure the caller passes in the right actor?

[ruisearch42]

(There may be more than one upstream_actors, so we need to pass them all in.)
is_nccl_write_to_upstream_actor is a short hand and makes the caller code more concise. We can change it to the following at caller side (dag_node_operation.py:792) if you find it easier to understand:

Original:
overlapped_schedule[j].is_nccl_write_to_upstream_actor
New:
overlapped_schedule[j].is_nccl_write_to_peer(overlapped_schedule[j].upstream_actors)

[stephanie-wang]

Why does downstream_nccl_actors check whether the current node requires NCCL while uspstream_nccl_actors checks whether the upstream node requires NCCL?

[ruisearch42]

Not sure if I get the underlying question, for a dag_node with NCCL type hint, downstream_nccl_actors means the actors reading from this NCCL channel, and upstream_nccl_actors means the actors writing to NCCL channels which are input for this dag_node. "downstream" and "upstream" are relative to the dag_node.

[ruisearch42]

Sorry, I didn't really understand the upstream/downstream NCCL actors change. Can you either clarify with some examples/comments or see if the suggested change could work instead?

Sure, for the test_torch_tensor_dag.py::test_torch_tensor_nccl_overlap2 example, we have the following overlapped schedule:

For the Actor_520f.send method, its downstream_nccl_actors: {Actor_e171}, its upstream_nccl_actors: {} (No actor is writing to this node using NCCL). Therefore we can overlap (see the red ovals).

[ruisearch42]

The relax algorithm would deadlock for 5d35b71
See #48953 (comment)