Skip to main content
  1. Posts/

从 DataLoader 到 Remote DataLoader:如何让 GPU 不再等待数据

·4999 words·10 mins
Table of Contents
Note: This article is available in Chinese only. 本文暂无英文版本。 View original

起因:GPU 没吃满,但不是算力问题
#

最近组里有同学训练模型时,最直接的感受是:GPU utilization 上不去,step time 里总有一段在等数据。

最开始的怀疑都很合理:

  • GPU compute 太快,导致数据侧跟不上。
  • DataLoader 的 num_workers 太小。
  • 宿主机带宽或者机房网络有问题。
  • 底层 storage backend 吞吐打满了。

profiling 做完之后,结论有点反直觉。底层 storage 集群还有余量,真正先撞墙的是宿主机侧能力:单个 training process 背后的 storage client,在 IOPS 或 request concurrency 上先到上限。

也就是说,数据路径里的短板并不是 “storage 总吞吐不够”,而是 “这台机器、这个 client 每秒能发和处理的请求数不够”。

当 sample/object 偏小、请求偏多时,系统很容易从 bandwidth-bound 变成 IOPS-bound(每秒请求数受限)。

问题就从这里展开。

DataLoader 到底是什么
#

先看最常见的 PyTorch 训练循环:

 1dataset = MyDataset(...)
 2
 3loader = DataLoader(
 4    dataset,
 5    batch_size=1024,
 6    num_workers=8,
 7)
 8
 9for x, y in loader:
10    x = x.cuda()
11    y = y.cuda()
12
13    output = model(x)
14    loss = criterion(output, y)
15
16    optimizer.zero_grad()
17    loss.backward()
18    optimizer.step()

很多人第一眼会把 DataLoader 理解成 “读磁盘文件”。

这个理解太窄了。真实训练里,DataLoader 往往至少包含这些阶段:

  • sample selection(选哪些样本)
  • storage read(从对象存储/分布式文件系统读取)
  • metadata lookup
  • network RPC
  • deserialize / decode
  • preprocessing / augmentation
  • batch / collate
  • CPU memory management
  • prefetch 与 pinned memory
  • Host-to-Device transfer(H2D)

所以更准确的定义是:DataLoader 不是一个文件读取 API,而是训练系统里的 input pipeline(输入流水线)。

DatasetSamplerDataLoader 可以这样分工理解:

  • Dataset:定义单个 sample 是什么、怎么取到。
  • Sampler:定义这一轮要取哪些 sample。
  • DataLoader:把这些 sample 高效生产成 GPU 可消费的 batch。
dataloader pipeline

用 Producer-Consumer 建一个简单心智模型
#

把训练流程抽象一下:

  • Producer:DataLoader / Data Workers
  • Consumer:GPU

定义两个时间:

  • \(T_{data}\):准备一个 batch 的平均时间
  • \(T_{gpu}\):GPU 完成一个 step 的平均时间

理想情况下:

$$ T_{step} \approx \max(T_{data}, T_{gpu}) $$

如果 \(T_{data} > T_{gpu}\),GPU 必然等数据。

DataLoader 优化目标不是把 I/O 吞吐无限做高,而是让数据生产速度稳定超过 GPU 消费速度。

也就是:

$$ Throughput_{producer} \ge Throughput_{consumer} $$

例如:

  • GPU compute:50ms / batch
  • Data preparation:80ms / batch

这时 step 被数据限制。

如果把 Data preparation 降到 30ms,而 GPU 仍是 50ms,那么数据阶段可以被隐藏在 compute 后面,GPU 基本不会因为取数据而 idle。

这就是 overlap、pipeline、prefetch 的意义:不是追求极限吞吐,而是追求端到端节拍匹配。

为什么 num_workers 总是第一层优化
#

DataLoader(num_workers=8) 本质上是在加 producer 并发。

多个 worker 可以并行做读取、解码、预处理、拼 batch,理论上会提高供给速度。

multi worker dataloader

这个优化很常见,也经常有效。

但它有一个隐含前提:瓶颈确实在 worker 并发本身,而不是更下游的共享资源。

为什么有时候把 num_workers 拉满也没用
#

真实系统里经常出现这种结构:

1Worker1 ┐
2Worker2 ├── Storage Client ── Storage
3...     │
4WorkerN ┘

64 个 worker,如果最后都复用同一个 storage client、同一套连接池、同一条 RPC 引擎和请求队列,那瓶颈还是那一个 client。

于是 num_workers: 8 -> 64 可能几乎不提升吞吐,甚至让队列更拥塞、尾延迟更差。

这里有个很通用的系统结论:只有当瓶颈资源本身也能扩展时,并发才有意义。

IOPS 和 Bandwidth 是两种完全不同的约束
#

很多 ML 工程实践里,大家习惯盯 GB/s,但训练数据路径常常先撞的是 IOPS

举个数字。假设平均 sample 只有 4KB,目标吞吐是 1GB/s:

$$ IOPS = \frac{1GB/s}{4KB} \approx 262,144 $$

也就是每秒大约 26 万次请求。

哪怕后端 storage 理论带宽是 100GB/s,你也可能因为 request rate、metadata、RPC、client queue、syscall、connection 并发上限,实际吞吐远低于预期。

小 I/O workload 的主矛盾常常不是 “每秒传多少字节”,而是 “每秒处理多少请求”。

这时常见的两个优化方向是:

  • 方向 A:减少请求数(tiny files -> large shards)
  • 方向 B:增加独立 client 数(scale-out readers)

需要把这两个方向区分清楚:Sharding 解决“请求数过多”,Remote DataLoader 解决“请求处理并行度不够”。 两者经常要叠加用,而不是二选一。

Remote DataLoader:input pipeline 的 scale-out
#

原始架构里,训练进程自己背所有事情:storage client、DataLoader、预处理、prefetch、GPU compute。

single client architecture

Remote DataLoader 的思路是把读取侧横向展开:

remote dataloader architecture

每个 remote worker 都可以有独立的 storage client、连接池、CPU、内存、请求队列。

理想情况下,聚合 IOPS 近似:

$$ IOPS_{aggregate} \approx N \times IOPS_{client} $$

直到遇到新的短板:

$$ Throughput = \min(Storage,\ Client,\ CPU,\ Network,\ RPC,\ Trainer) $$

所以本质上,Remote DataLoader 是把 input pipeline 从 scale-up 改成 scale-out。

Remote 的价值不只在 IOPS:它在做资源解耦
#

如果只从 IOPS 看 Remote DataLoader,会低估它的价值。

更重要的是它在做 resource disaggregation(资源解耦):把数据平面和 GPU 计算平面拆开。

GPU 节点贵,应该尽量干 forward、backward、optimizer,不该被文件系统 RPC、解码、预处理拖住。

resource disaggregation

拆开后,CPU/Data fleet 和 GPU fleet 可以独立扩缩容。这个思路在 tf.data service、Ray Data、Mosaic Streaming、NVIDIA DALI 这类系统里都能看到影子。

第一版做成 Job Scoped,其实非常合理
#

我们最早的 Remote DataLoader 设计很朴素:

1Job Start -> Remote DataLoader Start -> Load Dataset -> Train -> Job End -> DataLoader Exit

也就是 Remote DataLoader 生命周期和训练任务一一绑定。

这版并不 “落后”,对于 MVP 很合理:

  • ownership 清晰
  • 故障域简单
  • 不需要跨 job 隔离
  • 不需要 eviction 与版本治理
  • job 结束即释放内存

先把系统跑起来,很多时候就该这么做。

第二阶段问题:每个 Job 都重新 Cold Load 太贵
#

新的问题是后面才暴露出来的。

即使 Remote DataLoader 已经解决单 client 瓶颈,第一次把海量数据从 storage 拉进内存本身仍然很重。

典型情况:

1JobA: Load 1TB -> Train -> Release
25分钟后
3JobB: Load 同样 1TB -> Train -> Release

我们重复做了几乎完全相同的 I/O。

而且在真实业务里,短时间窗口(一天或一周)内,训练任务的数据集经常高度重叠:

  • JobA:dataset 2026-08-20
  • JobB:dataset 2026-08-20
  • JobC:dataset 2026-08-20 + small delta
  • JobD:dataset 2026-08-21

问题就从 “怎么读更快” 变成 “刚读过的数据为什么还要再读”。

从 Job Scoped 到 Service Scoped:Persistent Data Cache
#

下一步自然演化是把生命周期从 job 级别拉长到 service 级别。

persistent data service

数据一旦进入 remote worker 的内存,尽可能保留,让后续任务复用。

从这里开始,数据生命周期开始长于训练任务生命周期。

也是从 “Remote Reader” 走向 “Data Service” 的分界线。

这一步其实是系统抽象升级
#

一旦允许 data lifetime > job lifetime,系统就不再只是 DataLoader 了。

它会迅速长成一个 distributed data serving / caching system,因为你必须回答这些问题:

  • 数据放在哪些 worker
  • 如何定位 dataset/shard
  • 何时加载,谁触发加载
  • 多 job 如何共享与隔离
  • 内存满了怎么淘汰
  • dataset 新版本如何兼容
  • worker 挂了怎么办
  • 热点数据要不要复制
  • 权限和配额怎么做

这里最容易低估的是复杂度来源。很多复杂度不是来自 “remote”,而是来自 “persistent + shared”。

Cache 带来的关键系统问题
#

引入 cache 之后,必须先面对几个概念边界。

Cache Key 与 Dataset Identity
#

你必须能区分这些是否同一份数据:

  • dataset id
  • dataset version
  • shard id
  • preprocessing version

没有稳定 identity,就会出现 stale data 或重复缓存。

Eviction 之前,先想 Admission
#

内存不可能无限大。LRU/LFU/TTL/size-aware 都是常见 eviction 策略。

但工程上经常更关键的是 admission:哪些数据值得进昂贵内存,哪些不值得。

Cache 的核心不只是淘汰谁,而是接纳谁。

Locality 感知调度
#

如果 shard 已经在 Worker3,最好不要让 Worker7 再从 storage 拉一遍。

scheduler 需要 locality awareness,不然命中率上不去,I/O 白做。

Shared Cache 与隔离
#

共享提高 hit rate,但会带来 noisy neighbor、quota、fairness、permission 这些多租户问题。

Fault Tolerance:Reconstructable Cache
#

Remote worker 挂了,缓存丢失不等于数据丢失,原始数据还在 storage。

这类缓存天然是 reconstructable cache(可重建缓存)。

所以 replication 不一定强制。可以:

  • 默认不复制,挂了重建,省内存
  • 对热点数据 selective replication,换恢复速度

这是经典成本-可用性 trade-off,不存在单一最优解。

Prefetch 仍然是核心机制
#

Remote 不等于 “每 step 请求一次,然后同步等返回”。

正确模式依然是多阶段流水线重叠:

  • Remote 准备 Batch N+2、N+3
  • 网络传输 Batch N+1
  • CPU 准备 Batch N
  • GPU 计算 Batch N-1
prefetch pipeline

Prefetch 本质是 latency hiding(延迟隐藏)。

Backpressure:读得更快不代表系统更好
#

scale-out 以后很容易出现 producer 吞吐远大于 GPU 消费:

  • Data workers:100GB/s
  • GPU demand:20GB/s

如果没有 bounded queue 和 backpressure,后果很直接:

  • RAM 打爆
  • 网络浪费
  • cache pollution
  • 无效预处理

所以 Remote DataLoader 的目标不是 “无限快”,而是稳定匹配消费速率。

Scale-out 之后的调度问题:Sharding、Straggler、Shuffle
#

当规模变成 1000 shards / 32 readers / 256 GPUs,系统问题自然变成:

谁读哪个 shard?

Static Sharding
#

优点:简单、deterministic、可复现。
缺点:容易被慢 worker 拖住(straggler)。

Dynamic Scheduling
#

优点:负载均衡更好。
缺点:checkpoint、复现语义、exactly-once/at-least-once 处理复杂度更高。

再加上训练对 shuffle 的要求,事情会更复杂。海量数据下通常做不了真正全局随机,只能做近似:

  • shuffle shard order
  • shuffle within shard
  • shuffle buffer

这里的 trade-off 很典型:Randomness、Memory、I/O Efficiency 三者一起拉扯。

从缓存视角看性能:Cold Path 与 Hot Path
#

把路径拆开看更清楚。

Cold Path(cache miss)
#

cold path

Hot Path(cache hit)
#

hot path

对应延迟模型:

$$ T_{cold} = T_{storage} + T_{network} + T_{deserialize} + T_{preprocess} $$$$ T_{warm} \approx T_{memory} + T_{network} $$

如果 workload 的 dataset reuse rate 很高(例如 80%),那第二阶段收益常常不再来自 “多 client 并行读 storage”,而来自 “直接跳过 storage”。

Remote DataLoader 第一阶段优化的是 storage read throughput;Persistent Cache 第二阶段优化的是 storage read elimination。

说到底,最好的 I/O 优化,不是把 I/O 做得更快,而是根本不要做这次 I/O。

放进完整 ML Infra:Control Plane 与 Data Plane
#

当系统走到这一步,通常会形成下面这种分层。

control data plane

Control Plane 负责 metadata、placement、lifecycle、调度与 worker 管理。
Data Plane 负责读取、缓存、预处理、传输与 serving。

这个拆分和很多 distributed system 的经典模式是一致的。

一条常见演进路线
#

evolution stages

每一阶段主要在解决不同问题:

阶段主要解决的问题
Single Loader最基础的数据读取
Multi-workerCPU / 本地并发不足
Prefetchlatency hiding
Sharding降低 IOPS amplification
Remote DataLoader突破单 client / 单 host 限制
Persistent Cache消除重复 I/O
Data Serving System跨 job 共享与统一调度

这个表有一个隐含信息:后面的阶段不是替代前面,而是在前面基础上继续抽象。

不要过度设计:先 profiling,再做架构升级
#

不是所有训练都该上 Remote DataLoader,更不是都该上 persistent distributed cache。

如果你的场景是:

  • dataset 不大
  • 本地 NVMe 足够快
  • GPU 基本不等数据
  • storage client 没撞上瓶颈

那普通 PyTorch DataLoader 往往已经够好。

上分布式数据服务会带来额外代价:更多服务、更多故障点、网络依赖、调度复杂度、内存成本、运维成本。

更稳妥的路径是:

1GPU idle
2-> measure data wait
3-> measure storage latency / IOPS
4-> measure preprocessing CPU
5-> measure network
6-> locate bottleneck
7-> choose architecture

不要从 “我有 GPU 训练任务” 直接跳到 “我需要 distributed dataloader”。

Observability:怎么判断优化真的生效
#

Remote DataLoader 做得好不好,不看 worker 数量,看端到端指标是否改善。

建议至少盯这些指标:

Trainer
#

  • GPU utilization
  • GPU data wait time
  • step time
  • samples / second

DataLoader / Input Pipeline
#

  • batch ready latency
  • prefetch queue depth
  • producer throughput
  • consumer throughput

Storage
#

  • IOPS
  • read bandwidth
  • request latency P50 / P99
  • outstanding requests

Remote Workers
#

  • CPU utilization
  • memory utilization
  • network throughput
  • per-client throughput

Cache
#

  • hit rate / miss rate
  • eviction rate
  • cached bytes
  • cold-load latency

Scale-out 效率
#

你需要看到随着 worker 增加,GPU data wait time 真实下降,而不仅是 reader 侧吞吐上升。

举个典型观测:

11 worker  -> 10 GB/s
22 workers -> 19 GB/s
34 workers -> 35 GB/s
48 workers -> 42 GB/s

当曲线开始 flatten,就说明新的瓶颈已经出现,继续加 worker 的收益会快速递减。

问题最后都指向生命周期与边界
#

这件事最开始其实很简单:我们只看到 GPU 在等数据。

于是第一反应是多并发、多 client、并行读 storage。Remote DataLoader 确实解决了单机侧 request-rate 瓶颈。

再往后才发现,真正贵的不只是 “读得慢”,还有 “反复读同一批数据”。

当训练任务在同一个时间窗口里重复使用高度重叠的数据时,系统的最优边界往往已经不在 training process 内部,而在一个独立、持久、可共享的数据服务。

最后会回到几个基础问题:数据应该在哪里、由谁持有、存活多久,以及 GPU 是否真的需要感知这些细节。

很多 ML Infra 问题都是这样演化的:一开始像局部性能优化,往深处走其实是生命周期管理和系统边界设计。

Remote DataLoader 表面上在解决 IOPS,往深处看,是在回答 “数据面和计算面应该如何解耦”。

参考资料
#

Related