62 lines
2.3 KiB
Bash
Executable File
62 lines
2.3 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe}"
|
|
IMAGE="${IMAGE:-nvcr.io/nvidia/nemo:26.06}"
|
|
SOURCE_DATA_DIRS="${SOURCE_DATA_DIRS:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701:/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/logic_topup_proof_pile_17b_v1_20260701}"
|
|
WORK_DIR="${WORK_DIR:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_v1}"
|
|
JSONL="${JSONL:-$WORK_DIR/text.jsonl}"
|
|
OUTPUT_PREFIX="${OUTPUT_PREFIX:-$WORK_DIR/laoyao_2b_moe_8192_text_document}"
|
|
TOKENIZER_MODEL="${TOKENIZER_MODEL:-$REPO_ROOT/tokenizer/glm5.2}"
|
|
WORKERS="${WORKERS:-16}"
|
|
MAX_DOCS="${MAX_DOCS:-0}"
|
|
MIN_FREE_GB="${MIN_FREE_GB:-1000}"
|
|
KEEP_JSONL="${KEEP_JSONL:-0}"
|
|
|
|
mkdir -p "$WORK_DIR"
|
|
|
|
available_gb="$(df -BG "$WORK_DIR" | awk 'NR==2 {gsub("G", "", $4); print $4}')"
|
|
if [[ "$MAX_DOCS" == "0" && "$available_gb" -lt "$MIN_FREE_GB" ]]; then
|
|
cat >&2 <<EOF
|
|
Refusing full preprocessing: only ${available_gb}GB free at $WORK_DIR.
|
|
Full 200B-token Megatron indexed output is expected to require hundreds of GB,
|
|
and this script also stages a JSONL export. Set MIN_FREE_GB lower only if you
|
|
have confirmed an output filesystem with enough capacity, or run a bounded
|
|
probe with MAX_DOCS.
|
|
EOF
|
|
exit 2
|
|
fi
|
|
|
|
input_args=()
|
|
IFS=':' read -r -a source_dirs <<< "$SOURCE_DATA_DIRS"
|
|
for source_dir in "${source_dirs[@]}"; do
|
|
input_args+=(--input "${source_dir}/*.parquet")
|
|
done
|
|
|
|
docker run --rm --ipc=host --network=host \
|
|
--ulimit memlock=-1 --ulimit stack=67108864 \
|
|
-v /mnt/beegfs:/mnt/beegfs \
|
|
-w "$REPO_ROOT" \
|
|
"$IMAGE" \
|
|
bash -lc '
|
|
set -euo pipefail
|
|
python3 dataset/pretrain/scripts/export_pretrain_parquet_text_jsonl.py \
|
|
"$@" \
|
|
--output '"$JSONL"' \
|
|
--max-docs '"$MAX_DOCS"'
|
|
python3 /opt/Megatron-Bridge/3rdparty/Megatron-LM/tools/preprocess_data.py \
|
|
--input '"$JSONL"' \
|
|
--json-keys text \
|
|
--tokenizer-type HuggingFaceTokenizer \
|
|
--tokenizer-model '"$TOKENIZER_MODEL"' \
|
|
--append-eod \
|
|
--output-prefix '"$OUTPUT_PREFIX"' \
|
|
--workers '"$WORKERS"'
|
|
if [[ "'"$KEEP_JSONL"'" == "0" ]]; then
|
|
rm -f '"$JSONL"'
|
|
fi
|
|
ls -lh '"${OUTPUT_PREFIX}"'*
|
|
' bash "${input_args[@]}"
|
|
|
|
echo "Megatron indexed dataset prefix: $OUTPUT_PREFIX"
|