在生产场景中，从模型中获得最佳性能对于提供快速响应时间、低成本以及低资源需求至关重要。当计算资源需求较大时，尤其是处理大型模型和/或数据集时，或者当推理延迟和/或成本要求具有挑战性时，高性能建模变得尤为重要。
在本章中，我们将讨论如何通过数据和模型并行化来加速模型。我们还将探讨高性能建模技术，如分布式策略，以及高性能数据摄取管道，如TF Data。最后，我们将考虑巨型神经网络的兴起，以及在这种背景下应对由此产生的对高效、可扩展基础设施需求的方法。
分布式训练
当你开始原型设计时，训练模型可能是一个快速而简单的任务，尤其是在处理小数据集时。然而，完全训练一个模型可能会变得非常耗时。许多领域的数据集和模型架构都变得越来越大。随着训练数据集和模型的大小增加，模型的训练时间也变得越来越长。而且，不仅仅是每个周期的训练时间；通常，由于某些原因，模型的周期数也会增加。解决这种问题通常需要分布式训练。分布式训练允许我们在加速训练的同时，通过利用更多的计算资源来训练巨大的模型。
从高层次来看，分布式训练有两种主要方式：数据并行和模型并行。在数据并行中，这是实现起来相对简单的一种方法，你将数据划分为多个部分，并将完整的模型复制到所有工作节点。每个工作节点在不同的数据分区上进行操作，并且模型更新会在所有工作节点之间同步。这种类型的并行性与模型无关，可以应用于任何神经网络架构。通常，数据并行的规模与批处理大小相关。
在模型并行中，你将模型划分为不同的部分，在不同的工作节点上并行训练。每个工作节点在相同的数据上训练，工作节点只需要同步共享的参数，通常是在每次前向传播或反向传播步骤后进行一次同步。当你的模型较大，无法在加速器（如GPU或TPU）的内存中完全加载时，通常使用模型并行。与数据并行相比，模型并行的实现相对较为复杂。因此，我们关于分布式训练技术的讨论将集中在数据并行上。
数据并行
如前所述，在数据并行中，数据被分割成多个部分，通常分区的数量等于计算集群中可用的工作节点数量。如图7-1所示，你将模型复制到每个工作节点，每个工作节点在其自己的数据子集上进行训练。这要求每个工作节点有足够的内存来加载整个模型，而对于较大的模型，这可能成为一个问题。
每个工作节点独立地计算其对训练样本的预测与标签数据之间的误差，然后执行反向传播，基于误差更新其模型，并将所有更改传递给其他工作节点，以便它们可以更新自己的模型。这意味着，工作节点需要在每个批次结束时同步它们的梯度，以确保它们训练的是一致的模型。
同步训练与异步训练
数据并行中的两种基本训练方式是同步训练和异步训练。在同步训练中，每个工作节点在其当前的小批次数据上进行训练，应用自身的更新，向其他工作节点传递其更新，并在接收到并应用来自其他工作节点的所有更新后，才能进行下一个小批次的训练。全减少算法是其中的一个例子。
在异步训练中，所有工作节点独立地在其小批次数据上进行训练并异步更新变量。异步训练通常更高效，但实现起来也更为复杂。参数服务器算法是这种方法的一个例子。
异步训练的一个主要缺点是精度降低和收敛速度变慢，这意味着需要更多的步骤才能收敛。收敛速度变慢可能不会成为问题，因为异步训练的加速可能足以弥补这一点。然而，精度的损失可能会是一个问题，这取决于损失的精度有多大以及应用的要求。
分布式感知
要使用分布式训练，模型需要具备分布式感知。幸运的是，像Keras这样的高级API支持分布式训练。
你甚至可以创建自定义的训练循环，以提供更精确的控制。为了使模型能够以分布式方式进行训练或推理，你需要通过一些小的代码更改，使它们具有分布式感知。
tf.distribute：TensorFlow中的分布式训练
要在TensorFlow中执行分布式训练，您可以使用TensorFlow的tf.distribute.Strategy类。
该类支持多种分布式策略，用于高级API，并且还支持使用自定义训练循环进行训练。该类还支持在急切模式（eager mode）和图模式（graph mode）下执行TensorFlow代码，使用tf.function。除了训练模型外，还可以使用tf.distribute.Strategy在不同平台上以分布式方式执行模型评估和预测。
tf.distribute.Strategy类只需要少量额外的代码来将您的模型适配为分布式训练。您可以轻松地在不同策略之间切换，以进行实验并找到最适合您需求的策略。TensorFlow有许多不同的策略用于执行分布式训练。以下是最常用的几种策略：
OneDeviceStrategyMirroredStrategyParameterServerStrategyMultiWorkerMirroredStrategyTPUStrategyCentralStorageStrategy
在这里，我们将重点介绍前三种策略，以帮助您了解基本的相关问题和方法，因为后面三种策略都是前面三种的派生。TensorFlow网站上有更多关于这些策略的详细信息。
OneDeviceStrategy
OneDeviceStrategy 将在其作用域内创建的所有变量放置在指定设备上。通过此策略分发的输入将预先加载到指定的设备上。此外，任何通过 strategy.run 调用的函数也将被放置在指定设备上。
您可能会问：“如果只有一个设备，那它的意义何在？”此策略的典型使用场景可能是在切换到实际分布到多个设备/机器的其他策略之前，先使用 tf.distribute.Strategy API 测试您的代码。
MirroredStrategy
MirroredStrategy 支持在单台机器上的多个GPU上进行同步分布式训练。它为每个GPU设备创建一个副本，并且模型中的每个变量在所有副本之间进行镜像复制。这些变量通过应用相同的更新保持同步，形成一个单一的概念变量，称为镜像变量。
使用高效的全还原（all-reduce）算法来跨设备通信变量更新。全还原通过将所有设备的张量加起来进行汇总，然后使其在每个设备上可用。全还原是一种非常高效的融合算法，可以显著减少同步的开销。
MultiWorkerMirroredStrategy
MultiWorkerMirroredStrategy 将训练分布在多个工作节点上，每个节点可以有多个GPU。TPUStrategy 类似于 MirroredStrategy，但是将训练分布在多个TPU上，而不是GPU上。
最后，CentralStorageStrategy 不会镜像变量，而是将它们放置在CPU上，并在所有本地GPU上复制操作。
ParameterServerStrategy
ParameterServerStrategy 是一种常见的异步数据并行方法，用于在多台机器上扩展模型训练。参数服务器训练集群由工作节点和参数服务器组成。变量是在参数服务器上创建的，并由工作节点在每个步骤中读取和更新。默认情况下，工作节点独立地读取和更新这些变量，而不与其他工作节点进行同步。这就是为什么参数服务器风格的训练有时被称为异步训练的原因。
Fault Tolerance
通常，在同步训练中，如果一个或多个工作节点失败，整个工作节点集群都会失败。因此，在工作节点死亡或变得不稳定的情况下，考虑某种形式的容错机制是非常重要的。这使得可以从因预占工作节点而发生的故障中恢复过来。这可以通过在分布式文件系统中保留训练状态来完成。由于所有工作节点在训练的epoch和步骤上保持同步，其他工作节点需要等待失败或被预占的工作节点重新启动，才能继续训练。
例如，在 MultiWorkerMirroredStrategy 中，如果一个工作节点中断，整个集群会暂停，直到中断的工作节点重新启动。其他工作节点也会重新启动，中断的工作节点重新加入集群。然后，需要有一种方法让每个工作节点恢复其之前的状态，从而使集群重新同步，确保训练能够顺利进行。例如，Keras 提供了在 BackupAndRestore 回调中实现此功能。
高效的输入管道
加速器是高性能建模、训练和推理的关键部分。但加速器也很昂贵，因此使用它们时需要高效。这意味着要保持加速器的忙碌，这需要快速提供数据。因此，高效的输入管道在高性能建模中至关重要。
输入管道基础
输入管道是许多训练管道中的重要部分，但推理管道也常常有类似的需求。在训练管道的更大背景下，例如 TensorFlow Extended (TFX) 训练管道，高性能输入管道将是 Trainer 组件的一部分，可能还包括 Transform 等其他组件，这些组件通常需要对数据进行大量处理。
在提高输入管道效率时，了解输入管道摄取数据的基本步骤非常重要。你可以将输入管道视为一个提取、转换、加载（ETL）过程。该过程的第一步是从数据存储中提取数据，这些数据存储可能是本地的，也可能是远程的，如硬盘、固态硬盘（SSD）、云存储和 Hadoop 分布式文件系统（HDFS）。
第二步，数据通常需要进行预处理或转换。这包括数据的打乱、批处理、重复以及应用逐元素转换。如果这些转换耗时过长，加速器可能会在等待数据时被低效使用。此外，转换的顺序可能会影响管道的性能。在使用任何数据转换（如 map、batch、shuffle、repeat 等）时，你需要注意这一点。
输入管道的第三步是将预处理过的数据加载到模型中，模型可能在 CPU、GPU 或 TPU 上进行训练，并开始训练。高性能输入管道的关键要求是平行化数据处理，以尽量高效地使用可用的计算、I/O 和网络资源。特别是对于更昂贵的组件，如加速器，你希望尽可能让它们保持忙碌，而不是等待数据。
输入管道模式：提高效率
让我们来看一个典型的模式，这种模式很容易陷入，但你真的想避免它。
在图7-2中，关键硬件组件，包括CPU和加速器，处于空闲状态，等待前一步骤完成。如果你仔细想一想，ETL（提取、转换、加载）是一个很好的数据性能思维模型。为了帮助你理解管道如何执行，假设ETL的每个阶段都使用系统中的不同硬件组件。提取阶段利用的是磁盘，或者如果你从远程系统加载数据，则利用的是网络。转换阶段通常发生在CPU上，且可能非常消耗CPU资源。加载阶段则是利用直接内存访问（DMA）子系统和与加速器（通常是GPU或TPU）的连接。
图7-3中展示的方法比图7-2中的模式更加高效，尽管它仍然不是最优的。然而，实际上，在许多情况下，这种模式可能很难进一步优化。
如图7-3所示，通过并行化操作，你可以使用一种被称为软件管道（software pipelining）的技术，重叠ETL的不同部分。在软件管道中，你可以在执行第5步的数据提取的同时，进行第4步的数据转换，同时进行第3步的数据加载，并且进行第2步的训练，所有这些操作几乎是同时进行的。这样，你可以非常高效地利用计算资源。
因此，你的训练速度会更快，资源利用率也会更高。值得注意的是，现在只有少数几个时刻，你的硬盘和CPU会处于空闲状态。
优化输入管道与TensorFlow Data
那么，如何在实践中优化数据管道呢？有几种基本的方法可以加速你的管道。预取（Prefetch）是一种良好的实践，它可以在当前步骤完成之前开始加载下一步的数据。其他技术涉及并行化数据提取和转换。缓存数据集也是非常有效的，尤其是在你有足够缓存的情况下，这样当一个新的训练周期开始时，你可以立即开始训练。最后，你需要注意在管道中如何安排这些优化，以最大化管道的效率。
一个可以帮助实现这些方法的框架是TensorFlow Data（TF Data）。我们将以TF Data为例，说明如何设计一个高效的输入管道。
预取（Prefetch）
通过预取，你可以将“生产者”的工作与“消费者”的工作重叠。模型在执行步骤S时，输入管道正在读取步骤S+1的数据。这样可以减少一个步骤的总时间，无论是训练模型还是从磁盘提取数据（取决于哪个步骤花费的时间更长）。
TF Data API提供了tf.data.Dataset.prefetch转换。你可以使用它将数据生产的时间与数据消费的时间解耦。此转换使用后台线程和内部缓冲区，在请求元素之前提前从输入数据集中预取元素。理想情况下，预取的元素数量应该等于或可能大于单个训练步骤所消费的批次数量。你可以手动调整这个值，或者将其设置为tf.data.experimental.AUTOTUNE，这将配置TF Data运行时在运行时动态优化该值。
在实际应用中，输入数据可能存储在远程（例如，Google Cloud Storage或HDFS）。一个在本地读取数据时表现良好的数据集管道，在远程读取数据时可能会因I/O或网络带宽的瓶颈而受到限制，因为本地存储与远程存储之间存在以下差异：
- 首次字节时间（Time-to-first-byte）
从远程存储读取文件的第一个字节可能比从本地存储读取要慢几个数量级。 - 读取吞吐量（Read throughput）
虽然远程存储通常提供大的总带宽，但读取单个文件时可能只能利用这一带宽的小部分。
为了减少数据提取的开销，可以使用tf.data.Dataset.interleave转换来并行化数据加载步骤，包括交替加载其他数据集的内容。交替的数据集数量由cycle_length参数指定，而并行化级别则由num_parallel_calls参数设置。
与预取转换类似，交替转换支持tf.data.experimental.AUTOTUNE，它将决定使用何种并行化级别的任务委托给TF Data运行时。
并行化数据转换
在准备数据时，输入元素可能需要预处理。例如，TF Data API提供了tf.data.Dataset.map转换，它应用用户定义的函数来预处理输入数据集的每个元素。元素级预处理可以跨多个CPU核心并行化。与预取（prefetch）和交替（interleave）转换类似，map转换提供了num_parallel_calls参数，用于指定并行化的级别。
为num_parallel_calls参数选择最佳值取决于你的硬件、训练数据的特征（如大小和形状）、map函数的计算成本以及CPU同时处理的其他操作。一个简单的启发式方法是使用可用的CPU核心数量。然而，与预取和交替转换一样，map转换也支持tf.data.experimental.AUTOTUNE，它将决定使用何种并行化级别的任务委托给TF Data运行时。
缓存
tf.data.Dataset转换包括缓存数据集的功能，可以将数据集缓存到内存或本地存储中。在许多情况下，缓存是有利的，并且可以提高性能。缓存可以避免在每个周期中重复执行某些操作，例如打开文件和读取数据。当你缓存数据集时，缓存之前的转换（如打开文件和读取数据）只会在第一次迭代时执行，后续的周期将重用缓存的数据。
我们来考虑两种缓存的场景：
- 如果传递给
map转换的用户定义函数非常昂贵，最好在map转换之后应用缓存转换，只要结果数据集仍能适配内存或本地存储。 - 如果用户定义的函数增加了存储数据集所需的空间，超出了缓存容量，可以考虑在缓存转换之后应用它，或者考虑在训练作业开始之前预处理数据，以减少资源需求。
现在我们已经讨论了分布式训练和高效输入管道的基本概念，接下来我们将讨论巨型神经网络的崛起和高性能建模策略，这些策略有助于高效地训练此类模型。
训练大规模模型：巨型神经网络和并行化的崛起
近年来，机器学习（ML）数据集和模型的规模不断扩大，推动了在语音识别、视觉识别和语言处理等多个任务上取得更好的结果。特别是近期生成性人工智能（GenAI）模型的进展，如Gemini、GPT-4和Claude 3.5，展示了大规模模型的潜力。同时，硬件加速器，如GPU和TPU的计算能力也在不断增强，但增速显著较慢。模型规模的增长与硬件进步之间的差距，使得并行化变得越来越重要。
在这个背景下，并行化意味着在多个硬件设备上训练单个机器学习模型。一些模型架构，特别是小型模型，适合并行化，并且可以很容易地在硬件设备之间分割。然而，在超大规模模型中，同步成本会导致性能下降，阻碍它们的使用。
介绍开源库GPipe的博客文章（见《管道并行化的救援？》）强调了近年来模型规模的巨大增长，并取得了性能的提升。在这篇文章中，作者指出了ImageNet大规模视觉识别挑战赛的获胜者，特别是2014年和2017年获胜者之间，模型参数数量增长了36倍。
大量的权重和激活参数需要巨大的内存存储。仅凭硬件进步并未跟上模型规模的快速增长，巨型神经网络的崛起进一步加剧了有效解决内存约束的需求。但从某种程度上来说，这并不是一个新问题，接下来我们将对此进行讨论。
潜在解决方案及其局限性
在本节中，我们将探讨一些较早的应对巨型神经网络崛起所带来的需求的方式，并分析这些方法可能存在的局限性。最后，我们将讨论管道并行化及其如何解决这些局限性。
梯度累积
一种能够解决GPU内存不足问题的策略是梯度累积。梯度累积是一种将完整批次拆分为多个小批次的机制。在反向传播过程中，模型不会在每个小批次后立即更新，而是将梯度累积起来。当完整批次处理完成时，所有前面小批次的累积梯度会用于反向传播，从而更新模型。这一过程与使用完整批次训练网络的效果相同，因为模型参数更新的次数是相同的。
交换
第二种方法是交换（swapping）。在这种方法中，由于加速器上的存储不足，你将激活值复制回CPU或内存，然后再返回加速器。这个方法的问题在于它比较慢，且CPU或内存与加速器之间的通信成为瓶颈。
并行性：在巨型神经网络背景下的再探讨
回到我们对分布式训练的讨论，基本思想是将计算任务分配给多个工作节点。你已经看到了两种进行分布式训练的方法：数据并行和模型并行。数据并行是将输入数据分配给多个工作节点，而模型并行则是将模型分配给多个工作节点。
在数据并行中，不同的工作节点或GPU处理相同的模型，但处理不同的数据。模型会在多个工作节点之间复制，每个工作节点都进行前向和反向传播。当一个工作节点完成处理后，它会与其他设备同步更新后的模型权重，并计算整个小批量的更新权重。
在数据并行中，输入数据集被划分到多个GPU上。每个GPU都维护模型的完整副本，并在自己的数据分区上进行训练，同时定期与其他GPU同步权重，使用集体通信原语或参数服务器。参数同步的频率影响统计效率和硬件效率。
在每个小批量结束时进行同步，减少了用于计算梯度的权重的陈旧性，从而确保了良好的统计效率。不幸的是，这需要每个GPU等待其他GPU的梯度，这显著降低了硬件效率。由于神经网络的结构，数据并行训练中不可避免地会出现通信延迟，导致通信通常会占据总执行时间的主导地位。加速器速度的快速提升进一步将训练瓶颈转向了通信。
还有另一个问题。加速器具有有限的内存和与主机机器之间有限的通信带宽。这意味着，在加速器上训练更大的模型时，需要使用模型并行性，将模型划分为多个分区，并将不同的分区分配给不同的加速器。
在模型并行中，工作节点只需要同步共享参数，通常每次前向或反向传播后进行一次同步。此外，由于每个工作节点使用相同的训练数据操作模型的一部分，因此大模型不再是一个主要问题。当在训练中使用模型并行时，模型会被划分到K个工作节点上，每个工作节点持有模型的一部分。一个简单的模型并行方法是通过将N层神经网络分配到K个工作节点上，每个节点托管N/K层。然而，更复杂的方法会分析每一层的计算复杂度，确保每个工作节点的负载均衡。标准的模型并行性能够训练更大的神经网络，但它在性能上会遭遇较大的瓶颈，因为工作节点需要不断等待其他节点，且在任何时刻只有一个节点能进行更新。
总而言之，在神经网络的背景下，无论是数据并行还是模型并行，都面临着实现高性能的挑战，每种方法都有其自身的局限性。
管道并行来拯救？
数据并行和模型并行的问题促使了管道并行的发展。图7-4展示了使用四个加速器（设备0-3）进行管道并行的示例。训练模型的前向传播步骤用F0-3表示，梯度的反向传播用B0-3表示。正如图中所示，简单的模型并行策略由于模型的顺序特性，导致了严重的资源未充分利用：每次仅有一个加速器处于活动状态。
为了实现多加速器之间的高效训练，需要找到一种方法，将模型划分到不同的加速器上，并自动将训练样本的小批量拆分为更小的微批次，如图7-5所示。通过对微批次进行流水线化执行，加速器可以并行操作。此外，梯度会在所有微批次中一致地积累，这样一来，分区数量就不会影响模型的质量。
Google 的 GPipe 是一个开源库，用于使用管道并行高效训练大规模模型。在图7-5中，GPipe 将输入的小批量拆分为更小的微批次，使得不同的加速器能够同时处理各自的微批次。GPipe 本质上是一种新的模型并行方式，允许在多硬件设备上训练大型模型，并且几乎能实现一比一的性能提升。它还帮助模型包含显著更多的参数，从而在训练中取得更好的结果。微软的 PipeDream 也支持管道并行。GPipe 和 PipeDream 在很多方面相似。
像 GPipe 和 PipeDream 这样的管道并行框架整合了数据并行和模型并行，以实现高效率并保持模型的准确性。它们通过将小批量分割为更小的微批次，使得不同的工作器能够并行处理不同的微批次，从而在给定的加速器集上训练具有更多参数的模型。关于 GPipe 及其使用带来的内存和训练效率提升的更多信息，可以参考 Google Research 博客文章《Introducing GPipe, an Open Source Library for Efficiently Training Large-scale Neural Network Models》。
总结
本章为你介绍了一些涉及高性能建模的问题和技术。这是一个正在快速发展的领域，因为对像 GPT-4 和 Gemini 这样极大规模生成式 AI 模型的高效训练提出了巨大的计算资源和预算需求。本章提到的主要领域——分布式训练、高效输入管道和大型模型训练——几乎每周都有新的进展出现。通过本章的背景知识，你将能够在这个领域不断进步的过程中，更好地理解和评估这些新进展。
