在 Vertex AI 上使用 Tesla A100 GPU 和 Kubeflow 管道

Question

我正在使用以下代码行来指定所需的机器类型，并且 accelerator/GPU 在 Kubeflow 管道 (KFP) 上，我将通过 Vertex 运行以无服务器方式运行ning AI/Pipelines.

op().
set_cpu_limit(8).
set_memory_limit(50G).
add_node_selector_constraint('cloud.google.com/gke-accelerator', 'nvidia-tesla-k80').
set_gpu_limit(1)

它也适用于其他 GPU，即 Pascal、Tesla、Volta 卡。

但是，我不能对最新的加速器类型 Tesla A100 做同样的事情，因为它需要一种特殊的机器类型，至少是 a2-highgpu-1g.

当我在 Vertex 上运行时，如何确保这个特定组件运行在 a2-highgpu-1g 之上？

如果我只是按照旧 GPU 的方法：

op().
set_cpu_limit(12). # max for A2-highgpu-1g
set_memory_limit(85G). # max for A2-highgpu-1g
add_node_selector_constraint('cloud.google.com/gke-accelerator', 'nvidia-tesla-a100').
set_gpu_limit(1)

当 run/deployed 时抛出错误，因为生成的机器类型是通用类型，即 N1-Highmem-*

当我没有指定 cpu 和内存限制时发生了同样的事情，希望它会根据加速器约束自动 select 正确的机器类型。

    op().
    add_node_selector_constraint('cloud.google.com/gke-accelerator', 'nvidia-tesla-a100').
    set_gpu_limit(1)

错误： "NVIDIA_TESLA_A100" is not supported for machine type "n1-highmem-2",

Answer 1

目前，GCP 不支持普通 KF 组件的 A2 机器类型。现在一个潜在的解决方法是使用 GCP custom job component，您可以明确指定机器类型。

在 Vertex AI 上使用 Tesla A100 GPU 和 Kubeflow 管道

Using Tesla A100 GPU with Kubeflow Pipelines on Vertex AI

google-kubernetes-engine

kubeflow-pipelines

google-cloud-vertex-ai