Skip to content

perf: batched NLI inference + shared model instances (closes #330)#413

Open
mohitcek wants to merge 18 commits into
cvs-health:developfrom
mohitcek:mc/patch_nli
Open

perf: batched NLI inference + shared model instances (closes #330)#413
mohitcek wants to merge 18 commits into
cvs-health:developfrom
mohitcek:mc/patch_nli

Conversation

@mohitcek

@mohitcek mohitcek commented Jul 8, 2026

Copy link
Copy Markdown
Contributor

Description

NLI scoring is the main latency bottleneck for Semantic Entropy, Semantic Density,
Noncontradiction, Entailment, and Number of Semantic Sets (#330). This PR speeds it
up in three ways:

  1. Batched inference. NLI.predict_batch() scores many premise/hypothesis pairs
    per forward pass (configurable nli_batch_size, default 32) under
    torch.no_grad(). All callers now collect their pairs first and score them in
    bulk: clustering batches each round's comparisons, ConsistencyScorer and
    SemanticDensity batch across all prompts, and the longform claim/graph scorers
    batch their claim × response loops. predict(premise, hypothesis) remains as a
    single-pair wrapper for backward compatibility.

  2. One model copy instead of 2–4. Default BlackBoxUQ() loaded the ~1.4 GB NLI
    model twice; its scorers now share a single injected instance.
    SampledLogprobsScorer no longer re-instantiates SemanticEntropy /
    SemanticDensity / SentenceTransformer on every evaluate() call.

  3. Correct model inputs. Proper tokenizer(premise, hypothesis) pair encoding
    replaces the hand-concatenated " [SEP] " idiom, and real token-level truncation
    (512 tokens for the default model) replaces character slicing. Also fixed: the
    pair cache used collision-prone f"{r1}_{r2}" string keys (now tuples).

Extras: opt-in fp16 on CUDA/MPS, optional device_map="auto" via a new
uqlm[accelerate] extra, and ConsistencyScorer now honors device.

Benchmarks

Apple M-series, torch 2.10; before = main. Reproduce with python -m uqlm.nli.benchmark
(its sequential/pipeline modes also run on old releases). CUDA numbers welcome —
batching should gain more there.

Benchmark Device Before After
BlackBoxUQ noncontradiction+entailment, 64 pairs mps 15.5s 3.5s (4.4x)
BlackBoxUQ noncontradiction+entailment, 64 pairs cpu 15.4s 4.9s (3.1x)
BlackBoxUQ 3 NLI scorers, 8 prompts × 5 samples mps 8.6s 2.6s (3.3x)
Raw NLI throughput, short texts cpu 8.4 pairs/s 33.0 pairs/s (3.9x)
NLI model copies, default BlackBoxUQ() 2 (~2.8 GB) 1 (~1.4 GB)

Behavior notes — scores change slightly

The old " [SEP] " concatenation leaked a stray whitespace token into the model
input (and corrupted inputs for non-DeBERTa nli_model_name values). The new pair
encoding matches deberta-large-mnli's MNLI training format exactly, so NLI scores
are not bit-identical to previous releases: noncontradiction typically moves by
~1e-4..1e-3; entailment can move more on QA-style text. Texts beyond 512 tokens
were previously fed to the model over-length; they are now truncated correctly.
Test fixtures were re-scored with the same stored responses.

Batched vs sequential inference is a numeric no-op (parity test included), and
batched clustering was verified pair-for-pair identical to the old sequential loop.

Additional bug fixes

Five verified bugs found while working in these code paths, each with a regression
test that fails on main:

  1. Tuner optimized the wrong direction for log loss — a "logloss"/"log_loss"
    key mismatch made the ensemble tuner maximize log loss (uqlm/utils/tuner.py)
  2. Functional entropy assigned rotated cluster probabilities[j-1] index
    wraparound; token-probability-weighted code-UQ scores were wrong (uqlm/code/entropy.py)
  3. UQEnsemble.score() ignored its logprobs_results argument — fresh instances
    crashed, reused instances silently scored stale logprobs; adds
    sampled_logprobs_results (uqlm/scorers/shortform/ensemble.py)
  4. LongTextUQ(mode="matched_unit") always crashedfloat() applied to the
    list returned by the cosine scorer (uqlm/longform/luq/matched_unit.py)
  5. LLM-judge entailment retry path never worked — retried with DataFrame column
    names as prompts, then crashed indexing a list; case-sensitive yes/no parsing
    (uqlm/nli/entailment.py). Its test class was converted to
    IsolatedAsyncioTestCase — the async tests previously never executed.

Type of Change

  • Bug fix
  • New feature
  • Breaking change
  • Documentation update
  • Dependency update

Checklist

  • Tests added or updated for all changed behavior
  • Docstrings updated for any new or modified public API
  • Type annotations added for any new or modified functions
  • ruff check and ruff format pass locally

@mohitcek
mohitcek marked this pull request as ready for review July 8, 2026 12:49
@dylanbouchard
dylanbouchard changed the base branch from main to develop July 9, 2026 00:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants