Distributed Offline Packing

Purpose / 用途

Use this skill when packing a large SFT JSONL (hundreds of thousands to millions of samples) into Energon WebDataset shards at a fixed sequence length, using offline_packing/auto_pipe.sh parallelized across multiple nodes.

当需要把大规模 SFT JSONL（几十万到几百万样本）按固定序列长度打包成 Energon WebDataset shards，并通过 offline_packing/auto_pipe.sh 在多台机器上并行处理时，使用这个 skill。

Prerequisites / 前置条件

• All nodes share the same NFS mount (data + repo + output)
• Same docker image on every node (must contain transformers, energon, project repo)
• offline_packing/auto_pipe.sh and stage scripts s1_split_json_to_samples.py … s4_bins_to_webdataset.py
• Tokenizer + image processor available locally (HF format model dir)
• Source JSONL where each line has images, prompts, captions (multi-turn list-of-lists) and image paths are usable as-is

所有节点共享同一个 NFS（数据+代码+输出）。每台机器使用同一个 docker 镜像，里面要有 transformers、energon、项目代码。需要 offline_packing/auto_pipe.sh 和 s1–s4 四个 stage 脚本。Tokenizer 和 image processor 是本地 HF 格式目录。源 JSONL 每行包含 images / prompts / captions（多轮是 list-of-list），图片路径可直接使用。

Architecture / 架构

Pipeline Stages / 流水线阶段

JSONL (N samples)
  ├─ s1_split_json_to_samples.py   # validate + drop bad/missing-image samples
  │                                # output: per-sample serialized records
  ├─ s2_compute_token_lengths.py   # tokenize prompts/captions, compute image-patch tokens
  │                                # output: length array per sample
  ├─ s3_bin_packing.py             # BFD (Best-Fit-Decreasing) into bins of capacity L
  │                                # output: bin assignment
  └─ s4_bins_to_webdataset.py      # write tar shards + idx + .nv-meta/{dataset.yaml,split.yaml,sample_loader.py}

auto_pipe.sh runs all four stages sequentially on one node for one input file. To use N nodes, split the JSONL into N parts and run auto_pipe.sh independently on each — s3 BFD does NOT shard across nodes, so each node packs its own slice.

auto_pipe.sh 在一台机器上对一个输入文件顺序跑完四个 stage。要用 N 台机器就把 JSONL 切成 N 份，每台独立跑一次 auto_pipe.sh——s3 BFD 不能跨节点共享，所以每台只 pack 自己的那一份。

Key Design Decisions / 关键设计

• Per-node shard prefix: pass distinct --shard-prefix <name>_<a|b|...> so tar files don't collide on shared NFS
• Sample class: choose based on data shape — PackedCaptioningSample for image+text packed turns, MultiMixQASample for QA-style. The class controls the auto-generated sample_loader.py
• --no-npy flag: only set when JSONL has no precomputed patch_positions field. Without --no-npy s2 expects per-sample .npy files; missing files cause noisy warnings AND fall back to slow real-image tokenization
• Sequence length L: scan token lengths first to confirm drop rate is acceptable
• Output layout per node: <output_root>/node_<x>/webdataset/{*.tar, *.idx, .nv-meta/}
• Top-level Metadataset yaml combines all node outputs into one logical dataset

Step-by-step Workflow / 操作步骤

0. Pre-flight: choose L / 预检：选 L

Scan token lengths once on the full JSONL using s2_compute_token_lengths.py (or a quick standalone script). Pick L so that drop rate is acceptable (typically <0.1%).

# Inside container, on any single node
python offline_packing/s2_compute_token_lengths.py \
    --jsonl <path/to/full.jsonl> \
    --tokenizer <path/to/tokenizer> \
    --image-processor Qwen2_5_VLProcessor \
    --factor 48 --min-pixels 3136 --max-pixels 4000000 \
    --output <path/to/token_lens.txt>

# Then quickly inspect distribution (max, p99, count > L) before committing to L

1. Split JSONL across nodes / 切分 JSONL

TOTAL=$(wc -l < full.jsonl)
HALF=$(( (TOTAL + 1) / 2 ))
split -l $HALF -d --additional-suffix=.jsonl full.jsonl part_
# produces part_00.jsonl, part_01.jsonl

For >2 nodes, adjust -l accordingly.

2. Mount NFS on every node / 挂载 NFS

Make sure the same NFS is mounted on every node at the same path so paths in JSONL and outputs match.

3. Launch container on every node / 每台启动容器

Use the project's standard docker image with the repo bind-mounted. Working directory should be the repo root.

4. Run auto_pipe.sh in parallel on each node / 并行启动

On node A:

cd <repo_root>
bash offline_packing/auto_pipe.sh \
    --jsonl <data_root>/part_00.jsonl \
    --tokenizer <tokenizer_path> \
    --image-processor Qwen2_5_VLProcessor \
    --factor 48 --min-pixels 3136 --max-pixels 4000000 \
    --image-root / \
    --sample-class PackedCaptioningSample \
    --shard-prefix <dataset_name>_a \
    --output-dir <output_root>/node_a \
    --seq-len 4096 \
    --no-npy \
    2>&1 | tee <log_dir>/node_a.log

On node B (in parallel):

bash offline_packing/auto_pipe.sh \
    --jsonl <data_root>/part_01.jsonl \
    ... \
    --shard-prefix <dataset_name>_b \
    --output-dir <output_root>/node_b \
    ... \
    2>&1 | tee <log_dir>/node_b.log

[!IMPORTANT]

• --shard-prefix must differ between nodes so tar filenames don't collide
• --output-dir must differ between nodes
• --image-root / if JSONL paths are absolute; otherwise set it to the image root prefix
• Add --no-npy if the JSONL has no patch_positions field

5. Verify outputs / 验证产物

# Tar count per node should match s4 log
ls <output_root>/node_a/webdataset/*.tar | wc -l
ls <output_root>/node_b/webdataset/*.tar | wc -l

# Bin count + capacity utilization printed by s3
grep -E "(bins|util|efficiency)" <log_dir>/node_a.log

# Inspect one tar to confirm sample schema
mkdir -p /tmp/tar_inspect && cd /tmp/tar_inspect
tar -xf <output_root>/node_a/webdataset/<prefix>-000000.tar
ls | head
python -c "import json; d=json.load(open(open(__import__('glob').glob('*.json')[0]).name)); print(list(d.keys()))"

Expected JSON top-level keys for PackedCaptioningSample: images, prompts, captions, sample_count, patch_positions, timestamp_decimal.

patch_positions=[[""]] is normal under --no-npy — the auto-generated sample_loader.py handles it via sample.get(..., None).

6. Write top-level Metadataset yaml / 写顶层 Metadataset yaml

__module__: megatron.energon
__class__: Metadataset
splits:
  train:
    datasets:
      - weight: <num_samples_in_part_00>
        path: <output_root>/node_a/webdataset
        subflavors:
          augmentation: false
      - weight: <num_samples_in_part_01>
        path: <output_root>/node_b/webdataset
        subflavors:
          augmentation: false

Notes / 注意:

• weight is sample-count proportional (use the original input line count of each part, not the bin count)
• Use absolute paths — Energon resolves them as-is
• Add val split only if you actually need one; for SFT-only pipelines omit it
• subflavors is optional; here we mark augmentation: false since data is already packed

Common Pitfalls / 常见坑

Pitfall 1: forgetting --no-npy

Without it s2 hunts for per-sample .npy files and either spams warnings OR falls back to the slow real-image tokenization path. Always check the source JSONL for a patch_positions field first.