| #!/usr/bin/env bash |
| set -euo pipefail |
|
|
| PROFILE="${1:-today}" |
| WORK_ROOT="${NAYARA_WORK_ROOT:-/teamspace/studios/this_studio/nayara-work}" |
|
|
| case "$PROFILE" in |
| preview) |
| MAX_TOKENS=50000000 |
| MAX_MINUTES=30 |
| DATA_DIR="$WORK_ROOT/packed-preview" |
| TOKENIZER="$WORK_ROOT/tokenizer-preview.json" |
| ;; |
| today) |
| MAX_TOKENS=300000000 |
| MAX_MINUTES=90 |
| DATA_DIR="$WORK_ROOT/packed-today" |
| TOKENIZER="$WORK_ROOT/tokenizer-today.json" |
| ;; |
| overnight) |
| MAX_TOKENS=2000000000 |
| MAX_MINUTES=540 |
| DATA_DIR="$WORK_ROOT/packed-overnight" |
| TOKENIZER="$WORK_ROOT/tokenizer-overnight.json" |
| ;; |
| sprint) |
| MAX_TOKENS="${NAYARA_SPRINT_TOKENS:-2000000000}" |
| MAX_MINUTES="${NAYARA_SPRINT_MINUTES:-205}" |
| DATA_DIR="${NAYARA_PACKED_DATA:?Set NAYARA_PACKED_DATA to a pre-tokenized elite corpus}" |
| TOKENIZER="${NAYARA_TOKENIZER:?Set NAYARA_TOKENIZER to its tokenizer.json}" |
| ;; |
| *) |
| echo "usage: bash scripts/lightning_train_h100.sh [preview|today|overnight|sprint]" >&2 |
| exit 2 |
| ;; |
| esac |
|
|
| MAX_TOKENS="${NAYARA_MAX_TOKENS:-$MAX_TOKENS}" |
| MAX_MINUTES="${NAYARA_MAX_MINUTES:-$MAX_MINUTES}" |
| DATA_DIR="${NAYARA_PACKED_DATA:-$DATA_DIR}" |
| TOKENIZER="${NAYARA_TOKENIZER:-$TOKENIZER}" |
|
|
| |
| |
| REQUESTED_TOKENS="$MAX_TOKENS" |
| MAX_TOKENS="$(python - "$DATA_DIR/manifest.json" "$REQUESTED_TOKENS" "${NAYARA_RESUME:-}" <<'PY' |
| import json |
| import sys |
| from pathlib import Path |
| |
| manifest = json.loads(Path(sys.argv[1]).read_text()) |
| requested = int(sys.argv[2]) |
| resume_tokens = 0 |
| if sys.argv[3]: |
| resume_tokens = int(json.loads((Path(sys.argv[3]) / "progress.json").read_text())["tokens"]) |
| capacity = min( |
| int((item["tokens"] - resume_tokens * float(item["weight"])) / float(item["weight"])) |
| for item in manifest["sources"].values() |
| if float(item["weight"]) > 0 |
| ) |
| safe = max(1, int(capacity * 0.98)) |
| print(min(requested, safe)) |
| PY |
| )" |
| if [[ "$MAX_TOKENS" -lt "$REQUESTED_TOKENS" ]]; then |
| echo "Packed corpus supports $MAX_TOKENS unique mixed tokens; clamped from $REQUESTED_TOKENS." |
| fi |
|
|
| python - <<'PY' |
| import torch |
| assert torch.cuda.is_available(), "Switch the Studio to the H100 machine" |
| assert torch.cuda.is_bf16_supported(), "BF16 is unavailable" |
| name = torch.cuda.get_device_name(0) |
| memory = torch.cuda.get_device_properties(0).total_memory / 2**30 |
| print(f"GPU preflight: {name}, {memory:.1f} GiB, torch={torch.__version__}, CUDA={torch.version.cuda}") |
| assert memory >= 70, "This profile expects an 80GB H100" |
| PY |
|
|
| OUTPUT="$WORK_ROOT/checkpoints-${PROFILE}" |
| MICRO_BATCH="${NAYARA_MICRO_BATCH:-64}" |
| EXTRA_ARGS=() |
| if [[ -n "${NAYARA_RESUME:-}" ]]; then |
| EXTRA_ARGS+=(--resume "$NAYARA_RESUME") |
| fi |
| if [[ -n "${NAYARA_PEAK_LR:-}" ]]; then |
| EXTRA_ARGS+=(--peak-lr "$NAYARA_PEAK_LR") |
| fi |
| python -m nayara.train \ |
| --strategy single \ |
| --config configs/nayara_350m.json \ |
| --data "$DATA_DIR" \ |
| --output "$OUTPUT" \ |
| --micro-batch "$MICRO_BATCH" \ |
| --grad-accum 1 \ |
| --max-tokens "$MAX_TOKENS" \ |
| --max-minutes "$MAX_MINUTES" \ |
| --checkpoint-steps 2000 \ |
| --checkpoint-minutes 15 \ |
| --no-gradient-checkpointing \ |
| --peak-tflops 989 \ |
| "${EXTRA_ARGS[@]}" \ |
| 2>&1 | tee "$WORK_ROOT/train-${PROFILE}.log" |
|
|
| FINAL="$(find "$OUTPUT" -maxdepth 1 -type d -name 'step-*' | sort | tail -n 1)" |
| test -n "$FINAL" |
| cp "$TOKENIZER" "$FINAL/tokenizer.json" |
| python -m nayara.quantize "$FINAL" |
| printf '%s\n' "$FINAL" | tee "$WORK_ROOT/final-${PROFILE}.txt" |
| echo "Ready: $FINAL" |
|
|
