Merge f88153065f into f66a780f68

README.md

@@ -184,6 +184,7 @@ python -m pytest tests/test_rustbpe.py -v -s
 │   ├── smoltalk.py                 # Conglomerate dataset of SmolTalk from HF
 │   └── spellingbee.py              # Task teaching model to spell/count letters
 ├── tests
+│   └── test_engine.py
 │   └── test_rustbpe.py
 └── uv.lock
 ```

nanochat/engine.py

@@ -17,8 +17,9 @@ import signal
 import warnings
 from contextlib import contextmanager
 from collections import deque
-from nanochat.common import compute_init
+from nanochat.common import compute_init, autodetect_device_type
 from nanochat.checkpoint_manager import load_model
+from contextlib import nullcontext 
 
 # -----------------------------------------------------------------------------
 # Calculator tool helpers
@@ -328,6 +329,9 @@ if __name__ == "__main__":
     import time
     # init compute
     ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init()
+    device_type = autodetect_device_type()
+    autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=torch.bfloat16) if device_type == "cuda" else nullcontext()
+
     # load the model and tokenizer
     model, tokenizer, meta = load_model("base", device, phase="eval")
     bos_token_id = tokenizer.get_bos_token_id()
@@ -340,10 +344,11 @@ if __name__ == "__main__":
     torch.cuda.synchronize()
     t0 = time.time()
     stream = model.generate(prompt_tokens, **kwargs)
-    for token in stream:
-        generated_tokens.append(token)
-        chunk = tokenizer.decode([token])
-        print(chunk, end="", flush=True)
+    with autocast_ctx:
+        for token in stream:
+            generated_tokens.append(token)
+            chunk = tokenizer.decode([token])
+            print(chunk, end="", flush=True)
     print()
     torch.cuda.synchronize()
     t1 = time.time()
@@ -355,11 +360,12 @@ if __name__ == "__main__":
     stream = engine.generate(prompt_tokens, num_samples=1, **kwargs) # note: runs in fp32
     torch.cuda.synchronize()
     t0 = time.time()
-    for token_column, token_masks in stream:
-        token = token_column[0] # only print out the first row
-        generated_tokens.append(token)
-        chunk = tokenizer.decode([token])
-        print(chunk, end="", flush=True)
+    with autocast_ctx:
+        for token_column, token_masks in stream:
+            token = token_column[0] # only print out the first row
+            generated_tokens.append(token)
+            chunk = tokenizer.decode([token])
+            print(chunk, end="", flush=True)
     print()
     torch.cuda.synchronize()
     t1 = time.time()

nanochat/loss_eval.py

@@ -9,9 +9,9 @@ import torch.distributed as dist
 def evaluate_bpb(model, batches, steps, token_bytes):
     """
     Instead of the naive 'mean loss', this function returns the bits per byte (bpb),
-    which is a tokenization vocab size-indepedent metric, meaning you are still comparing
+    which is a tokenization vocab size-independent metric, meaning you are still comparing
     apples:apples if you change the vocab size. The way this works is that instead of just
-    calculating the average loss as usual, you calculate the sum loss, and indepependently
+    calculating the average loss as usual, you calculate the sum loss, and independently
     also the sum bytes (of all the target tokens), and divide. This normalizes the loss by
     the number of bytes that the target tokens represent.
 

rustbpe/src/lib.rs

@@ -27,6 +27,13 @@ pub struct Tokenizer {
 
 // ------------------------ internal helpers ------------------------
 
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[repr(i8)]
+enum Delta {
+    Rem = -1,
+    Ins = 1,
+}
+
 #[derive(Clone, Debug)]
 struct Word {
     ids: Vec<u32>,
@@ -48,7 +55,7 @@ impl Word {
     ///   -1 for removed pairs, +1 for newly created pairs.
     ///
     /// NOTE: this version deliberately avoids a HashMap in the hot loop.
-    fn merge_pair(&mut self, pair: Pair, new_id: u32) -> Vec<(Pair, i32)> {
+    fn merge_pair(&mut self, pair: Pair, new_id: u32) -> Vec<(Pair, Delta)> {
         let (a, b) = pair;
         let n = self.ids.len();
         if n < 2 {
@@ -56,7 +63,7 @@ impl Word {
         }
 
         let mut out: Vec<u32> = Vec::with_capacity(n);
-        let mut deltas: Vec<(Pair, i32)> = Vec::with_capacity(6);
+        let mut deltas: Vec<(Pair, Delta)> = Vec::with_capacity(6);
 
         let mut i = 0;
         while i < n {
@@ -66,13 +73,13 @@ impl Word {
 
                 // remove old pairs
                 if let Some(x) = left {
-                    deltas.push(((x, a), -1));
-                    deltas.push(((x, new_id), 1));
+                    deltas.push(((x, a), Delta::Rem));
+                    deltas.push(((x, new_id), Delta::Ins));
                 }
-                deltas.push(((a, b), -1));
+                deltas.push(((a, b), Delta::Rem));
                 if let Some(y) = right {
-                    deltas.push(((b, y), -1));
-                    deltas.push(((new_id, y), 1));
+                    deltas.push(((b, y), Delta::Rem));
+                    deltas.push(((new_id, y), Delta::Ins));
                 }
 
                 // write merged token
@@ -112,12 +119,10 @@ impl PartialOrd for MergeJob {
 impl Ord for MergeJob {
     fn cmp(&self, other: &Self) -> Ordering {
         // Max-heap by count; tie-break to ascending pair order (deterministic)
-        if self.count != other.count {
-            self.count.cmp(&other.count)
-        } else {
+        self.count.cmp(&other.count).then_with(||
             // ascending order on the pair when counts tie
             other.pair.cmp(&self.pair)
-        }
+        )
     }
 }
 
@@ -217,10 +222,10 @@ impl Tokenizer {
                 let changes = words[word_idx].merge_pair(top.pair, new_id);
                 // Update global pair counts based on this word's count
                 for (pair, delta) in changes {
-                    let delta_total = delta * counts[word_idx];
+                    let delta_total = (delta as i32) * counts[word_idx];
                     if delta_total != 0 {
                         *pair_counts.entry(pair).or_default() += delta_total;
-                        if delta > 0 {
+                        if delta == Delta::Ins {
                             local_pos_updates.entry(pair).or_default().insert(word_idx);
                         }
                     }

scripts/chat_eval.py

@@ -1,6 +1,6 @@
 """
 Evaluate the Chat model.
-All the generic code lives here, and all the evlauation-specific
+All the generic code lives here, and all the evaluation-specific
 code lives in nanochat directory and is imported from here.
 
 Example runs:

scripts/chat_sft.py

@@ -192,7 +192,7 @@ for step in range(num_iterations):
         })
         model.train()
 
-    # evlauate accuracy of the multiple choice tasks (which are quick to run)
+    # evaluate accuracy of the multiple choice tasks (which are quick to run)
     if last_step or (step > 0 and step % eval_metrics_every == 0):
         model.eval()
         metrics = {}