The HPC lab

#!/bin/bash
# Exercise 1: your first batch job (the work itself runs on the CPU)
#
# Alpha only accepts jobs that ask for a GPU, so this script asks for one even
# though it does not use it. On Barnard (no GPUs), delete the --gres line.
#
# Submit it from a login node:
#   mkdir -p logs && sbatch hello.sbatch
# Check on it:
#   squeue --me
# Read the result:
#   cat logs/hello-hpc-<jobid>.out
#
# Then try other core counts without editing this file:
#   sbatch --cpus-per-task=1 hello.sbatch
#   sbatch --cpus-per-task=6 hello.sbatch   (6 cores per GPU is the most Alpha allows)
#
#SBATCH --job-name=hello-hpc          # the name you see in squeue
#SBATCH --account=p_lv_llmsecrets      # our course project, it pays for the CPU time
#SBATCH --nodes=1                     # one machine (ZIH wants this line in every job)
#SBATCH --ntasks=1                    # one process
#SBATCH --cpus-per-task=4             # with 4 cores
#SBATCH --mem-per-cpu=1G              # memory per core
#SBATCH --gres=gpu:1                  # Alpha takes only GPU jobs; on Barnard delete this line
#SBATCH --time=00:05:00               # hard time limit, hh:mm:ss
#SBATCH --output=logs/%x-%j.out       # %x is the job name, %j the job ID
#SBATCH --error=logs/%x-%j.err        # the logs folder has to exist before you submit

echo "Job $SLURM_JOB_ID started on $(hostname) at $(date)"

# Make sure srun passes all our cores on to the program
# (newer Slurm versions don't do this by themselves).
export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK

# The system python3 is enough for this one.
srun python3 hello_hpc.py --samples 20000000

echo "Job finished at $(date)"

#!/usr/bin/env python3
"""
Exercise 1: where am I running, and what did Slurm give me?

Prints the node name, the Python version and the Slurm variables of the job.
Then it estimates pi with random numbers, split over all the CPU cores the job
got. Run it with --cpus-per-task=1 and with 4, and compare the time.

Only uses the standard library, so the system python3 is enough.
"""
import argparse
import multiprocessing as mp
import os
import platform
import random
import socket
import time


def slurm_info():
    keys = [
        "SLURM_JOB_ID", "SLURM_JOB_NAME", "SLURM_JOB_ACCOUNT", "SLURM_JOB_PARTITION",
        "SLURM_JOB_NODELIST", "SLURM_CPUS_PER_TASK", "SLURM_MEM_PER_CPU",
        "SLURM_ARRAY_TASK_ID", "SLURM_SUBMIT_DIR",
    ]
    return {k: os.environ.get(k, "-") for k in keys}


def available_cpus():
    """CPUs this process may use (respects Slurm's CPU binding, unlike os.cpu_count())."""
    try:
        return len(os.sched_getaffinity(0))
    except AttributeError:  # macOS / Windows
        return os.cpu_count() or 1


def count_hits(args):
    """Throw n random darts at the unit square; count those inside the quarter circle."""
    n, seed = args
    rng = random.Random(seed)
    hits = 0
    for _ in range(n):
        x, y = rng.random(), rng.random()
        if x * x + y * y <= 1.0:
            hits += 1
    return hits


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--samples", type=int, default=10_000_000, help="total number of darts")
    args = p.parse_args()

    print("=" * 60)
    print(f"Hello from        : {socket.gethostname()}")
    print(f"Python            : {platform.python_version()}  ({platform.python_implementation()})")
    print(f"os.cpu_count()    : {os.cpu_count()}   <- all cores of the node")
    n_cpus = available_cpus()
    print(f"CPUs I may use    : {n_cpus}   <- what Slurm allocated to me")
    print("-" * 60)
    for k, v in slurm_info().items():
        print(f"{k:<20}: {v}")
    print("=" * 60)

    # split the work into one chunk per CPU
    chunk = args.samples // n_cpus
    jobs = [(chunk, seed) for seed in range(n_cpus)]
    t0 = time.perf_counter()
    with mp.Pool(processes=n_cpus) as pool:
        hits = sum(pool.map(count_hits, jobs))
    dt = time.perf_counter() - t0

    pi = 4.0 * hits / (chunk * n_cpus)
    print(f"pi ~ {pi:.6f}   ({chunk * n_cpus:,} samples on {n_cpus} CPU(s) in {dt:.2f} s)")
    print(f"throughput        : {chunk * n_cpus / dt / 1e6:.2f} M samples/s")


if __name__ == "__main__":
    main()

#!/bin/bash
# Exercise 2: watch a job, then cancel it
#
# Alpha only accepts jobs that ask for a GPU, so this script asks for one even
# though it does not use it. On Barnard (no GPUs), delete the --gres line.
#
#   mkdir -p logs
#   sbatch sleepy.sbatch                 # prints "Submitted batch job <jobid>"
#   squeue --me                          # the ST column goes from PD (pending) to R (running)
#   scontrol show job <jobid>            # everything Slurm knows about the job
#   tail -f logs/sleepy-<jobid>.out      # live log, Ctrl+C stops tail but not the job
#   scancel <jobid>                      # stop it, then look for "received signal" in the log
#   sacct -j <jobid> --format=JobID,JobName,State,Elapsed,ExitCode
#
# Bonus, a job that waits on purpose:
#   sbatch --begin=now+5minutes sleepy.sbatch
#   squeue --me                          # the REASON column says (BeginTime)
#   whypending <jobid>                   # ZIH helper that explains why a job waits
#
#SBATCH --job-name=sleepy
#SBATCH --account=p_lv_llmsecrets
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=1
#SBATCH --mem-per-cpu=500M
#SBATCH --gres=gpu:1                  # Alpha takes only GPU jobs
#SBATCH --time=00:10:00
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err

# -u turns off output buffering, so tail -f shows every line right away
srun python3 -u sleepy.py --minutes 8 --every 10

#!/usr/bin/env python3
"""
Exercise 2: a job that is slow on purpose, so you can practise squeue,
scontrol and scancel.

It prints a line every few seconds. When you cancel the job, Slurm first sends
SIGTERM ("please stop") and only later SIGKILL. We catch SIGTERM and write a
small checkpoint file. A real training script should do the same.

Only uses the standard library, so the system python3 is enough.
"""
import argparse
import json
import os
import signal
import sys
import time

STATE = {"step": 0}


def save_checkpoint(reason):
    path = f"checkpoint-{os.environ.get('SLURM_JOB_ID', 'local')}.json"
    with open(path, "w") as f:
        json.dump({"step": STATE["step"], "reason": reason, "time": time.ctime()}, f)
    print(f"[sleepy] saved {path} at step {STATE['step']} ({reason})", flush=True)


def on_sigterm(signum, frame):
    print(f"[sleepy] received signal {signum} (SIGTERM) -> cleaning up", flush=True)
    save_checkpoint("cancelled")
    sys.exit(0)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--minutes", type=float, default=8.0)
    p.add_argument("--every", type=float, default=10.0, help="seconds between heartbeats")
    args = p.parse_args()

    signal.signal(signal.SIGTERM, on_sigterm)
    print(f"[sleepy] pid {os.getpid()} on {os.uname().nodename}, "
          f"running for {args.minutes} min. Cancel me with: scancel {os.environ.get('SLURM_JOB_ID', '<jobid>')}",
          flush=True)

    t_end = time.time() + 60 * args.minutes
    while time.time() < t_end:
        STATE["step"] += 1
        print(f"[sleepy] step {STATE['step']:4d}  {time.strftime('%H:%M:%S')}  still working...", flush=True)
        time.sleep(args.every)

    save_checkpoint("finished")


if __name__ == "__main__":
    main()

#!/bin/bash
# Exercise 3: ask for one GPU and see how fast it is
#
# Submit it from a login node of a GPU cluster:
#   ssh <zih-login>@login1.alpha.hpc.tu-dresden.de      (A100)  or
#   ssh <zih-login>@login1.capella.hpc.tu-dresden.de    (H100)
#   cd 03_gpu && mkdir -p logs && sbatch gpu.sbatch
#
# Every cluster has its own Slurm, so a job sent from Barnard never gets a GPU.
#
#SBATCH --job-name=gpu-check
#SBATCH --account=p_lv_llmsecrets
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=6             # CPU cores that come with the GPU
#SBATCH --gres=gpu:1                  # this is the line that gets you a GPU
#SBATCH --mem=32G
#SBATCH --time=00:10:00
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err

source ../00_setup/setup_env.sh       # workspace, Python module and venv with PyTorch

echo "---- nvidia-smi ----"
nvidia-smi
echo "--------------------"

export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK
srun python gpu_check.py --n 8192 --repeats 20

#!/usr/bin/env python3
"""
Exercise 3: did I really get a GPU, and how fast is it?

1. Prints the PyTorch and CUDA versions and the GPUs this job can see.
2. Multiplies two big matrices (the main work inside every LLM layer), first on
   the CPU and then on the GPU, and prints the speed in TFLOP/s.

Needs PyTorch, so run "source 00_setup/setup_env.sh" first.
"""
import argparse
import os
import time

import torch


def bench(device, n, dtype, repeats):
    a = torch.randn(n, n, device=device, dtype=dtype)
    b = torch.randn(n, n, device=device, dtype=dtype)
    _ = a @ b                                   # warm-up (CUDA kernels get loaded lazily)
    if device.type == "cuda":
        torch.cuda.synchronize()                # GPU work is asynchronous: wait for it!
    t0 = time.perf_counter()
    for _ in range(repeats):
        c = a @ b
    if device.type == "cuda":
        torch.cuda.synchronize()
    dt = (time.perf_counter() - t0) / repeats
    flops = 2 * n ** 3                          # n^3 multiply-adds = 2 n^3 floating point ops
    return dt, flops / dt / 1e12


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--n", type=int, default=4096, help="matrix size n x n")
    p.add_argument("--repeats", type=int, default=10)
    args = p.parse_args()

    print("=" * 60)
    print(f"torch version         : {torch.__version__}")
    print(f"CUDA available        : {torch.cuda.is_available()}")
    print(f"CUDA_VISIBLE_DEVICES  : {os.environ.get('CUDA_VISIBLE_DEVICES', '-')}")
    if torch.cuda.is_available():
        print(f"CUDA (torch build)    : {torch.version.cuda}")
        for i in range(torch.cuda.device_count()):
            prop = torch.cuda.get_device_properties(i)
            print(f"GPU {i}                 : {prop.name}, {prop.total_memory / 2**30:.1f} GiB")
    else:
        print("No GPU found. Did you ask for one with --gres=gpu:1 ?")
    print("=" * 60)

    n_cpu_threads = torch.get_num_threads()
    cpu_n = min(args.n, 2048)                   # keep the CPU run short
    dt, tflops = bench(torch.device("cpu"), cpu_n, torch.float32, max(1, args.repeats // 5))
    print(f"CPU  ({n_cpu_threads:2d} threads) fp32 {cpu_n}x{cpu_n}: {dt * 1e3:8.1f} ms  ->  {tflops:6.2f} TFLOP/s")

    if torch.cuda.is_available():
        gpu = torch.device("cuda")
        for dtype in (torch.float32, torch.bfloat16):
            dt, tflops = bench(gpu, args.n, dtype, args.repeats)
            name = str(dtype).replace("torch.", "")
            print(f"GPU  {name:>8s} {args.n}x{args.n}: {dt * 1e3:8.1f} ms  ->  {tflops:6.2f} TFLOP/s")
        print(f"peak GPU memory used  : {torch.cuda.max_memory_allocated() / 2**20:.0f} MiB")


if __name__ == "__main__":
    main()

#!/bin/bash
# Exercise 4: train a tiny GPT on one GPU
#
#   cd 04_tinygpt && mkdir -p logs && sbatch train.sbatch
#   tail -f logs/tinygpt-<jobid>.out          # watch the loss go down
#
# Afterwards:
#   python sample.py --ckpt $WS_DIR/runs/tinygpt-<jobid>/ckpt.pt --prompt "ROMEO:"
#
# Things to try:
#   A bigger model. Everything after the script name is passed on to Python:
#     sbatch train.sbatch --n-layer 6 --n-embd 384
#   A job that runs out of time. Is ckpt.pt still saved?
#     sbatch --time=00:02:00 train.sbatch
#
#SBATCH --job-name=tinygpt
#SBATCH --account=p_lv_llmsecrets
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=6
#SBATCH --gres=gpu:1
#SBATCH --mem=32G
#SBATCH --time=00:30:00
#SBATCH --signal=TERM@60              # send SIGTERM 60 s before the time limit, so we can save
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err

source ../00_setup/setup_env.sh

export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK
srun python -u train_tinygpt.py \
    --data "$WS_DIR/data/input.txt" \
    --out  "$WS_DIR/runs/tinygpt-$SLURM_JOB_ID" \
    --steps 5000 \
    "$@"

#!/usr/bin/env python3
"""
Exercise 4: train a tiny character-level GPT as a batch job. It is the same
model as in notebook "5 - Attention and TinyGPT".

    python train_tinygpt.py --data $WS_DIR/data/input.txt --steps 3000 --out $WS_DIR/runs/base

It writes three files into --out:
    ckpt.pt        the weights and the vocabulary (sample.py reads this)
    metrics.json   final train and val loss, run time, settings (used in exercise 5)
    samples.txt    some text written by the trained model

A few cluster habits are built in: it uses the GPU if Slurm gave you one, it
prints with flush=True so the log updates while the job runs, and it saves a
checkpoint when Slurm sends SIGTERM (scancel or time limit).
"""
import argparse
import json
import math
import os
import signal
import time

import torch
import torch.nn as nn
import torch.nn.functional as F


# ----------------------------------------------------------------------------- model
class CausalSelfAttention(nn.Module):
    """softmax(Q K^T / sqrt(d_head) + causal mask) V, for n_head heads in parallel."""

    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)  # Q, K, V in one matmul
        self.proj = nn.Linear(n_embd, n_embd)
        self.drop = nn.Dropout(dropout)
        mask = torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size)
        self.register_buffer("mask", mask)

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        # (B, T, C) -> (B, n_head, T, head_size)
        q, k, v = (t.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) for t in (q, k, v))
        att = (q @ k.transpose(-2, -1)) / math.sqrt(k.size(-1))       # (B, nh, T, T)
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))  # no peeking ahead
        att = self.drop(F.softmax(att, dim=-1))
        y = (att @ v).transpose(1, 2).contiguous().view(B, T, C)     # re-assemble heads
        return self.drop(self.proj(y))


class Block(nn.Module):
    """Transformer block: communicate (attention) then compute (MLP), with residuals."""

    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd), nn.GELU(), nn.Linear(4 * n_embd, n_embd), nn.Dropout(dropout)
        )

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x


class TinyGPT(nn.Module):
    def __init__(self, vocab_size, block_size=128, n_layer=4, n_head=4, n_embd=128, dropout=0.1):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab_size, n_embd)     # which character?
        self.pos_emb = nn.Embedding(block_size, n_embd)     # at which position?
        self.blocks = nn.Sequential(*[Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)        # scores for the next character

    def forward(self, idx, targets=None):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)
        x = self.ln_f(self.blocks(x))
        logits = self.lm_head(x)                            # (B, T, vocab)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            logits, _ = self(idx[:, -self.block_size:])     # crop to the context window
            logits = logits[:, -1, :] / temperature          # only the last position matters
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float("inf")
            probs = F.softmax(logits, dim=-1)
            idx = torch.cat([idx, torch.multinomial(probs, num_samples=1)], dim=1)
        return idx


# ----------------------------------------------------------------------------- data
def load_text(path):
    if not os.path.isfile(path):
        raise SystemExit(
            f"Could not find the data file '{path}'.\n"
            "   Download it (login node) with:\n"
            "   curl -sSL -o input.txt https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
        )
    with open(path, encoding="utf-8") as f:
        return f.read()


def get_batch(data, block_size, batch_size, device):
    ix = torch.randint(len(data) - block_size - 1, (batch_size,))
    x = torch.stack([data[i:i + block_size] for i in ix])
    y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])   # the same, shifted by one
    return x.to(device), y.to(device)


@torch.no_grad()
def estimate_loss(model, splits, args, device):
    model.eval()
    out = {}
    for name, data in splits.items():
        losses = torch.zeros(args.eval_iters)
        for k in range(args.eval_iters):
            x, y = get_batch(data, args.block_size, args.batch_size, device)
            losses[k] = model(x, y)[1].item()
        out[name] = losses.mean().item()
    model.train()
    return out


# ----------------------------------------------------------------------------- main
def main():
    default_data = os.path.join(os.environ.get("WS_DIR", "."), "data", "input.txt")
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--data", default=default_data)
    p.add_argument("--out", default="runs/tinygpt")
    p.add_argument("--steps", type=int, default=3000)
    p.add_argument("--lr", type=float, default=1e-3)
    p.add_argument("--batch-size", type=int, default=64)
    p.add_argument("--block-size", type=int, default=128)
    p.add_argument("--n-layer", type=int, default=4)
    p.add_argument("--n-head", type=int, default=4)
    p.add_argument("--n-embd", type=int, default=128)
    p.add_argument("--dropout", type=float, default=0.1)
    p.add_argument("--eval-interval", type=int, default=250)
    p.add_argument("--eval-iters", type=int, default=50)
    p.add_argument("--seed", type=int, default=1337)
    p.add_argument("--device", default="auto", help="auto | cpu | cuda")
    args = p.parse_args()

    torch.manual_seed(args.seed)
    device = torch.device("cuda" if (args.device == "auto" and torch.cuda.is_available()) else
                          ("cpu" if args.device == "auto" else args.device))
    if device.type == "cuda":
        torch.backends.cuda.matmul.allow_tf32 = True     # free speed-up on A100 / H100
    os.makedirs(args.out, exist_ok=True)

    # --- data: characters -> integers
    text = load_text(args.data)
    chars = sorted(set(text))
    stoi = {ch: i for i, ch in enumerate(chars)}
    data = torch.tensor([stoi[c] for c in text], dtype=torch.long)
    n = int(0.9 * len(data))
    splits = {"train": data[:n], "val": data[n:]}

    # --- model + optimiser
    model = TinyGPT(len(chars), args.block_size, args.n_layer, args.n_head, args.n_embd, args.dropout).to(device)
    optim = torch.optim.AdamW(model.parameters(), lr=args.lr)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"device={device}  vocab={len(chars)}  chars={len(text):,}  params={n_params / 1e6:.2f}M  "
          f"job={os.environ.get('SLURM_JOB_ID', 'local')}", flush=True)

    # --- save a checkpoint before Slurm kills us
    stop = {"flag": False}

    def on_sigterm(signum, frame):
        print("Got SIGTERM, saving a checkpoint and stopping", flush=True)
        stop["flag"] = True

    signal.signal(signal.SIGTERM, on_sigterm)

    def save(step, losses):
        torch.save({"model": model.state_dict(), "chars": chars, "args": vars(args), "step": step},
                   os.path.join(args.out, "ckpt.pt"))
        with open(os.path.join(args.out, "metrics.json"), "w") as f:
            json.dump({"step": step, **losses, "lr": args.lr, "params": n_params,
                       "seconds": round(time.time() - t0, 1), "device": str(device),
                       "job_id": os.environ.get("SLURM_JOB_ID", "local"),
                       "array_task": os.environ.get("SLURM_ARRAY_TASK_ID", "-")}, f, indent=2)

    # --- training loop: forward, loss, backward, update (the same as in notebook 3b)
    t0 = time.time()
    losses = {"train": float("nan"), "val": float("nan")}
    step = 0
    for step in range(1, args.steps + 1):
        x, y = get_batch(splits["train"], args.block_size, args.batch_size, device)
        _, loss = model(x, y)
        optim.zero_grad(set_to_none=True)
        loss.backward()
        optim.step()

        if step % args.eval_interval == 0 or step == args.steps or stop["flag"]:
            losses = estimate_loss(model, splits, args, device)
            tok_s = step * args.batch_size * args.block_size / (time.time() - t0)
            print(f"step {step:5d} | train {losses['train']:.4f} | val {losses['val']:.4f} | "
                  f"{tok_s / 1e3:7.1f}k tok/s | {time.time() - t0:6.1f}s", flush=True)
            save(step, losses)
        if stop["flag"]:
            break

    # --- show what it learned
    model.eval()
    start = torch.tensor([[stoi.get("\n", 0)]], device=device)
    sample = "".join(chars[i] for i in model.generate(start, 500, temperature=0.8)[0].tolist())
    with open(os.path.join(args.out, "samples.txt"), "w", encoding="utf-8") as f:
        f.write(sample)
    print("-" * 60 + "\n" + sample + "\n" + "-" * 60, flush=True)
    print(f"done: {step} steps in {time.time() - t0:.1f}s -> {args.out}", flush=True)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Let a trained TinyGPT write some text.

    python sample.py --ckpt $WS_DIR/runs/base/ckpt.pt --prompt "ROMEO:" --temperature 0.8

Nice to try in an interactive session: change --temperature and --top-k and see
what happens.
"""
import argparse

import torch

from train_tinygpt import TinyGPT


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--ckpt", required=True)
    p.add_argument("--prompt", default="\n")
    p.add_argument("--tokens", type=int, default=400)
    p.add_argument("--temperature", type=float, default=0.8)
    p.add_argument("--top-k", type=int, default=None)
    p.add_argument("--seed", type=int, default=0)
    args = p.parse_args()

    torch.manual_seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    ck = torch.load(args.ckpt, map_location=device)
    a, chars = ck["args"], ck["chars"]
    stoi = {c: i for i, c in enumerate(chars)}

    model = TinyGPT(len(chars), a["block_size"], a["n_layer"], a["n_head"], a["n_embd"], 0.0).to(device)
    model.load_state_dict(ck["model"])
    model.eval()

    unknown = sorted(set(args.prompt) - set(chars))
    if unknown:
        print(f"(ignoring characters the model never saw: {unknown})")
    ids = [stoi[c] for c in args.prompt if c in stoi] or [0]
    out = model.generate(torch.tensor([ids], device=device), args.tokens, args.temperature, args.top_k)
    print("".join(chars[i] for i in out[0].tolist()))


if __name__ == "__main__":
    main()

#!/bin/bash
# Exercise 5: send the array job, then a second job that waits for it.
#
#   source ../00_setup/setup_env.sh      # so that $WS_DIR is set
#   mkdir -p logs && ./submit_sweep.sh
set -e
: "${WS_DIR:?run 'source ../00_setup/setup_env.sh' first}"

# --parsable makes sbatch print only the job ID, so we can keep it in a variable
SWEEP_ID=$(sbatch --parsable sweep.sbatch)
echo "sweep (array) job : $SWEEP_ID"

COLLECT_ID=$(sbatch --parsable --dependency=afterok:"$SWEEP_ID" \
                    --export=ALL,SWEEP_ID="$SWEEP_ID" collect.sbatch)
echo "collect job       : $COLLECT_ID (waits for $SWEEP_ID, squeue shows it as Dependency)"
echo
echo "watch with:  squeue --me"
echo "results in:  logs/lr-collect-$COLLECT_ID.out"

#!/bin/bash
# Exercise 5: a learning-rate sweep as a job array
#
#   cd 05_sweep && mkdir -p logs && ./submit_sweep.sh     (sends the sweep and the summary job)
#   squeue --me                     # you will see <jobid>_0, <jobid>_1, ...
#   scancel <jobid>_3               # cancel just one task of the array
#
# Slurm starts this script 4 times. The only thing that changes is
# $SLURM_ARRAY_TASK_ID, which goes from 0 to 3.
#
#SBATCH --job-name=lr-sweep
#SBATCH --account=p_lv_llmsecrets
#SBATCH --array=0-3                   # 4 tasks; with 0-3%2 at most 2 run at the same time
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=6
#SBATCH --gres=gpu:1                  # every task gets its own GPU
#SBATCH --mem=32G
#SBATCH --time=00:20:00
#SBATCH --output=logs/%x-%A_%a.out    # %A is the array job ID, %a the task number
#SBATCH --error=logs/%x-%A_%a.err

LEARNING_RATES=(3e-3 1e-3 3e-4 1e-4)
LR=${LEARNING_RATES[$SLURM_ARRAY_TASK_ID]}
echo "array task $SLURM_ARRAY_TASK_ID of job $SLURM_ARRAY_JOB_ID uses lr=$LR on $(hostname)"

source ../00_setup/setup_env.sh

export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK
srun python -u ../04_tinygpt/train_tinygpt.py \
    --data "$WS_DIR/data/input.txt" \
    --out  "$WS_DIR/runs/sweep-$SLURM_ARRAY_JOB_ID/lr-$LR" \
    --steps 2000 \
    --lr "$LR"

#!/bin/bash
# Exercise 5: a tiny CPU job that prints the sweep results.
# submit_sweep.sh sends it with --dependency=afterok:<array job id>,
# so it only starts after all array tasks have finished without errors.
# It asks for a GPU only because Alpha accepts nothing else.
#
#SBATCH --job-name=lr-collect
#SBATCH --account=p_lv_llmsecrets
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=1
#SBATCH --mem-per-cpu=1G
#SBATCH --gres=gpu:1
#SBATCH --time=00:05:00
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err

# SWEEP_ID comes from submit_sweep.sh (sbatch --export=ALL,SWEEP_ID=...)
python3 collect_results.py "$WS_DIR/runs/sweep-$SWEEP_ID"

#!/usr/bin/env python3
"""
Exercise 5: print a small table of all runs in a sweep, best one first.
It reads every metrics.json below the folder you give it.

    python collect_results.py $WS_DIR/runs/sweep-1234567
"""
import glob
import json
import os
import sys


def main():
    root = sys.argv[1] if len(sys.argv) > 1 else "."
    rows = []
    for path in sorted(glob.glob(os.path.join(root, "**", "metrics.json"), recursive=True)):
        with open(path) as f:
            m = json.load(f)
        m["run"] = os.path.basename(os.path.dirname(path))
        rows.append(m)

    if not rows:
        sys.exit(f"no metrics.json found under {root}")

    rows.sort(key=lambda m: m.get("val", float("inf")))
    print(f"{'run':<14}{'lr':>10}{'steps':>8}{'train':>10}{'val':>10}{'secs':>8}  device")
    for m in rows:
        print(f"{m['run']:<14}{m['lr']:>10.0e}{m['step']:>8}{m['train']:>10.4f}{m['val']:>10.4f}"
              f"{m['seconds']:>8.0f}  {m['device']}")
    best = rows[0]
    print(f"\nbest: {best['run']}  (val loss {best['val']:.4f})")


if __name__ == "__main__":
    main()

#!/bin/bash
# Exercise 6: let a pretrained language model answer a file of prompts
#
#   source ../00_setup/setup_env.sh
#   python generate.py --download-only        # on the login node, fetches the model once
#   mkdir -p logs && sbatch generate.sbatch
#
# Things to try:
#   Another model:   sbatch generate.sbatch --model gpt2
#   CPU vs GPU:      copy this file, remove the --gres line, ask for more CPUs
#   Languages:       do the German and Hindi prompts get answers in the same language?
#
#SBATCH --job-name=hf-generate
#SBATCH --account=p_lv_llmsecrets
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=6
#SBATCH --gres=gpu:1
#SBATCH --mem=32G
#SBATCH --time=00:15:00
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err

source ../00_setup/setup_env.sh
export HF_HUB_OFFLINE=1               # use the downloaded copy, don't wait for the network

export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK
srun python -u generate.py --out "outputs-$SLURM_JOB_ID.jsonl" "$@"

#!/usr/bin/env python3
"""
Exercise 6: run a small pretrained language model from Hugging Face on a GPU node.

    # 1) on the login node: download the model once into $HF_HOME (in your workspace)
    python generate.py --download-only

    # 2) as a GPU job: answer every line of prompts.txt
    sbatch generate.sbatch

The default model is Qwen/Qwen2.5-0.5B-Instruct. It has 0.5 billion parameters,
speaks many languages and needs about 1 GB in bf16. You can also try --model gpt2,
which only knows English and has no chat template.
"""
import argparse
import json
import os
import time

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


def build_inputs(tok, prompts):
    """Instruct models expect their chat format; base models just get the raw text."""
    if getattr(tok, "chat_template", None):
        return [tok.apply_chat_template([{"role": "user", "content": p}], tokenize=False,
                                        add_generation_prompt=True) for p in prompts]
    return prompts


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="Qwen/Qwen2.5-0.5B-Instruct")
    p.add_argument("--prompts", default="prompts.txt")
    p.add_argument("--out", default="outputs.jsonl")
    p.add_argument("--max-new-tokens", type=int, default=120)
    p.add_argument("--temperature", type=float, default=0.7)
    p.add_argument("--top-p", type=float, default=0.9)
    p.add_argument("--batch-size", type=int, default=8)
    p.add_argument("--download-only", action="store_true")
    args = p.parse_args()

    print(f"HF_HOME = {os.environ.get('HF_HOME', '~/.cache/huggingface (careful: $HOME quota!)')}", flush=True)
    if args.download_only:                    # only fetch the files, do not load the model
        from huggingface_hub import snapshot_download
        path = snapshot_download(args.model)
        print(f"downloaded {args.model} to {path}. Now submit the GPU job.")
        return
    tok = AutoTokenizer.from_pretrained(args.model)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
    model = AutoModelForCausalLM.from_pretrained(args.model, dtype=dtype).to(device).eval()
    n_params = sum(p.numel() for p in model.parameters())
    print(f"model {args.model}: {n_params / 1e6:.0f}M params on {device} ({dtype})", flush=True)

    tok.padding_side = "left"                 # decoder-only models: pad on the left for batching
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token

    with open(args.prompts, encoding="utf-8") as f:
        prompts = [line.strip() for line in f if line.strip() and not line.startswith("#")]

    n_new_total, t0 = 0, time.time()
    with open(args.out, "w", encoding="utf-8") as fout:
        for i in range(0, len(prompts), args.batch_size):
            batch = prompts[i:i + args.batch_size]
            enc = tok(build_inputs(tok, batch), return_tensors="pt", padding=True).to(device)
            with torch.no_grad():
                out = model.generate(**enc, max_new_tokens=args.max_new_tokens, do_sample=True,
                                     temperature=args.temperature, top_p=args.top_p,
                                     pad_token_id=tok.pad_token_id)
            new_tokens = out[:, enc["input_ids"].shape[1]:]          # drop the prompt part
            for prompt, ids in zip(batch, new_tokens):
                answer = tok.decode(ids, skip_special_tokens=True).strip()
                n_new_total += int((ids != tok.pad_token_id).sum())
                fout.write(json.dumps({"prompt": prompt, "answer": answer}, ensure_ascii=False) + "\n")
                print(f"\n>>> {prompt}\n{answer}", flush=True)

    dt = time.time() - t0
    print(f"\n{len(prompts)} prompts, {n_new_total} new tokens in {dt:.1f}s "
          f"-> {n_new_total / dt:.1f} tokens/s  (written to {args.out})")


if __name__ == "__main__":
    main()

# one prompt per line; lines starting with # are ignored
Explain in two sentences what a GPU is good for.
What is the capital of Saxony, and what is it famous for?
Write a haiku about a supercomputer that is waiting in a queue.
Erkläre in zwei Sätzen, was ein Sprachmodell ist.
एक वाक्य में बताइए कि कृत्रिम बुद्धिमत्ता क्या है।
Translate to French: "The job was cancelled because it ran out of time."
A train leaves Dresden at 9:15 and arrives in Berlin at 11:05. How long is the trip?
List three risks of large language models.