跳到主要内容

数据截至 (上游 commit 669f534823b0)

01 · 工具箱骨架:标注与决策分离

这一章讲什么: Dolma 全部机制的地基——documents/attributes 分存的数据模型、并行执行框架 BaseParallelProcessor、以及 tagger 运行时如何在一遍扫描里跑完多个 tagger。读完你就理解了整个工具箱的「世界观」。

1. 它要解决的小问题​

清洗语料要跑很多「检查项」:数词数、判语言、查毒性、找 PII、查重复……每项都要把全量数据扫一遍。

天真的做法是每项检查直接改数据(删掉不合格的行)。这有两个致命伤:

  • 不可逆:阈值定错了、规则写错了,被删的数据就没了,只能从头再来。
  • 没法调配方:想对比「阈值 0.3 vs 0.5 哪个训出来好」,就得保存 N 份几十 TB 的数据副本。

Dolma 的答案是:检查只做检查,判决另算。

2. 思路:原文是只读的,检查报告单独存档​

核心设计一句话:

tagger 永远不修改 documents/,只往平行的 attributes/<实验名>/ 目录写「报告」;报告与原文逐行对齐,用 id 关联。

报告里写的是什么?不是「合格/不合格」的判决,而是测量值:一个 [start, end, score] 三元组的列表——text[start:end] 这段文本,在某个维度上得了 score 分。

判决(留/删/替换)被推迟到混合阶段,用声明式表达式去查这些测量值(那是第 03 章的事)。

目录长什么样​

documents/2019-09/0000.jsonl.gz        ← 原文:{id, text, source, metadata}
attributes/gopher_v1/2019-09/0000.jsonl.gz   ← 报告:{id, attributes:{…}}

docs/data-format.md 明确要求:attributes 文件必须和对应 documents 文件行数相同、顺序相同。这是一条强约定,第 03 章的 mixer 靠它做按行归并,并用 id 校验防错位。

属性名里的三段式命名​

多个实验、多个 tagger 的报告不能撞名,所以属性 key 是三段式:

<实验名>__<tagger名>__<信号类型>
gopher_rules__gopher_v1__word_count  →  [[0, 10234, 512]]

这个拼接发生在写出前(python/dolma/core/runtime.py:216,_write_sample_to_streams 内的 f-string);反向解析在 python/dolma/core/data_types.py:205(Span.from_spec,按 __ 切两段)。第 03 章你会在 JSONPath 表达式里反复见到这种长 key。

3. 图示:一遍扫描,多份报告​

documents/x.jsonl.gz
       │  逐行解码(每行只 decode 一次)
       ▼
 ┌──────────────── TaggerProcessor ────────────────┐
 │  row ──► tagger A(gopher)──► spans ─┐           │
 │     ──► tagger B(pii)─────► spans ─┤ 收集       │
 │     ──► tagger C(cld2)────► spans ─┘           │
 │                按输出文件聚合、写一行             │
 └──────┬──────────────────┬──────────────────────┘
        ▼                  ▼
 attributes/A/x.jsonl.gz   attributes/BC/x.jsonl.gz

怎么读: 左边一个输入流,右边 N 个输出流;同一行的 JSON 只被解码一次,被所有 tagger 轮流看。IO(这个尺度下唯一贵的资源)只付一遍。

4. 原理演示:最小 tagger 与并行骨架​

先演一个 tagger 的最小形态(官方文档 docs/taggers.md 里的同款例子,简化):

# 示意,非源码
@add_tagger("random_number_v1")          # 注册进 TaggerRegistry
class RandomNumberTagger(BaseTagger):
    def predict(self, doc: Document) -> DocResult:
        score = random.random()
        # span 覆盖全文,type 是信号名,score 是测量值
        return DocResult(doc, spans=[Span(0, len(doc.text), "random", score)])

重点看:返回值里没有「删」这个动作,只有测量值。BaseTagger.tag(python/dolma/core/taggers.py:53)负责把 predict 的 span 列表转成 {type: [[start,end,score], ...]} 字典(group_output,python/dolma/core/taggers.py:46,score 统一保留 5 位小数,TAGGER_SCORE_PRECISION = 5 在 python/dolma/core/taggers.py:22)。

再演并行框架怎么用:

# 示意,非源码
class MyProcessor(BaseParallelProcessor):
    @classmethod
    def increment_progressbar(cls, queue, /, files=0, documents=0):
        return super().increment_progressbar(queue, files=files, documents=documents)

    @classmethod
    def process_single(cls, source_path, destination_path, queue, **kwargs):
        ...  # 处理一个文件;成功返回后框架自动写 .done.txt

processor = MyProcessor(
    source_prefix="s3://bucket/documents/*.gz",
    destination_prefix="s3://bucket/out",
    metadata_prefix="/tmp/meta",
    num_processes=32,
)
processor()   # 断点续跑:已有 .done.txt 的文件自动跳过

5. 真实实现​

5.1 BaseParallelProcessor:一个文件 = 一个任务​

框架在 python/dolma/core/parallel.py:52。它把并行问题收敛成三件事:

① 枚举任务并对账。 _get_all_paths(python/dolma/core/parallel.py:456)展开 glob、打乱顺序(python/dolma/core/parallel.py:478,注释说明是为了进度条 ETA 更准),再查 metadata 目录里已有的 .done.txt(python/dolma/core/parallel.py:481-484),已完成的直接跳过。

② 多进程分发。 _multiprocessing_run_all(python/dolma/core/parallel.py:369)强制 spawn 方式起 multiprocessing.Pool,每个文件一个 apply_async;kwargs 先 pickle 再传(python/dolma/core/parallel.py:215 解包),避免子进程拿到不可序列化的对象。

③ 成功才盖章。 _process_single_and_save_status(python/dolma/core/parallel.py:202)先跑 process_single,内嵌一个只认 DolmaRetryableFailure 的重试循环(python/dolma/core/parallel.py:217-225),全部成功后才写 .done.txt(python/dolma/core/parallel.py:228-229)。半成品不算数——这就是断点续跑正确性的来源。

进度条走一个 Queue:子进程往队列里扔增量元组,主进程一个线程消费并更新 tqdm(_run_threaded_progressbar,python/dolma/core/parallel.py:251)。子类要声明自己跟踪哪些单位(files/documents/tokens),构造函数里用 inspect 校验签名(python/dolma/core/parallel.py:131-146),防写错。

5.2 TaggerProcessor:一遍扫描跑所有 tagger​

TaggerProcessor.process_single(python/dolma/core/runtime.py:241)是 tagger 命令的热路径:

  • 按 tagger 类型选解码器:任何一个 tagger 继承自 BaseTaggerWithMetadata 就用带 metadata 的 msgspec 解码器,否则用更快的无 metadata 版(python/dolma/core/runtime.py:290-293)。
  • 一行一喂:逐行 decode 后,把所有 tagger 依次调一遍(python/dolma/core/runtime.py:301-311),收集结果交给 _write_sample_to_streams 统一写。
  • 输出文件按路径去重:多个 tagger 写同一文件时共享一个句柄(_make_output_streams,python/dolma/core/runtime.py:158-186)。
  • 进度自适应降频:队列积压超过 CPU 数就把上报间隔翻倍(python/dolma/core/runtime.py:326-330),防止进度队列自己成为瓶颈。
  • 失败分级:S3 偶发的 IncompleteReadError 转成可重试异常;其余按 skip_on_failure 决定跳过整个文件还是致命(python/dolma/core/runtime.py:333-343)。

实验名缺省时,每个 tagger 自己成为实验名,输出路径用占位符 EXPERIMENT_PLACEHOLDER_NAME 逐 tagger 替换(python/dolma/core/runtime.py:135-156,_determine_output_paths_for_taggers)。

5.3 默认寻址:路径就是配置​

create_and_run_tagger 在用户不给 destination 时,直接把文档路径里的 documents 段替换成 attributes/<实验名>(python/dolma/core/runtime.py:41-68,_make_paths_from_substitution;调用处在 python/dolma/core/runtime.py:448)。这就是「目录约定」能够成立的另一半:寻址不靠注册表,靠字符串替换。

6. 关键细节与坑​

  • kwargs 要先 pickle 再进子进程(python/dolma/core/parallel.py:215):多进程边界上一切状态都要可序列化,自定义 tagger 的参数同理。
  • shuffle 的是文件顺序,不是内容(python/dolma/core/parallel.py:478):只为了让进度条估时更准;真正的数据混洗在 tokenize 阶段(第 05 章)。
  • 重试只认 DolmaRetryableFailure:其他异常一次就挂(转成 DolmaError,python/dolma/core/parallel.py:222-225)。写自定义 processor 时别指望泛异常重试。
  • increment_progressbar 签名有硬性约束:第一个参数必须是仅限位置的 queue,其余参数默认值必须为 0,且不能有 **kwargs——构造函数运行时用 inspect 强校验(python/dolma/core/parallel.py:131-146)。
  • span score 只有 5 位小数(python/dolma/core/taggers.py:22):做阈值实验时别把阈值定在小数点后第 6 位。
  • metadata 目录是易失品:不给 metadata_prefix 时会落在临时目录(python/dolma/core/runtime.py:457 的 tempfile.mkdtemp()),断点续跑能力随之消失——长跑任务务必显式指定。
  • 顺带一处代码里可见的蹊跷:create_and_run_tagger 里 TaggerProcessor 的构造缩进在 if isinstance(metadata, str): 分支内(python/dolma/core/runtime.py:458-464),即 metadata 以列表形式传入时这段不会执行。看起来是上游的缺陷,但本拆解不下定论,仅标注此处在该 commit 的原始形态。