Add training template: setup.sh

training-template/setup.sh

@@ -0,0 +1,359 @@
+#!/bin/bash
+# ============================================================================
+# Daimon Training — RunPod First-Boot Setup
+# Liberation Labs
+# ============================================================================
+#
+# Run this ONCE on a fresh RunPod pod before training.
+# It installs dependencies, pulls the model & data, and verifies the environment.
+#
+# Requirements:
+#   - 1x H200 SXM 141GB
+#   - 188GB+ system RAM (critical for CPU-offloaded optimizer)
+#   - 400GB+ disk on /workspace (persistent volume)
+#   - HF_TOKEN environment variable set (model is gated)
+#
+# Usage:
+#   export HF_TOKEN="hf_your_token_here"
+#   bash /workspace/runpod-template/setup.sh
+# ============================================================================
+
+set -e
+
+echo "============================================================"
+echo "  DAIMON FULL-PARAMETER SFT — POD SETUP"
+echo "  Liberation Labs"
+echo "  $(date)"
+echo "============================================================"
+
+# ── 1. Find Python ──────────────────────────────────────────────────────────
+export PATH=/opt/conda/bin:/usr/local/bin:$PATH
+PYTHON=$(which python3.11 2>/dev/null || which python3 2>/dev/null)
+echo "Python: $PYTHON ($($PYTHON --version 2>&1))"
+
+# ── 2. Check GPU ──────────────────────────────────────────────────────────
+echo ""
+echo "=== GPU Check ==="
+GPU_COUNT=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
+echo "GPUs detected: $GPU_COUNT"
+nvidia-smi --query-gpu=name,memory.total --format=csv,noheader
+
+if [ "$GPU_COUNT" -lt 1 ]; then
+    echo ""
+    echo "FATAL: No GPUs detected."
+    exit 1
+fi
+
+# Check VRAM (need >= 140GB for full SFT on single GPU)
+VRAM=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -1 | tr -d ' ')
+echo "VRAM: ${VRAM} MiB"
+if [ "$VRAM" -lt 140000 ]; then
+    echo ""
+    echo "FATAL: GPU has ${VRAM} MiB VRAM. Need >= 140,000 MiB."
+    echo ""
+    echo "Why: Full-parameter SFT memory budget:"
+    echo "  Model params (bf16):    ~70GB → GPU"
+    echo "  Activations (grad ckpt): ~20GB → GPU"
+    echo "  Total GPU:              ~90GB of 141GB"
+    echo ""
+    echo "Fix: Provision a pod with 1x H200 SXM 141GB."
+    exit 1
+fi
+
+# ── 3. Check system RAM ────────────────────────────────────────────────────
+echo ""
+echo "=== System RAM ==="
+TOTAL_RAM=$(free -g | grep Mem | awk '{print $2}')
+echo "Total: ${TOTAL_RAM} GB"
+
+# System RAM is CRITICAL for full SFT — optimizer states and gradients are CPU-offloaded
+if [ "$TOTAL_RAM" -lt 180 ]; then
+    echo "FATAL: System RAM is ${TOTAL_RAM}GB. Need >= 180GB."
+    echo ""
+    echo "Why: Full SFT CPU-offloaded memory budget:"
+    echo "  Gradients (bf16):        ~70GB → CPU"
+    echo "  Adafactor optimizer:     ~35GB → CPU"
+    echo "  Total CPU:              ~105GB"
+    echo "  Plus OS/data overhead:   ~30GB"
+    echo ""
+    echo "AdamW is NOT viable — its fp32 states would need ~280GB CPU RAM."
+    echo "Even Adafactor needs ~105GB + headroom."
+    echo ""
+    echo "Fix: Provision a pod with >= 188GB system RAM."
+    exit 1
+fi
+
+# ── 4. Check disk space ────────────────────────────────────────────────────
+echo ""
+echo "=== Disk Space ==="
+AVAIL_GB=$(df -BG /workspace | tail -1 | awk '{print $4}' | tr -d 'G')
+echo "Available on /workspace: ${AVAIL_GB} GB"
+
+if [ "$AVAIL_GB" -lt 300 ]; then
+    echo "FATAL: Less than 300GB on /workspace."
+    echo "Full SFT needs: model (~70GB) + data + checkpoints (~210GB for 3 × 70GB)."
+    echo "Mount a 400GB+ persistent volume."
+    exit 1
+elif [ "$AVAIL_GB" -lt 400 ]; then
+    echo "WARNING: Less than 400GB on /workspace."
+    echo "Full model checkpoints are ~70GB each. With save_total_limit=3, need ~210GB."
+    echo "Will be tight — consider a larger volume."
+fi
+
+# ── 5. Check HF_TOKEN ──────────────────────────────────────────────────────
+echo ""
+echo "=== HuggingFace Authentication ==="
+if [ -z "$HF_TOKEN" ]; then
+    echo "FATAL: HF_TOKEN environment variable not set."
+    echo "Qwen3.6-35B-A3B is a gated model. You need a HuggingFace token."
+    echo ""
+    echo "Fix: export HF_TOKEN='hf_your_token_here'"
+    echo "Or set it in the RunPod pod template environment variables."
+    exit 1
+else
+    echo "HF_TOKEN is set (${#HF_TOKEN} chars)"
+fi
+
+# ── 6. Install dependencies (pinned versions) ─────────────────────────────
+echo ""
+echo "=== Installing Dependencies ==="
+$PYTHON -m pip install --upgrade -q pip
+
+echo "Installing PyTorch..."
+$PYTHON -m pip install -q \
+    torch==2.7.1 \
+    torchvision==0.22.1 \
+    --index-url https://download.pytorch.org/whl/cu124 \
+    2>&1 | tail -2
+
+echo "Installing training stack (pinned versions)..."
+$PYTHON -m pip install -q \
+    transformers==5.12.1 \
+    trl==1.7.0 \
+    datasets==5.0.0 \
+    accelerate==1.14.0 \
+    deepspeed==0.16.7 \
+    safetensors==0.8.0 \
+    pyyaml==6.0.2 \
+    2>&1 | tail -3
+
+echo "Installing flash-attn (may take a few minutes to compile)..."
+$PYTHON -m pip install -q flash-attn --no-build-isolation 2>&1 | tail -3 || {
+    echo "WARNING: flash-attn failed to install. Will fall back to SDPA attention."
+    echo "This is fine — SDPA is only ~5% slower on H200."
+}
+
+echo "Dependencies installed."
+
+# ── 7. Verify critical packages ────────────────────────────────────────────
+echo ""
+echo "=== Package Verification ==="
+$PYTHON -c "
+import torch, transformers, trl, datasets, accelerate, deepspeed, safetensors
+print(f'torch:        {torch.__version__}')
+print(f'transformers: {transformers.__version__}')
+print(f'trl:          {trl.__version__}')
+print(f'datasets:     {datasets.__version__}')
+print(f'accelerate:   {accelerate.__version__}')
+print(f'deepspeed:    {deepspeed.__version__}')
+print(f'safetensors:  {safetensors.__version__}')
+print(f'CUDA:         {torch.version.cuda}')
+print(f'GPUs:         {torch.cuda.device_count()}')
+try:
+    import flash_attn
+    print(f'flash_attn:   {flash_attn.__version__}')
+except ImportError:
+    print('flash_attn:   not installed (using SDPA fallback)')
+"
+
+# ── 8. Verify Qwen3.6 architecture support ─────────────────────────────────
+echo ""
+echo "=== Model Architecture Check ==="
+MODEL_REVISION="995ad96eacd98c81ed38be0c5b274b04031597b0"
+$PYTHON -c "
+from transformers import AutoConfig
+c = AutoConfig.from_pretrained('Qwen/Qwen3.6-35B-A3B', revision='$MODEL_REVISION', trust_remote_code=True)
+print(f'Model type:      {c.model_type}')
+print(f'Hidden size:     {c.hidden_size}')
+print(f'Num layers:      {c.num_hidden_layers}')
+print(f'Num experts:     {getattr(c, \"num_experts\", \"N/A\")}')
+print(f'Vocab size:      {c.vocab_size}')
+print(f'Pinned revision: $MODEL_REVISION')
+print('Architecture supported: OK')
+" || {
+    echo "FATAL: Qwen3.6 architecture not supported by installed transformers."
+    echo "Upgrade: pip install --upgrade transformers"
+    exit 1
+}
+
+# ── 9. Pull model from HuggingFace ─────────────────────────────────────────
+echo ""
+echo "=== Model Download ==="
+MODEL_DIR="/workspace/models/Qwen3.6-35B-A3B"
+if [ -d "$MODEL_DIR" ] && [ -f "$MODEL_DIR/config.json" ]; then
+    echo "Model already downloaded at $MODEL_DIR"
+else
+    echo "Downloading Qwen3.6-35B-A3B (~70GB, this will take a while)..."
+    mkdir -p /workspace/models
+    $PYTHON -c "
+from huggingface_hub import snapshot_download
+import os
+snapshot_download(
+    'Qwen/Qwen3.6-35B-A3B',
+    revision='$MODEL_REVISION',
+    local_dir='$MODEL_DIR',
+    token=os.environ['HF_TOKEN'],
+)
+print('Model download complete.')
+"
+fi
+
+# ── 10. Pull training data ──────────────────────────────────────────────────
+echo ""
+echo "=== Training Data ==="
+DATA_DIR="/workspace/daimon-data"
+mkdir -p "$DATA_DIR"
+
+if [ -d "$DATA_DIR/train_arrow" ] && [ -d "$DATA_DIR/valid_arrow" ]; then
+    echo "Arrow data already present. Verifying..."
+    $PYTHON -c "
+from datasets import load_from_disk
+t = load_from_disk('$DATA_DIR/train_arrow')
+v = load_from_disk('$DATA_DIR/valid_arrow')
+print(f'Train: {len(t):,} samples | Valid: {len(v):,} samples — OK')
+"
+else
+    echo "Downloading and preparing training data..."
+    $PYTHON -c "
+import os, json, gzip, shutil
+from huggingface_hub import hf_hub_download, list_repo_files
+from datasets import Dataset, load_dataset
+
+DATA_DIR = '$DATA_DIR'
+REPO = 'HumboldtJoker/daimon-sft-data'
+token = os.environ.get('HF_TOKEN')
+
+try:
+    # Try loading as a HF dataset first
+    ds = load_dataset(REPO, token=token)
+    if 'train' in ds:
+        ds['train'].save_to_disk(f'{DATA_DIR}/train_arrow')
+        print(f'Train: {len(ds[\"train\"]):,} samples saved as Arrow')
+    if 'validation' in ds:
+        ds['validation'].save_to_disk(f'{DATA_DIR}/valid_arrow')
+        print(f'Valid: {len(ds[\"validation\"]):,} samples saved as Arrow')
+    elif 'test' in ds:
+        ds['test'].save_to_disk(f'{DATA_DIR}/valid_arrow')
+        print(f'Valid: {len(ds[\"test\"]):,} samples saved as Arrow')
+    else:
+        # Split train into train/valid
+        split = ds['train'].train_test_split(test_size=0.05, seed=42)
+        split['train'].save_to_disk(f'{DATA_DIR}/train_arrow')
+        split['test'].save_to_disk(f'{DATA_DIR}/valid_arrow')
+        print(f'Auto-split: Train {len(split[\"train\"]):,} | Valid {len(split[\"test\"]):,}')
+except Exception as e:
+    print(f'HF dataset load failed: {e}')
+    print('Trying file-based download...')
+
+    # Fall back to downloading individual files
+    try:
+        files = list_repo_files(REPO, repo_type='dataset', token=token)
+        for f in files:
+            if f.endswith(('.jsonl', '.jsonl.gz', '.json')):
+                print(f'Downloading {f}...')
+                hf_hub_download(REPO, f, repo_type='dataset', local_dir=DATA_DIR, token=token)
+    except Exception as e2:
+        print(f'File download also failed: {e2}')
+        print('DATA MUST BE UPLOADED MANUALLY to {DATA_DIR}/')
+        print('Expected format: JSONL with {\"messages\": [{\"role\": ..., \"content\": ...}, ...]}')
+
+    # Convert any JSONL files to Arrow
+    for split_name in ['train', 'valid']:
+        jsonl = f'{DATA_DIR}/{split_name}.jsonl'
+        gz = f'{DATA_DIR}/{split_name}.jsonl.gz'
+        arrow_dir = f'{DATA_DIR}/{split_name}_arrow'
+
+        if os.path.exists(gz) and not os.path.exists(jsonl):
+            with gzip.open(gz, 'rb') as fin, open(jsonl, 'wb') as fout:
+                shutil.copyfileobj(fin, fout)
+
+        if os.path.exists(jsonl) and not os.path.exists(arrow_dir):
+            data = []
+            with open(jsonl) as fh:
+                for line in fh:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        d = json.loads(line)
+                        if 'messages' in d and len(d['messages']) >= 2:
+                            data.append(d)
+                    except:
+                        pass
+            ds = Dataset.from_list(data)
+            ds.save_to_disk(arrow_dir)
+            print(f'{split_name}: {len(data):,} examples saved as Arrow')
+
+# Final verification
+try:
+    from datasets import load_from_disk
+    t = load_from_disk(f'{DATA_DIR}/train_arrow')
+    print(f'Verified train: {len(t):,} samples')
+    if os.path.isdir(f'{DATA_DIR}/valid_arrow'):
+        v = load_from_disk(f'{DATA_DIR}/valid_arrow')
+        print(f'Verified valid: {len(v):,} samples')
+except:
+    print('WARNING: Could not verify data. Check $DATA_DIR manually.')
+"
+fi
+
+# ── 11. Create persistent directories ──────────────────────────────────────
+echo ""
+echo "=== Creating Directories ==="
+mkdir -p /workspace/daimon-sft/logs
+mkdir -p /workspace/daimon-sft/checkpoints
+echo "Output directories created on persistent volume."
+
+# ── 12. Copy template files to /workspace ───────────────────────────────────
+echo ""
+echo "=== Copying Template Files ==="
+SCRIPT_DIR=$(dirname "$(readlink -f "$0")")
+cp "$SCRIPT_DIR/train_daimon.py" /workspace/runpod-template/train_daimon.py 2>/dev/null || true
+cp "$SCRIPT_DIR/train_daimon_config.yaml" /workspace/runpod-template/train_daimon_config.yaml 2>/dev/null || true
+cp "$SCRIPT_DIR/ds_config_zero2.json" /workspace/runpod-template/ds_config_zero2.json 2>/dev/null || true
+cp "$SCRIPT_DIR/launch.sh" /workspace/runpod-template/launch.sh 2>/dev/null || true
+cp "$SCRIPT_DIR/test_template.py" /workspace/runpod-template/test_template.py 2>/dev/null || true
+echo "Template files in /workspace/runpod-template/"
+
+# ── 13. Add SSH key for remote access ──────────────────────────────────────
+echo ""
+echo "=== SSH Key ==="
+mkdir -p ~/.ssh
+echo "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIOtjekz8l1s6xTAXlhZJg/A0N3d6mZAyF/EyrEMiBCDG thomas@coalition" >> ~/.ssh/authorized_keys
+chmod 700 ~/.ssh
+chmod 600 ~/.ssh/authorized_keys
+echo "SSH key added."
+
+# ── 14. Clean up HF token from disk cache ──────────────────────────────────
+echo ""
+echo "=== Security Cleanup ==="
+rm -f ~/.cache/huggingface/token 2>/dev/null || true
+echo "Cleared cached HF token from disk."
+
+# ── 15. Summary ─────────────────────────────────────────────────────────────
+echo ""
+echo "============================================================"
+echo "  SETUP COMPLETE — FULL-PARAMETER SFT"
+echo ""
+echo "  Memory budget:"
+echo "    GPU:  ~90GB of 141GB (model + activations)"
+echo "    CPU: ~105GB of ${TOTAL_RAM}GB (gradients + Adafactor)"
+echo ""
+echo "  Next steps:"
+echo "    1. Run validation:  python3 /workspace/runpod-template/test_template.py"
+echo "    2. Start training:  bash /workspace/runpod-template/launch.sh"
+echo ""
+echo "  Monitor:"
+echo "    watch -n 5 nvidia-smi"
+echo "    tail -f /workspace/daimon-sft/logs/training_*.log"
+echo "============================================================"