Skip to content

conformance: use real DistilBERT-QA checkpoint for batched_run test - #33

Merged
ausimian merged 1 commit into
mainfrom
fix-distilbert-qa-vocab-mismatch
Apr 18, 2026
Merged

conformance: use real DistilBERT-QA checkpoint for batched_run test#33
ausimian merged 1 commit into
mainfrom
fix-distilbert-qa-vocab-mismatch

Conversation

@ausimian

@ausimian ausimian commented Apr 18, 2026

Copy link
Copy Markdown
Owner

Summary

  • The Nx.Serving.batched_run conformance test paired distilbert-base-uncased's 30522-token tokenizer with the tiny-random DistilBertForQuestionAnswering model (1124-row embedding), so the tokenizer emitted token ids past the embedding bound. Emily.Backend.gather silently returned uninitialised memory (occasionally NaN) because MLX's gather does not bounds-check — the intermittent :nan score seen once on M4. Nx.BinaryBackend would have raised immediately.
  • Swap to distilbert-base-uncased-distilled-squad for both model and tokenizer (matching vocab, ~250 MB checkpoint). Tag the test :distilbert_full and exclude it from the default --only conformance run, mirroring :qwen3_full / :vit_full / :whisper_full.
  • Strengthen assertions: with a real SQuAD-fine-tuned model the returned spans should contain "Sarah" / "London", not just be arbitrary floats.
  • Underlying backend behaviour tracked in debug: compile-time debug flags for Emily.Backend (opt-in assertions) #32 (compile-time debug-flag proposal for opt-in bounds-checked gather and friends).

Test plan

  • mix precommit clean on the branch (364 tests, 0 failures; :distilbert_full correctly excluded from the default run)
  • mix test --only distilbert_full test/emily/conformance/distilbert_test.exs passes locally on M4 (250 MB checkpoint downloaded on first run)
  • CI runs the default suite — :distilbert_full stays opt-in, no CI regression expected

The old test paired distilbert-base-uncased's 30522-token tokenizer
with a tiny-random model whose 1124-row embedding can't accommodate
the ids the tokenizer produces — Emily.Backend silently returned
uninitialised rows (occasionally NaN) because MLX's gather does not
bounds-check, producing the intermittent `:nan` score seen once on
M4. Nx.BinaryBackend would have raised.

Swap to distilbert-base-uncased-distilled-squad for both model and
tokenizer (matching vocab, ~250 MB); tag `:distilbert_full` and
exclude from the default `--only conformance` run. Assertions
tightened to check the returned span text against the expected
entity.

Underlying backend behaviour (silent OOB gather vs. opt-in compile-
time bounds-check debug flag) tracked in #32.
@ausimian
ausimian merged commit 5b23df1 into main Apr 18, 2026
1 of 2 checks passed
@ausimian
ausimian deleted the fix-distilbert-qa-vocab-mismatch branch April 18, 2026 02:04
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.

1 participant