| #!/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 |
|
|
