AR / new_data.sh

Upload AR-4B pretraining codebase, configs, and cluster scripts

2ae6691 verified 18 days ago

2.4 kB

#!/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"

	#!/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"