数据流水线：indexed dataset 与 mmap

6.1 为什么 Megatron 不直接用 HuggingFace datasets

百亿 token 级别的预训练，数据 IO 不能马虎：

• HF datasets 默认从 Arrow / parquet 一条条读，随机访问慢；
• 万亿 token 不可能全部装内存，必须 mmap；
• 跨 GPU / 跨节点要保证每个 rank 读到不重叠的 sample；
• 多个数据源（书 / 网页 / 代码）需要按比例混合。

Megatron 自家造了一套"二进制 + 索引"的数据格式，叫 indexed_dataset。理解它是上手 Megatron 数据 pipeline 的第一步。

6.2 数据流水线全景

6.3 .bin / .idx 文件格式

每个 dataset 对应两个二进制文件：

.idx 文件本身很小（每文档 24 字节），加载时一次性读进内存；.bin 文件用 mmap 映射，不真正读盘，访问哪页才 page fault。这套做法对万亿 token 数据集是唯一可行的。

.idx 文件头格式

magic   8 bytes   b"MMIDIDX\x00\x00"
version 8 bytes   uint64 = 1
dtype   1 byte    uint8 (1=uint8, 2=int8, 3=int16, 4=int32, 5=int64, 6=float, 7=double, 8=uint16)
num_doc 8 bytes   uint64
num_seq 8 bytes   uint64
... 后跟 sizes / pointers / doc_idx 数组

用 Python 直接读：

from megatron.core.datasets.indexed_dataset import IndexedDataset
ds = IndexedDataset("path/to/my_dataset_text_document", multimodal=False)
print(len(ds))                  # 文档数
print(ds[0][:50])               # 第一个文档前 50 token

6.4 把原始数据转成 .bin / .idx

tools/preprocess_data.py 是官方工具。最小用法：

# 1. 把语料写成 jsonl，每行一条文档
echo '{"text": "你好，这是第一条文档。"}' > corpus.jsonl
echo '{"text": "第二条文档，含中文和英文 mixed."}' >> corpus.jsonl

# 2. 预处理（GPT-2 BPE 例子）
python tools/preprocess_data.py \
    --input corpus.jsonl \
    --output-prefix my_corpus \
    --tokenizer-type GPT2BPETokenizer \
    --vocab-file gpt2-vocab.json \
    --merge-file gpt2-merges.txt \
    --append-eod \
    --workers 32 \
    --chunk-size 25

产物：my_corpus_text_document.bin + my_corpus_text_document.idx。

关键 flag

6.5 在训练脚本里使用

--data-path /path/to/my_corpus_text_document     # ★ 注意：不带 .bin / .idx 后缀
--vocab-file gpt2-vocab.json
--merge-file gpt2-merges.txt
--tokenizer-type GPT2BPETokenizer
--split "949,50,1"

Megatron 会自动找 .bin 和 .idx 一对文件。--split 把整个 dataset 按比例切成训/验/测三段（按 sequence 切，不会拆文档）。

6.6 BlendedDataset：多源加权混合

真实预训练数据通常是 "30% 书 + 50% 网页 + 15% 代码 + 5% 中文"。Megatron 通过 BlendedMegatronDatasetBuilder 实现：

--data-path 0.3 /data/books_text_document \
            0.5 /data/web_text_document \
            0.15 /data/code_text_document \
            0.05 /data/zh_text_document

权重和会被 Megatron 归一化。底层是 BlendedDataset：

# megatron/core/datasets/blended_dataset.py
class BlendedDataset:
    def __init__(self, datasets, weights, size, ...):
        # 用 numpy.random 按 weights 抽 dataset，再从该 dataset 抽 sample
        self.dataset_idx, self.sample_idx = build_blending_indices(...)

6.7 GPTDatasetConfig：拆样本的"边界规则"

megatron/core/datasets/gpt_dataset.py 里的 GPTDataset 把"长 token 流"切成等长 sequence。关键配置：

"重置 position id + 遮断 attention" 是关键 trick：当一个 2048-seq 跨了多个短文档时，避免文档 A 的 token 看到文档 B 的 token，这能让长 seq 训练表现接近"每条样本独立"。

6.8 MMapIndexedDataset 内部

megatron/core/datasets/indexed_dataset.py 的 __getitem__ 大致是：

def __getitem__(self, idx):
    ptr     = self.pointers[idx]      # 该文档在 .bin 中的起点
    size    = self.sizes[idx]         # 该文档 token 数
    np_array = np.frombuffer(
        self._bin_buffer,             # mmap 出来的 buffer
        dtype=self._dtype,
        count=size,
        offset=ptr,
    )
    return np_array                    # 不发生真实 IO，直到访问元素

这套 mmap 路径的性能在第一次 epoch 会有一次"page fault 风暴"，之后 OS 把热数据放到 page cache，第二次 epoch 几乎全是内存访问。

6.9 Sampler：跨 rank 不重叠

Megatron 的 MegatronSampler 生成全局确定性的 sample 索引序列，然后按 (consumed_samples, micro_batch_size) 切分给每个 rank。重要后果：

• 同一 checkpoint 恢复时 sample 顺序完全一致（含 RNG）；
• 跨 rank 不重叠；
• 跨 dp_rank 同 step 看到不同 sample。

6.10 长文档训练：packing 思路

Megatron 的 indexed dataset 天然支持跨文档拼接——把若干短文档头尾连起来填满 --seq-length。但跨文档不希望 attention 串味，所以靠 reset_attention_mask=True 做内部 mask。

当你想严格"每条 sample 一个文档"，可以预处理时按 max-length 切：

python tools/preprocess_data.py \
    --input long_docs.jsonl \
    --output-prefix long \
    --split-sentences \
    --max-doc-length 8192 \
    ...

6.11 数据 cache

Sampler 第一次启动会算出 sample index 数组并存到 --data-cache-path。后续训练或恢复直接读 cache，省 10+ 分钟。如果换了 --seed、--global-batch-size、--train-iters，cache 失效，必须删掉重算。

rm -rf $DATA_CACHE_PATH       # 改完关键参数务必清

6.12 多模态数据

Megatron 0.10+ 支持多模态 indexed dataset（multimodal_dataset.py），把 image / video 特征也按二进制存。tools/preprocess_mmdata.py 是入口。多模态 dataset 训练时通过 --image-data-path 等 flag 引用。

6.13 常见错误

6.14 一份"从 jsonl 到训练"的完整流程

# 1. 假设你已有 corpus.jsonl
# 2. 下载 GPT-2 vocab
wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json
wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt

# 3. 预处理
python tools/preprocess_data.py \
    --input corpus.jsonl \
    --output-prefix my_corpus \
    --vocab-file gpt2-vocab.json \
    --merge-file gpt2-merges.txt \
    --tokenizer-type GPT2BPETokenizer \
    --append-eod \
    --workers 32

# 4. 训练
torchrun --nproc_per_node=8 pretrain_gpt.py \
    --num-layers 12 --hidden-size 768 --num-attention-heads 12 \
    --seq-length 2048 --max-position-embeddings 2048 \
    --micro-batch-size 4 --global-batch-size 32 \
    --train-iters 1000 \
    --data-path my_corpus_text_document \
    --vocab-file gpt2-vocab.json \
    --merge-file gpt2-merges.txt \
    --tokenizer-type GPT2BPETokenizer \
    --split 949,50,1 \
    --lr 1e-4 --min-lr 1e-5 --lr-warmup-fraction 0.01 \
    --bf16 \
    --save ./ckpt --save-interval 500