#!/usr/bin/env python3 """Work out a --tensor-split before you run one, and say what each host would hold. Purpose: turn a list of devices, each with the memory it can actually give the model and optionally the decode rate it managed on its own, into two candidate splits: one in proportion to memory, which is what llama.cpp does by default, and one in proportion to speed, which is what you want when the hosts are not equally fast. It prints the gigabytes and the layers each device would take under each plan, flags any device that would be asked to hold more than it has, and gives the split string in both the comma form llama-server takes and the slash form llama-bench takes. Platform: all (pure Python; nothing platform-specific) Minimum memory: 8 GB on whichever machine you run it on; it loads nothing Assumes: Python 3.9 or later; the usable-memory figures come from your own measurements, not from the box, because the memory a device will give a model is smaller than the memory it has. Usage: python3 plan-tensor-split.py --weights-gb 125.5 --layers 94 \ --device node-a:110 --device node-b:110 python3 plan-tensor-split.py --weights-gb 18.6 --layers 48 --kv-gb 1.6 \ --device desktop:22:38 --device strix:90:12 --plan speed A device given as name:memory_gb has no speed, so the speed plan falls back to equal shares for it and says so. Numbers here are arithmetic, not measurements: they tell you what will fit and roughly where the work will land, and the run tells you the rest. """ from __future__ import annotations import argparse import sys def parse_device(spec: str) -> dict: parts = spec.split(":") if len(parts) not in (2, 3): raise SystemExit(f"--device {spec!r}: expected name:memory_gb or name:memory_gb:tokens_per_s") name = parts[0] try: memory = float(parts[1]) rate = float(parts[2]) if len(parts) == 3 else None except ValueError as exc: raise SystemExit(f"--device {spec!r}: {exc}") from exc if memory <= 0: raise SystemExit(f"--device {spec!r}: usable memory must be greater than zero") return {"name": name, "memory_gb": memory, "rate": rate} def normalise(values: list[float]) -> list[float]: total = sum(values) return [v / total for v in values] if total > 0 else [1.0 / len(values)] * len(values) def whole_layers(shares: list[float], layers: int) -> list[int]: """Largest-remainder allocation, so the layers add up to exactly `layers`.""" exact = [s * layers for s in shares] base = [int(x) for x in exact] remaining = layers - sum(base) order = sorted(range(len(exact)), key=lambda i: exact[i] - base[i], reverse=True) for i in order[:remaining]: base[i] += 1 return base def report(plan: str, devices: list[dict], shares: list[float], args) -> bool: payload = args.weights_gb + args.kv_gb layers = whole_layers(shares, args.layers) print(f"\n== split by {plan} ==") print(f"{'device':<20}{'share':>9}{'GB held':>10}{'usable GB':>12}{'layers':>9}") over = False for device, share, count in zip(devices, shares, layers): held = payload * share flag = "" if held > device["memory_gb"]: flag = " <-- more than this device has" over = True print( f"{device['name']:<20}{share:>9.3f}{held:>10.1f}{device['memory_gb']:>12.1f}" f"{count:>9d}{flag}" ) proportions = [f"{s:.3f}".rstrip("0").rstrip(".") for s in shares] print(f"\n llama-server / llama-cli: -ts {','.join(proportions)}") print(f" llama-bench: -ts {'/'.join(proportions)}") return over def main() -> int: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--device", action="append", default=[], metavar="NAME:GB[:TOK_S]", help="one device, in the order the client will see it") parser.add_argument("--weights-gb", type=float, required=True, help="size of the weights on disk") parser.add_argument("--kv-gb", type=float, default=0.0, help="KV cache at your context length") parser.add_argument("--layers", type=int, required=True, help="the model's layer count") parser.add_argument("--plan", choices=["memory", "speed", "both"], default="both") args = parser.parse_args() if len(args.device) < 2: print("give at least two --device arguments; a split needs somewhere to split to", file=sys.stderr) return 1 devices = [parse_device(d) for d in args.device] payload = args.weights_gb + args.kv_gb capacity = sum(d["memory_gb"] for d in devices) print(f"model payload : {payload:.1f} GB ({args.weights_gb:.1f} GB weights" f" + {args.kv_gb:.1f} GB KV cache)") print(f"cluster memory: {capacity:.1f} GB across {len(devices)} device(s)") print(f"headroom : {capacity - payload:.1f} GB") if payload > capacity: print("\nThis model does not fit the cluster at all. Take a smaller quantisation," "\nshorten the context, or add a machine. No split rescues it.") return 2 over_anywhere = False if args.plan in ("memory", "both"): over_anywhere |= report("memory", devices, normalise([d["memory_gb"] for d in devices]), args) if args.plan in ("speed", "both"): rates = [d["rate"] for d in devices] if any(r is None for r in rates): mean = sum(r for r in rates if r is not None) / max(1, sum(1 for r in rates if r is not None)) \ if any(r is not None for r in rates) else 1.0 rates = [r if r is not None else mean for r in rates] print("\n(one or more devices had no rate; they were given the average of the others)") over_anywhere |= report("speed", devices, normalise(rates), args) print("\nRead this before using either split:") print(" * The order above must match the order the client sees: local devices first,") print(" then the RPC hosts in the order they appear in RPC_HOSTS. Confirm it with") print(" probe-rpc-devices.sh rather than assuming.") print(" * Usable memory is not installed memory. Leave room for the KV cache, the") print(" compute buffers and whatever else the machine is doing.") print(" * Splitting by speed only helps while every share still fits. A device flagged") print(" above will fall back to host memory or fail to allocate.") return 3 if over_anywhere else 0 if __name__ == "__main__": raise SystemExit(main())