训练一个大型机器学习模型,很多人首先想到的是GPU数量、显存容量和Tensor Core性能,却容易忽略一个非常现实的问题:GPU究竟有没有持续拿到足够的数据。
一块高端GPU执行矩阵运算时可以达到极高的计算吞吐,但训练数据不会凭空出现在显存里。数据首先要从本地SSD、网络文件系统、对象存储或者数据湖读取出来,然后经过解压、解码、解析、tokenization、数据增强、格式转换、batch组装等处理,最后才进入GPU。
如果这条数据流水线跟不上GPU,GPU即使拥有几百瓦甚至更高的功耗,也可能长时间等待输入。
因此,训练系统中的数据预处理并不是简单回答“CPU还是GPU”的问题,而是要判断每一个处理阶段的计算特征、数据位置、带宽需求和延迟要求,再决定把它放在哪里。
CPU通常是训练数据预处理的第一主力
CPU最适合承担大量控制流复杂、分支较多、依赖成熟软件库的数据处理任务。
文本数据的解析、JSON处理、文件读取、tokenization、样本过滤、数据格式转换以及batch组装,通常都可以由CPU完成。
图像训练中的JPEG、PNG解码、随机裁剪、尺寸调整、颜色转换等操作,也长期由CPU数据加载线程承担。
以PyTorch为例,DataLoader可以通过多个worker并行读取和处理数据。多个CPU worker同时从存储系统读取样本,再把准备好的batch交给训练进程。
这时候真正需要观察的并不是“CPU占用率有没有100%”,而是CPU数据加载是否能够稳定提供训练所需要的batch。
如果GPU每个step只需要几十毫秒,而CPU准备一个batch需要一百毫秒,那么即使GPU本身非常强,整体训练速度仍然会被数据流水线拖住。
因此数据预处理优化的目标不是让CPU跑满,而是让数据生产速度至少能够覆盖GPU消费速度,同时避免CPU线程数量、内存和I/O产生新的瓶颈。
GPU并不是不能做预处理
“数据预处理应该放CPU,GPU负责训练”也是一种过于简单的理解。
很多预处理操作本身就是高度并行的数值计算,例如图像resize、颜色变换、归一化、部分数据增强,以及某些张量变换。
如果这些操作已经成为明显瓶颈,就可以考虑把部分计算迁移到GPU。
现代深度学习训练框架也提供了越来越多的GPU数据处理能力。NVIDIA生态中还有专门针对数据加载和预处理的工具,可以利用GPU执行部分解码和增强任务。
但GPU预处理并不意味着“GPU越忙越好”。
训练模型本身和数据预处理需要竞争GPU计算、显存、内存带宽以及kernel调度资源。如果GPU已经被模型计算压得很满,把更多预处理塞进去反而可能降低训练吞吐。
所以判断是否应该GPU加速预处理,不能只看某一个操作的理论计算速度,而应该看它是否真的成为端到端训练的瓶颈。
真正重要的是把数据流水线做成异步流水线
现代训练系统通常不会等GPU完成一个batch以后,CPU才开始准备下一个batch。
更合理的结构是:
存储系统持续读取数据,CPU worker同时进行解码和预处理,下一批数据进入内存缓冲区,GPU正在训练当前batch时,CPU已经开始准备下一个batch。随后通过异步host-to-device拷贝,把下一批数据送入GPU。
这样才能形成真正的pipeline。
可以把训练过程理解成三个持续运行的阶段:数据读取和预处理、CPU到GPU的数据传输、GPU模型计算。
理想情况下,当GPU正在计算第N个batch时,CPU已经准备好第N+1个batch,而存储系统正在为第N+2个batch提供数据。
这比单纯增加CPU核心或者GPU数量更加重要。
如果三个阶段中任何一个阶段明显慢于另外两个阶段,整体吞吐就会被最慢的阶段限制。
不要把Redis和Kafka当成GPU前面的普通内存缓冲区
原稿把Redis或Kafka作为CPU处理后数据传给GPU的内存缓冲区,这在一般单机训练架构中并不是典型方案。
Redis是内存数据库,Kafka是分布式事件流平台,它们各自有明确的业务用途,但都不是训练进程与GPU之间天然的高速batch缓冲层。
在单机训练环境中,CPU预处理后的数据通常直接进入进程内存、pinned memory或者框架的数据队列,然后通过PCIe等路径进入GPU显存。
例如使用pinned memory,可以让CPU到GPU的数据传输更适合异步执行。训练程序可以在GPU计算当前batch时,同时准备和传输下一批数据。
因此,真正需要关注的是host memory、pinned memory、PCIe带宽、GPU显存以及数据加载队列,而不是简单地在CPU和GPU之间增加一个Redis。
Kafka在实时数据管道中当然非常有价值。例如推荐系统、事件流或者实时日志训练管道可以使用Kafka承接事件流,再由后端数据处理系统生成训练数据。
但这属于数据管道架构问题,并不意味着每一个batch训练数据都应该经过Kafka再送进GPU。
存储系统往往才是第一瓶颈
当训练数据达到TB甚至PB级别以后,CPU和GPU之外还必须考虑数据到底放在哪里。
如果每个训练worker都直接从远程对象存储读取大量小文件,性能可能非常差。
原因并不一定是对象存储的总带宽不足,而可能是大量小文件带来的请求数量、元数据操作、网络延迟和连接开销。
因此大规模训练通常需要考虑数据格式和布局。
把大量小文件重新组织成适合顺序读取的大文件、采用Parquet、WebDataset、TFRecord等适合训练的数据格式,或者建立本地缓存,都可能比单纯增加GPU更加有效。
如果一个8卡训练节点需要持续产生极高的数据吞吐,而所有GPU都共享一个很慢的网络文件系统,那么增加到16卡甚至32卡,很可能只是把存储瓶颈放大。
GPU数量增长以后,数据读取带宽必须同步增长。
GPU训练节点通常应该自己完成大部分轻量预处理
对于常见的分布式训练,一个合理的设计通常不是把所有数据预处理都搬到独立服务器。
训练节点本身一般已经拥有多核CPU、较大的系统内存、本地NVMe SSD以及高速PCIe总线。
如果数据处理量并不大,让训练节点上的CPU worker完成预处理往往最简单,也减少了额外网络传输。
例如一个训练节点拥有8张GPU,同时配备几十个CPU核心和数TB本地NVMe缓存,那么完全可以让CPU负责数据读取、解码、tokenization和batch组装,然后直接送到本机GPU。
这种架构最大的优势是数据路径短。
数据从本地SSD读取到系统内存,再进入GPU,避免了额外的数据处理服务器和网络跳转。
什么时候值得增加独立数据处理节点
独立数据处理节点并不是GPU训练系统的标准配置,而是当数据处理本身已经成为明显瓶颈时才值得考虑。
例如训练任务需要进行非常复杂的图像解码、视频解码、文本处理、数据过滤或者大规模特征计算,而训练GPU本身已经处于高利用率状态,这时把预处理任务放到独立计算集群可能有价值。
还有一种典型情况是多个训练任务需要使用相同的数据处理结果。
例如一个大型数据集需要经过数小时甚至数天的数据清洗、OCR、tokenization或者特征生成。如果每一个训练job启动以后都重新处理一次,那么大量CPU时间会被重复消耗。
这时候可以把重型预处理变成独立的数据准备阶段。
原始数据经过离线处理后生成标准化训练数据,之后多个训练任务直接读取这些已经准备好的数据。
这种设计通常比每次训练都重新处理一遍更加高效。
独立节点真正解决的是资源隔离
独立数据处理集群最大的价值之一并不是“拥有专用指令集”,而是资源隔离。
假设一个训练集群有大量GPU,而数据预处理非常吃CPU。如果训练任务和预处理任务共享同一组CPU资源,数据加载线程可能与训练框架、通信线程、系统进程发生竞争。
把预处理任务放到独立节点以后,可以让训练节点专注于模型计算,同时让数据处理集群按照CPU、内存和I/O需求独立扩展。
但代价也很明显。
数据必须从预处理节点传到训练节点。
如果处理后的数据体积巨大,那么网络就可能变成新的瓶颈。
因此独立节点架构只有在“减少的计算压力”大于“增加的数据传输成本”时才有意义。
RDMA和NVLink不能混为一谈
原稿把独立数据处理节点和NVLink联系起来,这里需要特别纠正。
NVLink主要用于GPU之间以及特定GPU系统内部的高速互联,它不是一种通用的跨服务器数据处理网络。
如果数据处理服务器与训练服务器之间进行高速网络传输,更常见的是高速以太网或者InfiniBand,并可以结合RDMA技术降低CPU参与和网络传输开销。
在多GPU训练中,NCCL可以利用高速互联进行GPU之间的数据通信。
因此应该区分三条不同的数据路径:存储到计算节点、计算节点CPU到GPU、GPU之间或者服务器之间的通信。
把NVLink、RDMA、PCIe和存储网络全部称为“低延迟数据传输”虽然听起来正确,但对于工程设计没有实际帮助。
真正需要做的是测量每一条路径的吞吐和延迟。
多GPU训练还需要考虑CPU和GPU之间的通信
当单机从一张GPU扩展到8张甚至更多GPU以后,数据加载问题也会发生变化。
如果每个GPU都拥有独立的数据加载worker,那么CPU线程数量、系统内存、PCIe拓扑以及GPU与CPU NUMA关系都会影响性能。
例如双路CPU服务器中,不同GPU可能挂载在不同CPU socket下面。如果所有数据加载线程都集中在一个NUMA节点,而GPU又位于另一个socket对应的PCIe域,那么跨NUMA访问可能增加延迟和内存带宽压力。
因此高性能训练服务器需要考虑CPU affinity、NUMA、PCIe拓扑以及GPU affinity,而不是简单地把DataLoader worker数量设置成CPU核心数量。
worker也不是越多越快。
线程和进程增加以后,会产生上下文切换、内存压力、文件描述符、锁竞争以及I/O争用。最终可能出现CPU利用率很高,但GPU吞吐反而下降的情况。
图像和文本的数据处理策略并不一样
不同数据类型对CPU和GPU的需求差异非常大。
文本模型通常涉及tokenization、Unicode处理、文本清洗、序列拼接和动态batch等工作。很多tokenizer本身已经高度优化,但如果训练数据规模巨大,tokenization仍然可能成为CPU瓶颈。
因此可以提前离线tokenize,把原始文本转换成适合训练的token ID格式,从而减少训练阶段CPU工作。
图像模型则可能面对JPEG解码、resize、随机裁剪、颜色增强、归一化等操作。
如果CPU解码成为瓶颈,可以考虑更高效的图片格式、预解码数据、GPU解码或者GPU augmentation。
视频模型的情况更加极端。
视频解码本身就可能消耗大量CPU资源和I/O带宽,训练时还需要随机访问帧、进行resize和增强。这类工作负载更容易值得建立专门的数据处理流水线。
因此不存在一个“CPU负责文本、GPU负责图像”的固定规则。真正的判断依据应该是profile结果。
先做Profile,再决定硬件放在哪里
数据流水线优化最容易犯的错误就是凭经验增加硬件。
如果GPU利用率只有40%,有人可能马上认为“CPU不够”;但真正原因可能是网络文件系统吞吐不足。
如果CPU达到90%,也不能马上得出“增加CPU”的结论,因为CPU可能是在等待磁盘I/O。
因此专业诊断应该同时观察多个指标。
GPU侧需要看SM和Tensor Core利用率、显存带宽、显存占用以及每个step的计算时间。
CPU侧需要看CPU utilization、每个worker处理时间、上下文切换、内存带宽以及NUMA访问。
存储侧需要看IOPS、吞吐、读延迟以及小文件比例。
网络侧需要看吞吐、丢包、延迟以及是否存在共享链路瓶颈。
训练框架还应该观察data loading time、GPU idle time、step time以及host-to-device copy时间。
只有把这些指标放在一起,才能判断究竟是CPU、GPU、存储、网络还是数据格式出了问题。
一个比较合理的分层架构
对于普通深度学习训练,可以采用相对简单的结构。
数据首先存放在对象存储、分布式文件系统或者本地存储中;训练节点负责读取数据;CPU worker完成解析、解码、tokenization和常规数据增强;经过pinned memory等机制进入异步数据传输阶段;GPU负责模型前向、反向传播和优化器计算。
当数据处理量明显增加时,可以增加本地NVMe缓存,减少重复访问远程存储。
如果预处理本身已经成为大型计算任务,则把重型数据处理独立出来,提前生成训练数据,再让GPU训练节点直接消费处理结果。
对于超大规模训练集群,则可以进一步建设独立的数据准备集群、缓存层、高速存储以及高速网络,让训练节点不再承担大量重复的数据处理工作。
这种架构的关键不是“把预处理全部搬到独立服务器”,而是把训练阶段真正需要的数据准备工作缩减到最低。
数据一致性比单纯追求吞吐更加重要
数据预处理系统还存在一个经常被忽视的问题:可复现性。
如果训练过程中随机增强、数据过滤或者采样策略发生变化,而系统没有记录对应版本,那么同一个模型重新训练时可能得到不同结果。
因此训练数据应该具备版本管理。
需要记录数据集版本、预处理代码版本、tokenizer版本、词表版本、数据过滤规则以及随机种子等关键参数。
对于分布式训练,还要避免不同worker产生重复样本或者遗漏样本。
这也是为什么现代训练系统越来越强调dataset versioning、checkpoint、manifest和可追踪的数据处理流程。
相比简单地使用“事务日志和校验码防止数据漂移”,这些机制更符合机器学习训练数据管理的实际需求。
最终目标不是让某一种硬件跑满
训练数据预处理的最终目标并不是让CPU达到100%,也不是让GPU一直保持99%利用率。
真正应该优化的是端到端训练吞吐。
如果GPU已经成为绝对瓶颈,那么继续增加CPU预处理能力没有意义。
如果GPU经常等待数据,那么继续购买更强的GPU同样没有意义。
如果CPU、GPU都没有明显问题,但远程存储已经达到吞吐上限,那么应该优化存储和缓存。
如果独立数据处理节点降低了训练节点CPU压力,却因为网络传输增加导致整体速度下降,那么这种架构同样没有达到目的。
所以,CPU、GPU和独立数据处理节点之间并不存在一个固定的“最佳分工表”。
CPU擅长通用数据处理和复杂控制流,GPU适合高度并行的数值处理,独立数据处理节点适合资源隔离和重型、可复用的数据准备任务,而存储和网络则决定这些组件之间能否真正形成高吞吐的数据流水线。
一个成熟的AI训练系统,应该让数据准备、数据传输和模型计算彼此重叠,让GPU尽可能少等待,同时避免为了追求理论峰值而把系统设计成一堆昂贵却互相争抢资源的硬件。
对于真正的大模型训练来说,决定训练速度的往往不是某一块GPU的峰值TFLOPS,而是整个数据路径能否持续、稳定地把正确的数据送到正确的计算单元。把这个问题解决以后,再谈增加GPU数量、升级高速网络或者部署独立预处理集群,才有工程意义。
训练数据预处理应该放在哪里,CPU、GPU 和独立数据处理节点如何分工
图片说明:示意图 图片来源:Public Domain(公有领域)
喜欢这篇报道?
使用下面的功能,方便以后继续阅读和分享 MNewsTV
关于文章收藏
收藏不需要注册帐号,收藏信息仅保存在当前浏览器中。
删除收藏请进入「我的收藏」进行管理。
捐助(Paypal): https://www.paypal.me/observeccp 订阅中国观察电报 Telegram : https://t.me/s/ObserveCCP