Add the Unsloth warm-fork demo (one golden, three concurrent fine-tunes)

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.

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)
+──────────────────────────────────────────────────────────────

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

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}")