File size: 5,091 Bytes
2d5c26a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | """CPU correctness checks using a tiny random Qwen and deterministic tokenizer."""
import unittest
from unittest.mock import patch
import torch
from transformers import Qwen2Config, Qwen2ForCausalLM
from decision_model import DecisionScorer
class CharacterTokenizer:
pad_token_id = 0
eos_token_id = 2
def encode(self, text, add_special_tokens=False):
assert add_special_tokens is False
if text in (" yes", " no"):
return [1 if text == " yes" else 2]
return [ord(character) + 3 for character in text]
class ScorerTests(unittest.TestCase):
@classmethod
def setUpClass(cls):
torch.set_num_threads(1)
def setUp(self):
torch.manual_seed(42)
config = Qwen2Config(
vocab_size=256, hidden_size=32, intermediate_size=48,
num_hidden_layers=3, num_attention_heads=4, num_key_value_heads=2,
max_position_embeddings=2048, attention_dropout=0.0,
)
config._attn_implementation = "sdpa"
with (
patch("decision_model.AutoTokenizer.from_pretrained", return_value=CharacterTokenizer()),
patch("decision_model.AutoModelForCausalLM.from_pretrained", return_value=Qwen2ForCausalLM(config)),
):
self.scorer = DecisionScorer("unused", train_layers=1, device="cpu").eval()
self.state = "The customer needs a refund."
self.questions = [
{"question": "Which team?", "choices": ["sales", "billing support", "tech"]},
{"question": "Is money involved?", "choices": ["yes", "no"]},
]
self.examples = [dict(state=self.state, **question) for question in self.questions]
def randomize_head(self):
# A zero head would make cache/permutation comparisons vacuous.
with torch.no_grad():
self.scorer.head.weight.normal_(std=0.1)
def test_zero_head_and_trainable_layers(self):
self.assertEqual(torch.count_nonzero(self.scorer.head.weight).item(), 0)
for index, layer in enumerate(self.scorer.lm.model.layers):
self.assertTrue(all(p.requires_grad == (index == 2) for p in layer.parameters()))
self.randomize_head()
scores = self.scorer.score_examples(self.examples)
sum(score.square().sum() for score in scores).backward()
self.assertGreater(self.scorer.head.weight.grad.abs().sum().item(), 0)
grads = [p.grad for p in self.scorer.lm.model.layers[-1].parameters() if p.grad is not None]
self.assertGreater(sum(grad.abs().sum().item() for grad in grads), 0)
self.assertIsNone(self.scorer.lm.model.embed_tokens.weight.grad)
def test_shared_parity_chunking_permutation_and_isolation(self):
self.randomize_head()
with torch.inference_mode():
full = self.scorer.score_examples(self.examples)
for batch_size in (1, 2, 4, 16):
shared = self.scorer.scores_shared(self.state, self.questions, batch_size)
for expected, actual in zip(full, shared):
torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
self.assertEqual(self.scorer.last_shared_timings["branches"], 5)
self.assertGreater(self.scorer.last_shared_timings["cache_bytes"], 0)
permuted = dict(self.questions[0], choices=list(reversed(self.questions[0]["choices"])))
actual = self.scorer.scores_shared(self.state, [permuted], 2)[0]
torch.testing.assert_close(actual.flip(0), full[0], rtol=1e-5, atol=1e-5)
altered = dict(self.questions[1], question="A completely different question?")
actual = self.scorer.scores_shared(self.state, [self.questions[0], altered], 3)[0]
torch.testing.assert_close(actual, full[0], rtol=1e-5, atol=1e-5)
def test_token_baseline_matches_two_logits(self):
with torch.inference_mode():
hidden, counts = self.scorer._full_hidden(self.examples)
logits = self.scorer.lm.lm_head(hidden)
expected = (logits[:, 1] - logits[:, 2]).split(counts)
actual = self.scorer.scores_token_baseline(self.examples)
for want, got in zip(expected, actual):
torch.testing.assert_close(got, want, rtol=1e-5, atol=1e-5)
def test_invalid_input_and_length(self):
with self.assertRaises(ValueError):
self.scorer.score_examples([])
with self.assertRaises(ValueError):
self.scorer.scores_shared("", self.questions)
with self.assertRaises(ValueError):
self.scorer.scores_shared(self.state, [dict(question="Q?", choices=["a", " a "])])
with self.assertRaises(ValueError):
self.scorer.scores_shared(self.state, self.questions, 0)
self.scorer.max_tokens = 10
with self.assertRaisesRegex(ValueError, "never silently truncated"):
self.scorer.score_examples(self.examples)
with self.assertRaisesRegex(ValueError, "never silently truncated"):
self.scorer.scores_shared(self.state, self.questions)
if __name__ == "__main__":
unittest.main()
|