[Train] Fix lightning checkpoint report callback

[woshiyyya]

Why are these changes needed?

In RayTrainReportCallback, we were using torch.distributed.barrier() to coordinate workers without specifying the device argument. If users do not set up the torch cuda device themselves, then the collective calls will be all bound to the default device (cuda:0). This could mess up the device map of the barrier call.

Why it hangs?

pytorch/pytorch#53658

Internally, torch.distributed.barrier calls allreduce on a dummy tensor link. The dummy tensor is created on the GPU specified by barrier(device_id=). NCCL will try to create a communicator for current process on device x if it doesn't exist, which is a blocking operation.

# rank 0 on node 0
collective1(tensor on dev 0) # creates new communicator, OK
collective2(tensor on dev 0) # doesn't try to create new communicator

# rank 1 on node 1
collective1(tensor on dev 0) # creates new communicator, OK
collective2(tensor on dev 1) # tries to create new communicator, hangs

When the users don't specify device_id, and the barrier will use the torch.cuda.current_device, which can be all cuda:0, if no torch.cuda.set_device was called before. The first call works, but the successive call hangs if it's binding processes with different device_id.

Therefore, to avoid deadlock, the rule-of-thumb is to always explicitly specify different device id for each workers for all the collective calls.

Solution

This PR switch from torch.distributed.barrier() to lightning's trainer.strategy.barrier(), which explicitly specified the cuda device for each barrier call to ensure the ordering.

Related issue number

Closes #42927

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 :(

[matthewdeng]

Awesome investigation!