HDF5 token 化语料

HDF5 token 化语料

本节摘要:下载的语料要落到训练器能按行速流式读取的布局里。JSONL 在磁盘撑不住 16 个 dataloader worker。HDF5 用可调整大小、分块的整数数据集可以。本节构建流式 token 化进可调整大小 HDF5 数据集、跨多文件分片写、训练时内存映射读、以及产出定长序列带正确打包规则的滑动窗 dataloader。建侧的诚实是关键:可调整数据集易误用——逐文档写则文件碎片到不可用、一次 resize 全写则进程死丢整分片;正解是「缓冲后扩展」,缓冲大小匹配分块大小,分片写把工作拆到多文件使崩溃至多丢一分片。

对应原课程:Phase 19 · Lesson 43 · hdf5-tokenized-corpus(原英文 phases/19-capstone-projects/43-hdf5-tokenized-corpus/docs/en.md)。本节属「预训练/分布式」赛道第二节。

学习目标

阅读完本节,你应当能够:

1. 把文档流式 token 化进带确定性分块的可调整大小 HDF5 整数数据集。
2. 跨多个 HDF5 文件分片写,使失败有界、并行可行。
3. 经 HDF5 页缓存支撑的分块布局把 token 读回,使 dataloader 只在批时拷进批缓冲。
4. 实现滑动窗 dataloader,产出定长训练序列与显式打包规则。

一、问题与直觉

现代语言模型训练跑以每秒几十万样本跨几十 worker 读 token。JSONL 在磁盘在第一次冷缓存页错误就死:JSON 解析器慢、文档边界不可寻址、「定位到样本 4,217,884」要扫整个文件。即便 Parquet(压缩好)也不合适,因为训练器不要列、要带 O(1) 随机访问的扁平 token 流。

HDF5 合适,因为它提供分块、可调整、纯整数的数据集,其分块在读时页缓存友好。训练器要 tokens[3,200,000:3,200,8192] 的切片,HDF5 从页缓存把请求的超平面拷进新分配的 NumPy 数组。成本是一个打开的文件句柄与每 worker 一个分块大小的页缓存足迹,与解码 JSONL 的成本相比可忽略。

可调整 HDF5 的正解

token 数据集用 maxshape=(None,) 与固定 chunks=(chunk_size,) 创建。写按把 token 缓冲到长 chunk_size 的 NumPy 数组进行;缓冲满时数据集恰 resize chunk_size,缓冲写进新范围;分片末尾残余缓冲写进最后部分范围。每次写连续且分块对齐,除最后一次——读者被告知在分片 HDF5 属性里记的 token_count 处截断。

分片写

单个 HDF5 文件是单点故障。流水线并行写分片:第 40 节每个输入分片产出一个 HDF5 输出分片。shards.json 索引每分片记文件路径、token 数、文档数、token 的 sha256。训练器读 shards.json 算全局偏移与校验语料。

二、从零实现

code/main.py 实现:

• tokenize_to_hdf5(jsonl_path, hdf5_path, tokenizer, chunk_size):缓冲后扩展写。
• shard_corpus(input_shards, out_dir):并行写多 HDF5 分片,发 shards.json。
• HDF5TokenDataset(shards_json, context_length):内存映射读,全局偏移。
• make_sliding_window_loader(dataset, batch_size, context_length):滑窗产定长序列。

写侧骨架:

def tokenize_to_hdf5(jsonl_path, hdf5_path, tokenizer, chunk_size=8192):
    with h5py.File(hdf5_path, "w") as f:
        ds = f.create_dataset("tokens", (0,), dtype="int32",
                              maxshape=(None,), chunks=(chunk_size,))
        buf = np.empty(chunk_size, dtype="int32"); n = 0; doc_count = 0
        for line in open(jsonl_path):
            ids = tokenizer.encode(json.loads(line)["text"])
            doc_count += 1
            for t in ids:
                buf[n] = t; n += 1
                if n == chunk_size:                              # 缓冲满 -> resize+写
                    old = ds.shape[0]
                    ds.resize((old + chunk_size,)); ds[old:] = buf
                    n = 0
        if n:                                                     # 残余
            old = ds.shape[0]; ds.resize((old + n,)); ds[old:old+n] = buf[:n]
        f.attrs["token_count"] = ds.shape[0]
        f.attrs["doc_count"] = doc_count
​

读侧滑窗:

class HDF5TokenDataset(IterableDataset):
    def __init__(self, shards_json, context_length):
        self.shards = json.load(open(shards_json))
        self.ctx = context_length
        # 累计每分片 token 数算全局偏移
        self.offsets = list(accumulate([0] + [s["token_count"] for s in self.shards]))
        self.total = self.offsets[-1]

    def __getitem__(self, idx):
        # 定位 idx 落在哪个分片
        si = bisect(self.offsets, idx) - 1
        local = idx - self.offsets[si]
        with h5py.File(self.shards[si]["path"], "r") as f:
            x = f["tokens"][local:local + self.ctx]              # 页缓存 -> NumPy
        y = f["tokens"][local+1:local+1+self.ctx]                # 错位一位目标
        return pad_or_truncate(x, self.ctx), pad_or_truncate(y, self.ctx)
​

设计要点:缓冲后扩展的「缓冲 = 分块大小」是关键——每次 resize 恰一个分块,写连续且对齐,文件无碎片。逐 token resize 会碎片到不可用,一次 resize 全写会让进程死丢整分片。滑窗的「错位一位」目标呼应第 34 节的损失对齐。内存映射读(HDF5 默认 chunked 即 mmap 友好)使 dataloader 只在批时拷贝,worker 间共享页缓存。

三、框架对比

HuggingFace datasets 的 .save_to_disk 用 Arrow 列存,适合表格与变长文本,但对纯 token 流的随机访问不如 HDF5 分块高效。MosaicML 的 StreamingDataset(MDS 格式)、NVIDIA 的 Megatron IndexedDataset(.bin + .idx)思路相同:扁平 token 流 + 索引 + 分片。本节手写让你看清「缓冲后扩展」的碎片控制、分片写的失败隔离、滑窗的打包。生产规模上,数据集常配预取(prefetch)与 worker 亲和性,但 token 数据集的核心布局与本节一致。LM Data Portal、RedPajama-Data 的 token 化器都用类似 HDF5/Arrow 分片。

四、可复用产物

code/main.py:tokenize_to_hdf5、shard_corpus、HDF5TokenDataset、make_sliding_window_loader 均可复用。demo 把小型 JSONL 流式 token 化进 HDF5 分片,发 shards.json,再用滑窗 dataloader 产出定长批,验证 token 数与文档数与清单一致。这套布局是第 42~47 节训练栈的数据底座——大语料训练器直接消费 HDF5TokenDataset。

五、练习

1. 碎片对比:逐 token resize vs 缓冲后扩展,对比产出 HDF5 文件大小与读速。
2. 打包规则:实现「跨文档打包」(把多短文档拼进一个 context_length 窗,加分隔 token)与「不跨文档」(短文档补填充),对比下游损失曲线。
3. 多 worker:把 dataloader 跑 8 worker,确认页缓存共享、无重复读取。
4. 校验往返:token 化后反 token 化,与原始文本比,确认无损(或量化损失)。
5. 动态上下文:支持训练中改变 context_length(短窗预热后切长窗),验证滑窗正确处理。

本节要点回顾

1. JSONL 撑不住 worker:解析慢、边界不可寻址,16 worker 冷缓存就死。
2. HDF5 分块可调整:maxshape=(None,) + chunks=(chunk_size,),页缓存友好。
3. 缓冲后扩展:缓冲 = 分块大小,resize 恰一分块,连续对齐无碎片。
4. 分片写隔离失败:每输入分片一 HDF5 文件,shards.json 记全局偏移与校验。
5. 滑窗错位目标:呼应损失对齐,内存映射读使 dataloader 只在批时拷贝。
6. 数据底座:这套布局是第 42~47 节大语料训练栈的底座。

下一节,我们做「Cosine 学习率与预热」——把第 34 节用过的调度单独抽出来,精确实现并钉住边界。