perf: batched NLI inference + shared model instances (closes #330)#413
Open
mohitcek wants to merge 18 commits into
Open
perf: batched NLI inference + shared model instances (closes #330)#413mohitcek wants to merge 18 commits into
mohitcek wants to merge 18 commits into
Conversation
Patch release: `v0.6.1`
Patch release: `v0.6.2`
mohitcek
marked this pull request as ready for review
July 8, 2026 12:49
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
Batched inference.
NLI.predict_batch()scores many premise/hypothesis pairsper forward pass (configurable
nli_batch_size, default 32) undertorch.no_grad(). All callers now collect their pairs first and score them inbulk: 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 asingle-pair wrapper for backward compatibility.
One model copy instead of 2–4. Default
BlackBoxUQ()loaded the ~1.4 GB NLImodel twice; its scorers now share a single injected instance.
SampledLogprobsScorerno longer re-instantiates SemanticEntropy /SemanticDensity / SentenceTransformer on every
evaluate()call.Correct model inputs. Proper
tokenizer(premise, hypothesis)pair encodingreplaces 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 newuqlm[accelerate]extra, andConsistencyScorernow honorsdevice.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.
BlackBoxUQ()Behavior notes — scores change slightly
The old
" [SEP] "concatenation leaked a stray whitespace token into the modelinput (and corrupted inputs for non-DeBERTa
nli_model_namevalues). The new pairencoding 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:"logloss"/"log_loss"key mismatch made the ensemble tuner maximize log loss (
uqlm/utils/tuner.py)[j-1]indexwraparound; token-probability-weighted code-UQ scores were wrong (
uqlm/code/entropy.py)UQEnsemble.score()ignored itslogprobs_resultsargument — fresh instancescrashed, reused instances silently scored stale logprobs; adds
sampled_logprobs_results(uqlm/scorers/shortform/ensemble.py)LongTextUQ(mode="matched_unit")always crashed —float()applied to thelist returned by the cosine scorer (
uqlm/longform/luq/matched_unit.py)names as prompts, then crashed indexing a list; case-sensitive yes/no parsing
(
uqlm/nli/entailment.py). Its test class was converted toIsolatedAsyncioTestCase— the async tests previously never executed.Type of Change
Checklist
ruff checkandruff formatpass locally