File size: 2,399 Bytes
2ae6691 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | #!/bin/bash
# One-time: fresh data for the longer Stage 1 (~300B tokens).
# English: fineweb-edu groups 007-009 -> new files fw007-fw009 (20B each)
# Code: 5 languages continue after the documents already used (appended)
# Math: finemath-3plus continues after the documents already used (appended)
# Run: nohup ./new_data.sh > logs/new_data.log 2>&1 &
cd /fsx/ajay-stu/AR
source .venv/bin/activate
D=/fsx/ajay-stu/ar-data/bin
P=sample/350BT
# 1) files we append to must end on a full document; if not, add one end-of-text id
# (adding is safe while training reads the file; cutting it would not be)
python - <<'PY'
import numpy as np
EOT = 151645
for name in ["stack_python", "stack_javascript", "stack_java", "stack_cpp", "stack_typescript", "finemath3"]:
p = f"/fsx/ajay-stu/ar-data/bin/{name}.bin"
last = int(np.memmap(p, dtype=np.uint32, mode="r")[-1])
if last != EOT:
with open(p, "ab") as f:
np.array([EOT], dtype=np.uint32).tofile(f)
print(f"{name}: added end-of-text after a half-written document")
else:
print(f"{name}: ends on a full document")
PY
# 2) the Python job crashed before printing its skip value: count its documents,
# plus 5% for the ~1% the code filter dropped, so the new part never overlaps
K=$(python -c "
import numpy as np
a = np.memmap('$D/stack_python.bin', dtype=np.uint32, mode='r'); n = 500_000_000
print(sum(int((a[i:i + n] == 151645).sum()) for i in range(0, len(a), n)))")
PYSKIP=$(( 20000 + K * 105 / 100 ))
echo "$(date -u) stack_python: $K documents in file, skip $PYSKIP"
# $1 source $2 out dir $3 tokens to add $4 docs to skip $5 data files (optional)
sub() {
[ -z "$4" ] && { echo "$(date -u) $1: no skip value, NOT submitted"; return; }
echo "$(date -u) $1 -> $2: add $3 tokens, skip $4 ${5:+files $5}"
sbatch --export="ALL,SOURCE=$1,TARGET=$3,OUT_DIR=$2,SKIP_DOCS=$4,DATA_FILES=${5:-}" tokenize.sbatch
sleep 300 # Hugging Face rate limit
}
for g in 007 008 009; do sub fineweb $D/fw$g 20000000000 0 "$P/${g}_*.parquet"; done
sub stack_python $D 10000000000 "$PYSKIP"
sub stack_javascript $D 4000000000 3917625
sub stack_java $D 4000000000 3062240
sub stack_cpp $D 3000000000 978851
sub stack_typescript $D 2000000000 3330360
sub finemath3 $D 16000000000 4601977
echo "$(date -u) all 9 jobs submitted"
|