CookLLM文档数据流水线
会员专享理解 Parquet shard 如何变成 input_ids、labels 和 attention_mask
这一节把前面几章串起来：预训练数据已经被切成 Parquet shard，tokenizer 已经能够把文本变成 token id，模型也已经定义好了。中间缺的就是数据流水线。
在 cookllm-bento 里，预训练 batch 的生成过程可以概括成：
Pretrain Data Pipeline
How a Parquet row becomes a training batch for next-token prediction.
1
Parquet shard
read text column from row groups
list[str]
2
Batch tokenization
prepend BOS, encode texts in parallel
token ids
3
Truncate to max_length
keep the first 512 tokens by default
sample
4
input_ids / labels
same tokens; shifting happens inside the model
tensors
5
Batch padding
pad input_ids, mask labels with -100
attention_mask
6
PretrainModule.training_step()
forward batch and log train/loss
loss
这条链路由两个文件负责：
| 层级 | 文件 | 作用 |
|---|---|---|
| DataModule | src/datamodule/pretrain_datamodule.py | 解析数据路径、加载 tokenizer、划分 train/val、创建 DataLoader |
| Dataset | src/dataset/pretrain.py | 流式读取 Parquet、批量 tokenization、生成单条训练样本 |
配置入口
预训练样例使用的数据配置核心字段如下：
data:
data_path: pretrain_data/fineweb_shards
tokenizer_path: tokens
max_length: 512
batch_size: 96
num_workers: 4
tokenizer_batch_size: 128
tokenizer_num_threads: 8
val_files: 10| 配置 | 含义 |
|---|---|
data_path | 预训练 Parquet shard 目录 |
tokenizer_path | tokenizer 文件目录，默认读取项目根目录下的 tokens/ |
max_length | 每条样本最多保留多少个 token |
batch_size | DataLoader 每次返回多少条样本 |
num_workers | DataLoader worker 数量 |
tokenizer_batch_size | 每次送进 tokenizer 的文本条数 |
tokenizer_num_threads | tokenizer 内部并行线程数 |
val_files | 随机保留多少个 Parquet 文件做验证集 |
登录以继续阅读
这是一篇付费内容，请登录您的账户以访问完整内容。