Container & Slurm Recipes: Bridging Single 4090 to Cloud Multi-Node

#!/bin/bash
#SBATCH --job-name=ftc-llama-70b-qlora
#SBATCH --nodes=2                            # 2 node × 8 GPU = 16 GPU
#SBATCH --ntasks-per-node=8                  # 1 task per GPU
#SBATCH --gres=gpu:h100:8
#SBATCH --cpus-per-task=16                   # 8 CPU per GPU
#SBATCH --mem=0                              # all available
#SBATCH --time=08:00:00                      # walltime limit
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err
#SBATCH --signal=B:SIGUSR1@90                # 90s before walltime → graceful shutdown
#SBATCH --requeue                            # preemption sonrası yeniden zamanla
 
set -euo pipefail
 
# === Repro environment ===
export PYTHONHASHSEED=42
export CUBLAS_WORKSPACE_CONFIG=:4096:8
export TOKENIZERS_PARALLELISM=false
export HF_HUB_ENABLE_HF_TRANSFER=1
export HF_HOME=/scratch/$USER/hf-cache
 
# === NCCL & distributed ===
export NCCL_DEBUG=WARN                       # INFO debug için, prod'da WARN
export NCCL_ASYNC_ERROR_HANDLING=1
export NCCL_IB_GID_INDEX=3                   # InfiniBand RoCE
export NCCL_IB_DISABLE=0                     # IB açık (yoksa TCP'ye düşer)
export NCCL_SOCKET_IFNAME=^lo,docker0        # exclude loopback & docker
export NCCL_NSOCKS_PERTHREAD=4
export NCCL_SOCKET_NTHREADS=2
export OMP_NUM_THREADS=4
 
# === Master addr — Slurm yardımıyla ===
MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n1)
MASTER_PORT=$((RANDOM % 10000 + 30000))
export MASTER_ADDR MASTER_PORT
echo "[$(date)] MASTER_ADDR=$MASTER_ADDR  MASTER_PORT=$MASTER_PORT"
 
# === Graceful preemption handler ===
trap 'echo "[$(date)] Got SIGUSR1, snapshotting checkpoint then requeue"; \
      scancel --signal=USR1 $SLURM_JOBID; \
      wait; \
      scontrol requeue $SLURM_JOBID; \
      exit 0' SIGUSR1
 
# === Run ===
srun --label \
  apptainer run --nv \
    --bind /scratch:/scratch \
    --bind /home/$USER:/home/$USER \
    ftc.sif \
    bash -c "\
      cd /workspace && \
      uv run torchrun \
        --nproc_per_node=8 \
        --nnodes=$SLURM_NNODES \
        --node_rank=$SLURM_NODEID \
        --master_addr=$MASTER_ADDR \
        --master_port=$MASTER_PORT \
        --rdzv_id=$SLURM_JOB_ID \
        --rdzv_backend=c10d \
        --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
        train.py \
        --config configs/llama3_70b_qlora_fsdp.yaml \
        --resume_from_checkpoint /scratch/$USER/ckpts/last \
    "
 

# nccl_test.py — Multi-node NCCL bandwidth & latency sanity check
import os, time
import torch
import torch.distributed as dist
 
def main():
    rank = int(os.environ["RANK"])
    world = int(os.environ["WORLD_SIZE"])
    local = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local)
    dist.init_process_group("nccl", rank=rank, world_size=world)
 
    if rank == 0:
        print(f"world={world}  device={torch.cuda.get_device_name(0)}")
 
    # 1) All-reduce bandwidth — eğitimin "gradient sync" benchmarki
    for size_mb in [1, 4, 16, 64, 256, 1024]:
        x = torch.randn(size_mb * 256 * 1024, device="cuda")  # MB → fp32 floats
        # Warmup
        for _ in range(5):
            dist.all_reduce(x)
        torch.cuda.synchronize()
        # Bench
        t0 = time.perf_counter()
        for _ in range(20):
            dist.all_reduce(x)
        torch.cuda.synchronize()
        t = (time.perf_counter() - t0) / 20
        bw = (size_mb * 1024 * 1024 * 4) / t / 1e9  # GB/s
        if rank == 0:
            print(f"all-reduce {size_mb:5d} MB  →  {t*1000:7.2f} ms   ~{bw:6.1f} GB/s")
 
    # 2) Send/recv ping
    if rank == 0:
        x = torch.zeros(1, device="cuda")
        dist.send(x, dst=1)
        dist.recv(x, src=1)
        print("ping/pong OK")
    elif rank == 1:
        x = torch.zeros(1, device="cuda")
        dist.recv(x, src=0)
        dist.send(x, dst=0)
 
    dist.barrier()
    dist.destroy_process_group()
 
if __name__ == "__main__":
    main()

# ckpt_manifest.py — Atomik, resumable checkpoint kaydetme
import os, json, shutil, hashlib
from pathlib import Path
import torch
from accelerate import Accelerator
 
class CheckpointManager:
    """
    Atomik checkpoint kaydı + manifest.
    Yarı-yazılmış ckpt'yi tespit eder ve önceki sağlamına döner.
    """
    def __init__(self, out_dir: str, accelerator: Accelerator, keep_last: int = 3):
        self.out = Path(out_dir)
        self.out.mkdir(parents=True, exist_ok=True)
        self.acc = accelerator
        self.keep_last = keep_last
        self.manifest_path = self.out / "manifest.json"
 
    def _hash_state(self, state: dict) -> str:
        h = hashlib.sha256()
        for k in sorted(state):
            v = state[k]
            if torch.is_tensor(v):
                h.update(v.cpu().numpy().tobytes())
            else:
                h.update(json.dumps(v, sort_keys=True, default=str).encode())
        return h.hexdigest()[:16]
 
    def save(self, step: int, model, optimizer, scheduler, scaler, metadata):
        if not self.acc.is_main_process:
            self.acc.wait_for_everyone()
            return
 
        tag = f"step-{step:08d}"
        staging = self.out / f"{tag}.partial"
        final = self.out / tag
 
        staging.mkdir(exist_ok=True)
        # accelerate state
        self.acc.save_state(staging)
        # extra
        torch.save({"scheduler": scheduler.state_dict(),
                    "scaler": scaler.state_dict() if scaler else None,
                    "metadata": metadata}, staging / "extras.pt")
 
        manifest = {
            "step": step,
            "tag": tag,
            "metadata": metadata,
            "files": sorted(p.name for p in staging.iterdir()),
        }
        with open(staging / "manifest.json", "w") as f:
            json.dump(manifest, f, indent=2)
 
        # Atomic rename
        if final.exists():
            shutil.rmtree(final)
        os.rename(staging, final)
 
        # Update top-level manifest (atomik write)
        tmp_top = self.manifest_path.with_suffix(".tmp")
        with open(tmp_top, "w") as f:
            json.dump({"last": tag, "step": step, "metadata": metadata}, f, indent=2)
        os.replace(tmp_top, self.manifest_path)
 
        # Cleanup old
        ckpts = sorted([d for d in self.out.iterdir() if d.is_dir() and d.name.startswith("step-")])
        for old in ckpts[:-self.keep_last]:
            shutil.rmtree(old)
 
        self.acc.wait_for_everyone()
 
    def load_latest(self):
        if not self.manifest_path.exists():
            return None
        with open(self.manifest_path) as f:
            mf = json.load(f)
        return self.out / mf["tag"]

# train.py içinde
import signal
 
class GracefulPreemption:
    def __init__(self):
        self.requested = False
        signal.signal(signal.SIGUSR1, self._handler)
        signal.signal(signal.SIGTERM, self._handler)
 
    def _handler(self, signum, frame):
        print(f"[preempt] received signal {signum} — flushing checkpoint")
        self.requested = True
 
preempt = GracefulPreemption()
 
# Training loop
for step, batch in enumerate(loader):
    loss = train_step(batch)
    if step % 50 == 0 or preempt.requested:
        ckpt_mgr.save(step, model, opt, sched, scaler,
                      metadata={"loss": loss.item(), "preempted": preempt.requested})
        if preempt.requested:
            print("[preempt] checkpoint saved, exiting cleanly")
            break