diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 852a4ed..12e5145 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -21,6 +21,8 @@ jobs: with: components: rustfmt, clippy - uses: Swatinem/rust-cache@v2 + - name: asset downloader tests (no network) + run: python3 -m unittest discover -s python -p 'test_*.py' - name: rustfmt run: cargo fmt --all -- --check - name: clippy diff --git a/docs/correctness.md b/docs/correctness.md index 3adf6aa..4ce17e1 100644 --- a/docs/correctness.md +++ b/docs/correctness.md @@ -93,3 +93,30 @@ exist: - No SSM `scan` op (not implemented). - No full transformer-LM end-to-end parity (no LM yet) — only the MLP is verified end-to-end, plus every individual operator above. + +## Invalid-input contracts and sampling + +CPU reference operators check dimension products before allocation. Normalization +requires a positive width and a finite non-negative epsilon; rotary embeddings +require a finite positive theta. Attention validates head widths and bounds +absolute causal offsets without wrapping. + +Use Sampler::sample_token for checked single-token sampling, and Sampler::validate +to check temperature. Empty or non-finite logits are rejected. Probabilities use +f64 intermediates; top-k ties prefer smaller token IDs and the retained candidates +have a stable ID order. Seeded sampled sequences may differ from earlier versions +because their candidate ordering and arithmetic were corrected. Greedy finite +argmax retains its earliest-ID tie behavior. + +Tokenizer merge parsing retains hash-prefixed merge tokens, rejects malformed or +duplicate pairs, and checks both operands. Repeated whitespace preserves GPT-2's +lookahead boundary before the next token. The additive try_decode and +try_decode_bytes APIs report unknown token IDs; existing decode APIs retain their +permissive behavior. + +The download script accepts --out, --timeout, and --force. Each transfer uses a +temporary file beside its destination, checks any advertised Content-Length, and +atomically replaces the destination only after completion. Existing nonempty +assets are skipped by default; --force refreshes them. Network-free downloader +regressions run in CI. These tests do not download or validate pretrained GPT-2 +weights and do not replace on-device Metal parity tests. diff --git a/python/fetch_gpt2.py b/python/fetch_gpt2.py index b099448..116987d 100644 --- a/python/fetch_gpt2.py +++ b/python/fetch_gpt2.py @@ -1,10 +1,10 @@ -"""Download GPT-2 (124M) weights + tokenizer into models/gpt2/. - -Pulls the public HuggingFace `openai-community/gpt2` files (~550 MB) with no -extra dependencies beyond the standard library. These files are gitignored. -""" - +"""Download GPT-2 weights and tokenizer with atomic, bounded transfers.""" +import argparse +import math import os +import shutil +import sys +import tempfile import urllib.request BASE = "https://huggingface.co/openai-community/gpt2/resolve/main" @@ -12,18 +12,58 @@ OUT = os.path.join("models", "gpt2") -def main(): - os.makedirs(OUT, exist_ok=True) - for name in FILES: - dst = os.path.join(OUT, name) - if os.path.exists(dst) and os.path.getsize(dst) > 0: - print(f" have {name}") - continue - print(f" downloading {name} …") - urllib.request.urlretrieve(f"{BASE}/{name}", dst) - print(f"Done. Weights in {OUT}/") - print('Try: cargo run --release --bin generate -- --prompt "Once upon a time"') +def positive_timeout(value): + try: + seconds = float(value) + except ValueError as error: + raise argparse.ArgumentTypeError("timeout must be a number") from error + if not math.isfinite(seconds) or seconds <= 0: + raise argparse.ArgumentTypeError("timeout must be finite and positive") + return seconds + + +def download_file(url, destination, timeout=60.0): + """Install a completed transfer atomically, preserving any previous asset.""" + directory = os.path.dirname(os.path.abspath(destination)) + descriptor, temporary = tempfile.mkstemp(prefix=".gpt2-", suffix=".part", dir=directory) + os.close(descriptor) + try: + with urllib.request.urlopen(url, timeout=timeout) as response, open(temporary, "wb") as output: + shutil.copyfileobj(response, output) + expected = response.headers.get("Content-Length") + size = os.path.getsize(temporary) + if size == 0: + raise OSError("downloaded file is empty") + if expected is not None and size != int(expected): + raise OSError(f"incomplete transfer: expected {expected} bytes, received {size}") + os.replace(temporary, destination) + finally: + if os.path.exists(temporary): + os.unlink(temporary) + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", default=OUT, help="asset directory (default: models/gpt2)") + parser.add_argument("--timeout", type=positive_timeout, default=60.0, help="socket timeout in seconds") + parser.add_argument("--force", action="store_true", help="replace existing assets") + args = parser.parse_args(argv) + try: + os.makedirs(args.out, exist_ok=True) + for name in FILES: + destination = os.path.join(args.out, name) + if not args.force and os.path.isfile(destination) and os.path.getsize(destination) > 0: + print(f" have {name}") + continue + print(f" downloading {name} …") + download_file(f"{BASE}/{name}", destination, args.timeout) + except (OSError, ValueError) as error: + print(f"Download failed: {error}", file=sys.stderr) + return 1 + print(f"Done. Weights in {args.out}/") + print(f'Try: cargo run --release --bin generate -- --model-dir "{args.out}" --prompt "Once upon a time"') + return 0 if __name__ == "__main__": - main() + raise SystemExit(main()) diff --git a/python/test_fetch_gpt2.py b/python/test_fetch_gpt2.py new file mode 100644 index 0000000..2185a7a --- /dev/null +++ b/python/test_fetch_gpt2.py @@ -0,0 +1,74 @@ +"""Network-free regression tests for GPT-2 asset installation.""" +import argparse +import io +import os +import tempfile +import unittest +from unittest.mock import patch + +import fetch_gpt2 + + +class Response(io.BytesIO): + def __init__(self, payload=b"complete", length=None): + super().__init__(payload) + self.headers = {} if length is None else {"Content-Length": str(length)} + + +class DownloadTests(unittest.TestCase): + def test_complete_download_replaces_asset(self): + with tempfile.TemporaryDirectory() as directory: + destination = os.path.join(directory, "model") + with patch("fetch_gpt2.urllib.request.urlopen", return_value=Response(length=8)) as opener: + fetch_gpt2.download_file("https://example.test/model", destination, 12.0) + opener.assert_called_once_with("https://example.test/model", timeout=12.0) + with open(destination, "rb") as stream: + self.assertEqual(stream.read(), b"complete") + self.assertEqual(os.listdir(directory), ["model"]) + + def test_failure_preserves_previous_asset(self): + with tempfile.TemporaryDirectory() as directory: + destination = os.path.join(directory, "model") + with open(destination, "wb") as stream: + stream.write(b"previous") + with patch("fetch_gpt2.urllib.request.urlopen", side_effect=OSError("interrupted")): + with self.assertRaises(OSError): + fetch_gpt2.download_file("https://example.test/model", destination) + with open(destination, "rb") as stream: + self.assertEqual(stream.read(), b"previous") + self.assertEqual(os.listdir(directory), ["model"]) + + def test_empty_and_truncated_downloads_never_become_final_assets(self): + for payload, length in [(b"", None), (b"short", 100)]: + with tempfile.TemporaryDirectory() as directory: + destination = os.path.join(directory, "model") + with patch("fetch_gpt2.urllib.request.urlopen", return_value=Response(payload, length)): + with self.assertRaises(OSError): + fetch_gpt2.download_file("https://example.test/model", destination) + self.assertEqual(os.listdir(directory), []) + + def test_timeout_validation(self): + for value in ["0", "-1", "NaN", "inf", "invalid"]: + with self.assertRaises(argparse.ArgumentTypeError): + fetch_gpt2.positive_timeout(value) + self.assertEqual(fetch_gpt2.positive_timeout("5"), 5.0) + + def test_existing_assets_are_skipped_unless_forced(self): + with tempfile.TemporaryDirectory() as directory: + for name in fetch_gpt2.FILES: + with open(os.path.join(directory, name), "wb") as stream: + stream.write(b"previous") + with patch("builtins.print"), patch("fetch_gpt2.download_file") as transfer: + self.assertEqual(fetch_gpt2.main(["--out", directory]), 0) + transfer.assert_not_called() + self.assertEqual(fetch_gpt2.main(["--out", directory, "--force", "--timeout", "3"]), 0) + self.assertEqual(transfer.call_count, len(fetch_gpt2.FILES)) + + def test_main_reports_network_failure_as_unsuccessful(self): + with tempfile.TemporaryDirectory() as directory: + with patch("builtins.print"), patch("fetch_gpt2.download_file", side_effect=OSError("network")): + self.assertEqual(fetch_gpt2.main(["--out", directory]), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/gpt2.rs b/src/gpt2.rs index 463b722..cb2f743 100644 --- a/src/gpt2.rs +++ b/src/gpt2.rs @@ -19,6 +19,8 @@ use crate::tensor::{Tensor, TensorError}; #[derive(Error, Debug)] pub enum Gpt2Error { + #[error("invalid sampling input: {0}")] + InvalidSampling(&'static str), #[error("tensor '{name}' has shape {found:?}; expected {expected:?}")] BadShape { name: String, @@ -366,6 +368,20 @@ impl Gpt2 { mut on_token: impl FnMut(usize) -> bool, ) -> Vec { let mut toks = prompt.to_vec(); + if max_new == 0 { + return toks; + } + self.config.validate().expect("invalid GPT-2 configuration"); + sampler.validate().expect("invalid generation sampler"); + assert!(!prompt.is_empty(), "generation requires a non-empty prompt"); + assert!( + eot_token < self.config.vocab_size, + "end-of-text ID exceeds GPT-2 vocabulary" + ); + assert!( + prompt.iter().all(|&token| token < self.config.vocab_size), + "prompt ID exceeds GPT-2 vocabulary" + ); let n_ctx = self.config.n_ctx; let mut rng = Rng::new(sampler.seed); for _ in 0..max_new { @@ -427,37 +443,78 @@ impl Sampler { } } + pub fn validate(&self) -> Result<(), Gpt2Error> { + if !self.temperature.is_finite() || self.temperature < 0.0 { + return Err(Gpt2Error::InvalidSampling( + "temperature must be finite and non-negative", + )); + } + Ok(()) + } + + /// Samples one token using this seed. Generation retains the RNG across tokens. + pub fn sample_token(&self, logits: &[f32]) -> Result { + self.sample_checked(logits, &mut Rng::new(self.seed)) + } + fn sample(&self, logits: &[f32], rng: &mut Rng) -> usize { - if self.temperature <= 0.0 { - return argmax(logits); + self.sample_checked(logits, rng) + .expect("invalid token sampling input") + } + + fn sample_checked(&self, logits: &[f32], rng: &mut Rng) -> Result { + self.validate()?; + if logits.is_empty() { + return Err(Gpt2Error::InvalidSampling("logits must not be empty")); + } + if logits.iter().any(|logit| !logit.is_finite()) { + return Err(Gpt2Error::InvalidSampling("logits must be finite")); + } + if self.temperature == 0.0 { + return Ok(argmax(logits)); } - // Optional top-k: keep only the k highest logits. let mut idx: Vec = (0..logits.len()).collect(); if self.top_k > 0 && self.top_k < logits.len() { idx.select_nth_unstable_by(self.top_k, |&a, &b| { - logits[b].partial_cmp(&logits[a]).unwrap() + if logits[a] == logits[b] { + a.cmp(&b) + } else { + logits[b].total_cmp(&logits[a]) + } }); idx.truncate(self.top_k); + idx.sort_unstable(); } - let max = idx.iter().map(|&i| logits[i]).fold(f32::MIN, f32::max); - let mut probs: Vec = idx + let max = idx + .iter() + .map(|&i| logits[i]) + .fold(f32::NEG_INFINITY, f32::max); + let mut probs: Vec = idx .iter() - .map(|&i| ((logits[i] - max) / self.temperature).exp()) + .map(|&i| ((f64::from(logits[i]) - f64::from(max)) / f64::from(self.temperature)).exp()) .collect(); - let sum: f32 = probs.iter().sum(); + let sum: f64 = probs.iter().sum(); for p in &mut probs { *p /= sum; } - let r = rng.next_f32(); - let mut acc = 0.0; - for (j, &p) in probs.iter().enumerate() { - acc += p; - if r <= acc { - return idx[j]; - } + Ok(select_probability(&idx, &probs, f64::from(rng.next_f32()))) + } +} + +fn select_probability(indices: &[usize], probabilities: &[f64], draw: f64) -> usize { + let mut cumulative = 0.0; + for (index, &probability) in probabilities.iter().enumerate() { + cumulative += probability; + if draw < cumulative { + return indices[index]; } - idx[idx.len() - 1] } + // Rounding near one must still select a candidate with positive weight. + let last = probabilities + .iter() + .rposition(|&p| p > 0.0) + .expect("sampling has no positive weight"); + indices[last] } fn argmax(v: &[f32]) -> usize { @@ -588,4 +645,97 @@ mod tests { assert_eq!(calls, 1); assert_eq!(tokens.len(), 2); } + + #[test] + fn checked_sampling_rejects_invalid_inputs() { + for temperature in [f32::NAN, f32::INFINITY, -1.0] { + assert!(Sampler { + temperature, + top_k: 1, + seed: 0 + } + .sample_token(&[0.0]) + .is_err()); + } + assert!(Sampler::greedy().sample_token(&[]).is_err()); + for value in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + assert!(Sampler::greedy().sample_token(&[value, 0.0]).is_err()); + } + } + + #[test] + fn extreme_finite_logits_sample_without_overflow() { + let sampler = Sampler { + temperature: f32::MIN_POSITIVE, + top_k: 0, + seed: 42, + }; + assert_eq!(sampler.sample_token(&[-f32::MAX, f32::MAX]).unwrap(), 1); + } + + #[test] + fn top_k_ties_choose_lowest_ids_and_repeat_from_a_seed() { + let sampler = Sampler { + temperature: 1.0, + top_k: 1, + seed: 42, + }; + assert_eq!(sampler.sample_token(&[0.0, 0.0, 0.0]).unwrap(), 0); + let sampler = Sampler { + top_k: 2, + ..sampler + }; + let first = sampler.sample_token(&[0.0, 0.0, 0.0]).unwrap(); + assert!(first < 2); + assert_eq!(first, sampler.sample_token(&[0.0, 0.0, 0.0]).unwrap()); + } + + #[test] + fn sampling_boundaries_never_choose_zero_weight() { + assert_eq!(select_probability(&[10, 20], &[0.0, 1.0], 0.0), 20); + assert_eq!( + select_probability( + &[10, 20, 30], + &[0.5, 0.4999999999999999, 0.0], + 0.9999999999999999 + ), + 20 + ); + } + + #[test] + fn zero_token_generation_returns_the_prompt_without_inference() { + let tokens = forward_test_model().generate( + &crate::model::CpuBackend, + &[], + 0, + &Sampler::greedy(), + 1, + |_| panic!("unexpected inference"), + ); + assert!(tokens.is_empty()); + } + + #[test] + #[should_panic(expected = "end-of-text ID exceeds")] + fn generation_rejects_unknown_stop_tokens() { + forward_test_model().generate( + &crate::model::CpuBackend, + &[0], + 1, + &Sampler::greedy(), + 999, + |_| {}, + ); + } + + #[test] + fn top_k_ties_treat_signed_zero_as_equal() { + let sampler = Sampler { + temperature: 1.0, + top_k: 1, + seed: 42, + }; + assert_eq!(sampler.sample_token(&[-0.0, 0.0]).unwrap(), 0); + } } diff --git a/src/model.rs b/src/model.rs index 1046f7a..de24a4a 100644 --- a/src/model.rs +++ b/src/model.rs @@ -102,7 +102,7 @@ impl Mlp { .ok_or_else(|| ModelError::MissingTensor(b_name.clone()))?; w.validate()?; b.validate()?; - if w.shape.len() != 2 { + if w.shape.len() != 2 || w.shape.contains(&0) { return Err(ModelError::BadShape { name: w_name, shape: w.shape.clone(), @@ -132,6 +132,40 @@ impl Mlp { Ok(Self { layers }) } + /// Revalidates public layers before inference. + pub fn validate(&self) -> Result<(), ModelError> { + if self.layers.is_empty() { + return Err(ModelError::NoLayers); + } + let mut previous_width = None; + for (index, (weight, bias)) in self.layers.iter().enumerate() { + weight.validate()?; + bias.validate()?; + if weight.shape.len() != 2 || weight.shape.contains(&0) { + return Err(ModelError::BadShape { + name: format!("layers.{index}.weight"), + shape: weight.shape.clone(), + }); + } + if bias.shape.as_slice() != [weight.shape[0]] { + return Err(ModelError::BadShape { + name: format!("layers.{index}.bias"), + shape: bias.shape.clone(), + }); + } + if let Some(expected) = previous_width { + if weight.shape[1] != expected { + return Err(ModelError::WidthMismatch { + expected, + got: weight.shape[1], + }); + } + } + previous_width = Some(weight.shape[0]); + } + Ok(()) + } + /// Input feature width expected by the first layer. pub fn in_features(&self) -> usize { self.layers[0].0.shape[1] @@ -151,9 +185,7 @@ impl Mlp { backend: &B, x: &Tensor, ) -> Result { - if self.layers.is_empty() { - return Err(ModelError::NoLayers); - } + self.validate()?; x.validate()?; let (_, in_f) = x.dims2().map_err(|_| ModelError::WidthMismatch { expected: self.in_features(), @@ -264,4 +296,32 @@ mod tests { Err(ModelError::NoLayers) )); } + + #[test] + fn checkpoint_rejects_zero_feature_width() { + let mut map = HashMap::new(); + map.insert("layers.0.weight".into(), Tensor::zeros(vec![2, 0])); + map.insert("layers.0.bias".into(), Tensor::zeros(vec![2])); + assert!(matches!( + Mlp::from_tensors(&map), + Err(ModelError::BadShape { .. }) + )); + } + + #[test] + fn forward_revalidates_mutated_public_weights_and_biases() { + let mut model = tiny_model(); + let input = Tensor::zeros(vec![1, 2]); + model.layers[0].0.data.pop(); + assert!(matches!( + model.forward(&CpuBackend, &input), + Err(ModelError::Tensor(_)) + )); + let mut model = tiny_model(); + model.layers[0].1 = Tensor::zeros(vec![1]); + assert!(matches!( + model.forward(&CpuBackend, &input), + Err(ModelError::BadShape { .. }) + )); + } } diff --git a/src/ops.rs b/src/ops.rs index 2b9f754..cd4907c 100644 --- a/src/ops.rs +++ b/src/ops.rs @@ -10,11 +10,19 @@ //! * `gelu` uses the tanh approximation, matching `jax.nn.gelu(approximate=True)`. //! * `rope` uses the rotate-half (GPT-NeoX / HF) convention. +fn checked_elements(rows: usize, cols: usize) -> usize { + rows.checked_mul(cols) + .expect("operator dimensions overflow") +} + /// Standard matrix multiply: `A`[m×k] · `B`[k×n] → `C`[m×n] (row-major). pub fn matmul(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec { - assert_eq!(a.len(), m * k, "A has wrong length"); - assert_eq!(b.len(), k * n, "B has wrong length"); - let mut c = vec![0.0f32; m * n]; + assert_eq!(a.len(), checked_elements(m, k), "A has wrong length"); + assert_eq!(b.len(), checked_elements(k, n), "B has wrong length"); + let mut c = vec![0.0f32; checked_elements(m, n)]; + if c.is_empty() { + return c; + } for row in 0..m { for i in 0..k { let a_ik = a[row * k + i]; @@ -34,10 +42,13 @@ pub fn matmul(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec { /// /// `x`[n×in], `w`[out×in], `b`[out] → `[n×out]`. pub fn linear(x: &[f32], w: &[f32], b: &[f32], n: usize, in_f: usize, out_f: usize) -> Vec { - assert_eq!(x.len(), n * in_f); - assert_eq!(w.len(), out_f * in_f); + assert_eq!(x.len(), checked_elements(n, in_f)); + assert_eq!(w.len(), checked_elements(out_f, in_f)); assert_eq!(b.len(), out_f); - let mut y = vec![0.0f32; n * out_f]; + let mut y = vec![0.0f32; checked_elements(n, out_f)]; + if y.is_empty() { + return y; + } for row in 0..n { for o in 0..out_f { let mut acc = b[o]; @@ -95,10 +106,15 @@ pub fn layernorm( d: usize, eps: f32, ) -> Vec { - assert_eq!(x.len(), rows * d); + assert!(d > 0, "normalization width must be positive"); + assert!( + eps.is_finite() && eps >= 0.0, + "normalization epsilon must be finite and non-negative" + ); + assert_eq!(x.len(), checked_elements(rows, d)); assert_eq!(gamma.len(), d); assert_eq!(beta.len(), d); - let mut out = vec![0.0f32; rows * d]; + let mut out = vec![0.0f32; checked_elements(rows, d)]; for r in 0..rows { let row = &x[r * d..r * d + d]; let mean = row.iter().sum::() / d as f32; @@ -113,9 +129,14 @@ pub fn layernorm( /// Row-wise RMSNorm over the last dimension of size `d`: `y = x / sqrt(mean(x²) + eps) * gamma`. pub fn rmsnorm(x: &[f32], gamma: &[f32], rows: usize, d: usize, eps: f32) -> Vec { - assert_eq!(x.len(), rows * d); + assert!(d > 0, "normalization width must be positive"); + assert!( + eps.is_finite() && eps >= 0.0, + "normalization epsilon must be finite and non-negative" + ); + assert_eq!(x.len(), checked_elements(rows, d)); assert_eq!(gamma.len(), d); - let mut out = vec![0.0f32; rows * d]; + let mut out = vec![0.0f32; checked_elements(rows, d)]; for r in 0..rows { let row = &x[r * d..r * d + d]; let ms = row.iter().map(|v| v * v).sum::() / d as f32; @@ -131,7 +152,12 @@ pub fn rmsnorm(x: &[f32], gamma: &[f32], rows: usize, d: usize, eps: f32) -> Vec /// `x`[rows×d] in place. `positions[r]` gives the absolute position of row `r`. /// `d` must be even. pub fn rope_inplace(x: &mut [f32], positions: &[usize], rows: usize, d: usize, theta: f32) { - assert_eq!(x.len(), rows * d); + assert!(d > 0, "rope head dim must be positive"); + assert!( + theta.is_finite() && theta > 0.0, + "rope theta must be finite and positive" + ); + assert_eq!(x.len(), checked_elements(rows, d)); assert_eq!(positions.len(), rows); assert_eq!(d % 2, 0, "rope head dim must be even"); let half = d / 2; @@ -166,15 +192,20 @@ pub fn attention( causal: bool, q_offset: usize, ) -> Vec { - assert_eq!(q.len(), m * d); - assert_eq!(k.len(), seq * d); - assert_eq!(v.len(), seq * d); + assert!(d > 0, "attention width must be positive"); + assert!( + seq > 0 || m == 0, + "non-empty attention queries require keys" + ); + assert_eq!(q.len(), checked_elements(m, d)); + assert_eq!(k.len(), checked_elements(seq, d)); + assert_eq!(v.len(), checked_elements(seq, d)); let scale = 1.0 / (d as f32).sqrt(); - let mut out = vec![0.0f32; m * d]; + let mut out = vec![0.0f32; checked_elements(m, d)]; let mut scores = vec![0.0f32; seq]; for qi in 0..m { let limit = if causal { - (qi + q_offset + 1).min(seq) + qi.saturating_add(q_offset).saturating_add(1).min(seq) } else { seq }; @@ -201,12 +232,16 @@ pub fn attention( /// heads laid out contiguously per row; output has the same shape. Query `i` /// attends causally over keys `j <= i`. This is the GPT-2 attention reference. pub fn mha(q: &[f32], k: &[f32], v: &[f32], seq: usize, heads: usize, head_dim: usize) -> Vec { - let hd = heads * head_dim; - assert_eq!(q.len(), seq * hd); - assert_eq!(k.len(), seq * hd); - assert_eq!(v.len(), seq * hd); + assert!( + heads > 0 && head_dim > 0, + "attention heads and head width must be positive" + ); + let hd = checked_elements(heads, head_dim); + assert_eq!(q.len(), checked_elements(seq, hd)); + assert_eq!(k.len(), checked_elements(seq, hd)); + assert_eq!(v.len(), checked_elements(seq, hd)); let scale = 1.0 / (head_dim as f32).sqrt(); - let mut out = vec![0.0f32; seq * hd]; + let mut out = vec![0.0f32; checked_elements(seq, hd)]; let mut scores = vec![0.0f32; seq]; for h in 0..heads { let base = h * head_dim; @@ -234,9 +269,12 @@ pub fn mha(q: &[f32], k: &[f32], v: &[f32], seq: usize, heads: usize, head_dim: /// Dequantizes a per-row INT8 weight matrix: `out[r,c] = q[r,c] · scale[r]`. pub fn dequantize_int8(q: &[i8], scales: &[f32], rows: usize, cols: usize) -> Vec { - assert_eq!(q.len(), rows * cols); + assert_eq!(q.len(), checked_elements(rows, cols)); assert_eq!(scales.len(), rows); - let mut out = vec![0.0f32; rows * cols]; + let mut out = vec![0.0f32; checked_elements(rows, cols)]; + if out.is_empty() { + return out; + } for r in 0..rows { let s = scales[r]; for c in 0..cols { @@ -369,4 +407,57 @@ mod tests { let out = dequantize_int8(&q, &scales, 1, 3); approx(&out, &[1.0, -1.0, 10.0], 1e-6); } + + #[test] + #[should_panic(expected = "operator dimensions overflow")] + fn matmul_rejects_overflow_before_allocation() { + matmul(&[], &[], usize::MAX, 2, 1); + } + + #[test] + fn normalization_rejects_invalid_epsilon() { + for eps in [f32::NAN, f32::INFINITY, -1.0] { + assert!( + std::panic::catch_unwind(|| layernorm(&[1.0], &[1.0], &[0.0], 1, 1, eps)).is_err() + ); + assert!(std::panic::catch_unwind(|| rmsnorm(&[1.0], &[1.0], 1, 1, eps)).is_err()); + } + } + + #[test] + fn normalization_distinguishes_empty_batches_from_zero_width() { + assert!(layernorm(&[], &[1.0], &[0.0], 0, 1, 1e-5).is_empty()); + assert!(rmsnorm(&[], &[1.0], 0, 1, 0.0).is_empty()); + assert!(std::panic::catch_unwind(|| layernorm(&[], &[], &[], 1, 0, 0.0)).is_err()); + } + + #[test] + fn rope_rejects_invalid_frequency_bases() { + for theta in [0.0, -1.0, f32::NAN, f32::INFINITY] { + assert!( + std::panic::catch_unwind(|| rope_inplace(&mut [1.0, 2.0], &[0], 1, 2, theta)) + .is_err() + ); + } + } + + #[test] + fn attention_large_offset_attends_to_all_available_keys() { + let out = attention(&[0.0], &[0.0, 0.0], &[2.0, 4.0], 1, 2, 1, true, usize::MAX); + assert_eq!(out, vec![3.0]); + } + + #[test] + fn mha_rejects_zero_and_overflowing_head_dimensions() { + assert!(std::panic::catch_unwind(|| mha(&[], &[], &[], 0, 0, 1)).is_err()); + assert!(std::panic::catch_unwind(|| mha(&[], &[], &[], 0, usize::MAX, 2)).is_err()); + assert!(mha(&[], &[], &[], 0, 1, 1).is_empty()); + } + + #[test] + fn empty_matrix_outputs_skip_unused_rows() { + assert!(matmul(&[], &[], usize::MAX, 0, 0).is_empty()); + assert!(linear(&[], &[], &[], usize::MAX, 0, 0).is_empty()); + assert!(dequantize_int8(&[], &[1.0; 3], 3, 0).is_empty()); + } } diff --git a/src/tokenizer.rs b/src/tokenizer.rs index 65350b1..56e28c0 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -13,6 +13,12 @@ use thiserror::Error; #[derive(Error, Debug)] pub enum TokenizerError { + #[error("unknown tokenizer ID: {0}")] + UnknownTokenId(usize), + #[error("duplicate BPE merge pair: {0:?}")] + DuplicateMerge(String), + #[error("invalid BPE merge line: {0:?}")] + InvalidMergeLine(String), #[error("IO error: {0}")] Io(#[from] std::io::Error), #[error("failed to parse vocab.json: {0}")] @@ -96,16 +102,29 @@ impl Tokenizer { let mut bpe_ranks = HashMap::new(); for (rank, line) in merges_raw .lines() - .filter(|l| !l.starts_with('#')) + .map(str::trim) + .filter(|l| !l.is_empty() && !l.starts_with("#version:")) .enumerate() { let mut it = line.split_whitespace(); - if let (Some(a), Some(b)) = (it.next(), it.next()) { + if let (Some(a), Some(b), None) = (it.next(), it.next(), it.next()) { + for operand in [a, b] { + if !vocab.contains_key(operand) { + return Err(TokenizerError::MissingVocabSymbol(operand.to_string())); + } + } let merged = format!("{a}{b}"); if !vocab.contains_key(&merged) { return Err(TokenizerError::MissingVocabSymbol(merged)); } - bpe_ranks.insert((a.to_string(), b.to_string()), rank); + if bpe_ranks + .insert((a.to_string(), b.to_string()), rank) + .is_some() + { + return Err(TokenizerError::DuplicateMerge(line.to_string())); + } + } else { + return Err(TokenizerError::InvalidMergeLine(line.to_string())); } } @@ -146,6 +165,28 @@ impl Tokenizer { ids } + /// Decodes token IDs, returning an error instead of skipping unknown IDs. + pub fn try_decode(&self, ids: &[usize]) -> Result { + Ok(String::from_utf8_lossy(&self.try_decode_bytes(ids)?).into_owned()) + } + + /// Decodes raw bytes while checking every token ID. + pub fn try_decode_bytes(&self, ids: &[usize]) -> Result, TokenizerError> { + let mut bytes = Vec::new(); + for id in ids { + let token = self + .decoder + .get(id) + .ok_or(TokenizerError::UnknownTokenId(*id))?; + bytes.extend( + token + .chars() + .filter_map(|c| self.byte_decoder.get(&c).copied()), + ); + } + Ok(bytes) + } + /// Decodes token ids back into text. pub fn decode(&self, ids: &[usize]) -> String { String::from_utf8_lossy(&self.decode_bytes(ids)).into_owned() @@ -249,6 +290,9 @@ fn pre_tokenize(text: &str) -> Vec { while e < n && chars[e].is_whitespace() { e += 1; } + if e < n && e - i > 1 { + e -= 1; + } out.push(chars[i..e].iter().collect()); i = e; continue; @@ -365,4 +409,63 @@ mod tests { assert_eq!(tokenizer.encode("ab"), vec![256]); assert_eq!(tokenizer.decode(&[256]), "ab"); } + + #[test] + fn whitespace_runs_match_gpt2_lookahead() { + assert_eq!(pre_tokenize("a b"), vec!["a", " ", " b"]); + assert_eq!(pre_tokenize("a b"), vec!["a", " ", " b"]); + assert_eq!(pre_tokenize("a "), vec!["a", " "]); + assert_eq!(pre_tokenize("a\n b"), vec!["a", "\n", " b"]); + } + + #[test] + fn hash_prefix_merges_are_not_comments() { + let mut vocab = byte_vocab(); + vocab.insert("##".into(), 256); + let tokenizer = Tokenizer::from_assets(vocab, "#version: 0.2\n\n# #\n").unwrap(); + assert_eq!(tokenizer.encode("##"), vec![256]); + } + + #[test] + fn malformed_merge_lines_are_reported() { + for line in ["a", "a b c"] { + assert!(matches!( + Tokenizer::from_assets(byte_vocab(), line), + Err(TokenizerError::InvalidMergeLine(_)) + )); + } + } + + #[test] + fn duplicate_merges_are_rejected() { + let mut vocab = byte_vocab(); + vocab.insert("ab".into(), 256); + assert!(matches!( + Tokenizer::from_assets(vocab, "a b\na b\n"), + Err(TokenizerError::DuplicateMerge(_)) + )); + } + + #[test] + fn missing_merge_operands_are_reported() { + let mut vocab = byte_vocab(); + vocab.insert("not_presenta".into(), 256); + assert!( + matches!(Tokenizer::from_assets(vocab, "not_present a"), Err(TokenizerError::MissingVocabSymbol(symbol)) if symbol == "not_present") + ); + } + + #[test] + fn checked_decode_reports_unknown_ids_and_preserves_bytes() { + let tokenizer = Tokenizer::from_assets(byte_vocab(), "").unwrap(); + assert!(matches!( + tokenizer.try_decode(&[999]), + Err(TokenizerError::UnknownTokenId(999)) + )); + assert_eq!( + tokenizer.try_decode_bytes(&[0xe2, 0x82, 0xac]).unwrap(), + vec![0xe2, 0x82, 0xac] + ); + assert_eq!(tokenizer.try_decode(&[0xe2, 0x82, 0xac]).unwrap(), "€"); + } }