PiD 是一个面向高分辨率图像的潜空间解码器（Fast and High-Resolution Latent Decoding with Pixel Diffusion）。要支撑 4K 级图像-文本对的训练，它的数据管道必须又快又稳。本文带你快速吃透 PiD 数据加载机制的两大核心：WebDataset 多 tar 同步迭代与Shardlist 分片——从数据如何分片、多模态如何按 key 对齐，到分片如何分发给每台机器、每张 GPU，并支持断点恢复。

PiD 的训练数据是海量 4K 级图像-文本对，单卡单进程根本喂不饱训练，整个加载机制就围绕三件事展开：

1. 怎么打包：多模态数据（image / caption / 嵌入）怎么分片成 WebDataset；
2. 怎么对齐：不同模态放在不同目录的同名 tar 里，怎么保证id_0.jpg一定配id_0.json；
3. 怎么分发：几百万个 tar 分片怎么分给多台机器、每张 GPU、每个 worker，还能断点续训。

一、数据长什么样：多模态同名 tar 布局

PiD 的数据遵循 WebDataset 约定：不同模态存放在不同目录，但 tar 文件名与内部样本 key 完全一致。比如：

/data/ ├── image/ │ └── part_000000/ │ ├── 000000.tar # 包含: id_0.jpg, id_1.jpg, id_2.jpg, ... │ ├── 000001.tar │ └── ... └── caption/ └── part_000000/ ├── 000000.tar # 包含: id_0.json, id_1.json, id_2.json, ... └── ...

image/与caption/下的000000.tar是"孪生"分片：文件名相同、内部样本 key 相同、顺序一致。caption 的 json 里保存{"prompt": str, "prompt_medium": str, "prompt_short": str}（后两者可选）。

这种"孪生 tar"布局是整个同步机制的地基——对齐不靠数据库 join，而靠文件名。

1.1 核心数据结构：TarSample 与 Wdinfo

每个 tar 分片由 TarSample 描述（定义在 config/schema.py）：

数据源目录下会放置wdinfo.json清单文件，记录data_list（全部 tar 相对路径）、total_key_count（样本总数）和chunk_size（每个 tar 的样本数）。它们被汇总进 Wdinfo：

@attrs.define class Wdinfo: tar_files: list[TarSample] # 所有 tar 文件列表 total_key_count: int # 总样本数 chunk_size: int # 每个 tar 内的样本数（决定 shuffle buffer）

💡 最终每个 tar 的真实路径拼为root/key/path，例如root/image/part_000000/000000.tar和root/caption/part_000000/000000.tar是一对。

二、加载链路总览：从 wdinfo 到训练 batch

训练从 scripts/train.py 启动，配置入口 configs/pid_training/config.py 通过 Hydra 的data_train选项组选择数据配置，整条链路如下：

data_source_local.py（源名 → wdinfo.json 搜索目录） → dataset_definition.py（逻辑数据集 → 一个或多个源） → create_dataset_infos_local()（递归找 wdinfo.json） → Dataset.parse_dataset_info()（解析成 Wdinfo / TarSample 列表） → WebDataset 管道（分片 → 同步展开 → 解码 → 增强）

其中数据源映射定义在 data_source_local.py（当前注册了MultiAspect_4K_1M源），逻辑数据集定义在 dataset_definition.py。

2.1 rank 0 解析 + 广播：避免存储打爆

解析逻辑在 webdataset.py 的parse_dataset_info()中。有个值得学习的工程细节：只让 rank 0 解析 wdinfo，再broadcast给其他 rank——

解析要下载所有 wdinfo 并实例化百万级 TarSample。如果所有 rank 同时做，会产生 40 倍并发连接打爆对象存储（GCS），触发 SSL-reset 风暴。所以 rank 0 解析一次，广播结果给所有人。

三、核心机制①：多 tar 同步迭代

同步的核心代码在 utils/iterators.py，由 tarfile_samples() 串起三步流水线：url_opener→tar_file_expander→group_by_keys。

3.1 url_opener：为每个模态开一条流

url_opener() 遍历TarSample.keys，为每个模态目录拼出完整路径并打开字节流：

for data_key in url.keys: # ["image", "caption"] url_path_full = os.path.join(url.root, data_key, url.path) stream.append(gopen(url_path_full)) # 本地 open / S3 流式读取 # 结果: stream = [image_tar_stream, caption_tar_stream]

本地文件直接open；对象存储则走 stream.py 的RetryingStream，带指数退避重试，单个坏分片不会卡死预取 worker。

3.2 tar_file_expander：两种同步模式

tar_file_expander() 是"孪生 tar"对齐的灵魂，提供两种模式：

简单模式（sample_keys_full_list = None）——直接zip()对齐：

tar_file_iterator_list = [iter(image_tar), iter(caption_tar)] for sample in zip(*tar_file_iterator_list): # sample[0] = id_0.jpg（来自 image tar） # sample[1] = id_0.json（来自 caption tar）

简单粗暴：各 tar 的样本按位置一一对齐，前提是所有 tar 内文件顺序和数量完全一致。sharding_wds.py打包时保证了这一点（见第五节）。

索引模式（sample_keys_full_list指向 parquet 索引文件）——容忍缺失与乱序：

1. 加载 parquet 索引得到合法 key 列表["id_0", "id_1", ...]；
2. 维护target_index，用 run_iterator_to_index() 推进每个 tar 迭代器到目标位置；
3. 只有所有 tar 都命中对应样本才 yield；缺失的 key 跳过并记录日志，重复/乱序 key 也能正确追平。

3.3 process_sample：改名标记模态来源

process_sample() 把模态目录名注入文件名，例如：

id_0.jpg → id_0.image.jpg id_0.json → id_0.caption.json

这样后续group_by_keys就能按__key__前缀（id_0）分组，把不同模态合并进同一个字典：

{"__key__": "id_0", "image.jpg": <bytes>, "caption.json": <bytes>}

3.4 后处理：shuffle → 解码 → 增强

WebDataset 类把上述流水线包装成DataPipeline，build_dataset() 再依次追加：

最终每个样本长这样：{"__key__": "id_0", "image": <PIL>, "caption": <str>}，经 collate 组成训练 batch。

四、核心机制②：Shardlist 分片分发

分片（shardlist）解决的问题是：上百万个 tar 怎么分给 N 台机器 × M 张 GPU × K 个 worker，互不重叠、可复现、可续训。

4.1 两级切分：按 rank，再按 worker

ShardlistBasic（distributors/basic.py）的 obtain_url_list() 是整个分发的核心，按顺序做四件事：

① 列表补齐（repeat_url）：用 repeat_list 把 tar 列表循环延长到能被world_size × num_workers整除，保证每个进程分到的分片数量相同，避免 DDP 因 batch 数不一致而挂起。

② 按节点切：urls = urls[rank::world_size]——间隔切片，天然打散。

③ 按 worker 切：urls = urls[worker_id::num_workers]。

④ 可复现 shuffle：random.Random(rank * num_workers + worker_id).shuffle(urls)，每个进程用自己的进程号作种子打乱自己那份——既保证随机性，又保证同一进程重启后拿到相同顺序。

全部 tar 文件 → 按 rank 切（node）→ 按 worker_id 切（进程内）→ 各自种子 shuffle

4.2 断点恢复：checkpoint 里记下"读到哪了"

ShardlistBasic支持从 checkpoint 恢复：通过环境变量WDS_START_INDEX/WDS_EPOCH_NUM还原 set_epoch() 的epoch与start_index，然后跳过每个 worker 已消费的分片：

self.start_index = int(os.environ.get("WDS_START_INDEX", 0)) // self.chunk_size start_index_per_worker = self.start_index // num_workers urls = urls[start_index_per_worker:] # 跳过已读分片

配合chunk_size（每个 tar 的样本数），"已读样本数"就能精确换算成"已读分片数"。

4.3 并行训练：同组 GPU 必须吃同一份数据

当启用 tensor / context parallelism 时，一个 batch 会被组内多张卡共同处理，组内所有 rank 必须看到相同数据。ShardlistBasicParallelSync（distributors/parallel_sync_basic.py）继承自ShardlistBasic，只改两点：

• 用数据并行组编号group_id替代rank做切分：urls = urls[group_id::num_groups]；
• 打乱种子改用group_id * num_workers + worker_id。

效果：同组 GPU 拿到相同分片 + 相同顺序，不同组之间仍然互斥。业务层在 get_image_dataset() 中自动选择分发器——检测到 CP/TP > 1 就切到ShardlistBasicParallelSync，并强制detshuffle=True保证组内顺序完全一致。

五、配套工具：如何把原始数据分片成 WebDataset

仓库提供了 scripts/sharding_wds.py 本地分片管道，它正是"孪生 tar"布局的生产者：

1. 自动发现：扫描--input-dir下的image/、caption/、umt5/等子文件夹；
2. basename 交集：以第一个文件夹为基准，与各文件夹按文件主名求交集，只保留各模态齐全的样本——这就是简单模式同步能成立的前提；
3. 智能分片：按--max-shard-size-mb（默认 500MB）按图像实际大小智能分组，或按固定--shard-size数量分片；
4. 同名打包：每个数据模态各写一个同名 tar，例如image/part_000000/000000.tar与caption/part_000000/000000.tar；
5. 宽高比分桶：可选--classify-aspect-ratio，按最近邻比例分桶（aspect_ratio_16_9等），让训练时按桶采样更均衡；
6. 生成清单：并行统计每个 tar 的样本数，写出wdinfo.json，供第二节所述的加载流程使用。

六、关键文件速查表

七、总结：这套设计好在哪？

🎯零 join 成本的多模态配对：同名 tar + key 前缀对齐，不需要任何数据库或索引服务，简单模式下连zip()都不用；

🎯两级互斥分片：rank × worker 间隔切片 + 进程号种子 shuffle，分片互不重叠、顺序可复现；

🎯生产级鲁棒性：rank 0 解析广播避免存储打爆、S3 指数退避重试、组内detshuffle同步、WDS_START_INDEX断点恢复——让 4K 级大规模训练既喂得饱又随时可续。

理解了这套机制，你不仅能读懂 PiD 的训练数据流，也能把它作为模板，套用到自己项目的 WebDataset 数据管道上。