AR / swap2.sh

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

2ae6691 verified 20 days ago

2.71 kB

#!/bin/bash
# One-time switch to the new loop job + Stage 1 extension to 200B tokens.
cd /fsx/ajay-stu/AR
source .venv/bin/activate

# 1) wait for the 4 tokenizing jobs, then check the files
while [ -n "$(squeue -u $USER -h -n ar-tok)" ]; do sleep 300; done
for g in 003 004 005 006; do
  [ -s /fsx/ajay-stu/ar-data/bin/fw$g/fineweb.bin ] || { echo "$(date) fw$g missing, aborting"; exit 1; }
done
echo "$(date) new data ready"

# 2) new settings: 200B tokens (end step 381,000) + new data mix
python - << 'PY'
import glob, os
d = "/fsx/ajay-stu/ar-data/bin/"
eng = [d + "fineweb.bin"] + sorted(glob.glob(d + "fw*/fineweb.bin"))
code = [f"{d}stack_{l}.bin" for l in
        ["python", "javascript", "java", "cpp", "typescript", "go", "rust", "shell", "sql"]]
math = [d + "finemath.bin", d + "finemath3.bin"]
def split(files, share):
    s = [os.path.getsize(f) for f in files]
    return [round(share * x / sum(s), 4) for x in s]
files = eng + code + math
w = split(eng, 0.65) + split(code, 0.20) + [0.09, 0.06]
out = []
for l in open("run_args.sh").read().splitlines():
    if l.startswith("BIN="): l = 'BIN="' + " ".join(files) + '"'
    elif l.startswith("WEIGHTS="): l = 'WEIGHTS="' + " ".join(map(str, w)) + '"'
    elif l.startswith("STOP_AT="): l = "STOP_AT=381000"
    elif l.startswith("LR_HORIZON="): l = "LR_HORIZON=381000"
    out.append(l)
open("run_args.sh", "w").write("\n".join(out) + "\n")
print("settings updated:", len(files), "files, stop at 381000")
PY
sed -i 's/ckpt_185999/ckpt_380999/' watchdog.sh     # watchdog's new "finished" marker

# 3) pause the watchdog so it can't launch a duplicate during the switch
pkill -f watchdog.sh
J=$(squeue -u $USER -h -n ar-real -o %i | head -1)
if [ -z "$J" ]; then
  echo "$(date) training not running; watchdog will start the new job"
  nohup ./watchdog.sh >> logs/watchdog.log 2>&1 &
  exit 0
fi
NODES=$(squeue -j $J -h -o %N)

# 4) after a checkpoint, switch only if no waiting job could take our nodes (max ~6 tries)
tries=0
while true; do
  N=$(grep -c "\[ckpt\] saved" logs/real_$J.log)
  while [ "$(grep -c "\[ckpt\] saved" logs/real_$J.log)" -eq "$N" ]; do
    [ -z "$(squeue -j $J -h 2>/dev/null)" ] && break 2    # old job died meanwhile
    sleep 30
  done
  tries=$((tries + 1))
  risky=$(squeue -p dev -t PD -h -o "%u %D" | awk -v me=$USER '$1 != me && $2 <= 4' | wc -l)
  if [ "$risky" -eq 0 ] || [ $tries -ge 6 ]; then break; fi
  echo "$(date) $risky waiting jobs could take our nodes; trying at the next checkpoint"
done
sleep 30
scancel $J 2>/dev/null
sleep 20
sbatch --nodelist=$NODES train_real.sbatch
echo "$(date) switched: old job $J stopped, new loop job submitted on $NODES"
nohup ./watchdog.sh >> logs/watchdog.log 2>&1 &

	#!/bin/bash
	# One-time switch to the new loop job + Stage 1 extension to 200B tokens.
	cd /fsx/ajay-stu/AR
	source .venv/bin/activate

	# 1) wait for the 4 tokenizing jobs, then check the files
	while [ -n "$(squeue -u $USER -h -n ar-tok)" ]; do sleep 300; done
	for g in 003 004 005 006; do
	[ -s /fsx/ajay-stu/ar-data/bin/fw$g/fineweb.bin ] \|\| { echo "$(date) fw$g missing, aborting"; exit 1; }
	done
	echo "$(date) new data ready"

	# 2) new settings: 200B tokens (end step 381,000) + new data mix
	python - << 'PY'
	import glob, os
	d = "/fsx/ajay-stu/ar-data/bin/"
	eng = [d + "fineweb.bin"] + sorted(glob.glob(d + "fw*/fineweb.bin"))
	code = [f"{d}stack_{l}.bin" for l in
	["python", "javascript", "java", "cpp", "typescript", "go", "rust", "shell", "sql"]]
	math = [d + "finemath.bin", d + "finemath3.bin"]
	def split(files, share):
	s = [os.path.getsize(f) for f in files]
	return [round(share * x / sum(s), 4) for x in s]
	files = eng + code + math
	w = split(eng, 0.65) + split(code, 0.20) + [0.09, 0.06]
	out = []
	for l in open("run_args.sh").read().splitlines():
	if l.startswith("BIN="): l = 'BIN="' + " ".join(files) + '"'
	elif l.startswith("WEIGHTS="): l = 'WEIGHTS="' + " ".join(map(str, w)) + '"'
	elif l.startswith("STOP_AT="): l = "STOP_AT=381000"
	elif l.startswith("LR_HORIZON="): l = "LR_HORIZON=381000"
	out.append(l)
	open("run_args.sh", "w").write("\n".join(out) + "\n")
	print("settings updated:", len(files), "files, stop at 381000")
	PY
	sed -i 's/ckpt_185999/ckpt_380999/' watchdog.sh # watchdog's new "finished" marker

	# 3) pause the watchdog so it can't launch a duplicate during the switch
	pkill -f watchdog.sh
	J=$(squeue -u $USER -h -n ar-real -o %i \| head -1)
	if [ -z "$J" ]; then
	echo "$(date) training not running; watchdog will start the new job"
	nohup ./watchdog.sh >> logs/watchdog.log 2>&1 &
	exit 0
	fi
	NODES=$(squeue -j $J -h -o %N)

	# 4) after a checkpoint, switch only if no waiting job could take our nodes (max ~6 tries)
	tries=0
	while true; do
	N=$(grep -c "\[ckpt\] saved" logs/real_$J.log)
	while [ "$(grep -c "\[ckpt\] saved" logs/real_$J.log)" -eq "$N" ]; do
	[ -z "$(squeue -j $J -h 2>/dev/null)" ] && break 2 # old job died meanwhile
	sleep 30
	done
	tries=$((tries + 1))
	risky=$(squeue -p dev -t PD -h -o "%u %D" \| awk -v me=$USER '$1 != me && $2 <= 4' \| wc -l)
	if [ "$risky" -eq 0 ] \|\| [ $tries -ge 6 ]; then break; fi
	echo "$(date) $risky waiting jobs could take our nodes; trying at the next checkpoint"
	done
	sleep 30
	scancel $J 2>/dev/null
	sleep 20
	sbatch --nodelist=$NODES train_real.sbatch
	echo "$(date) switched: old job $J stopped, new loop job submitted on $NODES"
	nohup ./watchdog.sh >> logs/watchdog.log 2>&1 &