31 lines
1.1 KiB
Bash
31 lines
1.1 KiB
Bash
|
#!/usr/bin/env bash
|
||
|
|
set -euo pipefail
|
||
|
|
|
||
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||
|
|
RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
|
||
|
|
GPU_AHA="${GPU_AHA:-0}"
|
||
|
|
GPU_L2A="${GPU_L2A:-1}"
|
||
|
|
cd "$RECIPE"
|
||
|
|
|
||
|
|
echo "AHA = per-token, per-KV-head gate; assigned visible GPU $GPU_AHA."
|
||
|
|
echo "L2A-style = per-token, head-shared gate; assigned visible GPU $GPU_L2A."
|
||
|
|
echo "Each variant is a separate world-size-1 run with global batch size 1."
|
||
|
|
|
||
|
|
python "$SCRIPT_DIR/validate_release.py"
|
||
|
|
python -m unittest -v tests.test_router_granularity
|
||
|
|
|
||
|
|
if [[ "$GPU_AHA" == "$GPU_L2A" ]]; then
|
||
|
|
"$SCRIPT_DIR/train_one.sh" aha "$GPU_AHA"
|
||
|
|
"$SCRIPT_DIR/train_one.sh" l2a_style "$GPU_L2A"
|
||
|
|
else
|
||
|
|
"$SCRIPT_DIR/train_one.sh" aha "$GPU_AHA" & aha_pid=$!
|
||
|
|
"$SCRIPT_DIR/train_one.sh" l2a_style "$GPU_L2A" & l2a_pid=$!
|
||
|
|
status=0
|
||
|
|
wait "$aha_pid" || status=1
|
||
|
|
wait "$l2a_pid" || status=1
|
||
|
|
[[ "$status" == 0 ]] || exit "$status"
|
||
|
|
fi
|
||
|
|
|
||
|
|
echo "AHA selected checkpoint: ${OUTPUT_ROOT:-$RECIPE/../outputs}/aha/stage2/checkpoint-25"
|
||
|
|
echo "L2A-style selected checkpoint: ${OUTPUT_ROOT:-$RECIPE/../outputs}/l2a_style/stage2/checkpoint-25"
|