File size: 6,766 Bytes
e79d917 | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """Speed up reranker inference. Zero effect on the scores - only the schedule changes.
MEASURED PROBLEM: GPU 1 oscillates 0% <-> 100% with CPU pinned at ~1.5 cores.
Tokenisation runs inline between GPU forward passes, so the card idles roughly half
the time: 988 pairs/s against a 4,033/s benchmark.
FIXES, in order of expected value:
1. prefetch tokenisation on THREADS, 3 batches ahead, so encoding overlaps the GPU.
Threads not processes: lexical_topk still forks every chunk, and fork + a warmed
fast tokenizer is exactly the deadlock that cost us an hour earlier. Each worker
thread gets its OWN tokenizer instance, so no object is shared across threads.
2. bigger eval batch (1024 -> 2048): we use 21 GB of a 98 GB card, and fewer, larger
launches amortise Python-loop and kernel-launch overhead.
3. torch.inference_mode() instead of no_grad() - skips version counter bookkeeping.
4. SDPA attention when available.
5. non_blocking host->device copies.
Ordering is preserved exactly: batches are consumed in the same sorted order, and the
resume checkpoint still records a prefix of that order.
"""
import io, re, sys
P = "src/cross_encoder.py"
s = io.open(P, encoding="utf-8").read()
if "_prefetch" in s:
print("= already patched")
sys.exit(0)
start = s.index("@torch.no_grad()\ndef predict_cross_encoder(")
new_fn = '''@torch.inference_mode()
def predict_cross_encoder(a, b, model_dir, cfg, cache=None):
"""P(match) per pair. If `cache` is given, scoring is resumable: the finished
prefix is reloaded and only the remainder is computed.
Tokenisation is prefetched on threads so it overlaps GPU compute - see the module
docstring. Scores and their order are identical to the serial version.
"""
import threading
from collections import deque
from concurrent.futures import ThreadPoolExecutor
dev = "cuda" if torch.cuda.is_available() else "cpu"
n = len(a)
fp = _texts_fp(a, b)
done = 0
out = np.empty(n, np.float32)
part_p = done_p = None
if cache:
os.makedirs(os.path.dirname(cache) or ".", exist_ok=True)
part_p, done_p = cache + ".part.npy", cache + ".done.json"
if os.path.exists(done_p) and os.path.exists(part_p):
try:
st = json.load(open(done_p))
if st.get("fp") == fp and st.get("n") == n:
prev = np.load(part_p)
done = int(st.get("done_sorted", st.get("done", 0)))
out[:len(prev)] = prev[:len(prev)]
print(f"[ce] resuming scoring: {done:,}/{n:,} already done")
else:
print("[ce] prediction cache mismatch - rescoring")
except Exception as e: # noqa: BLE001
print(f"[ce] unreadable prediction cache ({e}) - rescoring")
if done >= n:
return out
try:
model = AutoModelForSequenceClassification.from_pretrained(
model_dir, attn_implementation="sdpa").to(dev).eval()
except Exception: # noqa: BLE001
model = AutoModelForSequenceClassification.from_pretrained(model_dir).to(dev).eval()
ngpu = torch.cuda.device_count()
use_dp = ngpu > 1 and cfg.get("ce_data_parallel", False)
net = torch.nn.DataParallel(model) if use_dp else model
amp = _amp_dtype()
bs = int(cfg["ce_eval_batch"])
ck_rows = int(cfg.get("ce_pred_ckpt_rows", 2_000_000))
max_len = cfg["ce_max_len"]
# one tokenizer per worker thread: HF fast tokenizers are not documented as safe
# to share across threads, and a per-thread instance costs nothing.
_tl = threading.local()
def _tok():
t = getattr(_tl, "tok", None)
if t is None:
t = AutoTokenizer.from_pretrained(model_dir)
_tl.tok = t
return t
def _encode(idx):
tk = _tok()
enc = tk([a[i] for i in idx], [b[i] for i in idx], truncation="longest_first",
max_length=max_len, padding=True, return_tensors="pt")
return idx, enc
order = np.argsort([len(a[i]) + len(b[i]) for i in range(n)], kind="stable")
buf = np.empty(n, np.float32)
if done:
buf[:] = out[:]
batches = [order[s:s + bs] for s in range(done, len(order), bs)]
n_prefetch = int(cfg.get("ce_prefetch", 3))
pos, scored, since = done, 0, 0
t0 = time.time()
with ThreadPoolExecutor(max_workers=n_prefetch) as ex:
it = iter(batches)
futs = deque()
for _ in range(n_prefetch):
try:
futs.append(ex.submit(_encode, next(it)))
except StopIteration:
break
while futs:
idx, enc = futs.popleft().result()
try:
futs.append(ex.submit(_encode, next(it)))
except StopIteration:
pass
enc = {k: v.to(dev, non_blocking=True) for k, v in enc.items()}
with torch.autocast("cuda", dtype=amp, enabled=amp is not None):
lg = net(**enc).logits.squeeze(-1).float()
buf[idx] = torch.sigmoid(lg).cpu().numpy()
pos += len(idx); scored += len(idx); since += len(idx)
if scored % (bs * 1000) < bs:
rate = scored / max(time.time() - t0, 1e-9)
print(f"[ce] scored {pos:,}/{n:,} ({rate:,.0f}/s, "
f"{time.time()-t0:.0f}s)", flush=True)
if cache and since >= ck_rows:
np.save(part_p + ".tmp.npy", buf)
os.replace(part_p + ".tmp.npy", part_p)
_atomic_json(done_p, {"fp": fp, "n": n, "done_sorted": int(pos)})
print(f"[ce] scoring checkpoint: {pos:,}/{n:,} rows persisted "
f"({100*pos/max(n,1):.1f}%)", flush=True)
since = 0
if cache:
np.save(part_p + ".tmp.npy", buf)
os.replace(part_p + ".tmp.npy", part_p)
_atomic_json(done_p, {"fp": fp, "n": n, "done_sorted": n})
return buf
'''
s = s[:start] + new_fn
io.open(P, "w", encoding="utf-8").write(s)
print("+ cross_encoder.py: threaded tokenisation prefetch, sdpa, inference_mode")
C = "src/config.py"
c = io.open(C, encoding="utf-8").read()
old = ' "ce_eval_batch": 1024,'
new = (' "ce_eval_batch": 2048, # 21 GB used of 98 GB; larger batches amortise launch overhead\n'
' "ce_prefetch": 3, # tokenisation batches queued ahead of the GPU')
if '"ce_prefetch"' in c:
print("= config already patched")
else:
assert old in c, "ce_eval_batch anchor missing"
io.open(C, "w", encoding="utf-8").write(c.replace(old, new, 1))
print("+ config.py: ce_eval_batch 2048, ce_prefetch 3")
|