如何在Vertex AI Pipeline中正确启用GPU加速

以下是几个关键的排查和修正步骤：

• 检查组件镜像是否支持GPU
  确认generate_article_embeddings组件使用的镜像包含GPU适配的依赖（如CUDA、cuDNN，以及GPU版本的TensorFlow/PyTorch）。如果用的是默认CPU镜像，硬件配置再正确也无法调用GPU。可以替换为Vertex AI官方提供的GPU镜像，示例如下：

  from kfp.v2.dsl import container_component, ContainerSpec
  
  @container_component
  def generate_article_embeddings(transformer_model, article_dataset):
      return ContainerSpec(
          image='gcr.io/deeplearning-platform-release/tf2-gpu.2.13:latest',
          command=['python', 'embedding_script.py'],
          args=[
              '--transformer_model', transformer_model,
              '--article_dataset', article_dataset
          ]
      )
  

• 修正节点选择器约束的格式
  Vertex AI中节点选择器的正确键是cloud.google.com/gke-accelerator，值为小写的GPU型号标识，不是直接写NVIDIA_TESLA_T4。修正后的代码：

  embeddings_task = (
          generate_article_embeddings(
              transformer_model=TRANSFORMER_MODEL,
              article_dataset=articles.output,
          )
          .set_cpu_limit("32")
          .set_memory_limit("208G")
          .add_node_selector_constraint("cloud.google.com/gke-accelerator", "nvidia-tesla-t4")
          .set_gpu_limit("4")
      )
  

• 验证GPU配额与权限
  确认流水线运行的区域有足够的NVIDIA T4 GPU配额，同时提交流水线的服务账号需要拥有roles/aiplatform.user及以上权限，确保能申请到GPU资源。

• 代码中显式指定使用GPU
  即使硬件配置到位，代码也要主动指定GPU设备。比如：
  PyTorch示例：

  import torch
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  model = model.to(device)
  

  TensorFlow示例：

  import tensorflow as tf
  tf.debugging.set_log_device_placement(True)
  with tf.device('/GPU:0'):
      # 执行模型推理或训练操作
  

• 确认资源配置逻辑
  确保GPU相关配置（set_gpu_limit和节点选择器）的调用顺序不影响生效，一般建议在CPU、内存配置后设置GPU相关参数。

内容的提问来源于stack exchange，提问作者MMM