GPU 利用率从 15% 走向 90%：先修数据流水线，再谈模型优化

简要描述Brief Description

从 15% 到 90% 的 GPU 利用率：修复数据流水线，而不是模型From 15% to 90% GPU Utilization: Fix the Data Pipeline， Not the Model

目录Table of Contents

• 简介
  Intro
• 训练真的受 GPU 性能限制吗？
  Is training really GPU-bound?
• 四个瓶颈
  Four bottlenecks
• 为什么多模态处理非常消耗 CPU】【。
  Why multimodal is CPU-heavy
• 等待时间比率
  The wait-time ratio
• 我们的技术栈
  Our stack
• 数据
  The data
• 系统布局
  System layout
• 基线：利用率 15%
  Baseline: 15% utilization
• 修复 1：并发
  Fix 1: concurrency
• 异步、线程还是进程？
  Async, threads or processes?
• Asyncio 加 Ray actor
  Asyncio plus Ray actors
• 修复 2：预取
  Fix 2: prefetching
• 修复 3：使用 Ray 的对象存储实现零拷贝
  Fix 3: zero-copy with Ray's object store
• 回顾：等待时间从 85% 降至 20%
  Recap: 85% to 20% wait time
• 规模扩张：新的瓶颈
  Scaling up: a new bottleneck
• 网卡瓶颈
  The network card bottleneck
• 分散调度与零拷贝
  Spread scheduling and zero copy
• 课程
  Lessons

简介Intro

大家好。感谢大家来到这里。我今天要谈谈 Ray actors、视觉 token 和 GIL——如何构建一条能让 GPU 保持繁忙的多模态 SFT 数据流水线。关于如何优化训练范式，大家已经谈了很多。不过，我想从一个不同的角度，聊聊训练一个模型所涉及的整个生态系统。那么……Hello everyone. Thank you for joining me here. Um yeah, today I'm going to talk about ray actors, vision tokens, and the gill. um engineering a multimodal SFT data pipeline that keeps the GPUs busy. So there's a lot of talk about you know how to optimize training paradigms. Um but I want to give a different perspective on the whole ecosystem of what it takes to get a model trained. Um and so the

训练真的受 GPU 限制吗？Is training really GPU-bound?

我接下来要问的问题是：你的训练循环现在真的依赖 GPU 吗？它是被 GPU 卡住了吗？我认为我们会从这个实验中了解到，多模态训练吞吐量优化中最难的部分未必是优化内核，而是确保 G 个 GPU 始终有数据可处理。那么我们来概览一下……question I'm going to be asking here is does your training loop really depend on the GPUs right now? Is it blocked on the GPUs? And I think we're going to learn that we're going to see in this experiment that the hardest part of optimizing multimodal train throughput is not necessarily kernel optimization. It is making sure that the G GPUs are constantly fed with data. So let's go over the overview of what

四大瓶颈Four bottlenecks

这就是本次演讲要讨论的主题。总之，我们会先搭建一个基础的多模态训练流水线，然后找出瓶颈，并逐一解决。大致来说，瓶颈有以下几个：第一，并发性——在把图像输入模型之前，处理一张图像要多快？第二，预取——如何确保我们能按需获取数据？第三，如何尽量减少不同进程之间的数据拷贝？第四，额外的框架优化，也就是调整各种细节。this topic what this talk will be about. Anyway, so we will establish a baseline multimodal training pipeline um and then we will encounter bottlenecks and then address them one at a time. So loosely they are one concurrency how fast can you process one image before feeding it to the model two pre-fetching how can we make sure that we have data on demand three um how can we minimize data copies across different processes and four f additional framework optimizations um so just tweaking your bells and whistles

为什么多模态处理会占用大量 CPU 资源Why multimodal is CPU-heavy

嗯，是的，那我们来设定一下前提。我们都知道，GPU 成本高昂，因此通常希望在拥有 GPU 时尽可能提高其利用率。多模态处理与纯文本处理不同，因为在真正把数据输入模型之前，往往需要进行大量 CPU 密集型操作。比如，加载、调整图像大小或对图像进行归一化，可以使用像 PIL 这样的库；处理音频时，你会进行很多加载和裁剪操作，可以使用 NumPy 数组；而视频则可以看作两者的某种结合。还有一点也很重要：训练集群通常配备的 CPU 远多于 GPU，但在这个处理流程中，真正供不应求的却是 GPU。今天我们要介绍一个叫作……的公式。um yeah so let's set up our premise We all know that GPUs incur expensive costs and are generally want to maximize their utilization while we have them at hand. Multimodal processing is different than text only um because there's a lot of CPU inensive operations that have to happen before you actually feed the model um the data. So whether that is um you know loading resizing or normalizing images using like a library like pill for audio you do a lot of loading and cropping you can use like numpy arrays and then videos are some sort of combination of both um and then it's also true that usually your training clusters have a lot more CPUs than GPUs at hand yet the GPU is what is starving in this pipeline. So today we're going to introduce this equation called like

等待时间比率The wait-time ratio

等待时间比率通常是这样定义的：准备数据并将其传输到 GPU 所需的时间，除以完成这两项工作以及训练所需的总时间。嗯，也就是说，分母是整个流水线的总耗时。嗯，我们会看到，在运行过程中，利用率会从 15% 提升到 90%。因此，我们的基线方案有大约 85% 的时间都花在数据获取上；而完成这些优化后，这一比例会降到 10% 以下。所以，就本次演讲的范围而言，GPU 并不是我们的目标；我们的目标是确保数据永远不会成为瓶颈。就是这样。weight time ratio which is generally I formulated as like the time it takes to prep your data and transfer it to the GPUs. Um and then we divide all of that by the time it takes to do both of those and the train as well. So basically the denominator is like your total pipeline. Um and we're going to see that during our run we will go from a utilization of 15% to utilization of 90%. So our baseline has like 85% of its spends 85% of its time um backed on data acquisition whereas after we do these optimizations less than 10 and so for the scope of this talk GPUs are not the goal. It's making sure that the data is never the bottleneck is the goal. And so

我们的技术栈Our stack

这里简单介绍一下我们的技术栈大致会是什么样。嗯，我们会采用一个基于 Quen 3VL 的模型，尝试用聊天式 VQA 和多轮 VLM 对话范式来训练它。简单回顾一下，在 GPU 开始进行任何计算之前，我们必须先加载数据、转换数据，然后还得把数据传输到 GPU 上。好，这是一个……这是一个概览，介绍了……here's a little bit of background as to like what our stack will look like. um we kind of take like a Quen 3VL based model um kind of try to train it on the paradigms of chat VQA and like multi-term VLM uh conversations. So to recap before the GPUs do any math um we have to load the data we have to transform the data and then we have to we have to transport it to the GPUs. Okay, this is a this is a overview of

数据The data

我们的数据实际是什么样的。因此，我们会使用 JSON L 文件。嗯，图片会通过相对路径链接到示例所在的位置。然后我们假设数据存储在 S3 中。嗯，比如在这个例子中，33471.jpeg 这张图片看起来像是一辆蓝色公交车。what our data actually looks like. So, we will be using JSON L files. Um, and then images are linked relatively uh to where the example lives. And then we're going to assume that our data is just stored in S3. Um, in this example, for example, it seems like that image 33471.jpeg seems to be a picture of like a blue bus.

系统布局System layout

嗯，是的。这就是我们的系统布局。左侧是一个数据生成器，负责数据加载、数据转换和数据传输。嗯，它被标记为向 Megatron 的多个数据并行 rank 扇出。嗯，各个并行 rank 会接收由协调器决定的数据分区。它会对这些数据进行汇总，并执行前向传播和反向传播。因此，Megatron 会负责在各个 DP rank 之间优化模型权重。我们将假设它像一个黑盒一样工作。Um, yeah. So this is our system layout. On the left we have a data generator which is responsible for the data loading, data transforming and data transporting. Um and the that's like labeled as like fan out to the mega megatron data parallel ranks. Um so the parallel ranks receive partitions of the data as decided by the coordinator. Um it collates it performs forward pass and backwards pass on this data. So, Megatron takes care of optimizing the model weights across the DP ranks. And we're going to assume that it's like a it's kind of it is kind of working as a black box.

基线：利用率 15%Baseline: 15% utilization

嗯，那我们先从最简单的基线流程开始。对于每个样本，我们都从对象存储中获取图像。嗯，我们会对图像进行分词、调整大小、写入调整后的图像、归一化，然后将它们输入模型。嗯，对，这样做时我们可以看到，红色部分代表串行操作，也就是逐张获取所有图像，然后逐张处理。因此，如果训练器要求的批次大小是 256，那么你就必须串行处理所有这些图像。我们可以看到，训练时间大约在 15 到 20 秒之间。然而，实际用于获取数据的时间超过了 100 秒。这就是为什么 GPU 的利用率只有 15%。嗯，你甚至可以从性能分析火焰图中看到，我们大部分时间都花在从 S3 加载数据上；然后在末尾还有一小段很长的尾部，表示通过 PIL 加载实际图像的过程。Um, so let's start with the naive baseline path. So for each example, we fetch the images from object store. Um, we kind of like we tokenize it, we resize it, we write resize the images, normalize them, and then feed them to the model. Um yeah and then when we do that we can see that there's serial operations in the red which is namely getting all the images one by one and then processing them one by one. So if your if your trainer is asking for a batch size of like 256 then you have to serially process all these images and then we can see that our train time is somewhere between 15 to 20 seconds. uh however the data the time spent to actually fetch the data is more than a 100 seconds. So that's why we result in only 15% utilization of our GPUs. Um and then you can even see a profiling flame graph that we spend most of the time doing S3 loads and then you can see a long a little tail at the end about like loading the actual image on the pill uh

修复 1：并发性Fix 1: concurrency

库。所以好，解决方案一，我们来试着为它增加一些并发。与其使用那些红色的单一串行进程，不如通过某种方式设置一个工作线程池来尝试并行处理。这样，当你请求 256 个示例时，也许我们可以直接一次性处理完。嗯，比如说，我们有两项开销：一是获取图像，二是进行处理，包括解码、视觉预处理，然后对文本进行分词。嗯，样本彼此独立。这就意味着它们非常容易并行化。嗯，所以，是的，你有一个设计。library. So okay, solution one, let's try to add some concurrency to this. Instead of those red singular uh you know serial processes, uh let's try to parallelize it by kind of having like a worker pool. So when you request 256 examples, maybe let's just process all of those at once. Um yeah, so for example, for example, we have two costs. One is fetching the images, two is processing. So that's like decoding, vision pre-processing, and then tokenizing your text. Um, samples are very are independent. So this just means it's like embarrassingly parallelizable. Um, and so yeah, you have a design

异步、线程还是进程？Async, threads or processes?

这里需要做出一些决策。你想采用哪种并发方式？有异步编程，嗯，它适合处理 I/O 工作，比如执行 SQL 查询或发送网络请求。还有线程，嗯，你可以用它来处理 CPU 密集型任务，前提是这些任务会释放 Python 的 GIL，也就是全局解释器锁。不过，哪些操作实际上会释放这个锁并不容易判断，所以你需要多做一些调查。常见的用法包括使用 Hugging Face Transformers 的分词器、PIL transforms API，或者执行某些 pandas 操作。进程的开销更大，但也能让各个工作负载完全隔离，因此你可以在完成任务的过程中利用机器上的多个核心。嗯，深入探讨如何选择合适的方案，足以单独做一场分享；不过网上有很多资源可以帮你根据自己的使用场景做出决定。decisions here. What forms of concurrency do you want? So there's asynchronous programming, uh, which is suitable for like IO work such as making SQL queries or like performing network requests. There's threads uh which you can use for like CPU intensive uh workloads that relinquish the Python gill like the global interpreter lock. Um so this it's not straightforward what operations actually release the lock. So you have to do a little more of your investigation but common uses are like using hugging face transformers tokenizers using the pill transforms API or like some pandas operations processes um are more heavyweight but also provide total isolation between your workload so you can leverage multiple cores on your machine as you finish your work. Um, yeah, doing a deep dive on the right solution would warrant its own talk, but like there's plenty of online resources that can help you decide for your use case.

对我们的解决方案来说，重要的是要针对我们的瓶颈采取措施，而不是简单地沿用以前对我们有效的方法。嗯，所以我们会……What's important for our solution is that we need to match our bottleneck. Not simply going with a lever that's worked for us before. Um, so we'll

Asyncio 加 Ray actorAsyncio plus Ray actors

在这里选择异步 IO 来获取图像，然后使用 Ray actors 作为处理器，也就是我们的多进程库。嗯，Ray actors 本质上和普通进程差不多，但它们提供了跨 worker 通信的 API，这样我们就不必重复造轮子了。所以，没错，你的进程速度只取决于最慢的环节。嗯，我们可以看到，等待时间占比已经降低了。嗯，在第二个例子中，也就是粉色的那条，你可以看到，每个批次花在获取数据上的时间基本上降到了大约 40 秒。嗯，所以这让速度提升了 25%。不过，正如你所看到的，随着我们经历越来越多的步骤，会出现很大的间隔。嗯，那么，我们能不能以某种方式避免瓶颈——也就是等待 trainer 发出请求之后才配置 worker 池来获取数据？choose async IO for fetching images here and then use ray actors as our processors uh as our process multipprocessing library. Um, ray actors are essentially the same as like general processes, but they offer us APIs to communicate across workers that just make it so that we don't have to reinvent the wheel. So, yeah, your process is only as slow as your slowest link. Um, and we can see that our weight time ratio has gone down. Um, on the second example, which is in the pink, uh, you can see that the time spent fetching data per batch has gone basically, yeah, around like around like 40 40 seconds. Um, so this is this gives us a 25% speed up. However, you yeah that's the as you can see the there's a huge gap that starts generating especially as we're going through many many steps. Um, so can we somehow limit our bottleneck of waiting on the trainer to request before provisioning our worker pool to fetch the data?

所以我们能不能提前获取数据？嗯……So can we just like preemptively fetch our data? Um

修复方案 2：预取Fix 2: prefetching

于是我们提出了方案二：使用一个后台线程提前填充队列，并始终让队列的进度领先于消费者。这意味着，生产者生成样本的速度是消费者速度的若干倍，因此训练器不会因为请求数据而受阻。它请求数据时，数据就已经可以按需提供了。 预取也有一个需要注意的地方：它会让模型和数据加载器的检查点保存与恢复过程变得稍微复杂一些；此外，还可能遇到冷启动阶段，因为队列需要先填充起来，才能为训练做好准备。不过，如果能妥善处理这些细节，我认为预取是个很不错的主意。总之，确保生产者始终领先于消费者，会为你带来很大的收益。so we came up with solution two which was that we have a background thread that fills up a queue ahead of time and it's always stays ahead of the consumer which means that you your producer is producing examples at a multiplier of your tra of your consumer and so the trainer is never bottlenecked on actually requesting the data. When it requests data it's it is available on demand. There's a caveat to prefetching. It makes your model and data loader checkpointing and resumption like a little intricate of a process and it can also suffer from a cold start period in which your queue is actually getting filled up uh so it can be ready for training. Um but uh if you work around those those details I think pre-fetching is a is a great idea. Um yeah, so I think making sure your producer is always ahead of your consumer will be will pay great dividends for you here.

嗯，不出所料，每个批次的数据获取和处理时间接近 0 秒，因为我们早在训练器需要数据之前，就已经获取并处理好了数据。嗯，因此我们实际上解决了数据加载和转换这两个瓶颈。采取这两项缓解措施后，我们以为已经大功告成了：吞吐量提升，GPU 得到了数据供给；然而，我们观察到数据……Um unsurprisingly, we achieved close to 0 seconds of data fetching and processing per batch because we've already fetched data and processed it way before the trainer even needs it. Um and so we've effectively solved both the data load and the transform bottlenecks. So after these two mitigations, we thought we were done. throughput up GPUs fed however we observe that the data

修复 3：使用 Ray 的对象存储实现零拷贝Fix 3: zero-copy with Ray's object store

传输成本仍然存在，事实上还增加了。因此，每个样本都会在多个进程边界之间多次传输大量数据，比如图像数组。每个样本都会在 worker 上序列化，复制到 coordinator，然后再序列化一次传给 trainer。每个样本、每个步骤都会产生一次没人真正需要的内存复制。那么，这会如何体现出来呢？你的数据生成器经常会……嗯，多模态张量非常庞大，每个样本就有好几兆字节。嗯，默认情况下，每个样本都是按值返回的，这意味着这些开销很大的数组会在进程之间复制。解决这个问题的一种方法是使用对象存储。嗯，这是 Ray 提供的一种很方便的方式。每个样本都带有一个对象引用，它实际上只是指向数据的指针。transport costs still linger in fact they've gone up so every sample transfers a large amount of like image arrays across process boundaries multiple times every sample is pickled on the worker copied to the coordinator and then repickled onto the trainer per sample per step we have a mem copy that nobody really asked for so what h how this manifest is that your data generator can often oo. Um, and so multimodal tensors are huge. They're megabytes per sample. Um, and so by default, every sample is returned by value, which means this that these expensive arrays are copied over from process to process. So a fix for this, you can come up with an object store. Um, this is a nice convenience method of that array provides you. Each sample carries a object reference which is just a pointer to the data.

嗯，然后 driver 就只是持有这个指针。张量不会往返传输到 head 节点。嗯，Megatron 的数据并行 rank 可以直接调用 Ray store，按需获取数据。因此，无论选择哪种并行和并发方案，都必须考虑额外增加的通信成本。要谨慎规划。好。是的，我们利用 Ray store，消除了预取器方案中一个多余的、类似 put object 的调用，这是我们原本需要处理的。我们可以看到，通信成本降低了，等待时间所占的比例也……Um and then the driver just holds the pointer. The tensor never round trips to the head node. Um and the Megatron data parallel ranks they can just directly call the race store to fetch the data as they need. So with every choice of your parallelism and concurrency solution, you always have to think about the commu added communication costs. So plan wisely around that. Okay. Yeah. And we were able to eliminate an extraneous like put object uh kind of call that our prefetcher solution had that we did not have to deal with by using the ray store. And we can see that our communication costs go down and the ratio of the time spent waiting also

回顾：等待时间从 85% 降至 20%Recap: 85% to 20% wait time

下降了。总结一下，我们先采用了串行基线方案，然后为数据加载和处理引入并发。我们还添加了预取功能，以便赶在训练请求之前准备好数据，并添加了 race store 来尽量减少数据传输。我们将基线等待比率从 85% 降到了 20%。你可以看到，顶部图表中的锯齿状模式在很大程度上得到了缓解，这意味着 GPU 的利用率提高了。现在，当我们扩大规模，并且已经拥有良好的多模态训练……goes down. So to recap, we took a serial baseline. We added concurrency for data loading and processing. We added pre-fetching to stay ahead of training requests and we added the race store to minimize data transfer. We went from a baseline weight ratio of 85% to 20%. And you can see that the saw to uh pattern that you see in the top graph is alleviated uh to a large extent which means that our utilization of the GPUs has gone up. So when we're scaling up and now that we have like a good multimodal training

扩大规模：新的瓶颈Scaling up: a new bottleneck

我们来试着把这个流水线扩大到生产规模。在这个过程中，我们发现，对于同一条流水线，仅仅是扩大了一些数值，等待时间占比就上升了。所以我们把数据并行环路的数量增加到了四倍，增加了大约 200 条数据流，然后将批次大小扩大了 4 倍。嗯，结果发现，由于数据生成器和工作池共置，可以实现快速数据传输……pipeline, let's try to scale it up to like a production run. While doing that, we observed that our weight time ratio for the same pipeline went up when all we did was scale some numbers. So we just added four times the data parallel rings. We added like 200 data streams and then batch size is x by4. Um so it turns out that since the data generator and the worker pool are colllocated for fast data transfer

网卡瓶颈The network card bottleneck

实际情况是，该节点的网络接口卡出现了瓶颈。也就是说，当我们把所有这些工作和处理任务都集中到特定节点上时，这些节点的网络吞吐量无法满足训练流水线的需求。于是，我们对正在使用的框架做了一些调查，发现 Ray 默认会优先在同一节点上调度并放置 worker，这意味着一台机器很容易先出现瓶颈，然后才会用到另一台机器。因此，我们考虑将整个流程拆分开来，让 worker 分散到多个节点上。这样它们就可以各自处理图像并将其存储起来，再以惰性的按需方式检索图像，供我们的 DP 环使用。what happens is that there's a bottleneck on the network interface card of that of that node. Uh which means that while we're packing all of this working and processing into specific nodes that their network throughput is not able to keep up with the demand of the training pipeline. Um so at this point we did a little bit of digging into the frameworks that we are using and we saw that by default Ray schedules workers on this it prefers to schedule workers on the same nodes and colllocate workers which means that it's easy to bottleneck one machine before you move to another machine. So what we did was consider a possibility where we disagregate this whole process where workers are spread across multiple nodes. Um so they can do their own image processing and storing and in a lazy just in time manner are retrieved and used by our DP rings.

所以，没错，重新说一下，然后这就是我们调整过的两个参数、两个标志位：其中一个是零拷贝检索，另一个是……So yeah to re and then so these are the two parameter two flags that we kind of tuned one is a zero copy retrieval and

价差调度与零拷贝Spread scheduling and zero copy

第二个是分散式调度机制。嗯，这两种机制在小规模实验中都没有实际效果，但结合起来，在大规模实验中却产生了协同效应。嗯，是的，所以某种配置在小规模下可能表现不佳，但在更大规模下却可能表现得非常好。因此，你应该始终考虑如何在不同的扩展阶梯上进行消融实验。最终，正是这两个标志对我们的数据处理流水线速度产生了巨大影响。现在，这就是我们最终实现的流水线：一个并发、急切执行的系统，它能够感知数据访问模式，在多个负载间共享训练，并让 GAPU 保持活跃，同时确保训练过程平稳稳定。总的来说，仅靠这两个标志，我们就在大规模场景下实现了额外 50% 的吞吐量。second is a spread scheduling mechanism. Um neither of these actually helped in a small scale experiment but together they had um a synergistic effect in a large scale experiment. Um yeah so config can lose a config can do poorly at a small scale but can do really good at a larger scale. So you should always think about how to ablate across your your scaling ladders. And these were the two flags that ended up making a huge difference uh in how how fast our data processing pipeline was. And so now this is the final pipeline that we've accomplished a concurrent eager um and a data access pattern aware system that shares the training across loads and enables that GAPU stay warm while providing smoothness and stability in training. Overall we achieved a 50% additional throughput at scale just from these two flags.

同一个旋钮，结论却相反。It's the same knob opposite verdict

课程Lessons

由瓶颈决定。因此，即便没有别的收获，这些就是我们从这次分享中得到的经验。嗯，GPU 很贵，但有时候最大的难题却发生在 GPU 之外。所以，先问问自己：你是受 CPU 限制，还是受 GPU 限制？是受 IO 限制，还是受计算能力限制？而且，解决一个瓶颈往往会暴露出另一个瓶颈。对我们来说，起初是吞吐量问题，接着变成了等待时间问题，然后是内存问题，最后是检索成本问题。嗯，经典方法并非万能。你知道，并行处理和预取很重要，我觉得很多内容都讲过这一点。不过，你所使用的框架的语义，以及你面临的约束，同样重要；而且，规模总会改变答案。零拷贝浪费之类的调优参数，以及分散工作进程，只有在规模更大时才会变得重要，在较小规模下可能看不出明显的改进。set by the bottleneck. So these are the learnings we get from this talk if nothing else. Um GPUs are expensive but sometimes the biggest battles happen outside of the GPU. So first ask yourself are you CPU bound or GPU bound? Are you IO bound or computebound? And then solving one bottleneck often exposes another. So first it was throughput for us then it became weight time then it became memory then it became retrieval cost. Um, classic moves are not the end. You know, parallelism and prefetching matter and I think a lot of content has like covered this. Uh, but so do the semantics of the frameworks that you're using and the and the constraints that you're dealing with and scale always changes the answer. Tuning parameters like zero copy waste um and spreading workers become important only at larger scales and they may not manifest as obvious improvements at smaller levels.

总的来说，对基线进行性能分析有助于取得进展。嗯，是的，感谢大家参加这次演讲。如果有任何问题，我很乐意回答。谢谢。Overall, profiling your baseline helps you move the needle. Um, yeah, thank you everyone for attending the talk. Uh, if you have any questions, I'm happy to answer them. Thanks.