Enabling GPU and TPU

Out of date

This guide contains outdated information pertaining to Kubeflow 1.0. This guide needs to be updated for Kubeflow 1.1.

This page describes how to enable GPU or TPU for a pipeline on GKE by using the Pipelines DSL language.

To enable GPU and TPU on your Kubeflow cluster, follow the instructions on how to customize the GKE cluster for Kubeflow before setting up the cluster.

The code above will be compiled into Kubernetes Pod spec:

  2.   ...
  3.   resources:
  4.       nvidia.com/gpu: "2"

If the cluster has multiple node pools with different GPU types, you can specify the GPU type by the following code.

The code above will be compiled into Kubernetes Pod spec:

  1. container:
  3.   resources:
  4.     limits:
  5.       nvidia.com/gpu: "2"
  6. nodeSelector:

Use the following code to configure ContainerOp to consume TPUs on GKE:

The above code uses 8 v2 TPUs with TF version to be 1.12. The code above will be compiled into Kubernetes Pod spec:

  2.   ...
  3.   resources:
  4.     limits:
  5.       cloud-tpus.google.com/v2: "8"
  6.   metadata:

See the to learn more about TPU settings.