在 Vertex AI 上使用 Tesla A100 GPU 和 Kubeflow Pipelines
Posted
技术标签:
【中文标题】在 Vertex AI 上使用 Tesla A100 GPU 和 Kubeflow Pipelines【英文标题】:Using Tesla A100 GPU with Kubeflow Pipelines on Vertex AI 【发布时间】:2021-11-11 03:13:57 【问题描述】:我正在使用以下代码行在 Kubeflow 管道 (KFP) 上指定所需的机器类型和加速器/GPU,我将通过 Vertex AI/管道以无服务器方式运行。
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)
运行/部署时会引发错误,因为生成的机器类型是通用类型,即 N1-Highmem-*
当我没有指定cpu和内存限制时也发生了同样的事情,希望它会根据加速器约束自动选择正确的机器类型。
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",
【问题讨论】:
您在哪个区域运行 GCP 服务?Tesla A100
对可以找到的区域非常有限 here 你能在没有私人信息的情况下分享你的整个代码吗?我猜它看起来类似于this
管道在/locations/us-central1/
和A2 Machine Type
上执行,A100 GPUs
应该对us-central1
区域普遍可用,如here 所述
【参考方案1】:
目前,GCP 不支持普通 KF 组件的 A2 机器类型。目前一个潜在的解决方法是使用GCP custom job component,您可以明确指定机器类型。
【讨论】:
以上是关于在 Vertex AI 上使用 Tesla A100 GPU 和 Kubeflow Pipelines的主要内容,如果未能解决你的问题,请参考以下文章