195 lines
6.3 KiB
Bash
195 lines
6.3 KiB
Bash
#!/usr/bin/env bash
|
|
# =============================================================================
|
|
# apply_optimizations.sh — Apply v2 optimizations and restart training
|
|
#
|
|
# Optimizations applied:
|
|
# 1. QKV Projection Fusion (+8-12% throughput)
|
|
# 2. NUMA CPU Affinity (fix 69% cross-NUMA workers)
|
|
# 3. Batch size 4→5 (11h saved over full run)
|
|
# 4. NCCL NVLS algorithm + 256MB buffers
|
|
# 5. DDP bucket_cap_mb 400→800
|
|
# 6. DataLoader num_workers 4→6, prefetch_factor 3→4
|
|
# 7. MADV_RANDOM + WILLNEED for PackedDataset
|
|
# 8. numactl --interleave=all on torchrun
|
|
#
|
|
# Usage:
|
|
# bash scripts/apply_optimizations.sh # full migration
|
|
# bash scripts/apply_optimizations.sh --test-only # just validate, don't restart
|
|
# bash scripts/apply_optimizations.sh --skip-stop # don't stop current training
|
|
# =============================================================================
|
|
set -u
|
|
|
|
cd "$(dirname "$0")/.."
|
|
|
|
RUN_NAME="korean_3b_fp8_run1"
|
|
CKPT_DIR="checkpoints/${RUN_NAME}"
|
|
PID_FILE="${CKPT_DIR}/train.pid"
|
|
LOG_FILE="${CKPT_DIR}/train.log"
|
|
|
|
TEST_ONLY=false
|
|
SKIP_STOP=false
|
|
for arg in "$@"; do
|
|
case "$arg" in
|
|
--test-only) TEST_ONLY=true ;;
|
|
--skip-stop) SKIP_STOP=true ;;
|
|
esac
|
|
done
|
|
|
|
echo "=================================================================="
|
|
echo " FRANKENSTALLM 3B — Optimization Migration v2"
|
|
echo " $(date)"
|
|
echo "=================================================================="
|
|
|
|
# ---- Step 1: Validate all modified files --------------------------------
|
|
echo ""
|
|
echo "[1/6] Validating modified files..."
|
|
ERRORS=0
|
|
|
|
for pyfile in model/attention.py train/pretrain.py data/dataset.py scripts/migrate_qkv_checkpoint.py; do
|
|
if python3 -c "import ast; ast.parse(open('$pyfile').read())" 2>/dev/null; then
|
|
echo " ✓ $pyfile — syntax OK"
|
|
else
|
|
echo " ✗ $pyfile — SYNTAX ERROR"
|
|
((ERRORS++))
|
|
fi
|
|
done
|
|
|
|
if bash -n scripts/launch_3b_pretrain.sh 2>/dev/null; then
|
|
echo " ✓ scripts/launch_3b_pretrain.sh — syntax OK"
|
|
else
|
|
echo " ✗ scripts/launch_3b_pretrain.sh — SYNTAX ERROR"
|
|
((ERRORS++))
|
|
fi
|
|
|
|
# Check YAML
|
|
python3 -c "
|
|
import yaml
|
|
with open('configs/korean_3b_fp8.yaml') as f:
|
|
cfg = yaml.safe_load(f)
|
|
assert cfg['train']['batch_size'] == 5, f'batch_size should be 5, got {cfg[\"train\"][\"batch_size\"]}'
|
|
print(' ✓ configs/korean_3b_fp8.yaml — valid, batch_size=5')
|
|
" 2>/dev/null || { echo " ✗ configs/korean_3b_fp8.yaml — INVALID"; ((ERRORS++)); }
|
|
|
|
if [[ $ERRORS -gt 0 ]]; then
|
|
echo ""
|
|
echo "[ERROR] $ERRORS file(s) failed validation. Aborting."
|
|
exit 1
|
|
fi
|
|
echo " All files validated successfully."
|
|
|
|
if $TEST_ONLY; then
|
|
echo ""
|
|
echo "[INFO] --test-only mode. Exiting without restart."
|
|
exit 0
|
|
fi
|
|
|
|
# ---- Step 2: Stop current training (graceful) ---------------------------
|
|
if ! $SKIP_STOP; then
|
|
echo ""
|
|
echo "[2/6] Stopping current training (SIGTERM → emergency checkpoint)..."
|
|
if [[ -f "$PID_FILE" ]]; then
|
|
PID=$(cat "$PID_FILE")
|
|
if kill -0 "$PID" 2>/dev/null; then
|
|
echo " Sending SIGTERM to PID $PID..."
|
|
kill "$PID"
|
|
echo " Waiting for graceful shutdown (up to 120s)..."
|
|
for i in $(seq 1 120); do
|
|
if ! kill -0 "$PID" 2>/dev/null; then
|
|
echo " Process stopped after ${i}s"
|
|
break
|
|
fi
|
|
sleep 1
|
|
done
|
|
if kill -0 "$PID" 2>/dev/null; then
|
|
echo " [WARN] Process still running after 120s. Force killing..."
|
|
kill -9 "$PID" 2>/dev/null || true
|
|
sleep 2
|
|
fi
|
|
else
|
|
echo " Process $PID not running."
|
|
fi
|
|
else
|
|
echo " No PID file found."
|
|
fi
|
|
|
|
# Wait for all GPU processes to clear
|
|
echo " Waiting for GPU processes to terminate..."
|
|
for i in $(seq 1 30); do
|
|
if ! pgrep -f "pretrain.py.*korean_3b" >/dev/null 2>&1; then
|
|
echo " All GPU processes cleared."
|
|
break
|
|
fi
|
|
sleep 1
|
|
done
|
|
fi
|
|
|
|
# ---- Step 3: Find and migrate latest checkpoint -------------------------
|
|
echo ""
|
|
echo "[3/6] Migrating latest checkpoint (QKV fusion)..."
|
|
LATEST_CKPT=$(ls -d "${CKPT_DIR}"/checkpoint-* 2>/dev/null | sort -V | tail -1)
|
|
if [[ -z "$LATEST_CKPT" ]]; then
|
|
echo " [ERROR] No checkpoint found!"
|
|
exit 1
|
|
fi
|
|
echo " Latest checkpoint: $LATEST_CKPT"
|
|
|
|
# Backup original model.pt
|
|
cp "${LATEST_CKPT}/model.pt" "${LATEST_CKPT}/model.pt.backup_pre_qkv"
|
|
echo " Backup created: model.pt.backup_pre_qkv"
|
|
|
|
# Run migration
|
|
python3 scripts/migrate_qkv_checkpoint.py "$LATEST_CKPT"
|
|
echo " QKV fusion migration complete."
|
|
|
|
# ---- Step 4: Quick validation test (5 steps) ----------------------------
|
|
echo ""
|
|
echo "[4/6] Running 5-step validation test..."
|
|
# Use single GPU for fast test
|
|
timeout 120 python3 train/pretrain.py \
|
|
--config configs/korean_3b_fp8.yaml \
|
|
--train_data data/3b_train.bin \
|
|
--checkpoint_dir /tmp/frankenstallm_test \
|
|
--max_steps 5 \
|
|
--batch_size 5 \
|
|
--resume "$LATEST_CKPT" \
|
|
2>&1 | tail -10
|
|
|
|
TEST_EXIT=$?
|
|
if [[ $TEST_EXIT -eq 0 ]]; then
|
|
echo " ✓ 5-step test passed!"
|
|
else
|
|
echo " ✗ 5-step test FAILED (exit code $TEST_EXIT)"
|
|
echo " [WARN] Restoring original checkpoint..."
|
|
cp "${LATEST_CKPT}/model.pt.backup_pre_qkv" "${LATEST_CKPT}/model.pt"
|
|
echo " Original checkpoint restored. Aborting."
|
|
exit 1
|
|
fi
|
|
|
|
# ---- Step 5: Clean up test artifacts ------------------------------------
|
|
echo ""
|
|
echo "[5/6] Cleaning up test artifacts..."
|
|
rm -rf /tmp/frankenstallm_test
|
|
|
|
# ---- Step 6: Launch full training with optimizations --------------------
|
|
echo ""
|
|
echo "[6/6] Launching optimized training..."
|
|
echo ""
|
|
echo " Changes applied:"
|
|
echo " • QKV Projection Fusion (single GEMM)"
|
|
echo " • NUMA CPU Affinity (cores 0-35→GPU0-3, 36-71→GPU4-7)"
|
|
echo " • Batch size: 4 → 5"
|
|
echo " • NCCL: NVLS,Ring algorithm, 256MB buffers"
|
|
echo " • DDP: bucket_cap_mb 400 → 800"
|
|
echo " • DataLoader: 4→6 workers, prefetch 3→4"
|
|
echo " • MADV_RANDOM + WILLNEED for dataset mmap"
|
|
echo " • numactl --interleave=all on torchrun"
|
|
echo ""
|
|
|
|
bash scripts/launch_3b_pretrain.sh
|
|
|
|
echo ""
|
|
echo "=================================================================="
|
|
echo " Migration complete! Monitor with:"
|
|
echo " tail -f ${LOG_FILE}"
|
|
echo "=================================================================="
|