File size: 1,147 Bytes
47ea602 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | #!/usr/bin/env bash
# Runs on the main pod (all GPUs): setup -> text cache -> [teacher samples] -> training.
# DISC_REAL=adm|teacher TEACHER_COUNT=2400 bash scripts/main_run.sh [--set key=value ...]
# Expects code at /workspace/code and prompts at /workspace/prompts/{train,eval}.jsonl.
set -euo pipefail
cd /workspace/code
N=$(nvidia-smi -L | wc -l)
M=/workspace/models/MOSS-Image-SFT
echo "== setup $(date -u +%T)"
bash scripts/pod_setup.sh
echo "== text cache $(date -u +%T)"
torchrun --nproc-per-node "$N" -m rapid_moss.textenc --model $M \
--prompts /workspace/prompts/train.jsonl /workspace/prompts/eval.jsonl --out /workspace/cache
EXTRA=(disc_real=adm)
if [[ "${DISC_REAL:-teacher}" == teacher ]]; then
echo "== teacher samples $(date -u +%T)"
torchrun --nproc-per-node "$N" -m rapid_moss.teachergen --model $M --cache /workspace/cache \
--prompts /workspace/prompts/train.jsonl --out /workspace/teacher --count "${TEACHER_COUNT:-2000}"
EXTRA=()
fi
echo "== train $(date -u +%T)"
torchrun --nproc-per-node "$N" -m rapid_moss.train --config configs/main.json --set "${EXTRA[@]}" "$@"
echo "== done $(date -u +%T)"
|