[CUDA graphs] Sync after replay (#57556)

Summary:
Right now** there's a bug in libcuda.so that triggers sometimes when graphs with certain topologies are replayed back to back without a sync in between. Replays that hit this bug turn into spaghetti: kernels reordered ignoring dependencies, kernels elided, corrupted results. Currently, the only workaround I know that fixes all our repros is a manual sync between replays.

I'll remove the sync (or special case it based on cuda version) in a later PR, as soon as a fixed libcuda.so is available.

The only substantive change is the cudaDeviceSynchronize, other lines changed are de-indenting an unneeded scope.

** The bug is in current and semi-recent public versions of libcuda.so. We discovered the bug recently and we're not sure yet which public release was first affected. The version that ships with 11.3 is definitely affected, versions that shipped with 11.1 and earlier are likely not affected.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/57556

Reviewed By: mruberry

Differential Revision: D28343043

Pulled By: ngimel

fbshipit-source-id: 3b907241aebdb8ad47ae96a6314a8b02de7bfa77
This commit is contained in:
Michael Carilli 2021-05-11 09:37:29 -07:00 committed by Facebook GitHub Bot
parent 565550d89a
commit dbedb1fa1c
2 changed files with 20 additions and 14 deletions

View file

@ -165,22 +165,26 @@ void CUDAGraph::replay() {
TORCH_CHECK(has_graph_exec_,
"Called CUDAGraph::replay without a preceding successful capture.");
c10::OptionalDeviceGuard device_guard{capture_stream_.device()};
// Just like any RNG consumer kernel!
auto* gen = get_generator_or_default<CUDAGeneratorImpl>(
c10::nullopt, cuda::detail::getDefaultCUDAGenerator());
PhiloxCudaState rng_engine_inputs;
{
c10::OptionalDeviceGuard device_guard{capture_stream_.device()};
// Just like any RNG consumer kernel!
auto* gen = get_generator_or_default<CUDAGeneratorImpl>(
c10::nullopt, cuda::detail::getDefaultCUDAGenerator());
PhiloxCudaState rng_engine_inputs;
{
std::lock_guard<std::mutex> lock(gen->mutex_);
rng_engine_inputs = gen->philox_cuda_state(wholegraph_increment_);
}
offset_extragraph_.fill_(int64_t(rng_engine_inputs.offset_.val));
// graph_exec_ may be replayed in any stream.
AT_CUDA_CHECK(cudaGraphLaunch(graph_exec_, at::cuda::getCurrentCUDAStream()));
std::lock_guard<std::mutex> lock(gen->mutex_);
rng_engine_inputs = gen->philox_cuda_state(wholegraph_increment_);
}
offset_extragraph_.fill_(int64_t(rng_engine_inputs.offset_.val));
// graph_exec_ may be replayed in any stream.
AT_CUDA_CHECK(cudaGraphLaunch(graph_exec_, at::cuda::getCurrentCUDAStream()));
// Temporary workaround for bug in libcuda.so that causes replayed graphs
// with certain topologies to be corrupted (kernels elided, internal syncs
// ignored) when replayed back to back without a sync in between.
// I hate to use a hard sync, but it's the only surefire workaround at the moment.
cudaDeviceSynchronize();
#else
TORCH_CHECK(false, "CUDA graphs may only be used in Pytorch built with CUDA >= 11.0");
#endif

View file

@ -3187,6 +3187,8 @@ torch.cuda.synchronize()
torch.cuda.synchronize()
torch.cuda.empty_cache()
@unittest.skip("Temporarily disabled due to a graphs bug in libcuda.so, " +
"see https://github.com/pytorch/pytorch/pull/57556")
@unittest.skipIf((not TEST_CUDA) or
TEST_WITH_ROCM or
int(torch.version.cuda.split(".")[0]) < 11, "CUDA >= 11.0 required for graphs")