| # 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)" | |