使用 Habana Gaudi 进行深度学习训练入门

随着基于图像和手势识别、语音识别、自然语言处理、推荐系统等的应用程序和服务数量的不断增长，对高性能深度学习 (DL) 训练的需求正在加速。 Habana® Gaudi® AI 处理器旨在最大化训练吞吐量和效率，同时为开发人员提供优化的软件和工具，这些工具可以扩展到多种工作负载和系统。Habana Gaudi 软件的设计充分考虑了最终用户，提供了多功能性和易于编程性，以满足用户专有模型的独特需求，同时允许将现有模型无缝、轻松地迁移到 Gaudi。我们希望帮助您开始使用 Habana 的 Gaudi HPU（Habana 处理器单元）。许多新老用户可能都有在 GPU 或 CPU 上运行工作负载的经验，本文的目标是引导您完成将现有 AI 工作负载和模型迁移到 Gaudi HPU 的过程。我们将介绍硬件架构、软件堆栈以及让模型在 Gaudi 上运行所需的工具的基础知识。

Gaudi 架构和 Habana SynapseAI® 软件套件

我们将从概述基本架构开始。Gaudi 从头开始设计，用于加速 DL 训练工作负载。其异构架构包含一个完全可编程的张量处理核心 (TPC) 集群，以及相关的开发工具和库，还有一个可配置的矩阵数学引擎。 

TPC 是一个 VLIW SIMD 处理器，其指令集和硬件针对训练工作负载进行了定制。它是可编程的，为用户提供了最大的创新灵活性，并具有许多面向工作负载的功能，例如：

• GEMM 操作加速
• 张量寻址
• 延迟隐藏能力
• 随机数生成
• 特殊函数的先进实现

TPC 本地支持以下数据类型：FP32、BF16、INT32、INT16、INT8、UINT32、UINT16 和 UINT8。 

Gaudi 内存架构包括片上 SRAM 和每个 TPC 中的本地内存。此外，芯片封装集成了四个 HBM 设备，提供 32 GB 容量和 1 TB/s 带宽。PCIe 接口提供主机接口，并支持第三代和第四代模式。 

Gaudi 是首款在芯片上集成了 RoCE v2（通过以太网的 RDMA）引擎的 DL 训练处理器。这些引擎具有高达 2 Tb/s 的双向吞吐量，在训练过程中所需的处理器间通信中起着至关重要的作用。RoCE 的这种原生集成允许客户在服务器和机架内部（纵向扩展）以及跨机架（横向扩展）使用相同的扩展技术。这些可以直接连接到 Gaudi 处理器之间，或通过任何标准的以太网交换机连接。图 1 显示了硬件框图。

SynapseAI 软件套件旨在促进 Habana Gaudi 加速器上的高性能 DL 训练，它能够有效地将神经网络拓扑映射到 Gaudi 硬件上。该软件堆栈包括 Habana 的图编译器和运行时、TPC 内核库、固件和驱动程序，以及开发人员工具，如用于自定义内核开发的 TPC SDK 和 SynapseAI Profiler。SynapseAI 集成了流行的框架 TensorFlow 和 PyTorch，并针对 Gaudi 进行了性能优化。图 2 显示了 SynapseAI 软件套件的组成部分。为了轻松地将 SynapseAI 软件集成到您的工作环境中，Habana 提供了一组用于 TensorFlow 和 PyTorch 的 Docker 镜像，其中包含创建运行模型所需环境的所有组件。我们将探讨如何将这些库集成到您的模型中。

设置环境

在本节中，我们将介绍如何设置环境，然后添加确保框架能够识别 Gaudi HPU 并开始执行模型中的操作所需的简单步骤。除了 Docker 镜像之外，Habana 还在我们的 ModelReferences GitHub 存储库中提供了一系列参考模型和示例，可用于指导如何将适当的组件添加到您的模型中。 

入门的第一步是确保您拥有完整的构建环境，包括 Habana 驱动程序、SynapseAI 软件堆栈和框架。在大多数情况下，确保创建此完整环境的最佳方法是使用 Habana 提供的预构建 Docker 镜像，这些镜像可在我们的 Software Vault 中找到。这些 Docker 镜像同时包含单节点和横向扩展的二进制文件，无需额外的安装步骤。如果您在云环境中使用了 Gaudi，请务必选择您的云服务提供商提供的 Habana Gaudi 镜像，其中包含完整的驱动程序和相应的框架。对于需要独立安装完整驱动程序和 SynapseAI 软件堆栈（作为本地安装）的用户，您可以参考专用的 Setup and Install GitHub 存储库以获取详细说明。第二步是加载 Habana 库并定位 Gaudi HPU 设备。我们将在下一节中介绍 TensorFlow 和 PyTorch 的实现方法。

TensorFlow 入门

对于 TensorFlow，Habana 在插件中使用 `tf.load_library` 和 `tf.load_op_library` 调用库模块和自定义操作/内核，将 TensorFlow 框架与 SynapseAI 集成。框架集成包含三个主要组件：

• SynapseAI 辅助函数
• Device
• 图传递

TensorFlow 框架控制着图构建或图执行所需的大部分对象。SynapseAI 允许用户在设备上创建、编译和启动图。Graph passes 库使用模式匹配、标记、分段和封装 (PAMSEN) 操作来优化 TensorFlow 图。它旨在操作 TensorFlow 图以充分利用 Gaudi 的硬件资源。给定一组具有 Gaudi 实现的图节点，PAMSEN 会尝试合并尽可能多的图节点，同时保持图的正确性。通过保留图语义并自动发现可以融合为一个实体的子图，PAMSEN 提供的性能应与原生 TensorFlow 相当（或更好）。

要准备您的模型，您必须加载 Habana 模块库。调用位于 _library_loader.py_ 下的 `load_habana_module()`。此函数加载在 TensorFlow 级别使用 Gaudi HPU 所需的 Habana 库。

import tensorflow as tf
from habana_frameworks.tensorflow import load_habana_module
load_habana_module()
tf.compat.v1.disable_eager_execution() # if using tf.session

使用旧版 (TF1.x) 模型运行时有一些特定要求，使用 Horovod 进行多卡和多节点训练的要求也有所不同。如果您在单卡上运行 TF2.x 模型，则不需要最后一条语句。 

加载后，Gaudi HPU 将在 TensorFlow 中注册，并优先于 CPU。这意味着当给定操作同时适用于 CPU 和 Gaudi HPU 时，该操作将被分配给 Gaudi HPU。在完成上述步骤的模型初始迁移后，您可以开始查看 Habana 操作和 `habana_ops` 对象中定义的自定义 TensorFlow 操作的集成。它可以通过以下命令导入：`from habana_frameworks.tensorflow import habana_ops`。调用 `load_habana_module()` 后才能使用它。自定义操作用于与标准 TensorFlow 操作进行模式匹配。当模型被移植到 Gaudi HPU 时，软件堆栈会决定哪些操作放置在 CPU 上，哪些放置在 Gaudi HPU 上。优化传递会自动将不支持的操作放置在 CPU 上。 

完成模型中的这些步骤后，您就可以在 Gaudi HPU 实例上运行它了。下面是一个简单的示例，展示了如何为 Gaudi 启用 MNIST 模型。（所需的最小更改以粗体突出显示。）

import tensorflow as tf
from TensorFlow.common.library_loader import load_habana_module
load_habana_module()
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
model = tf.keras.models.Sequential([
tf.keras.layers.Flatten(input_shape=(28, 28)),
tf.keras.layers.Dense(10),
])
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)
model.compile(optimizer=optimizer, loss=loss, metrics=['accuracy'])
model.fit(x_train, y_train, epochs=5, batch_size=128)
model.evaluate(x_test, y_test)

PyTorch 入门

PyTorch Habana 桥梁在框架和 SynapseAI 软件堆栈之间进行接口，以驱动深度学习模型在 Habana Gaudi 设备上的执行（图 3）。Habana 提供的安装包在标准 PyTorch 版本的基础上进行了修改，并包含在 Habana 提供的 Docker 镜像中。需要使用此安装包中的自定义框架来集成 PyTorch 与 Habana 桥梁。在这种情况下，您将通过加载 PyTorch Habana 插件库并 `import habana_frameworks.torch.core` 模块来修改 PyTorch 深度学习模型训练脚本，以与 Habana Bridge 集成。 

在 Gaudi HPU 上执行 PyTorch 模型有两种主要模式，由 Habana PyTorch 支持：

• Eager 模式 – 按标准 PyTorch eager 模式脚本中定义的逐操作执行。
• Lazy 模式 – 图的延迟执行，由脚本逐操作交付，类似于 Eager 模式。它提供了 Eager 模式的体验，并在 Gaudi 上实现了高性能。

通常，建议在 Eager 模式下进行初始模型开发以确定功能，然后切换到 Lazy 模式以获得最佳性能。可以通过设置运行时标志（如下面的第 5 步所示）来选择。

要在 Habana 上运行模型，需要进行以下代码添加。以下步骤涵盖了 Eager 和 Lazy 执行模式。

1. 加载 Habana PyTorch 插件库 `libhabana_pytorch_plugin.so`。

import torch
from habana_frameworks.torch.utils.library_loader import load_habana_module
load_habana_module()

2. 定位 Gaudi HPU 设备

device = torch.device("hpu")

3. 将模型移动到设备

model.to(device)

4. 导入 Habana Torch 库

import habana_frameworks.torch.core as htcore

5. 通过设置下面的环境变量启用 Lazy 执行模式。如果您想在 Eager 模式下执行代码，请不要设置此环境变量。

os.environ["PT_HPU_LAZY_MODE"] = "1"

6. 在 Lazy 模式下，当数据从 Habana 设备读取回主机时，执行就会被触发。例如，如果您正在运行一个拓扑并将损失值从设备获取到主机，使用 `loss.item()`，就会触发执行。添加 `mark_step()` 也可以触发执行。

htcore.mark_step()

在以下位置的代码中需要放置 `mark_step()`：

• 在 `optimizer.step()` 之后，以清晰地划分训练迭代；
• 在 `loss.backward` 和 `optimizer.step()` 之间，如果使用的优化器是 Habana 自定义优化器。

当模型被移植到 Gaudi HPU 上运行时，软件堆栈会决定哪些操作放置在 CPU 上，哪些放置在 HPU 上。此决定基于操作是否以 HPU 作为后端在 PyTorch 中注册。如果操作没有以其后端注册为 HPU，则操作的执行会自动回退到主机 CPU。

Habana 提供了针对 Habana 设备定制的一些复杂 PyTorch 操作的实现。在给定模型中，用自定义 Habana 版本替换这些复杂操作将获得更好的性能。 

以下是 Habana 设备当前支持的自定义优化器列表：

• FusedAdagrad – 参考 torch.optim.Adagrad
• FusedAdamW – 参考 torch.optim.AdamW
• FusedLamb – 参考 LAMB 优化器论文
• FusedSGD – 参考 torch.optim.SGD

结论

本文的目的是引导您完成将现有深度学习模型迁移到 Gaudi 的过程，并展示让模型准备好运行的基本步骤。请继续关注主网站 Habana.ai 的产品发布公告，并注册我们的兴趣列表，以接收即将举行的活动的通知。

有用资源

1. 开发者网站是入门的最佳去处。该网站将提供我们的模型、文档和其他资源的链接。
2. Habana 文档包含所有详细资料，特别是迁移指南。
3. Habana Vault 是下载最新驱动程序和 Docker 镜像的地方。
4. 用户论坛是社区讨论所有与 Habana 和 Gaudi 相关事宜的场所。

注册以获取《并行宇宙 6》的未来发行版。