做过大模型训练的人多半都被"喂数据"折磨过。GPU 明明很贵、很快,却常常在那儿干等——因为数据还堵在从对象存储到本地磁盘、再到显存的那条又长又绕的路上。TOS PyTorch Connector 想解决的,正是这个"最后一公里"的问题。

它的思路并不玄乎,一句话就能说清:别再把数据先整包搬到本地,而是让训练直接从对象存储里"边读边训";同时把访问路径压到最短,再用 PyTorch 原生的抽象把用起来的门槛降到最低。 剩下的性能优势,基本都是从这个思路里长出来的。

一、传统方式的三笔"隐形账单"

在没有专用连接器之前,把 TOS 里的数据接进训练,通常有两条路,而且都不便宜。

第一条是预下载(预热):先把整个数据集从对象存储拉到本地或某个缓存层,再开始训练。数据量动辄几百 GB 甚至上 TB,这一步既费时间、又占存储,还得有人管理这份副本什么时候删、怎么同步。这是第一笔账:成本和管理负担。

第二条是走类似 FSX 这样的文件系统挂载:好处是"看起来像本地文件",但代价是请求要多经过一次 FUSE 协议的转换,链路更长、复杂度更高。这是第二笔账:多余的协议开销。

还有藏得更深的第三笔账:开发成本。无论走哪条路,你往往还得自己写代码,把数据封装成 PyTorch 认识的 Dataset,把 checkpoint 的读写也自己接一遍。这些"胶水代码"不产生业务价值,却要人维护。

二、核心思想:流式 + 最短路径 + 原生抽象

TOS PyTorch Connector 把上面三笔账一次性抹掉,靠的是三个相互咬合的设计。

第一,流式加载,不再预热。 数据在训练需要它的那一刻,才通过 HTTP 直接从 TOS 读出来。没有整包搬运,也就没有本地副本要管,成本和管理负担自然消失。

第二,直连 TOS,路径最短。 它通过 HTTP 协议(配合长连接)直接读写 TOS,不经过 FUSE 那一层转换。链路越短,延迟越低、吞吐越稳。

第三,面向 PyTorch 开箱即用。 它把最烦人的两件事——数据集和 checkpoint——都做成了现成的类,配置几行参数就能用,不用再自己造轮子。

三、落到 API 上的三个抽象

理念要好用,得体现在接口上。连接器主要给了三样东西:

映射式数据集 TosMapDataset,适合随机访问,训练时想按索引取哪条就取哪条;

可迭代式数据集 TosIterableDataset,适合流式顺序读取,处理连续的大数据流很高效;

检查点接口 TosCheckpoint,训练时从 TOS 加载 checkpoint,周期性存盘时又能直接写回 TOS,读写对称、语义清晰。

两种数据集都提供了 from_prefix()(按统一前缀构建,适合路径有规律的场景)和 from_urls()(按 URL 列表构建,适合路径明确但分散的场景)两种入口。这个区分后面会解释为什么重要。

用起来大致就是这个样子——上下文管理器一包,凭证直接传参,干净利落:

from tostorchconnector import TosMapDataset, CredentialProvider

with TosMapDataset.from_prefix(
        'tos://bucket/prefix',
        region='cn-beijing',
        endpoint='http://tos-cn-beijing.ivolces.com',
        cred=CredentialProvider(ak=ak, sk=sk)) as dataset:
    item = dataset[0]          # 随机取一条
    content = item.read()      # 真正读取数据

四、数字说话:快到什么程度

理念再漂亮,也得拿吞吐量来验证。官方在 200 万张图片(平均 120 KB,总量约 220 GB)、96 核 / 384G / 144Gbps 内网带宽的环境下,把它和 TOS Python SDK、s3torchconnector 做了横向对比。结论很直接:

• 在内网域名下,from_urls 场景里 TosMapDataset 的 QPS 约为 Python SDK 对应方案的 20 倍,TosIterableDataset 约 12 倍;即便是 from_prefix 也有约 8 倍 提升。
• 和 s3torchconnector 相比,TosIterableDataset.from_urls 的 QPS 约为对方的 10 倍。
• 叠加 TOS 加速器 后还能再上一个台阶,其中 from_urls 提升最明显(约 2 倍),而 from_prefix 约 1 倍——说明加速器对"URL 列表式精确并发"的优化收益更大。

这里也顺带解释了 from_urls 和 from_prefix 的差别:前者路径明确、并发调度更精准,天花板更高;后者靠前缀去列举对象,省事但要多一层枚举开销。追求极致吞吐时,from_urls 往往是更优解。

五、正本清源:它比本地 SSD 还快吗?

看到"快 20 倍",很多人第一反应是——那我干脆别用本地盘了?这里必须泼盆冷水:这组基准测试从头到尾没跟本地 SSD 比过。 它比的是三个"连接器"读同一个对象存储时谁更高效,说的是"连接器之间的差距",不是"网络存储比本地盘快"。这是两码事,别混淆。

那到底能不能比 SSD 快?得分维度看。

先说吞吐和延迟这一层,本地 NVMe 通常更强。 拿它自己的数字算一下:220 GB / 123 秒 ≈ 1.8 GB/s,开加速器 49 秒 ≈ 4.5 GB/s;而一块像样的 NVMe SSD 顺序读大概 3.5~7 GB/s,PCIe 5 能到十几 GB/s。换句话说,数据一旦已经躺在本地 SSD 上,本地盘的原始吞吐和延迟基本都赢——网络请求单次延迟是毫秒级,SSD 是微秒到亚毫秒级,随机小读差距尤其明显。连接器之所以拼命堆并发、预取、长连接,本质是在"用吞吐掩盖网络延迟",让 GPU 不至于等数据。这是补短板,不是天生比 SSD 快。

但换个视角,它就赢回来了,靠的是这三点:

一是省掉了"搬进来"这一步。要用本地 SSD,你得先把数据从对象存储下载到盘上(预热)。几百 GB 到几 TB 的下载时间如果算进端到端,流式直读往往整体更快,因为它压根没有这个搬运阶段。

二是容量根本装不下。测试的 220 GB 能塞进 SSD,但真实训练集常常是几 TB 到 PB 级,本地盘要么放不下,要么得反复换批下载。这时"比 SSD 快"这个问题本身就不成立——SSD 根本不是一个可选项。

三是云上的"本地盘"未必真快。很多云主机的所谓本地盘是网络挂载或被限速的,单块 SATA SSD 也就 500 MB/s 量级。配合几十上百 Gbps 内网带宽和分布式对象存储的横向扩展,聚合吞吐反而能压过它。

所以务实的结论是:它的定位从来不是"比 SSD 快",而是"让你不用先搬到 SSD"。 数据集不大、能整份放进本地 NVMe、又只训一轮,那老实用本地盘、甚至配上它的带缓存 Reader(SEQUENTIAL 把整对象缓存进内存)会更省心。它真正的主场,是大规模、装不下、每轮都要重下的场景——省掉搬运、绕开容量墙,同时把吞吐做到"足够喂饱 GPU"。

六、想再压榨一点性能

如果默认配置还不够,有几个很实在的调优方向,思路都围绕"并发"和"缓存":

• 大对象读写:调高 shared_prefetch_tasks(顺序读并发)和 shared_upload_part_tasks(分片上传并发),几百的量级都很常见。
• 数据集加载:调大预加载并发 prefetch_concurrency,让读取尽量跑在训练前面。
• 同一对象反复读:选带缓存的 Reader。SEQUENTIAL 把整份数据缓存在内存,适合反复读整对象;RANGED 只缓存一定大小,适合小块反复读又想控内存;DIRECT 不缓存,适合一次性读取。

工程实践上还有两个细节值得记住:一是三个核心类都实现了上下文管理,推荐用 with 加统一的 TosException 捕获;二是多进程加载(DataLoader 多 worker)时,日志需要在每个进程里各自初始化,用 worker_init_fn 来完成。

结语

把这套东西的价值抽干净,其实就一句话:它让对象存储从"训练前要先搬运的仓库",变成了"训练时可以直接读取的数据源"。

省掉预热,你少了一份副本和一堆管理成本;直连 HTTP,你少了一层协议转换的开销;原生抽象,你少写了一大坨胶水代码。至于和本地 SSD 的关系——拼单盘极限速度它不一定赢,但拼"端到端不用预热 + 装得下 PB 级数据",它才是为这个而生的。对于数据规模越来越大、GPU 越来越贵的今天,把数据加载这一环做薄、做快,本身就是一种很划算的优化。