diff --git a/demos/unsloth-fork-sweep/README.md b/demos/unsloth-fork-sweep/README.md new file mode 100644 index 00000000..ac0fcc67 --- /dev/null +++ b/demos/unsloth-fork-sweep/README.md @@ -0,0 +1,39 @@ +# Concurrent Unsloth fine-tunes from one loaded model + +`smolvm machine fork` clones a running CUDA VM. This demo boots a golden VM +that loads Qwen2.5-1.5B (4-bit) with unmodified Unsloth and freezes; three +forked clones then fine-tune different tasks concurrently. + +``` +golden (frozen) ──fork──► uns-add ┐ + ──fork──► uns-mul │ train concurrently, isolated + ──fork──► uns-sub ┘ +``` + +Measured on an RTX 3070, 8 GB (`demo-output.txt`): cold start to first +training step ~54 s; fork to first training step ~4 s; peak VRAM 3.9 GB for +the golden plus all three clones, with base weights shared across them. + +## Run it + +```sh +./demo.sh # defaults to Qwen2.5-1.5B 4-bit +SMOLVM_DEMO_VENV=~/ptwork ./demo.sh unsloth/Qwen2.5-0.5B-Instruct-bnb-4bit +``` + +Requirements: linux + NVIDIA driver, a smolvm build with CUDA forking, a host +venv containing unsloth/torch (mounted read-only into the VM), and the smolvm +CUDA shim libs in `./drvlib` (`libcudart.so.12`, `libcuda.so.1`). The first +run downloads the model into the coord mount's HF cache; later runs are +offline. + +## How it works + +- The golden loads the model + LoRA, runs `fix_untrained_tokens` once + (recommended for 3+ clones with weight sharing), writes a READY marker, and + blocks on a GO file. `machine fork` snapshots it in that state. +- Each clone claims a distinct task via `O_CREAT|O_EXCL` files on the shared + mount, builds its trainer, trains, and writes its result. +- Weight sharing is opt-in (`SMOLVM_CUDA_FORK_SHARE_WEIGHTS=1`): a chunk is + shared only if fork-time verification shows its device content still matches + the uploaded weights; everything else is copied per clone. diff --git a/demos/unsloth-fork-sweep/demo-output.txt b/demos/unsloth-fork-sweep/demo-output.txt new file mode 100644 index 00000000..f44b8dc8 --- /dev/null +++ b/demos/unsloth-fork-sweep/demo-output.txt @@ -0,0 +1,22 @@ +────────────────────────────────────────────────────────────── + smolvm CUDA fork demo + model: unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit GPU baseline: 609 MiB +────────────────────────────────────────────────────────────── +[1/3] boot golden VM, load model + golden ready in 55s (2595 MiB); frozen as the fork base + +[2/3] fork three clones + uns-add forked in 0.4s + uns-mul forked in 0.4s + uns-sub forked in 0.3s + gate released; clones train concurrently (add / mul / sub) + +[3/3] results +────────────────────────────────────────────────────────────── + task GO→training loss ask it verdict + add 5.7s 2.60 → 1.39 6+7 = 13 PASS (want 13) + mul 6.7s 2.68 → 1.40 6*7 = 42 PASS (want 42) + sub 7.6s 3.02 → 1.61 6-7 = -1 PASS (want -1) +────────────────────────────────────────────────────────────── + peak VRAM: 3794 MiB (base weights shared across the clones) +────────────────────────────────────────────────────────────── diff --git a/demos/unsloth-fork-sweep/demo.sh b/demos/unsloth-fork-sweep/demo.sh new file mode 100755 index 00000000..3b146929 --- /dev/null +++ b/demos/unsloth-fork-sweep/demo.sh @@ -0,0 +1,68 @@ +#!/bin/bash +# Boots a golden VM that loads a model with Unsloth and freezes, then forks +# three clones that each fine-tune a different task concurrently. +# +# ./demo.sh [model] (default unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit) +# +# Requirements: see README.md. `machine start --forkable` enables CUDA forking; +# the one env var below opts into cross-clone weight sharing. +set -u +MODEL=${1:-unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit} +S=${SMOLVM_BIN:-smolvm} +VENV=${SMOLVM_DEMO_VENV:-$HOME/ptwork} # host venv with unsloth +DRV=${SMOLVM_DEMO_DRV:-$(dirname "$0")/drvlib} # CUDA shim libs +COORD=${SMOLVM_DEMO_COORD:-$(mktemp -d)} # shared clone<->host mount +WORKLOAD="$(cd "$(dirname "$0")" && pwd)/workload.py" +export SMOLVM_CUDA_FORK_SHARE_WEIGHTS=1 +gpumem() { nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits; } +line() { printf '%s\n' "──────────────────────────────────────────────────────────────"; } + +cp "$WORKLOAD" "$COORD/workload.py" +rm -f "$COORD"/GO "$COORD"/claim_* "$COORD"/result_* "$COORD"/marks.txt +for m in add mul sub golden; do $S machine stop --name uns-$m >/dev/null 2>&1; $S machine rm --name uns-$m --force >/dev/null 2>&1; done + +line +echo " smolvm CUDA fork demo" +echo " model: $MODEL GPU baseline: $(gpumem) MiB" +line +echo "[1/3] boot golden VM, load model" +T0=$(date +%s) +$S machine create --name uns-golden --cuda --net \ + -v "$VENV:$VENV:ro" -v "$DRV:/opt/drvlib:ro" -v "$COORD:/opt/coord:rw" \ + --image debian:bookworm-slim --storage 20 --overlay 15 -- sh -c " + export LD_PRELOAD='/opt/drvlib/libcudart.so.12 /opt/drvlib/libcuda.so.1' \ + HF_HOME=/opt/coord/hf CC=gcc HF_HUB_DISABLE_TELEMETRY=1 \ + TRITON_CACHE_DIR=/root/.triton SMOLVM_MODEL='$MODEL' + apt-get update -qq >/dev/null 2>&1 && apt-get install -y -qq gcc ca-certificates >/dev/null 2>&1 + ln -sf /opt/drvlib/libcuda.so.1 /usr/lib/x86_64-linux-gnu/libcuda.so + $VENV/venv/bin/python /opt/coord/workload.py + " >/dev/null 2>&1 +$S machine start --forkable --name uns-golden >/dev/null 2>&1 +until grep -q READY "$COORD/marks.txt" 2>/dev/null; do sleep 2; done +echo " golden ready in $(( $(date +%s) - T0 ))s ($(gpumem) MiB); frozen as the fork base" +echo +echo "[2/3] fork three clones" +for m in add mul sub; do + TF=$(date +%s.%N) + $S machine fork --golden uns-golden --name uns-$m >/dev/null 2>&1 + printf " uns-%s forked in %.1fs\n" "$m" "$(echo "$(date +%s.%N) $TF" | awk '{print $1-$2}')" +done +echo go > "$COORD/GO" +echo " gate released; clones train concurrently (add / mul / sub)" +echo +PEAK=0 +until [ "$(ls "$COORD"/result_* 2>/dev/null | wc -l)" -ge 3 ]; do + M=$(gpumem); [ "$M" -gt "$PEAK" ] && PEAK=$M + sleep 2 +done +echo "[3/3] results" +line +printf " %-5s %-14s %-14s %-12s %s\n" task "GO→training" "loss" "ask it" verdict +for f in "$COORD"/result_*.txt; do + IFS='|' read -r name t l0 l1 q gen expect verdict < "$f" + printf " %-5s %-14s %-14s %-12s %s\n" "$name" "${t}s" "$l0 → $l1" "$q = $gen" "$verdict (want $expect)" +done +line +echo " peak VRAM: ${PEAK} MiB (base weights shared across the clones)" +line +for m in add mul sub golden; do $S machine stop --name uns-$m >/dev/null 2>&1; $S machine rm --name uns-$m --force >/dev/null 2>&1; done diff --git a/demos/unsloth-fork-sweep/workload.py b/demos/unsloth-fork-sweep/workload.py new file mode 100644 index 00000000..68e63fd8 --- /dev/null +++ b/demos/unsloth-fork-sweep/workload.py @@ -0,0 +1,72 @@ +# Demo workload: the golden loads the model and blocks on a GO file; each fork +# claims one task over the shared mount and fine-tunes it. See README.md. +import os, time + +os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1" +GO = "/opt/coord/GO" +MODEL = os.environ.get("SMOLVM_MODEL", "unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit") +def mark(m): + open("/opt/coord/marks.txt", "a").write(f"{time.time():.2f} {m}\n") + +t0 = time.time() +from unsloth import FastLanguageModel +model, tok = FastLanguageModel.from_pretrained(MODEL, max_seq_length=256, load_in_4bit=True) +model = FastLanguageModel.get_peft_model( + model, r=8, lora_alpha=16, + target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], + use_gradient_checkpointing="unsloth", random_state=42) + +# Run fix_untrained_tokens in the golden so clones don't write the embedding +# at trainer setup (recommended for 3+ clones with weight sharing). +from unsloth_zoo.tokenizer_utils import fix_untrained_tokens +from datasets import Dataset +_ds = Dataset.from_list([{"text": f"### Q: {a}+{b}?\n### A: {a+b}"} + for a in range(1, 13) for b in range(1, 13)]) +FastLanguageModel.for_training(model) +fix_untrained_tokens(model, tok, _ds, [], eps=1e-16) +import torch; torch.cuda.synchronize() +FastLanguageModel.for_inference(model) +mark(f"READY load={time.time()-t0:.1f}s model={MODEL}") + +while not os.path.exists(GO): + time.sleep(0.2) # golden freezes here; forks are taken during this wait +t_go = time.time() + +# Claim one task per fork (O_EXCL on the shared mount). +TASKS = [("add", "+", lambda a, b: a + b), + ("mul", "*", lambda a, b: a * b), + ("sub", "-", lambda a, b: a - b)] +idx = None +for i in range(len(TASKS)): + try: + fd = os.open(f"/opt/coord/claim_{i}", os.O_CREAT | os.O_EXCL | os.O_WRONLY) + os.close(fd); idx = i; break + except FileExistsError: + continue +if idx is None: + mark("golden idle (all tasks claimed)") +else: + name, op, fn = TASKS[idx] + ds = Dataset.from_list([{"text": f"### Q: {a}{op}{b}?\n### A: {fn(a,b)}"} + for a in range(1, 13) for b in range(1, 13)]) + from trl import SFTTrainer, SFTConfig + FastLanguageModel.for_training(model) + tr = SFTTrainer(model=model, tokenizer=tok, train_dataset=ds, args=SFTConfig( + per_device_train_batch_size=2, max_steps=int(os.environ.get("SMOLVM_STEPS", "20")), + learning_rate=3e-4, logging_steps=10, optim="adamw_8bit", seed=42, + output_dir=f"/root/o{idx}", report_to=[], + dataset_text_field="text", max_seq_length=256, warmup_steps=2)) + t_train = time.time() + mark(f"clone[{name}] training after {t_train-t_go:.1f}s") + tr.train() + losses = [h["loss"] for h in tr.state.log_history if "loss" in h] + FastLanguageModel.for_inference(model) + q = f"### Q: 6{op}7?\n### A:" + ids = tok(q, return_tensors="pt").to("cuda") + gen = tok.decode(model.generate(**ids, max_new_tokens=5, do_sample=False)[0], + skip_special_tokens=True)[len(q):].strip() + gen = (gen.split() or [""])[0][:8] # first token; keep the result file one-line + ok = gen == str(fn(6, 7)) + open(f"/opt/coord/result_{idx}.txt", "w").write( + f"{name}|{t_train-t_go:.1f}|{losses[0]:.2f}|{losses[-1]:.2f}|6{op}7|{gen}|{fn(6,7)}|{'PASS' if ok else 'FAIL'}\n") + mark(f"clone[{name}] done ok={ok}")