File size: 3,881 Bytes
4397e12 | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | #!/bin/bash
# Pause the main run just before its LR decay, generate Qwen3.8-27B teacher trajectories on the
# freed GPU, build them into the synth_teacher source, then resume the run (which then decays with
# the teacher data in its mix). If vLLM fails to come up, the run is resumed anyway.
# The pre-decay checkpoint is kept, so pretraining can be continued later from it; on resume the
# decay is forced right away and lasts --decay_frac of the steps done (no wall-clock re-fit).
# PAUSE_AT=14:00 GEN_HOURS=10 setsid nohup bash scripts/teacher_window.sh &
set -u
cd "$(dirname "$0")/.."
source env.sh
RUN=$TA_DATA/runs/main
PAUSE_AT=${PAUSE_AT:-14:00}
GEN_HOURS=${GEN_HOURS:-10}
L=$TA_DATA/logs/teacher_window.log
log() { echo "[$(date '+%F %T')] $*" | tee -a "$L"; }
COOKBOOK=${COOKBOOK:?set COOKBOOK to a dir with the vLLM-XPU patches}
MODEL_DIR=${MODEL_DIR:?set MODEL_DIR to a local Qwen3.8-27B GPTQ copy}
IMAGE='vllm/vllm-openai-xpu@sha256:f01e24f6c7ff01f1e0662234255a1372297d1dbd89d003cf13c8fad3eab1ba4f'
TRAIN_CMD=($TA_PY scripts/train.py --size L --engram 1 --lr 3e-3 --tokens 1e10 --max_minutes 0 --decay_T 4096
--eval_every 250 --snapshot_tokens 5e8 --out "$RUN")
target=$(date -d "today $PAUSE_AT" +%s)
now=$(date +%s)
[ "$target" -gt "$now" ] && { log "waiting until $PAUSE_AT"; sleep $((target - now)); }
PID=$(ps -eo pid,args | awk '$3=="scripts/train.py" && $0 ~ /runs\/main/ {print $1}' | head -1)
if [ -z "$PID" ]; then log "main run not running; nothing to pause"; exit 1; fi
log "pausing main run (pid $PID)"
touch "$RUN/STOP"
while kill -0 "$PID" 2>/dev/null; do sleep 10; done
log "main run stopped: $(tail -1 "$RUN/log.jsonl" | cut -c1-120)"
TOK=$($TA_PY -c "import json;print(f\"{json.loads(open('$RUN/log.jsonl').readlines()[-1])['tokens']/1e9:.2f}B\")")
cp "$RUN/ckpt.pt" "$RUN/ckpt_predecay_$TOK.pt" && log "kept pre-decay checkpoint ckpt_predecay_$TOK.pt"
if [ ! -f "$MODEL_DIR/quantize_config.json" ]; then
log "teacher model not downloaded; skipping generation"
else
docker rm -f tinyagent-teacher >/dev/null 2>&1 || true
RGID="$(stat -c '%g' /dev/dri/render* | sort -u | sed -n '1p')"
docker run -d --name tinyagent-teacher -p 8000:8000 --device /dev/dri --group-add "$RGID" \
-v /dev/dri:/dev/dri:ro -v "$MODEL_DIR:/model:ro" \
-v "$COOKBOOK/patches/patch_mtp_nightly.py:/patch_mtp.py:ro" \
-v "$COOKBOOK/patches/patch_mtp_boundary.py:/patch_boundary.py:ro" \
-e VLLM_TARGET_DEVICE=xpu -e ZE_FLAT_DEVICE_HIERARCHY=COMPOSITE -e ZE_AFFINITY_MASK=0 \
-e B70_MTP_BF16_DRAFT=1 -e VLLM_XPU_ENABLE_XPU_GRAPH=1 -e PYTORCH_ALLOC_CONF=expandable_segments:True \
--entrypoint bash "$IMAGE" -lc \
"set -e; python /patch_mtp.py; python /patch_boundary.py; exec vllm serve /model --quantization gptq --dtype float16 --max-model-len 32768 --gpu-memory-utilization 0.90 --kv-cache-dtype fp8 --port 8000 --max-num-seqs 64 --max-num-batched-tokens 8192 --no-enable-prefix-caching --served-model-name qwen38 --language-model-only" \
>> "$L" 2>&1
ok=0
for i in $(seq 1 120); do
if curl -sf http://127.0.0.1:8000/health >/dev/null; then ok=1; break; fi
sleep 10
done
if [ "$ok" = 1 ]; then
log "vLLM up; generating teacher trajectories for $GEN_HOURS h"
$TA_PY scripts/teacher.py --hours "$GEN_HOURS" --concurrency 48 \
--out "$TA_DATA/teacher/qwen38_$(date +%Y%m%d).jsonl" >> "$TA_DATA/logs/teacher.log" 2>&1
log "teacher done: $(tail -1 "$TA_DATA/logs/teacher.log" | cut -c1-300)"
else
log "vLLM did not become healthy; last docker logs:"
docker logs --tail 40 tinyagent-teacher >> "$L" 2>&1
fi
docker rm -f tinyagent-teacher >/dev/null 2>&1 || true
$TA_PY scripts/gen_synth.py --only teacher >> "$L" 2>&1
fi
log "resuming main run with the decay forced now"
touch "$RUN/DECAY"
nohup "${TRAIN_CMD[@]}" >> "$TA_DATA/runs/main.out" 2>&1 &
log "resumed (pid $!)"
|