Repository: pytorch/pytorch

Issue:
https://github.com/pytorch/pytorch/pull/74944 breaks CUDA graph capture
### 🐛 Describe the bug

Prior to https://github.com/pytorch/pytorch/pull/74944, pytorch's native sync batch norm was safe to use with CUDA graph capture.

After https://github.com/pytorch/pytorch/pull/74944, Thorsten Kurth (@azrael417, one of our engineers using Pytorch for mlperf models) found that graph-capturing syncbn fails:
```
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/batchnorm.py", line 748, in forward
    return sync_batch_norm.apply(
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/_functions.py", line 72, in forward
    count_all = count_all[mask]
RuntimeError: CUDA error: operation not permitted when stream is capturing
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Exception raised from memcpy_and_sync at /opt/pytorch/pytorch/c10/cuda/CUDAFunctions.h:76 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0x6c (0x7fbcd6293cec in /opt/conda/lib/python3.8/site-packages/torch/lib/libc10.so)
frame #1: void at::native::nonzero_cuda_out_impl<bool>(at::Tensor const&, at::Tensor&) + 0x1136 (0x7fbcd8239286 in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cuda.so)
frame #2: at::native::nonzero_out_cuda(at::Tensor const&, at::Tensor&) + 0x607 (0x7fbcd821ec97 in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cuda.so)
frame #3: at::native::nonzero_cuda(at::Tensor const&) + 0x252 (0x7fbcd821f102 in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cuda.so)
frame #4: <unknown function> + 0x2d2ebfd (0x7fbcd9053bfd in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cuda.so)
frame #5: <unknown function> + 0x2d2ec85 (0x7fbcd9053c85 in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cuda.so)
frame #6: at::_ops::nonzero::call(at::Tensor const&) + 0x139 (0x7fbd029d4a09 in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so)
frame #7: <unknown function> + 0x1421904 (0x7fbd02651904 in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so)
frame #8: at::native::index(at::Tensor const&, c10::List<c10::optional<at::Tensor> > const&) + 0x8c (0x7fbd02652d1c in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so)
frame #9: <unknown function> + 0x2d1ec29 (0x7fbcd9043c29 in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cuda.so)
frame #10: <unknown function> + 0x2d1ecb8 (0x7fbcd9043cb8 in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cuda.so)
frame #11: at::_ops::index_Tensor::redispatch(c10::DispatchKeySet, at::Tensor const&, c10::List<c10::optional<at::Tensor> > const&) + 0x88 (0x7fbd02ac4b08 in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so)
frame #12: <unknown function> + 0x2cac7e0 (0x7fbd03edc7e0 in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so)
frame #13: <unknown function> + 0x2cacf5b (0x7fbd03edcf5b in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so)
frame #14: at::_ops::index_Tensor::call(at::Tensor const&, c10::List<c10::optional<at::Tensor> > const&) + 0x27e (0x7fbd02b1b0ae in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so)
frame #15: torch::autograd::THPVariable_getitem(_object*, _object*) + 0x51d (0x7fbd0af8b7ed in /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_python.so)
frame #16: _PyEval_EvalFrameDefault + 0xacd (0x5622e7e261ed in /opt/conda/bin/python)
...python frames
```

We believe the failure is a CPU sync incurred by [`count_all = count_all[mask]`](https://github.com/pytorch/pytorch/pull/74944/files#diff-6083dbd5d169b44f41c2fd5e9638c88d246aff4cd86d2a839c72972ed198546bR72).

How hard would it be to remove the sync, or at least reenable the old syncfree code if the user wants graph capture?

cc @ezyang @gchanan @zou3519 @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @kwen2501 @mcarilli


Discussion:
Synced with @mcarilli offline, for v1.12, we are going to use [is_current_stream_capturing](https://pytorch.org/docs/master/generated/torch.cuda.is_current_stream_capturing.html?highlight=is_current_stream_capturing#torch.cuda.is_current_stream_capturing) to guard the code below as a temporary fix. 

https://github.com/pytorch/pytorch/blob/cd4ffc865b337f3e0f51d6e757ac2a2ab83e9f20/torch/nn/modules/_functions.py#L70-L74

For longer term, we might need to update the CUDA kernel for `batch_norm_gather_stats_with_counts` to deal with 0 counts. 

cc @datumbox 
Created #78656 to track longer-term solution

Resolving patch:
diff --git a/torch/nn/modules/_functions.py b/torch/nn/modules/_functions.py
--- a/torch/nn/modules/_functions.py
+++ b/torch/nn/modules/_functions.py
@@ -67,11 +67,19 @@ def forward(self, input, weight, bias, running_mean, running_var, eps, momentum,
             # world_size * (2C + 1) -> world_size * C, world_size * C, world_size * 1
             mean_all, invstd_all, count_all = torch.split(combined, num_channels, dim=1)
 
-        # remove stats from empty inputs
-        mask = count_all.squeeze(-1) >= 1
-        count_all = count_all[mask]
-        mean_all = mean_all[mask]
-        invstd_all = invstd_all[mask]
+        if not torch.cuda.is_current_stream_capturing():
+            # The lines below force a synchronization between CUDA and CPU, because
+            # the shape of the result count_all depends on the values in mask tensor.
+            # Such synchronizations break CUDA Graph capturing.
+            # See https://github.com/pytorch/pytorch/issues/78549
+            # FIXME: https://github.com/pytorch/pytorch/issues/78656 describes
+            # a better longer-term solution.
+
+            # remove stats from empty inputs
+            mask = count_all.squeeze(-1) >= 1
+            count_all = count_all[mask]
+            mean_all = mean_all[mask]
+            invstd_all = invstd_all[mask]
 
         # calculate global mean & invstd
         mean, invstd = torch.batch_norm_gather_stats_with_counts(


Test patch: