pre-train-ckp/launch_train.sh

#!/bin/bash
set -uo pipefail
cd /ephemeral
source /ephemeral/venv/bin/activate

export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export NCCL_SHM_DISABLE=1
export TOKENIZERS_PARALLELISM=false
export OMP_NUM_THREADS=8
export MKL_NUM_THREADS=8

# Use CUDA 12.4 toolkit so DeepSpeed op JIT compilation matches torch
export CUDA_HOME=/usr/local/cuda-12.4
export PATH="/usr/local/cuda-12.4/bin:${PATH}"
export LD_LIBRARY_PATH="/usr/local/cuda-12.4/lib64:${LD_LIBRARY_PATH:-}"

# Kill stale training
pkill -9 -f torchrun 2>/dev/null || true
pkill -9 -f train_scrapegoat_sft_fast 2>/dev/null || true
sleep 3

# Clear logs
rm -f /ephemeral/scrapegoat-lora/train.log

# By default preserve load_state.json so resuming after a crash skips shard reload.
# To force a fresh launch (reload all shards), pass FRESH=1.
if [[ "${FRESH:-0}" == "1" ]]; then
  rm -f /ephemeral/scrapegoat-lora/load_state.json
  echo "[launch] FRESH=1: clearing shard load state"
fi

# Only clear stale DeepSpeed op build cache on a fresh launch; once CPUAdam
# compiles successfully with CUDA 12.4 we want to reuse it.
if [[ "${FRESH:-0}" == "1" ]]; then
  rm -rf ~/.cache/torch_extensions/py312_cu124/cpu_adam 2>/dev/null || true
fi

# Write DeepSpeed ZeRO-3 config (must be parsed by TrainingArguments.__post_init__)
cat > /ephemeral/ds_config_zero3.json <<'EOF'
{
  "bf16": {"enabled": true},
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {"device": "cpu", "pin_memory": false},
    "offload_param": {"device": "cpu", "pin_memory": false}
  },
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": 1.0,
  "steps_per_print": 10,
  "wall_clock_breakdown": false
}
EOF

PORT=29555
exec torchrun --nproc_per_node=8 --master_port=$PORT /ephemeral/train_scrapegoat_sft_fast.py \
  --output_dir /ephemeral/scrapegoat-lora \
  --train_data_path /ephemeral/scrapegoat-fp8 \
  --max_seq_length 4096 \
  --max_steps 200 \
  --save_steps 50 \
  --learning_rate 2e-4 \
  --warmup_steps 10 \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 8 \
  --deepspeed /ephemeral/ds_config_zero3.json \
  > /ephemeral/scrapegoat-lora/train.log 2>&1