Debugging a slow PyTorch GPU Inference Pipeline on Beam/Google Cloud Dataflow

Question

We are attempting to use Google Cloud Dataflow to build a simple GPU-based classification pipeline that looks like this: Pub/Sub request comes in with link to a file on GCS → Read data from GCS → Chop up and batch data → Run inference in PyTorch.

Background

We deploy our pipeline on Dataflow with a custom Docker image adapted from the pytorch-minimal sample.

We ingest Pub/Sub messages and download data audio files from GCS using pathy, then chop audio into chunks for classification.

We've adapted Beam's relatively new RunInference function. Currently, there is no GPU support for RunInference on Dataflow (see open issue https://issues.apache.org/jira/browse/BEAM-13986). Upon building the Beam pipeline locally before deploy to Dataflow, the model initialization step doesn't recognize a CUDA environment and defaults to a CPU device for inference. This configuration gets propagated to the Dataflow execution environment that is properly GPU-enabled. So, we force a GPU device if requested without a CUDA device check. Other than that, the code is the same as the general RunInference code: A BatchElements operation followed by a ParDo that calls the model.

Problem

Everything is working-ish, but GPU inference is very slow – much slower than what we can time the same GPU-instance on processing batches on Google Cloud Compute Engine.

We're looking for advice on how to debug and speed up the pipeline. We suspect that the issue might have to do with threading as well as how Beam/Dataflow manages load across the pipeline stages. We kept running into CUDA OOM issues with multiple threads trying to access the GPU in the ParDo function. We launch our jobs with --num_workers=1 --experiment="use_runner_v2" --experiment="no_use_multiple_sdk_containers" to avoid multi-processing altogether. We saw that this 2021 beam summit talk on using Dataflow for local ML batch inference recommended going even further and just using a single worker thread --number_of_worker_harness_threads=1. However, we ideally don't want to do this: it's pretty common practice in ML pipelines like these to have multiple threads doing the I/O work of downloading data from the bucket and preparing batches so that the GPU never sits idle. Unfortunately, it seems that there is no way to tell beam to use a certain max number of threads per stage (?), so the best solution we could come up with is to protect the GPU with a Semaphore like so:

class _RunInferenceDoFn(beam.DoFn, Generic[ExampleT, PredictionT]):
    ...

    def _get_semaphore(self):
        def get_semaphore():
            logging.info('intializing semaphore...')
            return Semaphore(1)

        return self._shared_semaphore.acquire(get_semaphore)

    def setup(self):
        ...
        self._model = self._load_model()
        self._semaphore = self._get_semaphore()

    def process(self, batch, inference_args):
        ...
        logging.info('trying to acquire semaphore...')
        self._semaphore.acquire()
        logging.info('semaphore acquired')
        start_time = _to_microseconds(self._clock.time_ns())
        result_generator = self._model_handler.run_inference(
            batch, self._model, inference_args)
        end_time = _to_microseconds(self._clock.time_ns())
        self._semaphore.release()
        ...

We make three odd observations in that setup:

1. Beam always uses the minimum possible batch size we allow; if we specify a batch size of min 8 max 32, it'll always choose a batch size of at most 8, sometimes lower.
2. The inference timed here is still much much slower when allow multiple threads (--number_of_worker_harness_threads=10) than when we single-thread (--number_of_worker_harness_threads=1). 2.7s per batch vs. 0.4s per batch, both of which are a bit slower than running on compute engine directly.
3. In the multi-threaded setup, we keep seeing occasional CUDA OOM errors despite using a conservative batch size.

Would appreciate any and all debugging guidance for how to make this work! Right now, the whole pipeline is so slow that we have resorted to just running things in batches on Compute Engine again :/ – but there must be a way to make this work on Dataflow, right?

For reference:

• Single-threaded job: catalin-debug-classifier-test-1660143139 (Job ID: 2022-08-10_07_53_06-5898402459767488826)
• Multi-threaded job: catalin-debug-classifier-10threads-32batch-1660156741 (Job ID: 2022-08-10_11_39_50-2452382118954657386)

Answer 1

Thanks for trying RunInference!

I believe the problems you were encountering have been documented in the following issues. Can you please confirm if this is the case, or if it's different, explain the errors you were getting? We intend to work on those soon.

1. Map state_dict to the correct device during loading in PytorchModelHandler
2. Warn user about automatic GPU to CPU conversion.

---

Beam always uses the minimum possible batch size we allow; if we specify a batch size of min 8 max 32, it'll always choose a batch size of at most 8, sometimes lower.

The way that BatchElements decides on the size is by "profiling the time taken by (fused) downstream operations". Please see here and here for more information. It could be due to the specific size/nature of your data that creates a certain type of pattern of historical timings that causes this. Curious: are the data in your pipeline very similar?

---

The inference timed here is still much much slower when allow multiple threads (--number_of_worker_harness_threads=1) than when we single-thread (--number_of_worker_harness_threads=10). 2.7s per batch vs. 0.4s per batch, both of which are a bit slower than running on compute engine directly.

Just to clarify (probably a typo): did you mean "single threads (--number_of_worker_harness_threads=1) than when we multi-thread (--number_of_worker_harness_threads=10)" take 2.7s per batch vs. 0.4s per batch, respectively?

---

Some other questions:

• How big is the model that you are using?
• What types of GPUs have you tried this on?

---

We are currently looking into a similar issue with respect to GPUs and multithreading in the PR of our TensorRT RunInference. Here's a thread on that discussion. Something we are actively looking into to manage threads is to use start_bundle and finish_bundle. Please stay tuned on an update on this.

And let me dig into the logs of the job to see what is going on, and get back to you as soon as I can.