#!/usr/bin/env bash # Purpose: train a draft head for one target model with whichever of the three published # recipes fits the machine - Medusa heads on a frozen backbone, EAGLE-3 through # the EAGLE repository, or EAGLE-3 through SpecForge - by fetching the upstream # repository at a pinned commit and running its own documented training command # against the data make-draft-data.py wrote # Platform: spark, strix, nvidia (CUDA, or ROCm which also presents as cuda to PyTorch). # Track M is not supported here: neither upstream trainer documents an Apple # silicon path, and the lab page gives Track M a reduced path that needs no # training at all. # Minimum memory: 24 GB for a Medusa head on an 8B target at bfloat16; 12-16 GB is enough # for a Medusa head on a 4B target. EAGLE-3 training wants considerably more and # the lab page states what each track can realistically attempt. # Assumes: git, python3 and a Python environment with torch installed; the recipe's own # requirements are installed by this script into that environment; the data file # from make-draft-data.py exists; draftlog.py sits next to this script. Nothing # here deletes anything: an existing checkout is reused, not replaced. # # Usage: bash train-draft.sh RECIPE TARGET_MODEL DATA_FILE [OUT_DIR] # RECIPE medusa | eagle3 | specforge # TARGET_MODEL Hugging Face id or local directory of the model the draft will serve # DATA_FILE the JSONL make-draft-data.py generate wrote # OUT_DIR where the trained draft is written (default drafts/RECIPE) # # Environment: WORK_DIR (default $HOME/draft-training) for the upstream checkouts, # GPUS (default 1), EPOCHS, LR, MEDUSA_NUM_HEADS, MEDUSA_NUM_LAYERS, # MAX_LENGTH, SPECFORGE_CONFIG, LABBOOK, DRY_RUN=1 to print the training # command and stop without running it. set -euo pipefail RECIPE="${1:-}" TARGET="${2:-}" DATA="${3:-}" OUT_DIR="${4:-drafts/${RECIPE:-unset}}" WORK_DIR="${WORK_DIR:-$HOME/draft-training}" GPUS="${GPUS:-1}" EPOCHS="${EPOCHS:-2}" LR="${LR:-1e-3}" MEDUSA_NUM_HEADS="${MEDUSA_NUM_HEADS:-3}" MEDUSA_NUM_LAYERS="${MEDUSA_NUM_LAYERS:-1}" MAX_LENGTH="${MAX_LENGTH:-1024}" LABBOOK="${LABBOOK:-labbook.md}" HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" # Pinned so that a rerun trains the same thing. Replace these with the commit you actually # used and record it in the notebook; both repositories move. MEDUSA_REPO="https://github.com/FasterDecoding/Medusa" EAGLE_REPO="https://github.com/SafeAILab/EAGLE" SPECFORGE_REPO="https://github.com/sgl-project/SpecForge" die() { echo "train-draft: $*" >&2; exit 1; } [[ -n "$RECIPE" && -n "$TARGET" && -n "$DATA" ]] || die "usage: bash train-draft.sh RECIPE TARGET_MODEL DATA_FILE [OUT_DIR]" [[ "$RECIPE" == "medusa" || "$RECIPE" == "eagle3" || "$RECIPE" == "specforge" ]] || die "RECIPE must be medusa, eagle3 or specforge (got '$RECIPE')" [[ -f "$DATA" ]] || die "$DATA does not exist; run make-draft-data.py generate first" [[ -f "$HERE/draftlog.py" ]] || die "draftlog.py is not next to this script" command -v git >/dev/null || die "git is not on PATH" command -v python3 >/dev/null || die "python3 is not on PATH" python3 -c "import torch" 2>/dev/null || die "torch is not importable in this environment" EXAMPLES="$(wc -l < "$DATA" | tr -d ' ')" [[ "$EXAMPLES" -gt 0 ]] || die "$DATA has no rows" mkdir -p "$WORK_DIR" "$OUT_DIR" echo "==> recipe $RECIPE" echo " target $TARGET" echo " data $DATA ($EXAMPLES example(s))" echo " output $OUT_DIR" echo " checkouts in $WORK_DIR" echo # Clones once and leaves an existing checkout alone, so a second run does not throw away # local edits or re-download hundreds of megabytes. fetch_repo() { local url="$1" dir="$2" if [[ -d "$dir/.git" ]]; then echo " reusing existing checkout at $dir" else echo " cloning $url into $dir" git clone --depth 1 "$url" "$dir" fi git -C "$dir" rev-parse --short HEAD } run_or_print() { echo echo " the command about to run:" printf ' %s\n' "$*" echo if [[ "${DRY_RUN:-0}" == "1" ]]; then echo " DRY_RUN=1, so stopping here." return 0 fi "$@" } START="$(date -u +%Y-%m-%dT%H:%M:%SZ)" SECONDS=0 COMMIT="" case "$RECIPE" in medusa) DIR="$WORK_DIR/Medusa" COMMIT="$(fetch_repo "$MEDUSA_REPO" "$DIR")" echo " Medusa at commit $COMMIT (Apache-2.0)" echo " Installing the repository into the active environment" python3 -m pip install --quiet -e "$DIR" command -v torchrun >/dev/null || die "torchrun is not on PATH; it ships with torch" # The option names are Medusa's own, from the training command in its README. run_or_print torchrun --nproc_per_node="$GPUS" "$DIR/medusa/train/train_legacy.py" \ --model_name_or_path "$TARGET" \ --data_path "$DATA" \ --bf16 True \ --output_dir "$OUT_DIR" \ --num_train_epochs "$EPOCHS" \ --per_device_train_batch_size 1 \ --gradient_accumulation_steps 4 \ --save_strategy "no" \ --learning_rate "$LR" \ --weight_decay 0.0 \ --warmup_ratio 0.1 \ --lr_scheduler_type "cosine" \ --logging_steps 1 \ --model_max_length "$MAX_LENGTH" \ --lazy_preprocess True \ --medusa_num_heads "$MEDUSA_NUM_HEADS" \ --medusa_num_layers "$MEDUSA_NUM_LAYERS" ;; eagle3) DIR="$WORK_DIR/EAGLE" COMMIT="$(fetch_repo "$EAGLE_REPO" "$DIR")" echo " EAGLE at commit $COMMIT (Apache-2.0)" echo " Its README recommends SpecForge for out-of-the-box EAGLE-3 training;" echo " this branch runs the repository's own trainer for readers who want it." command -v deepspeed >/dev/null || die "deepspeed is not on PATH; install it or use RECIPE=specforge" [[ -f "$DIR/eagle/traineagle3/main.py" ]] || die "eagle/traineagle3/main.py is missing from the checkout; the layout has changed, read the README" [[ -f "$DIR/eagle/traineagle3/ds_config.json" ]] || die "eagle/traineagle3/ds_config.json is missing; the layout has changed, read the README" echo " Set the dataset and target paths inside the repository's own config before this runs;" echo " the trainer reads them from there and not from this script's arguments." ( cd "$DIR/eagle/traineagle3" && run_or_print deepspeed main.py --deepspeed_config ds_config.json ) ;; specforge) DIR="$WORK_DIR/SpecForge" COMMIT="$(fetch_repo "$SPECFORGE_REPO" "$DIR")" echo " SpecForge at commit $COMMIT (MIT)" python3 -m pip install --quiet -e "$DIR" command -v specforge >/dev/null || die "the specforge command is not on PATH after installation" CONFIG="${SPECFORGE_CONFIG:-$DIR/examples/configs/online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml}" [[ -f "$CONFIG" ]] || die "config $CONFIG not found; list $DIR/examples/configs and set SPECFORGE_CONFIG" # SpecForge takes every setting from one YAML file and accepts dotted overrides on the # command line, so the target and the data are passed as overrides rather than flags. run_or_print specforge train --config "$CONFIG" \ "data.train_data_path=$DATA" \ "output_dir=$OUT_DIR" ;; esac ELAPSED="$SECONDS" if [[ "${DRY_RUN:-0}" == "1" ]]; then echo "==> DRY_RUN=1: nothing was trained and nothing was recorded." exit 0 fi echo echo "==> training finished in $((ELAPSED / 60)) minute(s); recording the run" python3 "$HERE/draftlog.py" --record --labbook "$LABBOOK" <