《PytorchConference2023翻译系列》25 数据及加载技术的演进

我们推出了一个新的系列，对PytorchConference2023 的博客进行中文编译，会陆续在公众号发表。建议PC端阅读本文。或者访问作者博客链接：https://www.aispacewalk.cn/docs/ai/framework/pytorch/pytorch-guides

大纲¶

dataloading概述
dataloading任务
行业趋势影响
当前dataloading生态
dataloading难点
合作展望

详细要点¶

1. dataloading概述¶

将数据提供给训练循环
从存储fetch数据,transform为张量

2. dataloading任务¶

fetch数据:从存储系统fetch样例
transform数据:预处理样例为张量

3. 行业趋势影响¶

数据集和模型规模增长
训练硬件和数据类型多样化
引起存储、计算资源需求变化

4. 当前dataloading生态¶

PyTorch提供基础API
用户选项和方式多样化
分散性难以解决问题

5. dataloading难点¶

异质性:训练场景复杂多样
性能瓶颈:吞吐量难满足要求
系统设计难兼容各种情景

6. 合作展望¶

共同讨论改进方案
打造通用高性能dataloading系统
面临挑战需广泛合作

我叫劳伦斯·拉斯内尔，是Meta公司的工程经理，我负责PyTorch团队的一部分。在过去的两年半里，我一直专注于PyTorch库，例如Torch vision,audio,multimodel。我们在生成式人工智能、大规模内容理解模型和大规模推荐系统等方面做了大量的工作。今天我将讲述PyTorch中dataloading的发展现状。

What is dataloading?
Conceptual model of dataloading
The "formula" for dataloading
Industry trends
Ecosystem today
Why it's a hard problem to solve
Let's work together!

这将是一个对dataloading中的挑战进行高层次概述的演讲，并介绍随着模型变得更大、系统变得更快，这些问题的性质如何发生变化。首先，我将从dataloading的某些概述开始介绍，这些方面使它从一般方式解决起来很棘手。以及如何解决这些问题以及一些与行业趋势的关系。我们致力于构建更好的PyTorch数据加载的抽象和工具。
那么，什么是数据加载呢？让我们从一个过于简化的模型开始。

数据加载器（dataloading）的目的是将数据批量提供给训练循环。实质上，我们调用next来获取一组张量，用于前向-反向传递。在本次讲座中，我将使用数据加载系统和数据加载器这两个术语来表达同一含义。有许多不同设计和权衡的实现方法。数据加载器（dataloading）实际上有两个关键功能。

首先，我们需要从存储系统中获取一些潜在的、可能是未经转换的训练样例。然后，我们需要对数据进行预处理，使其准备好供我们的模型使用。数据加载器通常还会执行一些其他准备工作，例如构建正确大小的批次，并将张量传送到正确的设备上。我还想强调一下，这个图示只是一个概念性模型。PyTorch中实际的数据加载器稍有不同，但我们稍后会涉及到。

所以，让我们从讨论获取数据开始。获取数据是我们访问数据存储的过程。这可能是一个云对象存储器。在本地附加的SSD上，一个高性能的网络文件系统，类似NFS或Lustre的东西，一个数据库，甚至可能是一个分布式队列。对于每个数据加载系统，获取数据需要我们指定如何定位正确的数据集，并实际调用API来完成此操作。例如，用户可能在文件系统上指定一个路径，可能附带一些过滤器或glob模式。在其他情况下，他们可能使用索引或元数据文件来直接查找正确的对象进行加载。例如，一个包含对象存储路径的CSV文件。一旦我们从存储系统中获取了原始字节，我们就需要将它们转换为适用于训练循环的张量。现在，转换几乎总是特定于你的用例的。如本例所示，计算机视觉模型可能需要解码JPEG图像，调整大小或裁剪它们，应用随机翻转等转换，最终将它们转换为张量。这里我要提醒一点，通常变换或预处理是我们做的不可微分的操作。它与反向传播无关。这一区别所以重要，是因为我们可以将这些操作分别和独立地扩展或并行化到模型中。正是这一方面使得我们能够启动多个工作进程，并在训练循环之外执行transform操作。然而，在某些情况下，图像、视频或音频或分词器也可能作为训练循环的一部分进行训练。因此，在这种情况下，它们实际上只是模型的一部分。上述的简单模型看起来很简单，但在实践中，实现起来要复杂得多。这就是本次演讲的主题。通常来说，每个步骤都是并行化的。我们可能会有多个Python进程来获取example或执行转换操作。有些用户可能会选择使用专用服务进行水平扩展，以提高性能。大多数数据加载器都支持某种形式的预取或流水线处理，可以与训练循环同时进行获取和转换操作。你可以看到复杂性已经开始积累起来了。现在我们可能要处理线程或进程管理，管理共享队列和内存缓冲区，并且还需要处理潜在的资源争用。你可能也能意识到，这里有很多决策都要依赖于具体情况、可用资源和加载的数据类型。这种异质性将贯穿整个演讲。因此，在思考数据加载时，请记住以下两个方程式。

在几乎所有情况下，作为机器学习工程师，我们希望尝试提高模型的 QPS ，以便我们的训练任务能更快地完成，从而能够训练出更大、更好的模型。只有当数据加载成为瓶颈时，数据加载才真正重要。

这给我们带来了第一个公式。只要我们的获取和转换时间比训练步骤时间更快，我们就不需要太关心改进我们的数据加载器。减少获取和转换时间的一种方法是通过并行化利用我们可用的计算资源。这与我们之前讨论的转换不可训练的问题有关。我们可以在前后传递过程中并行且独立地执行它们。这可以在训练主机的多个进程上执行，或者可能在一个单独的工作队列上执行，这就是分母的compute部分。第二个公式确保我们不会因资源争用而降低整体训练速度。我们需要注意训练主机上数据学习所使用的CPU和内存利用率。创建更多的工作进程可以帮助提高吞吐量，但如果与训练器存在竞争，很快就会导致性能退化。当资源竞争带来训练延迟时，这可能特别隐蔽，由于每个训练器在分布式训练中通常是同步进行的，训练步骤时间取决于最慢的工作进程。如果数据加载导致一个工作进程暂停，因为CPU使用率飙升，它会导致每个工作进程变慢。

让回到演讲标题，机器学习领域的发展如何演变？这对数据加载生态系统有何影响？现在让我们来看一些我们最近观察到的关键趋势，并预计这些趋势将继续发展。首先要注意的趋势是数据集的大小在不断增大。我们现在谈论的是数万亿个标记、数百万甚至数十亿个图像。数据集规模的增加实际上不像你一开始可能期望的那样有那么大的影响。

这是因为我们的系统首先会受到模型训练的瓶颈限制，只要这个方程作为一个起点仍然成立。虽然更多的训练样本可能会涉及更多的metadata跟踪，但对于已经具备可伸缩性的系统来说，这不会产生太大的差异。然而，如果你的训练数据集适合单个节点（无论是内存还是磁盘），那么你可能需要做一些改变，以适应分布式计算的模式。实际的挑战在于，我们通常会在增加数据集大小时尽量保持训练时间的一致。我们可以通过增加训练器的数量或使用更快的硬件来实现这一点。这就引出了我们的第二个趋势，即training硬件的速度越来越快。我们将其视为摩尔定律的延续。我们的CPU核心数增多，GPU的计算能力更强，我们的系统变得更好。更快的training硬件的真正影响是我们的training step 时间缩短了。这使得我们要么增加用于数据加载的计算量，要么提高我们的提取和转换性能来弥补这一差距。不过，GPU计算、CPU计算和内存带宽没有以相同的速度在加速，带来了一些新的问题。我们的第三个趋势与数据的速度相关。增加模型的及时程度，使用更近期的训练数据变得越来越重要和可行。虽然现在并非每个应用案例都如此，因为数据获取依赖于大量的人工input，比如为llm或者图像生成微调数据集，我预计我们将继续看到对于训练数据的快速提供以及机器学习工程师希望能够迅速使用数据的压力增加。这会对我们稍后讨论权衡的一些设计产生影响。例如，通过预先打乱和打包数据进行训练的方式正在变得普遍。这与相对静态的数据集相比，步骤本质上会减慢数据保持最新的更新速度，并且在相对稳定的数据集中增加了成本。更重要的变化是，随着多模式模型变得越来越普遍，模型训练所使用的数据类型正在增加。语言模型的上下文长度正在扩大到数十万个标记。图像的训练规模也逐渐增大，并且更丰富的视频等格式的使用也变得更加常见。这给存储和处理系统带来了压力，因为它们需要处理更大的对象。检索图像与视频之间的差异是数量级的。

现在，随着数据变得更加复杂，通常我们的模型也变得更大、更需要计算资源来发挥优势。这实际上意味着training step time也会面临上升压力。总体来说，情况有点复杂。

随着示例的尺寸增加，计算成本和转化形式也在发生变化。例如，如今使用视频时是通过采样少量帧来降低帧率的。然而，在生成式人工智能中，对于匹配生成输出的高分辨率输入和更高帧率的训练模型很重要。以每秒 24 帧的 512 像素宽度的帧训练，相对于采样 224 像素的 10 帧，内存带宽、存储和计算资源的增加接近千倍。与此同时，我们还面临着转化时间增加和训练步骤时间可能也在增加的不断压力。所以，影响因素是多样的。

在计算方面存在一个更大的问题是资源使用。加载和转换数据所需的内存和CPU计算量可能会超出我们在训练主机上可用的硬件限制。随着这些模型变得越来越复杂，资源的压力也会增加。例如，视频解码在CPU利用率方面需求特别高，我们将使用更多的内存来缓冲转换训练example。话虽如此，我们也看到了硬件方面的变化。在过去的十年中，我们看到了内存带宽与网络带宽到计算比率的变化。作为这一趋势的一个例子，现在可用的云实例接口的吞吐量达到每秒600GB。这是几年前的数量级增加。随着约束条件的变化，构建数据加载系统的最佳方式也会发生变化。接下来，让我们谈谈可能是最重要的变化。过去，在单个训练主机上只有一个GPU，

在单台机器上使用了多个GPU，现在利用更大的集群来扩展成千上万个GPU的情况也并不罕见。这对数据加载（dataloading）系统产生了重大影响。如前所述，目前的默认范式是完全同步的训练，大多数模型的前向和后向路径在延迟方面相当稳定。权重和梯度通过无争用的互连进行共享，方差较低。而数据加载则具有较高的方差。数据从共享网络上的存储中拉取。物理存储设备可能存在竞争请求。hot spot可能形成("hot spot" 通常指的是一个程序中的那部分代码，它占用了大量的计算资源或者是执行时间的瓶颈。这意味着程序的这一部分是最频繁执行的，可能是一个循环、一个频繁调用的函数或者是一块计算密集型的代码。)。transform延迟可能取决于数据的形状，如图像的大小等。所有这些因素意味着数据加载的延迟分布要比训练步骤时间的分布要宽。

回到之前的方程式。最大fetch时间和transform延迟时间严格小于训练步骤时间。这里使用最大值是因为训练是同步进行的。一个更慢的训练者会拖慢其他训练者的速度。我们的平均、中位数，甚至可能是P99.9（99.9%）的获取时间可能会小于训练步骤的时间。但要确保最大步骤时间更短则更难实现。一次慢速的远程过程调用可能导致数据获取时间延长，甚至延续几秒钟。被数据存储限制、图片大小过大均会导致这种情况发生。随着数据加载器数量的增加，我们更有可能频繁遇到延迟分布的长尾部分这最终会成为训练性能的瓶颈——曾经是千分之一的事件现在开始影响每个训练步骤，导致训练速度变得缓慢。类似地，我们在增加训练器数量时可能还会看到分布的变化。存储系统支持的总吞吐量和QPS也需要相应地扩展。最终，当达到它们支持的输入/输出操作数或吞吐量上限时，你可能需要切换数据存储。有一些方法可以通过使用预取等手段来减轻这种延迟差异——使用队列或缓冲区以及增加并行性。但是，这并非免费，这需要更多的内存且通常需要手动调优。随着数据规模越来越大，变换越来越复杂，这也可能导致成本的增加。在我看来，这是当前数据加载系统面临的最大挑战，因为机器学习领域正在发生变化。与将训练集群扩展不同，模型变得更大对于负责数据加载系统的人来说实际上更有利。

随着模型的扩展，训练步骤通常会变长，我们会有更多的余地。当某人从训练一个70亿参数的模型转变为训练一个130亿甚至600亿参数的模型时，延迟会增加。所以这对我们是有利的。

我们讨论了机器学习中的许多行业趋势以及它们如何改变数据加载的性质。这种变化的最大推动力是训练数据变得越来越复杂，使用和transform的数据类型越来越丰富、从架构的角度看硬件约束如何转变以及改变可能性的状况、扩展到数百或数千个数据加载器时所涉及的挑战。那么，数据加载生态系统现在处于什么地步呢？简而言之，它是零散的。

大多数机器学习工程师使用构建在数据集之上的某个东西，以及作为PyTorch库的一部分提供的数据加载器API。根据您的用例，您有不同的选择，包括数据存储位置、存储格式、预处理的数量和类型等等。我不会详细介绍每个选项，这需要花费多个小时来讨论设计、权衡和推荐。不过，我会简要介绍内置的数据加载器，因为它是大多数人熟悉的。下面是一个相当基本的示例，展示了一些不同的API。用户创建或使用一个数据集，其中指定了数据的位置和获取方式。数据加载器处理将示例整理成批次等。

通过Python的multiprocessing进行并行化、提取和转换。它可以进行预取，并且可以进行shuffle，如果你提供了map式数据集还可以排序。不过，目前我们和很多PyTorch用户在使用这些API时都遇到了问题。它们很难组合使用，很难避免Python解释器的开销。没有一种一致的方式来指定下推到存储系统的过滤器。当出现问题时，调试起来也很困难，并且性能需要大量手动调优和对内部的了解。我们试图通过Torch Data来解决这个问题，我们尝试建立一个通用的数据加载库来解决这种分散的问题。但事实证明，这比我们想象的要困难。我们在今年七月暂停了目前的工作。我们真的很想给用户提供更好的东西，但我们意识到我们正在努力解决的设计并不一定能解决我们所见到的所有问题。目前，我们仍在致力于解决这个问题。而且正如我之前提到的，我们正在寻找合作伙伴、兴奋的用户和有创意的人进行合作。所以如果你有想法或者想要帮忙，请在演讲结束后与我联系。那么接下来我要谈谈构建数据加载库的挑战。如果我必须用一个词来总结其复杂性的话——heterogeneity（异质性），两个词——shuffling。

为什么它很困难有很多原因。大多数情况下，需求在不同的用例中是不同的，并且系统设置差异如此之大，没有明确的通用解决方案。你的设计将不得不做出一种妥协，这对一部分用户来说是无法接受的。从用户那里听到的最重要和最关心的问题一直都是性能。我们拥有庞大昂贵的GPU集群，我们希望最大限度地利用它们。我们不希望我们的训练QPS受到数据加载的阻碍。为了达到最佳性能，用户通常会花费大量时间手动调整他们现有解决方案，来适应他们最重要的模型。因此，当有人将你的库与他们的进行比较时，这就已经提出了一个很高的基准与之竞争。如果用户不是被数据加载限制，那么他们可能专注于其他问题，比如提高训练查询每秒请求数(QPS)、扩展模型或调试数据问题。他们需要抽出时间来采用你的产品，而这可能会影响到更高优先级的事情。是的，你可以提供更好的API，是的，你可以消除维护的麻烦，但他们只会在当你的库在性能上匹敌或超越时才会选择切换。现在，关于性能的另一件事，这与shuffle有关，随机访问很难在许多系统上进行优化。总体而言，对于数据访问更加可预测的连续性数据，构建高性能系统会更加容易。所以，这是一个很好的过渡，来谈谈吞吐量。我们方程中的一个关键变量是训练步骤时间（training step time）。对于某些模型来说，这个时间非常短。由于这些模型具有很高的查询每秒请求数和很大的批次大小，它们更可能受到数据加载和预处理的限制。对于其他用例，比如训练大型语言模型，训练步骤的时间可能会相当长。在等待前向和后向传递完成的时候，我们可以隐藏很多计算和潜在的低效率。这意味着对于一个使用情况来说可以接受的额外开销——比如通过Python解释器运行一切，对于另一个使用情况：你要计算mallocs并且试图消除任何解释器开销——这可能是无法接受。

吞吐量的差异和系统优化的影响也对现有数据加载器提供的功能和用户对其的期望有很大影响。例如，如果你专注于低计算强度模型的高吞吐量，可能愿意权衡放弃许多功能。例如，你可以进行多个并行请求，并平衡地从数据存储中获取数据，以避免会减慢训练速度的延迟峰值。然而，这意味着放弃确定性排序，在调试时确定性排序可以非常有用，当你训练视觉模型或者语言模型的时候，我们看到这些结果是在预期中的。你可能会无意识地为某个特定目标中的一部分构建系统，这可能会带来一些权衡的考虑。另一个复杂性的巨大来源是我们在数据存储方面所见到的多样性。首先，你需要提供许多不同的集成点。尽管POSIX文件系统API相当常见，但有许多不同的数据存储需要定制集成。然而，lower到API层下面，数据存储具有完全不同的性能特征。有些可以实现快速的随机访问和相对较高的吞吐量，例如HPC风格的网络文件系统或本地连接的SSD。而其他一些则可以在处理大文件时提供高顺序吞吐量，但通常延迟较高，这在云对象存储中很常见。这很重要，因为它改变了ML工程师设计整个数据加载系统的方式。如果你有快速的随机访问，事情就变得简单得多，因为你可以忽略随机洗牌的成本。如果你在云存储上访问像JPEG这样的小文件，那么你就需要在缓存和预取方面要聪明一些。你可能还需要一个高度优化的适配器，以便能够与对象存储进行高并发的请求交互。因此，影响数据加载设计的关键因素是从数据存储中获取数据的速度有多快，包括延迟和吞吐量方面的考虑，以及特别要注意的是在给定访问模式下的分布情况。其次，你是进行顺序访问还是随机访问，还是介于两者之间。

对于延迟高、吞吐量高的系统，把数据打包成捆绑格式是非常常见的做法，甚至可以使用tar压缩文件。如果你正在处理捆绑文件，那么一切就变得更加复杂。你需要确保你的API支持不同种类或级别的随机性。你需要预先获取和下载分片。你需要进行内存管理或缓存清理。对于API和实现中的这个用例来说，有很多复杂性。

此外，不同的访问模式也有很多变化。这在PyTorch中表现为map数据集和可迭代样式数据集之间的差异，以及如何支持采样。对于吞吐量非常高的系统，您可以通过多个并行请求进行负载均衡。您的排序是系统定义的，无论如何返回的是最快请求返回的结果。当您的数据存储具有一些限制时，您需要更加小心地采样或访问数据。例如，对于捆绑包，您需要确保您的访问有可能命中已加载和缓存的捆绑包，这个结果包括准随机的排序中。（准随机的排序（Quasi-random or Pseudo-random ordering）是一种数据访问顺序，它介于完全随机和完全顺序之间。这种排序方法的目的是在保持一定程度的随机性的同时，还能利用数据结构或存储系统的特性来提高效率，特别是在缓存利用方面。）如果由于数据存储而具有快速的随机读取，您可以使用map样式的API，进行高效的随机读取。您甚至可以根据类别定制您的采样策略，例如按权重进行采样。对于所有这些不同的访问模式，API设计相当困难。如果抽象概念不清晰，用户很容易自废武功。我们遇到过用户仅在一个包中排序的问题，他们创建了大型的内存缓冲区，这可能导致内存不足的错误或意外的缓慢。因此，在这个领域，正确设计是棘手的，依赖于工作负载、数据存储和文件格式，而最后两个领域相辅相成但又有所不同。当你使用PyTorch的原生数据集和数据加载器时，计算拓扑如下图所示。

对于每个训练器和主机，通常每个GPU一个训练器，我们启动多个独立的进程来获取数据并进行transform。这是一个合理的起点，但如果我们在主机上造成资源争用，它就开始失效。这可能有很多原因，与每个训练环境从硬件角度设置的差异以及我们实际运行的transform相关。用户可能最大化使用CPU解码视频。在这种情况下，将fetch和transform任务转移到另一组可以独立扩展的主机可能是有意义的，只要它们具备网络带宽。另一个主机可能没有问题，因为他们所访问的集群的CPU和GPU的比例是不同的。对于可用的内存容量、网络带宽以及其他资源，同样存在类似这样的紧缺情况。因此，在许多情况下，将我们的数据fetch、transform、training loop分开，并独立地扩展它们是有意义的。以上是一个示例的拓扑结构，但你可以想象其他的配置。例如，我们可以将transform保留在我们的训练主机上，但扩展我们的数据fetch。因此，在计算拓扑中，确定在哪里运行数据获取和转换的最佳位置取决于获取和转换的成本，包括CPU计算成本、所需内存和所需的入口和出口网络带宽。这可能取决于你的模型。如果你的模型受到CPU限制，那么你可能很容易遇到这种竞争。它还取决于可用的硬件，这会因用户而异。因此，设计一个既可以正常工作又能提供高性能且不会遇到任何瓶颈的系统是很难的，很多复杂和抽象需要考虑。我们还有很多的问题需要解决，需要一起合作，欢迎联系我们思考解决方案，谢谢。

本文总阅读量次