Spaces:

ASTERIZER
/

LUNA

Running

LUNA / setup_and_train_300m.sh

BHARGAV REDDY

Upload setup_and_train_300m.sh with huggingface_hub

e0d3d11 verified 7 months ago

4.96 kB

#!/usr/bin/env bash
# =============================================================================
#  LUNA 300M — Cloud Setup & Train Entrypoint (One Command)
#  Downloads dataset, pretrains 300M model, validates every 100 steps,
#  uploads all checkpoints to HuggingFace at the end.
#
#  USAGE:
#    bash setup_and_train_300m.sh huggingface ASTERIZER/Luna_Dataset
#    bash setup_and_train_300m.sh huggingface ASTERIZER/Luna_Dataset 10000000  # smoke test
#    bash setup_and_train_300m.sh local /workspace/data/litdata_pretrain_final
#
#  REQUIRED ENV:
#    HF_TOKEN=hf_xxx   (for uploading checkpoints to HuggingFace)
# =============================================================================

set -e

DATA_SOURCE="${1:-huggingface}"
DATA_ID="${2:-ASTERIZER/Luna_Dataset}"
MAX_TOKENS="${3:-4515286950}"
DATA_DIR="/workspace/data/litdata_pretrain_final"
OUT_DIR="/workspace/out/pretrain/luna-300m"

echo "=========================================="
echo "  LUNA 300M — Cloud Setup & Pretrain"
echo "  Source  : $DATA_SOURCE"
echo "  ID/Path : $DATA_ID"
echo "  Tokens  : $MAX_TOKENS"
echo "  Output  : $OUT_DIR"
echo "=========================================="

# ── 1. Python packages ────────────────────────────────────────────────────────
echo ""
echo "[1/5] Installing dependencies..."

pip install -q --upgrade pip
pip install -q \
    torch torchvision \
    psutil \
    pyyaml \
    huggingface_hub \
    gdown \
    tensorboard \
    transformers \
    datasets \
    pyarrow 2>/dev/null || true

echo "      Done."

# ── 2. Download dataset ───────────────────────────────────────────────────────
echo ""
echo "[2/5] Fetching dataset..."

if [ "$DATA_SOURCE" = "gdrive" ]; then
    python fetch_data.py --source gdrive --gdrive_id "$DATA_ID" --out_dir "$DATA_DIR"
elif [ "$DATA_SOURCE" = "huggingface" ]; then
    HF_TOKEN="${HF_TOKEN:-}"
    python fetch_data.py --source huggingface --hf_repo "$DATA_ID" --out_dir "$DATA_DIR" --hf_token "$HF_TOKEN"
elif [ "$DATA_SOURCE" = "local" ]; then
    python fetch_data.py --source local --local_path "$DATA_ID" --out_dir "$DATA_DIR"
else
    echo "Unknown source: $DATA_SOURCE (use: gdrive | huggingface | local)"
    exit 1
fi

DATA_DIR_ENV="$DATA_DIR" python - <<'PYEOF'
from pathlib import Path
import json
import os

data_dir = Path(os.environ["DATA_DIR_ENV"])
index_path = data_dir / "index.json"
if not index_path.exists():
    candidates = sorted(data_dir.glob("**/index.json"))
    if candidates:
        print("ERROR: index.json still not at dataset root after fetch.")
        print("Found candidates:")
        for candidate in candidates[:10]:
            print(f"  - {candidate}")
    else:
        print(f"ERROR: No index.json found under {data_dir}")
    raise SystemExit(1)

with open(index_path, encoding="utf-8") as f:
    idx = json.load(f)

chunks = idx.get("chunks", [])
total = sum(c.get("dim", 0) for c in chunks)
print(f"  Dataset ready: {len(chunks)} chunks, {total:,} tokens at {index_path}")
PYEOF

# ── 3. System probe ──────────────────────────────────────────────────────────
echo ""
echo "[3/5] System probe..."
python - <<'PYEOF'
import os
import psutil
import torch

props = torch.cuda.get_device_properties(0) if torch.cuda.is_available() else None
if props:
    print(f"  GPU  : {props.name} ({props.total_memory / 1024**3:.1f} GB)")
else:
    print("  GPU: None")
print(f"  RAM  : {psutil.virtual_memory().total / 1024**3:.1f} GB")
print(f"  CPUs : {os.cpu_count()}")
PYEOF

# ── 4. Pretrain (with validation every 100 steps) ─────────────────────────────
echo ""
echo "[4/5] Starting LUNA 300M pretraining..."
echo "      - Validation: every 100 steps (10 generation prompts)"
echo "      - Checkpoints: every 1000 steps"
echo "      - Auto-detects VRAM and fills to maximum utilization"
echo ""

python train_300m.py \
    --config train_config_300m.yaml \
    --data_path "$DATA_DIR" \
    --out_dir "$OUT_DIR" \
    --max_tokens "$MAX_TOKENS"

# ── 5. Summary ────────────────────────────────────────────────────────────────
echo ""
echo "=========================================="
echo "  LUNA 300M — COMPLETE"
echo "  Output      : $OUT_DIR"
echo "  Final model : $OUT_DIR/final/lit_model.pth"
if [ -n "$HF_TOKEN" ]; then
echo "  HF Upload   : https://huggingface.co/ASTERIZER/LUNA-300M"
else
echo "  HF Upload   : SKIPPED (set HF_TOKEN to enable)"
fi
echo "=========================================="

	#!/usr/bin/env bash
	# =============================================================================
	# LUNA 300M — Cloud Setup & Train Entrypoint (One Command)
	# Downloads dataset, pretrains 300M model, validates every 100 steps,
	# uploads all checkpoints to HuggingFace at the end.
	#
	# USAGE:
	# bash setup_and_train_300m.sh huggingface ASTERIZER/Luna_Dataset
	# bash setup_and_train_300m.sh huggingface ASTERIZER/Luna_Dataset 10000000 # smoke test
	# bash setup_and_train_300m.sh local /workspace/data/litdata_pretrain_final
	#
	# REQUIRED ENV:
	# HF_TOKEN=hf_xxx (for uploading checkpoints to HuggingFace)
	# =============================================================================

	set -e

	DATA_SOURCE="${1:-huggingface}"
	DATA_ID="${2:-ASTERIZER/Luna_Dataset}"
	MAX_TOKENS="${3:-4515286950}"
	DATA_DIR="/workspace/data/litdata_pretrain_final"
	OUT_DIR="/workspace/out/pretrain/luna-300m"

	echo "=========================================="
	echo " LUNA 300M — Cloud Setup & Pretrain"
	echo " Source : $DATA_SOURCE"
	echo " ID/Path : $DATA_ID"
	echo " Tokens : $MAX_TOKENS"
	echo " Output : $OUT_DIR"
	echo "=========================================="

	# ── 1. Python packages ────────────────────────────────────────────────────────
	echo ""
	echo "[1/5] Installing dependencies..."

	pip install -q --upgrade pip
	pip install -q \
	torch torchvision \
	psutil \
	pyyaml \
	huggingface_hub \
	gdown \
	tensorboard \
	transformers \
	datasets \
	pyarrow 2>/dev/null \|\| true

	echo " Done."

	# ── 2. Download dataset ───────────────────────────────────────────────────────
	echo ""
	echo "[2/5] Fetching dataset..."

	if [ "$DATA_SOURCE" = "gdrive" ]; then
	python fetch_data.py --source gdrive --gdrive_id "$DATA_ID" --out_dir "$DATA_DIR"
	elif [ "$DATA_SOURCE" = "huggingface" ]; then
	HF_TOKEN="${HF_TOKEN:-}"
	python fetch_data.py --source huggingface --hf_repo "$DATA_ID" --out_dir "$DATA_DIR" --hf_token "$HF_TOKEN"
	elif [ "$DATA_SOURCE" = "local" ]; then
	python fetch_data.py --source local --local_path "$DATA_ID" --out_dir "$DATA_DIR"
	else
	echo "Unknown source: $DATA_SOURCE (use: gdrive \| huggingface \| local)"
	exit 1
	fi

	DATA_DIR_ENV="$DATA_DIR" python - <<'PYEOF'
	from pathlib import Path
	import json
	import os

	data_dir = Path(os.environ["DATA_DIR_ENV"])
	index_path = data_dir / "index.json"
	if not index_path.exists():
	candidates = sorted(data_dir.glob("**/index.json"))
	if candidates:
	print("ERROR: index.json still not at dataset root after fetch.")
	print("Found candidates:")
	for candidate in candidates[:10]:
	print(f" - {candidate}")
	else:
	print(f"ERROR: No index.json found under {data_dir}")
	raise SystemExit(1)

	with open(index_path, encoding="utf-8") as f:
	idx = json.load(f)

	chunks = idx.get("chunks", [])
	total = sum(c.get("dim", 0) for c in chunks)
	print(f" Dataset ready: {len(chunks)} chunks, {total:,} tokens at {index_path}")
	PYEOF

	# ── 3. System probe ──────────────────────────────────────────────────────────
	echo ""
	echo "[3/5] System probe..."
	python - <<'PYEOF'
	import os
	import psutil
	import torch

	props = torch.cuda.get_device_properties(0) if torch.cuda.is_available() else None
	if props:
	print(f" GPU : {props.name} ({props.total_memory / 1024**3:.1f} GB)")
	else:
	print(" GPU: None")
	print(f" RAM : {psutil.virtual_memory().total / 1024**3:.1f} GB")
	print(f" CPUs : {os.cpu_count()}")
	PYEOF

	# ── 4. Pretrain (with validation every 100 steps) ─────────────────────────────
	echo ""
	echo "[4/5] Starting LUNA 300M pretraining..."
	echo " - Validation: every 100 steps (10 generation prompts)"
	echo " - Checkpoints: every 1000 steps"
	echo " - Auto-detects VRAM and fills to maximum utilization"
	echo ""

	python train_300m.py \
	--config train_config_300m.yaml \
	--data_path "$DATA_DIR" \
	--out_dir "$OUT_DIR" \
	--max_tokens "$MAX_TOKENS"

	# ── 5. Summary ────────────────────────────────────────────────────────────────
	echo ""
	echo "=========================================="
	echo " LUNA 300M — COMPLETE"
	echo " Output : $OUT_DIR"
	echo " Final model : $OUT_DIR/final/lit_model.pth"
	if [ -n "$HF_TOKEN" ]; then
	echo " HF Upload : https://huggingface.co/ASTERIZER/LUNA-300M"
	else
	echo " HF Upload : SKIPPED (set HF_TOKEN to enable)"
	fi
	echo "=========================================="