QNNPACK convolution race condition

[dbalchev]

🐛 Bug

We're hitting a QNNPACK race condition, when executing quantized convolutions concurrently in more than 1 thread.

To Reproduce

Steps to reproduce the behavior:

1. Run the following script

   import random
   import threading
   
   import torch
   
   torch.set_grad_enabled(False)
   torch.backends.quantized.engine = "qnnpack"
   
   
   def _make_quant_layer():
       layer = torch.nn.intrinsic.ConvReLU2d(
           torch.nn.Conv2d(64, 64, 3, padding=1), torch.nn.ReLU()
       )
       layer.qconfig = torch.quantization.default_per_channel_qconfig
       fake_model = torch.nn.Sequential(layer)
       torch.quantization.prepare(fake_model, inplace=True)
       fake_model(torch.zeros(1, 64, 1, 1))
       layer = torch.nn.intrinsic.quantized.ConvReLU2d.from_float(layer)
       return layer
   
   
   class _RaceExperiment:
       def __init__(self, num_threads):
           self.module = torch.jit.script(
               torch.nn.Sequential(*[_make_quant_layer() for _ in range(50)])
           )
           self.module.eval()
   
           self.enter_barrier = threading.Barrier(num_threads)
           self.exit_barrier = threading.Barrier(num_threads)
           self.num_threads = num_threads
   
       def thread_run(self):
           data = torch.zeros((1, 64, 50, random.choice([50, 100, 200, 400])))
           data = torch.quantize_per_tensor(data, 1.0, 0, dtype=torch.quint8)
           self.enter_barrier.wait()
           self.module(data)
           self.exit_barrier.wait()
           self.module(data)
   
       def run_experiment(self):
           threads = [
               threading.Thread(target=self.thread_run) for _ in range(self.num_threads)
           ]
           for t in threads:
               t.start()
           for t in threads:
               t.join()
           print("done")
   
   
   _RaceExperiment(1).run_experiment()
   
   for _ in range(10):
       _RaceExperiment(3).run_experiment()

2. Pretty consistently I get either a segmentation fault, another memory error (e.g. double free) or
   the following message:

   RuntimeError: The following operation failed in the TorchScript interpreter.
    Traceback of TorchScript (most recent call last):
      File "/home/ubuntu/repro-venv/lib/python3.7/site-packages/torch/nn/modules/container.py", line 119, in forward
        def forward(self, input):
            for module in self:
                input = module(input)
                        ~~~~~~ <--- HERE
            return input
      File "/home/ubuntu/repro-venv/lib/python3.7/site-packages/torch/nn/intrinsic/quantized/modules/conv_relu.py", line 85, in forward
                input = F.pad(input, _reversed_padding_repeated_twice,
                              mode=self.padding_mode)
            return torch.ops.quantized.conv2d_relu(
                   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
                input, self._packed_params, self.scale, self.zero_point)
    RuntimeError: pack_w != nullptr INTERNAL ASSERT FAILED at "/pytorch/aten/src/ATen/native/quantized/cpu/qconv.cpp":678, please report a bug to PyTorch. Packed Weights are NULL
   

Expected behavior

The script should finish without any issues. This can be achieved either by using the FBGEMM
backend or setting num_threads = 1 in _RaceExperiment.

Environment

Please copy and paste the output from our
environment collection script
(or fill out the checklist below manually).

Collecting environment information...
PyTorch version: 1.8.1+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: Could not collect

Python version: 3.7 (64-bit runtime)
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.2
[pip3] torch==1.8.1+cpu
[pip3] torchaudio==0.8.1
[pip3] torchvision==0.9.1+cpu
[conda] Could not collect

Additional context

We've ran into this issue on one of our client's machines, so the bug is not theoretical. The actual
issue we observed there was the error message:

RuntimeError: pack_w != nullptr INTERNAL ASSERT FAILED at "/pytorch/aten/src/ATen/native/quantized/cpu/qconv.cpp":678, please report a bug to PyTorch. Packed Weights are NULL

(this is adapted to 1.8.1 since we run on 1.7.1, but I'm able to reproduce with the above script on
1.8.1)

Unfortunately we can't access AVX2 on that machine, so we cannot use FBGEMM.

My hypothesis of the cause of the race condition is the following:

1. Before the first execution of the convolution, PackedConvWeightsQnnp::w is nullptr and
   PackedConvWeightsQnnp::input_scale doesn't contain a value (as outlined in https://github.com/pytorch/pytorch/blob/v1.8.1/aten/src/ATen/native/quantized/cpu/qconv_prepack.cpp#L254-L257)
2. One of the threads starts executing enters https://github.com/pytorch/pytorch/blob/v1.8.1/aten/src/ATen/native/quantized/cpu/qconv.cpp#L604
   and stores a value in PackedConvWeightsQnnp::input_scale, but hasn't yet assigned
   PackedConvWeightsQnnp::w
3. Another thread sees that value of PackedConvWeightsQnnp::input_scale, skips the branch on https://github.com/pytorch/pytorch/blob/v1.8.1/aten/src/ATen/native/quantized/cpu/qconv.cpp#L604,
   reaches https://github.com/pytorch/pytorch/blob/v1.8.1/aten/src/ATen/native/quantized/cpu/qconv.cpp#L678 and exits
   with the error we see above.

cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo

[raghuramank100]

cc @kimishpatel

[kimishpatel]

Trouble here is that you have one model that is being run on different threads. So each thread is doing its own inference using the same underlying model. This usecase has its own trouble at the moment and it can rightfully run into race condition. Your hypothesis is pretty much on the spot, most likely.

Is your usease what I described? That is you have one model and you are doing inference on different threads using the same model? One way around, and it is certainly cringworthy, is to do dummy inference outside of multithreaded setup and then do multithreaded inference. This has not been quite tested though.

We can certainly look into this to make it more robust if this is crucial to your project.

[dbalchev]

@kimishpatel Thanks for the response.
I confirm our use-case involves running inference of the same torch.nn.Module concurrently in multiple threads. We were even happily surprised when we saw that TorchScript doesn't touch the GIL.

One way around, and it is certainly cringworthy, is to do dummy inference outside of multithreaded setup and then do multithreaded inference.

It's really cringeworthy yeah. But we tried that and that didn't work too for some reason - we observe a segmentation fault. It can be reproduced by modifying the script I posted.

Our current workaround involves thread local models. We observed torch.jit.script is not thread-safe, so we have an additional lock for initialization (but this is a relatively minor problem).

Nevertheless thanks for the suggestion (:

We can certainly look into this to make it more robust if this is crucial to your project.

That would be great (:. To be honest we prefer using FBGEMM, since it's significantly faster. But as it turns out some of our clients' machines don't have access to AVX2 for some reason, so we're stuck with QNNPACK for them.

[kimishpatel]

@dbalchev thanks for all the inputs. We will keep this in our planning. I dont have a date yet by which we can do this but once I hav a better idea I will post the timeline here.

[kimishpatel]

@dbalchev, I just realized why you saw the seg fault with work around.

This is due to stateful ops. In general pytorch ops are stateless, such that no mutable state persist, and thus, for inference, likely thread safe. With quantized conv/linear ops we introduce objects that can contain state. In the case of fbgemm I think it does not maintain any mutable state, but for qnnpack this is different.
Qnnpack based conv/linear op contains two pieces of information in the state.

1. Immutable packed weights and this is what you ran into and described.
2. Mutable state that contains some information about activation tensor. This state is tracked from one inference to next and is used for performance optimization. However this piece is not thread safe and I highly suspect that is what is leading to seg fault. Now to make this thread safe will require some effort without sacrificing performance. We will certainly look into this, but just wanted to provide some context. For source code pointer, here is where the state is tracked: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/quantized/cpu/qnnpack/src/conv-run.cc#L331-L335

[dbalchev]

@kimishpatel Thank you so much for the detailed explanation (:

Regarding the code you linked. This might be a really uninformed question, since I have little to no experience with torch internals, but isn't input always different? At least when used in the python torch.nn.Sequential interfance, since we pass a new tensor each time?

[kimishpatel]

@kimishpatel Thank you so much for the detailed explanation (:

Regarding the code you linked. This might be a really uninformed question, since I have little to no experience with torch internals, but isn't input always different? At least when used in the python torch.nn.Sequential interfance, since we pass a new tensor each time?

@dbalchev sorry took me a while to get back to this thread. So no input does not have to be different every time. Especially if say your tensor allocation pattern, when running the model, does not change. This can be the case for many conv models. Plus if you have an allocator that can cache allocations and give you the same pointer back. We have a CPUCachingAllocator that caches allocation which can, not guaranteed, give you the same memory pointer back.

[dbalchev]

@kimishpatel thank you for the explanation (:

[jshen]

This issue has been fixed @dbalchev