跳到主要内容

数据截至 (上游 commit a649de79c14a)

05 · Token 化与混洗写出

这一章讲什么： 产线的最后一公里——把清洗后的文本变成训练框架能直接读的 token 文件。.ds 二进制格式长什么样、loss mask 怎么存、以及最关键的设计：混洗文档不重算 token，只搬字节。

1. 它要解决的小问题​

训练框架（GPT 系）要的不是一堆文本文件，而是：

  • 一个巨大的 token 数组（训练时按偏移直接切序列）；
  • 每个 token 定长编码（uint16 够用就别用 uint32，省一半盘）；
  • 文档之间要打乱——同一网站/同一主题的文档连续出现会伤害训练（inferred：混洗是预训练数据的标准实践，代码注释未解释动机）；
  • 有些 token 不该算 loss（比如 prompt 部分），需要逐 token 的掩码。

朴素做法是「文本全部读进内存 → token 化 → shuffle → 写出」，TB 级数据直接爆内存。DataTrove 的做法是两遍扫描：第一遍顺序 token 化落盘并记下每篇文档的边界，第二遍按边界随机 seek、搬字节。

2. .ds 文件格式​

TokenizedFile（src/datatrove/pipeline/tokens/tokenizer.py:23）一次写出最多四个伴生文件：

文件内容写入处
xxx.dstoken 本体，定长小端整数流write（:133-143）
xxx.ds.index文档边界：uint64 数组，第 i 项 = 第 i 篇文档的结束 token 偏移close（:72-85）
xxx.ds.loss逐 token 的 loss mask，每 token 1 字节布尔write_loss_bytes（:121-131）
xxx.ds.metadata文本小卡片：tokenizer 名、token 总数（人类可读版）write_final_metadata（:247-271）

token 宽度按词表大小自动选：超过 uint16 上限用 4 字节，否则 2 字节（PipelineStepWithTokenizer.token_size，src/datatrove/utils/tokenization.py:60-67）。EOS token 通过 tokenizer 的 post-processor 模板追加（TemplateProcessing(single="$A <EOS>")，:91-98），所以「每篇文档结尾有 EOS」是编码器层面保证的，不靠手写拼接。

index 文件是整套混洗机制的关键：有了「每篇文档的起止偏移」，文档就变成可随机访问的字节段。

3. DocumentTokenizer：token 化 + 文档内混洗​

3.1 直觉演示​

# 示意,非源码
# 第一遍:顺序编码落盘,记边界
doc_ends = []
with open("unshuffled.ds", "wb") as f:
    for batch in batched(docs, 10000):
        for doc, enc in zip(batch, tokenizer.encode_batch(batch)):
            f.write(pack(enc.ids))
            doc_ends.append(f.tell() // token_size)

# 第二遍:按边界随机 seek,搬字节
with open("unshuffled.ds", "rb") as fin, open("shuffled.ds", "wb") as fout:
    for doc_id in rng.permutation(len(doc_ends)):
        start, end = doc_ends[doc_id - 1], doc_ends[doc_id]
        fin.seek(start * token_size)
        fout.write(fin.read((end - start) * token_size))

重点：第二遍没有任何 token 化，只是字节搬运 + 重建 index。

3.2 真实实现​

DocumentTokenizer（tokenizer.py:281）的 run（:417-476）正是这个两遍结构：

  1. write_unshuffled（:378-415）：按 batch_size=10000 批量 encode_batch（Rust 实现的 HF tokenizers，批量才能吃满），逐篇写 TokenizedFile，同时 stat_update("tokens", ...) 记总数；
  2. shuffle_documents=True（默认）时调 outputfile.copy(...)（TokenizedFile.copy，:145-245）按 rand.permutation 的顺序搬文档，写完删原文件（cleanup，:454）。

copy 里有两个为远端文件系统做的细节：

  • 读原文件时显式关掉缓存、块大小只设 50KB（SHUFFLING_READ_BLOCK_SIZE / SHUFFLING_CACHE_TYPE，:15-17）——随机跳读场景下，缓存和预读纯属浪费，注释说明 50KB 约为「过滤后 CC 文档的均值 + 2σ」；
  • max_tokens_per_file 可以把混洗输出切成多个小文件（文件名加 00000 序号，:227-241），S3 上分片上传失败率更低（upload_block_size 参数注释，:325-326）。

3.3 loss mask​

文档 metadata 里带 no_loss_ranges（字符级区间列表）时，get_loss_values（:356-376）把区间映射到 token 区间（encoded.char_to_token），对应位置写 0——训练时这些 token 不参与 loss。mask 跟着 token 一起被 copy 搬运（:223-225），混洗不会错位。写 .ds 时若 loss 段被裁短，token 也同步裁短（:407-409）。

4. 三级混洗与 merger​

文档级混洗只在单个 rank 的输出文件内打乱。要全局打乱，还有两个层次：

层类粒度并行度
文档级DocumentTokenizer（内建）单文档，rank 内多任务
跨文件级DocumentTokenizerMerger单文档或定长块，全数据集单任务
上下文窗口级DocumentTokenizerContextShuffler定长 token 窗口（默认 2049）多任务（按文件切 shard）

DocumentTokenizerMerger（src/datatrove/pipeline/tokens/merger.py:15）：读入所有 .ds + .index，把「文件 id + 文档」展平成一个全局序列做 permutation（get_ordering，:75-86），然后按新顺序从各文件读字节段写进输出，写满 max_tokens_per_file（默认 100G token）就开新文件（:183-194）。它只能单任务（assert world_size == 1，:100），且 docstring 明确警告在 S3 上随机读会慢、建议本地盘（:19-21）。shuffle_chunk_size 参数把粒度从「单文档」换成「定长 token 块」：chunk_doc_ends（src/datatrove/utils/tokenization.py:25-36）把文档边界按块归组，同一块内的文档一起搬——打乱的随机性换顺序读的性能。

DocumentTokenizerContextShuffler（src/datatrove/pipeline/tokens/context_shuffler.py:13）更激进：完全无视文档边界，按 window_size（默认 2049 token，即训练上下文长度 +1）把整个文件切成窗口、窗口级 permutation 后搬运，本地文件用 mmap 做随机访问（:76-83，代码里 TODO 注明 mmap 只支持本地）。产出就是「训练时连续 token 大概率来自不同来源」的最终形态。

如果只想数 token 不写文件，用 TokensCounter（src/datatrove/pipeline/tokens/counter.py:7）——批量编码后只把 token_count 写进 metadata 并记 stats，FineWeb 在 dedup 前后各挂一个来对比（examples/fineweb.py:163-166）。

5. 关键细节 / 坑​

  • 混洗是两遍 I/O，中间文件最好落本地盘。 local_working_dir 参数（tokenizer.py:314）把未混洗的中间文件写到本地；不给且输出是远端时会警告「会很慢」（:337-344）。
  • 块级混洗会丢掉不足一块的尾巴：chunk_doc_ends 按 range(shuffle_chunk_size, doc_ends[-1] + 1, ...) 切（tokenization.py:28），shuffle_chunk_size 的 docstring 明说「最后一个不满 chunk_size 的块会被丢弃」（tokenizer.py:330-331）。
  • token 宽度在混洗/合并链路里必须一致。merger 从输入的 .metadata 文件读回 token_size（merger.py:114-121）；缺 metadata 时默认 2 字节——混用 tokenizer 会导致错位解读。
  • merger 是单点：全局混洗的代价是一个任务扫完全数据集的字节（虽然是流式的），超大语料这一步要单独给足时间和本地盘。
  • 代码里看不出的：.ds 格式没有自封的描述头（版本、tokenizer 指纹在 .metadata 文本里而非二进制头），跨团队/跨版本混用这批二进制文件时的兼容性约定，代码层面没有强制。

6. 本章代码地图​

主题文件路径符号名
token 文件读写src/datatrove/pipeline/tokens/tokenizer.pyTokenizedFile.write_bytes、TokenizedFile.copy、get_output_filename
token 化主步骤同上DocumentTokenizer.run、write_unshuffled、get_loss_values
tokenizer 加载与 EOSsrc/datatrove/utils/tokenization.pyPipelineStepWithTokenizer.tokenizer、token_size、chunk_doc_ends
跨文件合并混洗src/datatrove/pipeline/tokens/merger.pyDocumentTokenizerMerger.run、get_ordering、load_doc_ends
窗口级混洗src/datatrove/pipeline/tokens/context_shuffler.pyDocumentTokenizerContextShuffler.run
token 计数src/datatrove/pipeline/tokens/counter.pyTokensCounter.run