mirror of
https://github.com/karpathy/nanochat.git
synced 2026-02-19 18:10:24 +00:00
101 lines
4.8 KiB
Bash
101 lines
4.8 KiB
Bash
#!/bin/bash
|
|
|
|
# This script is configured to train your own GPT-2 grade LLM (pretraining + finetuning)
|
|
# It is designed to run on a blank 8XH100 GPU node and takes approximately 3 hours to complete.
|
|
|
|
# 1) Example launch (simplest):
|
|
# bash runs/speedrun.sh
|
|
# 2) Example launch in a screen session (because the run takes ~3 hours):
|
|
# screen -L -Logfile runs/speedrun.log -S speedrun bash runs/speedrun.sh
|
|
# 3) Example launch with wandb logging, but see below for setting up wandb first:
|
|
# WANDB_RUN=speedrun screen -L -Logfile runs/speedrun.log -S speedrun bash runs/speedrun.sh
|
|
|
|
# Default intermediate artifacts directory is in ~/.cache/nanochat
|
|
export OMP_NUM_THREADS=1
|
|
export NANOCHAT_BASE_DIR="$HOME/.cache/nanochat"
|
|
mkdir -p $NANOCHAT_BASE_DIR
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Python venv setup with uv
|
|
|
|
# install uv (if not already installed)
|
|
command -v uv &> /dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh
|
|
# create a .venv local virtual environment (if it doesn't exist)
|
|
[ -d ".venv" ] || uv venv
|
|
# install the repo dependencies
|
|
uv sync --extra gpu
|
|
# activate venv so that `python` uses the project's venv instead of system python
|
|
source .venv/bin/activate
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# wandb setup
|
|
# If you wish to use wandb for logging (it's nice!, recommended).
|
|
# 1) Make sure to first log in to wandb, e.g. run:
|
|
# `wandb login`
|
|
# 2) Set the WANDB_RUN environment variable when running this script, e.g.:
|
|
# `WANDB_RUN=d26 bash speedrun.sh`
|
|
if [ -z "$WANDB_RUN" ]; then
|
|
# by default use "dummy" : it's handled as a special case, skips logging to wandb
|
|
WANDB_RUN=dummy
|
|
fi
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# During the course of the run, we will be writing markdown reports to the report/
|
|
# directory in the base dir. This command clears it out and writes a header section
|
|
# with a bunch of system info and a timestamp that marks the start of the run.
|
|
python -m nanochat.report reset
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Tokenizer
|
|
|
|
# Download the first ~2B characters of pretraining dataset
|
|
# each data shard is ~250M chars
|
|
# so we download 2e9 / 250e6 = 8 data shards at this point
|
|
# each shard is ~100MB of text (compressed), so this is about ~800MB of data on disk
|
|
# look at dev/repackage_data_reference.py for details on how this data was prepared
|
|
python -m nanochat.dataset -n 8
|
|
# Immediately also kick off downloading more shards in the background while tokenizer trains
|
|
# Approximately 350 shards are needed for 10B tokens of data for pretraining.
|
|
# The maximum total number of shards available in the entire dataset is 1822.
|
|
python -m nanochat.dataset -n 370 &
|
|
DATASET_DOWNLOAD_PID=$!
|
|
# train the tokenizer with vocab size 2**15 = 32768 on ~2B characters of data
|
|
python -m scripts.tok_train
|
|
# evaluate the tokenizer (report compression ratio etc.)
|
|
python -m scripts.tok_eval
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Base model (pretraining)
|
|
echo "Waiting for dataset download to complete..."
|
|
wait $DATASET_DOWNLOAD_PID
|
|
|
|
# Number of processes/GPUs to use
|
|
NPROC_PER_NODE=8
|
|
|
|
# d24 model (slightly overtrained is enough to beat GPT-2 => increase data:params ratio from compute optimal 10.5 (default) to 12)
|
|
torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_train -- --depth=24 --target-param-data-ratio=12 --run=$WANDB_RUN
|
|
# evaluate the model: CORE metric, BPB on train/val, and draw samples
|
|
torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_eval
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# SFT (teach the model conversation special tokens, tool use, multiple choice)
|
|
|
|
# download 2.3MB of synthetic identity conversations to impart a personality to nanochat
|
|
# see dev/gen_synthetic_data.py for details on how this data was prepared and to get a sense of how you can easily tune it
|
|
curl -L -o $NANOCHAT_BASE_DIR/identity_conversations.jsonl https://karpathy-public.s3.us-west-2.amazonaws.com/identity_conversations.jsonl
|
|
|
|
# run SFT and eval the model
|
|
torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.chat_sft -- --run=$WANDB_RUN
|
|
torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.chat_eval -- -i sft
|
|
|
|
# chat with the model over CLI! Leave out the -p to chat interactively
|
|
# python -m scripts.chat_cli -p "Why is the sky blue?"
|
|
|
|
# even better, chat with your model over a pretty WebUI ChatGPT style
|
|
# python -m scripts.chat_web
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Generate the full report by putting together all the sections
|
|
# report.md is the output and will be copied to current directory for convenience
|
|
python -m nanochat.report generate
|