如何在 Vertex AI Pipeline 中启用 GPU

Question

我对 Vertex AI Pipeline 比较陌生，我面临着无法在管道中启用 GPU 使用的问题。尽管为特定组件指定了硬件类型，但它仍然在 CPU 上运行，如下图所示。我怀疑我的管道设置中可能缺少配置。

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("NVIDIA_TESLA_T4")
        .set_gpu_limit("4")
    )

有人可以提供有关如何在 Vertex AI Pipeline 中正确启用 GPU 使用的指导吗？

任何帮助将不胜感激。

我已经安装了所有必需的 GPU 库来启用 GPU 加速，但是当我运行代码时，它仍然默认使用 CPU。

Answer 1

您需要确保已在基础容器中安装了 cuda 依赖项。

我制作这个组件是为了测试 GPU 是否可以使用张量：

@dsl.component
def simple(
    word: str,
    number: int,
) -> None:
    log.info("Running Simple")


    if torch.cuda.is_available():
        device = torch.device("cuda")  # set the device to be the first cuda device
        log.info("PyTorch is using the GPU")
        log.info("Current device: %s", torch.cuda.current_device())
        log.info("Device name: %s", torch.cuda.get_device_name(device))
    else:
        log.info("PyTorch is using the CPU")

    # Create a tensor on the CPU
    x = torch.randn(100, number)

    if torch.cuda.is_available():
        # Move the tensor to the GPU
        x = x.to(device)

    # Perform some computations on the tensor on the GPU
    y = x * 10**4

    # Move the tensor back to the CPU for printing
    y = y.to("cpu")