128 lines
3.5 KiB
Bash
128 lines
3.5 KiB
Bash
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
if [[ $# -lt 1 || $# -gt 2 ]]; then
|
|
echo "Usage: $0 aha|l2a_style [GPU_ID]" >&2
|
|
exit 2
|
|
fi
|
|
|
|
ARM="$1"
|
|
GPU="${2:-0}"
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
|
|
REPO="$(cd "$RECIPE/.." && pwd)"
|
|
VANILLA_DIR="${VANILLA_DIR:-$REPO}"
|
|
DATA_DIR="${DATA_DIR:-$RECIPE/data/am_distilled_long_mix}"
|
|
OUTPUT_ROOT="${OUTPUT_ROOT:-$REPO/outputs}"
|
|
|
|
case "$ARM" in
|
|
aha) GRANULARITY=token_kv_head ;;
|
|
l2a_style) GRANULARITY=token ;;
|
|
*) echo "Unknown arm: $ARM (expected aha or l2a_style)" >&2; exit 2 ;;
|
|
esac
|
|
|
|
ARM_OUT="$OUTPUT_ROOT/$ARM"
|
|
HOTSTART="$ARM_OUT/hotstart"
|
|
STAGE1="$ARM_OUT/stage1"
|
|
STAGE2="$ARM_OUT/stage2"
|
|
LOG_DIR="$ARM_OUT/logs"
|
|
mkdir -p "$LOG_DIR"
|
|
|
|
export CUDA_VISIBLE_DEVICES="$GPU"
|
|
export PYTHONUNBUFFERED=1
|
|
export TOKENIZERS_PARALLELISM=false
|
|
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
|
|
|
|
if [[ ! -f "$VANILLA_DIR/model.safetensors" ]]; then
|
|
echo "Missing tuned-vanilla checkpoint: $VANILLA_DIR/model.safetensors" >&2
|
|
exit 2
|
|
fi
|
|
|
|
if [[ ! -f "$HOTSTART/config.json" ]]; then
|
|
python "$SCRIPT_DIR/export_tuned_vanilla_aha_hotstart.py" \
|
|
--vanilla-path "$VANILLA_DIR" \
|
|
--output-path "$HOTSTART" \
|
|
--window-size 128 \
|
|
--local-kind sink_recent \
|
|
--gate-init-full-prob 0.90 \
|
|
--router-granularity "$GRANULARITY" \
|
|
2>&1 | tee "$LOG_DIR/hotstart.log"
|
|
fi
|
|
|
|
if [[ ! -f "$STAGE1/checkpoint-300/config.json" ]]; then
|
|
AHA_TRAIN_GATE_HARD_THRESHOLD=0.50 \
|
|
python "$RECIPE/dynamic_duo_train.py" \
|
|
--aha_checkpoint "$HOTSTART" \
|
|
--model_path "$VANILLA_DIR" \
|
|
--output_dir "$STAGE1" \
|
|
--data_source am_distilled \
|
|
--am_dataset_path "$DATA_DIR" \
|
|
--am_dataset_split train \
|
|
--am_label_mode full \
|
|
--max_length 8192 \
|
|
--num_steps 300 \
|
|
--warmup_ratio 0.10 \
|
|
--lr 3e-5 \
|
|
--reg_weight 0.1 \
|
|
--ce_weight 0.0 \
|
|
--batch_size 1 \
|
|
--grad_accum 1 \
|
|
--save_steps 100 \
|
|
--log_steps 10 \
|
|
--seed 42 \
|
|
--dtype bfloat16 \
|
|
--attn_impl sdpa \
|
|
--aha_local_kind sink_recent \
|
|
--router_granularity "$GRANULARITY" \
|
|
2>&1 | tee "$LOG_DIR/stage1.log"
|
|
fi
|
|
|
|
if [[ ! -f "$STAGE2/checkpoint-75/config.json" ]]; then
|
|
MODEL_PATH="$VANILLA_DIR" \
|
|
AHA_CHECKPOINT_PATH="$STAGE1/checkpoint-300" \
|
|
DATASET_PATH="$DATA_DIR" \
|
|
DATASET_SPLIT=train \
|
|
OUTPUT_DIR="$STAGE2" \
|
|
AHA_MODE=dynamic \
|
|
AHA_ROUTER_GRANULARITY="$GRANULARITY" \
|
|
AHA_LOCAL_KIND=sink_recent \
|
|
AHA_CE_WEIGHT=1.0 \
|
|
AHA_DISTILL_WEIGHT=0.5 \
|
|
AHA_REG_WEIGHT=0.01 \
|
|
AHA_TRAIN_GATE_HARD_THRESHOLD=0.58 \
|
|
GROUPED_LR=1 \
|
|
GATE_ONLY=0 \
|
|
LEARNING_RATE=3e-6 \
|
|
GATE_LEARNING_RATE=3e-6 \
|
|
BACKBONE_LEARNING_RATE=3e-7 \
|
|
LR_SCHEDULER_TYPE=constant_with_warmup \
|
|
WARMUP_RATIO=0.10 \
|
|
WEIGHT_DECAY=0.0 \
|
|
MAX_SEQ_LENGTH=8192 \
|
|
PER_DEVICE_TRAIN_BATCH_SIZE=1 \
|
|
GRADIENT_ACCUMULATION_STEPS=1 \
|
|
MAX_STEPS=75 \
|
|
LOGGING_STEPS=5 \
|
|
SAVE_STEPS=25 \
|
|
SAVE_TOTAL_LIMIT=3 \
|
|
FREEZE_EMBEDDINGS_LM_HEAD=1 \
|
|
REPORT_TO=none \
|
|
SEED=47 \
|
|
python "$RECIPE/sft.py" \
|
|
2>&1 | tee "$LOG_DIR/stage2.log"
|
|
fi
|
|
|
|
python - "$STAGE2/checkpoint-25" "$GRANULARITY" <<'PY'
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
checkpoint = Path(sys.argv[1])
|
|
expected = sys.argv[2]
|
|
config = json.loads((checkpoint / "config.json").read_text())
|
|
actual = config.get("aha_router_granularity", "token_kv_head")
|
|
if actual != expected:
|
|
raise RuntimeError(f"checkpoint granularity {actual!r} != {expected!r}")
|
|
print(f"selected_checkpoint={checkpoint} router_granularity={actual}")
|
|
PY
|