73 lines
2.7 KiB
Bash
Executable File
73 lines
2.7 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_direct_smoke_v1}"
|
|
TOKENIZER_MODEL="${TOKENIZER_MODEL:-$REPO_ROOT/tokenizer/glm5.2}"
|
|
OUTPUT_PREFIX_PREFIX="${OUTPUT_PREFIX_PREFIX:-laoyao_2b_moe_8192_direct}"
|
|
TEXT_KEY="${TEXT_KEY:-text}"
|
|
PARALLEL_FILES="${PARALLEL_FILES:-1}"
|
|
WORKERS_PER_FILE="${WORKERS_PER_FILE:-8}"
|
|
BATCH_SIZE="${BATCH_SIZE:-8192}"
|
|
CHUNKSIZE="${CHUNKSIZE:-128}"
|
|
MAX_FILES="${MAX_FILES:-1}"
|
|
MAX_DOCS="${MAX_DOCS:-100000}"
|
|
MAX_SEQ_LEN="${MAX_SEQ_LEN:-65536}"
|
|
MIN_FREE_GB="${MIN_FREE_GB:-20}"
|
|
OVERWRITE="${OVERWRITE:-1}"
|
|
|
|
mkdir -p "$WORK_DIR"
|
|
|
|
available_gb="$(df -BG "$WORK_DIR" | awk 'NR==2 {gsub("G", "", $4); print $4}')"
|
|
if [[ "$MAX_FILES" == "0" && "$MAX_DOCS" == "0" && "$available_gb" -lt "$MIN_FREE_GB" ]]; then
|
|
cat >&2 <<EOF
|
|
Refusing full direct preprocessing: only ${available_gb}GB free at $WORK_DIR.
|
|
Direct Megatron indexed output for 200B GLM5.2 tokens is expected to require
|
|
roughly 800GB plus index/cache overhead. Increase free space or lower MAX_FILES/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
|
|
|
|
max_file_args=()
|
|
if [[ "$MAX_FILES" != "0" ]]; then
|
|
max_file_args+=(--max-files "$MAX_FILES")
|
|
fi
|
|
|
|
overwrite_args=()
|
|
if [[ "$OVERWRITE" == "1" ]]; then
|
|
overwrite_args+=(--overwrite)
|
|
fi
|
|
|
|
docker run --rm --ipc=host --network=host \
|
|
--ulimit memlock=-1 --ulimit stack=67108864 \
|
|
-v /mnt/beegfs:/mnt/beegfs \
|
|
-w "$REPO_ROOT" \
|
|
"$IMAGE" \
|
|
python3 dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py \
|
|
"${input_args[@]}" \
|
|
--output-dir "$WORK_DIR" \
|
|
--manifest "$WORK_DIR/manifest.json" \
|
|
--megatron-dir /opt/Megatron-Bridge/3rdparty/Megatron-LM \
|
|
--tokenizer-type HuggingFaceTokenizer \
|
|
--tokenizer-model "$TOKENIZER_MODEL" \
|
|
--text-key "$TEXT_KEY" \
|
|
--output-prefix-prefix "$OUTPUT_PREFIX_PREFIX" \
|
|
--parallel-files "$PARALLEL_FILES" \
|
|
--workers-per-file "$WORKERS_PER_FILE" \
|
|
--batch-size "$BATCH_SIZE" \
|
|
--chunksize "$CHUNKSIZE" \
|
|
--max-docs "$MAX_DOCS" \
|
|
--max-seq-len "$MAX_SEQ_LEN" \
|
|
"${max_file_args[@]}" \
|
|
"${overwrite_args[@]}"
|
|
|
|
echo "Direct Megatron indexed dataset manifest: $WORK_DIR/manifest.json"
|