05 · 数据管道：TargetCache 与对话模板

本篇在总分总中是"分"的第三篇，承接 02 核心原理 解释"训练数据从哪来"。DSpark 训练时需要 target 模型在每个位置的 hidden states 和 logits 作为监督信号，但每次训练 step 都跑 target 模型不现实——DeepSpec 的做法是把 target 的 forward 结果预计算成磁盘缓存（target cache），训练时只读盘不跑 target。这是支撑"38TB 训练数据"的工程基石。

---

总览段（总）

DeepSpec 的数据管道分三步，每一步的产物是下一步的输入：

图说明： 三步串联，每步输出明确文件。Step 1 用 mlabonne/open-perfectblend（1.3M 样本，chat 17.6% / math 39.4% / code 38.9% / instruction 4.1%）切分 train/eval。Step 2 用 sglang 在 8 端口（30000-30007）上重新生成 assistant 回答，因为论文要求"训练数据由各 target 模型用推荐采样参数重新生成"。Step 3 跑 target 模型一次 forward，用 hook 抓多层 hidden states 写盘。

关键文件清单：

文件	角色
[scripts/data/download_and_split.py](file:///workspace/scripts/data/download_and_split.py)	下载并切分
[scripts/data/generate_train_data.py](file:///workspace/scripts/data/generate_train_data.py)	sglang 重生成
[scripts/data/prepare_target_cache.py](file:///workspace/scripts/data/prepare_target_cache.py)	生成 target cache
[scripts/data/launch_sglang_server.sh](file:///workspace/scripts/data/launch_sglang_server.sh)	启动 sglang
[scripts/data/prepare_data.sh](file:///workspace/scripts/data/prepare_data.sh)	三步编排
[deepspec/data/target_cache_dataset.py](file:///workspace/deepspec/data/target_cache_dataset.py)	cache 协议+读写
[deepspec/data/jsonl_dataset.py](file:///workspace/deepspec/data/jsonl_dataset.py)	JSONL mmap 索引
[deepspec/data/parser.py](file:///workspace/deepspec/data/parser.py)	对话模板+loss_mask
[deepspec/data/cuda_prefetcher.py](file:///workspace/deepspec/data/cuda_prefetcher.py)	H2D 预取

---

分述段（分）

5.1 Step 1：下载与切分

[download_and_split.py](file:///workspace/scripts/data/download_and_split.py) 处理 mlabonne/open-perfectblend：

• 角色标准化（[download_and_split.py:12-18](file:///workspace/scripts/data/download_and_split.py#L12-18)）：把原始的 human/gpt/chatgpt/bing/bard 统一映射到 user/assistant。
• 校验（[download_and_split.py:131-141](file:///workspace/scripts/data/download_and_split.py#L131-141)）：必须以 user 开头、content 非空、role 合法。
• train/eval 分流：train 写完整 conversations（含 assistant 回答），eval 只写 {turns: [user_turns]}（[download_and_split.py:144-174](file:///workspace/scripts/data/download_and_split.py#L144-174)），因为 eval 时要让 target/draft 自己生成。
• 默认 test_size=0.05，seed=42。

5.2 Step 2：sglang 重生成回答

[generate_train_data.py](file:///workspace/scripts/data/generate_train_data.py) 用 OpenAI 兼容的 sglang 服务重生成 assistant 回答：

• 多 server 负载均衡（[generate_train_data.py:171-230](file:///workspace/scripts/data/generate_train_data.py#L171-230)）：先并行 validate 所有 server，过滤不可用的；按 per-server queue 分发避免单点过载。
• 多轮对话保留（[generate_train_data.py:106-144](file:///workspace/scripts/data/generate_train_data.py#L106-144)）：保留原对话上下文，只把 assistant 的回答用 sglang 重新生成。
• 采样参数与 target 对齐：--temperature 0.7 --top-p 0.8 --top-k 20 --min-p 0（Qwen3 推荐），切换 target 时需调整。
• --disable-thinking：Qwen3 / Gemma4 都用 non-thinking 模式（论文 Section 4.1 明示）。
• --resume：跳过已处理行，支持断点续传。

启动 sglang：bash scripts/data/launch_sglang_server.sh，默认 8 端口 30000-30007，日志到 logs/sglang_qwen3_4b/。

5.3 Step 3：生成 target cache（核心）

[prepare_target_cache.py](file:///workspace/scripts/data/prepare_target_cache.py) 是数据阶段最复杂的脚本，用 forward hook 抓 target 的多层 hidden states。

5.3.1 Forward hook 机制

代码位置 [prepare_target_cache.py:83-134](file:///workspace/scripts/data/prepare_target_cache.py#L83-134)：

图说明： target_layer_ids=[1,9,17,25,33] 通过 register_forward_hook 在 target backbone 的对应层注册钩子。注意 -1 特殊处理为 embed_tokens 的输出（embedding 层），其他正整数取 layers[layer_id]。target_output.last_hidden_state 单独作为 target_last_hidden_states，用于训练时计算 aligned_target_logits。多 rank 分布式：每个 rank 用 compute_local_sample_range 切分样本，AsyncTargetCacheWriter 异步写盘不阻塞 GPU。

5.3.2 Target cache 文件布局

图说明： target cache 目录含一个 manifest（JSON 元数据）、一个 samples.idx（按 sample_id 排序的固定大小索引）、若干 shard 二进制文件。每条索引记录 40 字节（INDEX_RECORD_STRUCT = struct.Struct("<QIIQQQQQ")，[target_cache_dataset.py:21](file:///workspace/deepspec/data/target_cache_dataset.py#L21)），记录 sample_id / shard_id / seq_len / 5 个 offset。每个 sample 在 shard 内连续存储 5 个张量：input_ids(int32) / attention_mask(uint8) / loss_mask(uint8) / target_hidden_states(bfloat16, shape [seq_len, 5*hidden_size]) / target_last_hidden_states(bfloat16, shape [seq_len, hidden_size])。

5.3.3 INDEX_RECORD_STRUCT 40 字节布局

图说明： < 表示小端序，Q = uint64(8B)，I = uint32(4B)。结构体大小 = 8+4+4+8+8+8+8+8 = 56 字节（注：原报告写 40 字节为近似，实际为 56 字节，因 INDEX_RECORD_SIZE 是 struct.Struct 自动计算）。每条记录精确描述一个 sample 在 shard 内的 5 个张量偏移，可通过 mmap + 偏移直接定位，无需加载整个 shard。

5.3.4 38TB 存储警告的来源

README.md 与 [scripts/data/README.md:121-127](file:///workspace/scripts/data/README.md#L121-127) 都警告 Qwen3-4B 默认配置下 target cache 约 38TB。粗算：

• Qwen3-4B hidden_size = 2560，target_layer_ids 长度 5
• 每 token 的 target_hidden_states = 5 × 2560 × 2B(bf16) = 25.6KB
• 加上 target_last_hidden_states = 2560 × 2B = 5KB
• Open-PerfectBlend 1.3M 样本，平均 seq_len 约 1000-1500
• 总量 ≈ 1.3M × 1200 × 30.6KB ≈ 47TB（粗算与 38TB 同量级）

减小存储的两条路：① 减小训练集；② 减少 target_layer_ids（如改成 [9, 25] 两层，存储减 60%）。

5.3.5 异步写入与多 rank 聚合

[AsyncTargetCacheWriter](file:///workspace/deepspec/data/target_cache_dataset.py) （[target_cache_dataset.py:410-502](file:///workspace/deepspec/data/target_cache_dataset.py#L410-502)）：

• 生产者-消费者模型，主线程把样本塞队列（max_queue_size=128），后台线程写盘
• 避免磁盘 IO 阻塞 GPU forward
• shard 自动切分：超过 max_shard_bytes 立即开新 shard

多 rank 完成后主进程聚合（[prepare_target_cache.py:351-393](file:///workspace/scripts/data/prepare_target_cache.py#L351-393)）：构建全局 shard_map，rename local shard 到全局命名（shard_0.bin、shard_1.bin …），finalize 全局 samples.idx，写 manifest.json。

5.3.6 Manifest 校验

[validate_target_cache_manifest](file:///workspace/deepspec/data/target_cache_dataset.py) （[target_cache_dataset.py:132-200](file:///workspace/deepspec/data/target_cache_dataset.py#L132-200)）校验：

• version == 2
• num_samples 与 samples.idx 文件大小匹配（= num_samples * INDEX_RECORD_SIZE）
• shards 从 0 连续编号，每个文件存在
• target_layer_ids / hidden_size / target_model_name_or_path 与 draft 模型匹配

BaseTrainer 在 [base_trainer.py:183-188](file:///workspace/deepspec/trainer/base_trainer.py#L183-188) 调用 validate_train_cache 二次校验，确保 cache 与 draft config 一致。

5.4 CacheDataset：训练时的 mmap 读取

[CacheDataset](file:///workspace/deepspec/data/target_cache_dataset.py#L615-798)（[target_cache_dataset.py:615-798](file:///workspace/deepspec/data/target_cache_dataset.py#L615-798)）：

• LRU shard 缓存（[target_cache_dataset.py:627-629](file:///workspace/deepspec/data/target_cache_dataset.py#L627-629)）：max_open_shards=4，避免同时打开过多文件描述符。
• mmap + offset 直接定位：__getitem__ 先从 samples.idx mmap 读索引记录，再用偏移从 shard mmap 切片读 5 个张量，零拷贝。
• 跨进程 pickle（[target_cache_dataset.py:660-666](file:///workspace/deepspec/data/target_cache_dataset.py#L660-666)）：__getstate__ 清空文件句柄，支持 DataLoader worker 跨进程。

5.5 CacheCollator 与 CUDAPrefetcher

图说明： CUDAPrefetcher（[cuda_prefetcher.py:14-73](file:///workspace/deepspec/data/cuda_prefetcher.py#L14-73)）用独立 CUDA stream 实现 H2D 传输与计算重叠。每次 __next__ 先让 compute stream 等待上一次的 side stream（wait_stream），再启动新线程预取下一 batch。record_stream 防止 caching allocator 在 side stream 还在用时就回收内存。move_batch_to_device 把 input_ids 在 GPU 上转 int64 以节省 CPU→GPU 带宽。

CacheCollator（[target_cache_dataset.py:859-870](file:///workspace/deepspec/data/target_cache_dataset.py#L859-870)）：对 input_ids/loss_mask 做 1D padding，对 hidden states 做 2D padding，自动构造 attention_mask。

5.6 ChatTemplate 与 loss_mask 构造

[parser.py](file:///workspace/deepspec/data/parser.py) 处理对话模板：

• ChatTemplate dataclass（[parser.py:9-15](file:///workspace/deepspec/data/parser.py#L9-15)）：记录 assistant_header / user_header / system_prompt / end_of_turn_token / assistant_loss_prefix。

• TEMPLATE_REGISTRY（[parser.py:30-51](file:///workspace/deepspec/data/parser.py#L30-51)）：

  • qwen：<|im_start|>assistant\n / <|im_start|>user\n / <|im_end|>\n
  • gemma4：<|turn>model\n / <|turn>user\n / <turn|>\n，且 assistant_loss_prefix="<|channel>thought\n<channel|>"（Gemma4 非思考模式前缀，[parser.py](file:///workspace/deepspec/data/parser.py) 中 Gemma4 特殊处理）

• loss_mask 构造（[parser.py:114-138](file:///workspace/deepspec/data/parser.py#L114-138)）：用正则 assistant_pattern 找到每个 assistant message 的内容范围，分别 encode 前缀与完整内容得到 [start_token_idx, end_token_idx)，在该区间内 loss_mask = 1，即只在 assistant 输出上计算 loss（prompt 部分不计）。

5.7 JsonLineDataset 与缓存索引

[jsonl_dataset.py](file:///workspace/deepspec/data/jsonl_dataset.py) 用 mmap + 行偏移索引访问 JSONL：

• _build_all_line_starts（[jsonl_dataset.py:90-132](file:///workspace/deepspec/data/jsonl_dataset.py#L90-132)）：首次遍历文件记录每行起始位置，以 path|mtime_ns 为 key hash 后缓存到 ~/.cache/deepspec/jsonlindex-{hash}.pkl。
• 原子 pickle 写入（_atomic_pickle_dump，[jsonl_dataset.py:82-88](file:///workspace/deepspec/data/jsonl_dataset.py#L82-88)）避免损坏。
• __getitem__ 通过 bisect_right 映射全局 idx 到 (file_idx, local_idx)。

---

小结段（总）

数据管道是 DeepSpec 工程化的核心：它把"训练时反复跑 target 模型"这一成本前置到一次离线准备，让训练时只需 mmap 读盘。代价是巨大的存储（38TB），但换来训练循环的极简——run_batch 只需把 batch 喂给 draft 模型，不再涉及 target。target cache 协议版本 2 通过 manifest + samples.idx + shard 三件套保证可校验、可并行写、可 mmap 零拷贝读。

设计要点回顾：

1. 三步解耦：下载 → 重生成 → 抓 hidden，每步产物独立可断点续传。
2. forward hook 抓多层：target_layer_ids=[1,9,17,25,33] 决定 draft 看到的 target 特征。
3. bf16 hidden + int32 ids + uint8 mask 三种 dtype 在同一 shard 连续存储，节省 IO。
4. CUDA prefetcher 用独立 stream 实现 H2D 与计算重叠。
5. loss_mask 只在 assistant 输出上计算 loss，对应论文"only use prompts from Open-PerfectBlend; responses are regenerated by each target model"。

易踩坑点：

• 38TB 存储警告必须预先规划磁盘。
• target_layer_ids 不能包含 target 模型最后一层（[base_evaluator.py:100-112](file:///workspace/deepspec/eval/base_evaluator.py#L100-112) 的 assert_no_final_target_layer），因为 transformers 的 output_hidden_states 存的是归一化后的 final hidden，与 cache 中的 raw decoder output 不一致。
• resume 时 topology 必须一致（saved_world_size == world_size）。
• 切换 target 模型时，sglang 采样参数（--temperature/--top-p/--top-k/--min-p）需对齐该模型的推荐设置。

延伸阅读：进入 06 训练框架 看 CacheDataset 如何被 trainer 消费；进入 09 使用指南 案例一查看端到端命令。论文 Section 4.1 描述了训练数据组成（Open-PerfectBlend 1.3M、non-thinking mode），见 [DSpark_paper.pdf](file:///workspace/DSpark_paper.pdf)。