File size: 3,577 Bytes
b890fab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """causalscale LowRankGNN benchmark: measures VRAM scaling (the key 100M claim).
Uses train_lowrank_gnn's actual correlation-reconstruction objective.
"""
import time, json, numpy as np, sys
sys.stdout.reconfigure(encoding='utf-8')
# causalscale installed via pip install -e .
import torch, causalscale as cs
print(f"GPU: {torch.cuda.get_device_name(0)}, VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB")
print("=== causalscale LowRankGNN memory/timing benchmark ===\n")
results = []
for d, r in [(500, 32), (1000, 32), (2000, 16), (5000, 8), (10000, 4)]:
n = min(2000, 2*d)
print(f"d={d}, r={r}, n={n}... ", end='', flush=True)
np.random.seed(42)
X = np.random.randn(n, d).astype(np.float32)
X = (X - X.mean(0)) / (X.std(0) + 1e-8)
torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats()
t0 = time.time()
try:
result_d = cs.train_lowrank_gnn(X, rank=r, epochs=200, lr=3e-4, device='cuda', verbose=False)
t1 = time.time()
mem = torch.cuda.max_memory_allocated()/1e9
f1 = result_d['f1']
nedges = result_d['gnn_edges']
gt_n = result_d['gt_edges']
rv = {'d':d,'r':r,'n':n,'time_s':round(t1-t0,1),'vram_gb':round(mem,2),
'edges':nedges,'f1':f1,'gt_edges':gt_n}
results.append(rv)
print(f"F1={f1:.3f}, {nedges}/{gt_n} edges, {t1-t0:.1f}s, {mem:.2f}GB")
except torch.cuda.OutOfMemoryError:
print("OOM"); results.append({'d':d,'r':r,'vram_gb':'OOM'}); break
out = r'results/memory_scaling.json'
json.dump(results, open(out,'w'), indent=2)
valid = [r for r in results if isinstance(r['vram_gb'],(int,float)) and r['vram_gb']>0]
if len(valid)>=2:
# Model params: U (d*r) + V (d*r) = 2dr float32 params = 8dr bytes
# Adam states: 2x params = 16dr bytes
# Total model: ~24dr bytes + X data (n*d*4 bytes) + intermediates
# Dominant: model params = 24*d*r bytes
ds=[r['d'] for r in valid]; rs=[r['r'] for r in valid]
model_bytes = [24*ds[i]*rs[i]/1e9 for i in range(len(valid))]
data_bytes = [valid[i]['n']*ds[i]*4/1e9 for i in range(len(valid))]
vrams = [r['vram_gb'] for r in valid]
for i in range(len(valid)):
print(f" d={ds[i]}: model={model_bytes[i]:.2f}GB, data={data_bytes[i]:.2f}GB, actual={vrams[i]:.2f}GB")
# VRAM = model + data + overhead. Overhead ~ 20-40% for intermediates
# Extrapolation: d=1e8, r=4
est_model = 24*1e8*4/1e9 # 9.6 GB
est_data_1k = 1000*1e8*4/1e9 # 400 GB for n=1000
est_data_500 = 500*1e8*4/1e9 # 200 GB
print(f"\n=== Extrapolation to d=1e8 ===")
print(f"d=1e8, r=4: model params = {est_model:.1f} GB")
print(f"d=1e8, n=500: data = {est_data_500:.1f} GB")
print(f"d=1e8, n=1000: data = {est_data_1000:.1f} GB")
print(f"FP16 model: {est_model/2:.1f} GB")
print(f"FP16 with gradient checkpoint: {est_model/4:.1f} GB")
print(f"\nKey insight: data (n*d) dominates at d=1e8, not model params (d*r)")
print(f" With n=500 (streaming batches): feasible on A100 80GB")
print(f" With n=100: model+data ~ {est_model+100*1e8*4/1e9:.0f} GB on FP16")
print(f" Consumer 8GB GPU: requires batched processing + CPU offloading")
time_slope = np.polyfit(ds, [r['time_s'] for r in valid], 1)[0] if len(valid)>=2 else 0
if time_slope > 0:
print(f"Time ~ {time_slope*1000:.1f} ms per 1000d")
print(f"d=1e8: ~{time_slope*1e8/3600:.0f} hours (linear scaling)")
print(f"\nSaved {out}")
|