Merge 8cfa0451f4 into 0aaca56805

Add setuptools for CPU run

.gitignore

@@ -11,3 +11,4 @@ eval_bundle/
 # Local setup
 CLAUDE.md
 wandb/
+*.egg-info/

nanochat/core_eval.py

@@ -201,6 +201,9 @@ def evaluate_example(idx, model, tokenizer, data, device, task_meta):
         for t, s, e in zip(tokens, start_idxs, end_idxs):
             if len(t) > max_tokens:
                 num_to_crop = len(t) - max_tokens
+                # Take the last max_tokens tokens instead of the first ones.
+                # The overly long questions are usually the few-shot contexts. They are placed
+                # at the beginning of the sequence, so cropping from the start should be ok.
                 new_tokens.append(t[-max_tokens:]) # take the last max_tokens tokens
                 new_start_idxs.append(s - num_to_crop) # shift the indices down
                 new_end_idxs.append(e - num_to_crop)
@@ -228,7 +231,11 @@ def evaluate_example(idx, model, tokenizer, data, device, task_meta):
         # predictions[i] predict input_ids[i+1] autoregressively
         predicted_tokens = predictions[0, si-1:ei-1]
         actual_tokens = input_ids[0, si:ei]
-        is_correct = torch.all(predicted_tokens == actual_tokens).item()
+        # Make the matching case-insensitive for LM tasks
+        predicted_text = tokenizer.decode(predicted_tokens.cpu().tolist()).lower()
+        actual_text = tokenizer.decode(actual_tokens.cpu().tolist()).lower()
+        # is_correct = torch.all(predicted_tokens == actual_tokens).item()
+        is_correct = (predicted_text == actual_text)
     elif task_type in ['multiple_choice', 'schema']:
         # For MC/schema: find the option with lowest average loss
         mean_losses = [losses[i, si-1:ei-1].mean().item()

nanochat/gpt.py

@@ -238,6 +238,11 @@ class GPT(nn.Module):
         for i in range(n_layer):
             self.x0_lambdas.data[i] = 0.20 - (0.15 * i / max(n_layer - 1, 1))
 
+        # Smear/backout scalars and smear gate must be explicitly initialized 
+        torch.nn.init.zeros_(self.smear_lambda)
+        torch.nn.init.constant_(self.backout_lambda, 0.2)
+        torch.nn.init.uniform_(self.smear_gate.weight, 0.0, 0.02)
+
         # Value embeddings (init like c_v: uniform with same std)
         for ve in self.value_embeds.values():
             torch.nn.init.uniform_(ve.weight, -s, s)

pyproject.toml

@@ -54,6 +54,7 @@ explicit = true
 
 [project.optional-dependencies]
 cpu = [
+    "setuptools>=65.0.0",
     "torch==2.9.1",
 ]
 gpu = [

uv.lock

@@ -1507,6 +1507,7 @@ dependencies = [
 
 [package.optional-dependencies]
 cpu = [
+    { name = "setuptools" },
     { name = "torch", version = "2.9.1", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" },
     { name = "torch", version = "2.9.1+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" },
 ]
@@ -1530,6 +1531,7 @@ requires-dist = [
     { name = "kernels", specifier = ">=0.11.7" },
     { name = "psutil", specifier = ">=7.1.0" },
     { name = "rustbpe", specifier = ">=0.1.0" },
+    { name = "setuptools", marker = "extra == 'cpu'", specifier = ">=65.0.0" },
     { name = "tiktoken", specifier = ">=0.11.0" },
     { name = "tokenizers", specifier = ">=0.22.0" },
     { name = "torch", specifier = "==2.9.1" },