File size: 3,561 Bytes
46b9eea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | #!/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
|