Use the NVIDIA® L4 GPU type

The page explains how to run your Dataflow pipeline with the NVIDIA® L4 GPU type. The L4 GPU type is useful for running machine learning inference pipelines.

Requirements

  • Use the Apache Beam SDK version 2.46 or later. Apache Beam 2.50 or later is recommended.
  • You need L4 GPU quota (NVIDIA_L4_GPUS) in the region that your job runs in. For more information, see GPU quotas.
  • The L4 GPU type is available only with the G2 accelerator-optimized machine type. For more information, see The G2 machine series. Pipelines that use the L4 GPU type are subject to the G2 standard limitations.
  • The NVIDIA L4 GPU type uses the NVIDIA driver version 525.0 or later and the CUDA toolkit version 12.0 or later. Any code that you use in your pipeline must be compatible with the NVIDIA driver version and CUDA toolkit version. For example, if you use PyTorch, you need to use PyTorch version 23.01 or later.

Run pipelines with the NVIDIA® L4 GPU type

To use the NVIDIA L4 GPU type, you need to include the following pipeline options and service options in your pipeline code.

Java 

--workerMachineType=G2_MACHINE_TYPE
--dataflowServiceOptions="worker_accelerator=type:nvidia-l4;count:GPU_COUNT;install-nvidia-driver"

Python 

--machine_type=G2_MACHINE_TYPE
--dataflow_service_options="worker_accelerator=type:nvidia-l4;count:GPU_COUNT;install-nvidia-driver"

Go 

--worker_machine_type=G2_MACHINE_TYPE
--dataflow_service_options="worker_accelerator=type:nvidia-l4;count:GPU_COUNT;install-nvidia-driver"

Replace the following values:

  • G2_MACHINE_TYPE: the G2 machine type to use
  • GPU_COUNT: The number of GPUs to use. Each G2 machine type has a fixed number of NVIDIA L4 GPUs. To find the correct number of GPUs for your machine type, see the GPU count column in the G2 standard machine types table.

For more information about running pipelines with GPUs, see Run a pipeline with GPUs.

Manage dependencies

To manage dependencies, use a custom container. For more information, see Use custom containers in Dataflow.

The following Dockerfile example contains compatible dependencies for a pipeline that uses the NVIDIA L4 GPU type.

RUN apt-get -y update
RUN apt-get install [system packages]

# Install the SDK.
RUN pip install --no-cache-dir apache-beam[gcp]==2.51.0
# Install the machine learning dependencies.
RUN pip install --no-cache-dir tensorflow[and-cuda]
RUN pip install xgboost
RUN pip install transformers accelerate
(etc…..)
# Verify that the image doesn't have conflicting dependencies.
RUN pip check

# Copy files from official SDK image, including the script and dependencies.
COPY --from=apache/beam_python3.10_sdk:2.51.0 /opt/apache/beam /opt/apache/beam

# Set the entrypoint to Apache Beam SDK launcher.
ENTRYPOINT ["/opt/apache/beam/boot"]

What's next