Skip to content

Commit 26cd2a5

Browse files
committed
feat(lumen-model): smoke probe in EmbeddingModel::load — catch loader corruption at load time
Runs a single embed forward over a tiny multilingual probe ("hello 안녕") at the end of `EmbeddingModel::load()` and verifies every output component is finite + L2 norm ≈ 1.0. Catches silent NaN-class loader bugs (the v0.4.2 BF16-as-F16 scales dtype mismatch that made Qwen3-Embedding-4B-4bit-DWQ emit all-NaN embeddings) BEFORE the model is exposed to user requests. Failure path: load() returns Err → EmbeddingService::load propagates → main.rs in lumen-server logs the diagnostic + sets /v1/embeddings to 503; chat endpoint untouched. Overhead: ~50ms on 0.6B-8bit, ~150ms on 4B-DWQ. Cheap relative to load (~1-4s) so always-on with no opt-out.
1 parent f9f8e1b commit 26cd2a5

1 file changed

Lines changed: 53 additions & 2 deletions

File tree

crates/lumen-model/src/embedding.rs

Lines changed: 53 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -167,14 +167,65 @@ impl EmbeddingModel {
167167
t0.elapsed().as_secs_f32()
168168
);
169169

170-
Ok(Self {
170+
let mut model = Self {
171171
model,
172172
tokenizer,
173173
device,
174174
dim,
175175
max_seq_len,
176176
model_id: model_id_or_path.to_string(),
177-
})
177+
};
178+
// Smoke check: catch silent loader corruption (e.g. dtype mismatch on
179+
// scales/biases — the `Qwen3-Embedding-4B-4bit-DWQ` NaN-emission bug
180+
// that motivated this guard) BEFORE the model is exposed to requests.
181+
// One forward over a tiny multilingual probe; ~50ms for 0.6B-8bit,
182+
// ~1s for 4B variants — cheap relative to load itself.
183+
model
184+
.smoke_check()
185+
.with_context(|| format!("post-load smoke check failed for {model_id_or_path}"))?;
186+
187+
Ok(model)
188+
}
189+
190+
/// Run one forward pass on a tiny probe and verify the output embedding
191+
/// is finite. Surfaces silent loader corruption (NaN-class bugs from
192+
/// scales/biases dtype mismatch, weight shard misalignment, etc.) at
193+
/// load time rather than letting it propagate into user requests.
194+
///
195+
/// Probe = `"hello 안녕"` — 2-3 tokens in any BPE we ship, exercises
196+
/// both Latin and Hangul codepoints so a tokenizer mis-mapping also
197+
/// surfaces here rather than as garbage embeddings.
198+
fn smoke_check(&mut self) -> Result<()> {
199+
let probe = vec!["hello 안녕".to_string()];
200+
let batch = self.embed(&probe).context("smoke probe embed failed")?;
201+
let v = batch.embeddings.first().ok_or_else(|| {
202+
anyhow!("smoke probe: embed returned 0 vectors (expected 1)")
203+
})?;
204+
let nan = v.iter().filter(|x| !x.is_finite()).count();
205+
if nan != 0 {
206+
return Err(anyhow!(
207+
"smoke probe: {nan}/{} embedding components are NaN/Inf — \
208+
model output corrupted (likely loader bug: scales/biases \
209+
dtype mismatch, truncated shard, or quantization mode \
210+
misdetected)",
211+
v.len()
212+
));
213+
}
214+
let l2: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
215+
if !(0.5..=2.0).contains(&l2) {
216+
// EmbeddingModel L2-normalizes — output norm should be ≈ 1.0.
217+
// A wildly different norm would indicate a normalization
218+
// path bug; allow a generous band to avoid false positives.
219+
return Err(anyhow!(
220+
"smoke probe: embedding L2 norm {l2:.3} far from unit \
221+
(expected ≈ 1.0 since EmbeddingModel L2-normalizes its output)"
222+
));
223+
}
224+
eprintln!(
225+
"[embedding] smoke probe OK (dim={}, l2={l2:.3})",
226+
v.len()
227+
);
228+
Ok(())
178229
}
179230

180231
pub fn model_id(&self) -> &str {

0 commit comments

Comments
 (0)