使用GPU和其他加速器

Metaflow 使得在使用 AWS Batch 或 Kubernetes 时能够访问硬件加速的计算，特别是 GPU。您可以利用

• 单一加速器 - 例如 @resources(gpu=1)
• 具有多个加速器的单个实例 - 例如 @resources(gpu=4)
• 多个实例与多个加速器

您可以找到许多使用Metaflow微调LLMs和其他生成性AI模型的示例，以及如何训练计算机视觉和其他深度学习模型的内容 在这些文章中。

使用加速器​

在您开始利用硬件加速步骤之前，您需要满足两个前提条件：

1. 在 您的 Metaflow 堆栈 中添加硬件加速的实例。 查看 AWS Batch 和 Kubernetes 的具体提示。

2. 配置您的流程以包含必要的驱动程序和框架。

之后，使用下面解释的加速器。

GPU​

要在您的计算环境中使用GPU，请使用该 @resources 装饰器快速访问一个或多个GPU，如下面的示例所示：

from metaflow import FlowSpec, step, resources

class GPUFlow(FlowSpec):

    @resources(memory=32000, cpu=4, gpu=1)
    @step
    def start(self):
        from my_script import my_gpu_routine
        my_gpu_routine()
        self.next(self.end)

    @step
    def end(self):
        pass

if __name__ == '__main__':
    GPUFlow()

与@resources一样，该装饰器在本地运行时会被忽略。这允许您在本地开发代码，例如，在本地工作站上使用GPU资源。要在云中获取请求的资源，请运行流程--with kubernetes 或 --with batch。

如果您需要更细粒度的控制所使用的GPU，请使用特定于计算环境的装饰器：例如，@kubernetes 允许您指定一个 gpu_vendor，而@batch 允许您指定一个 queue，以针对包含特定GPU的计算环境。如需更多信息，请参阅AWS Batch和Kubernetes的指南。

使用 AWS Trainium 和 Inferentia​

在AWS上，您可以使用AWS特定的硬件加速器，Trainium和Inferentia。有关更多详细信息，请参阅一篇在Metaflow上下文中概述它们的博客文章。

在使用AWS Batch时，您只需在@batch中定义Trainium或Inferentia核心的数量即可请求加速器：

• @batch(trainium=16)
• @batch(inferentia=16)

请注意，Metaflow 支持在多个 Trainium 实例上进行 分布式训练。有关详细说明，请访问 metaflow-trainium 存储库。

使用谷歌的张量处理单元（TPUs）​

如您有兴趣在Google云中使用TPU与Metaflow Slack进行联系。

监控GPU利用率​

要监视GPU设备及其利用率，请在您的GPU步骤中添加 a custom card @gpu_profile。