/
llama2.mojo
937 lines (780 loc) · 31.4 KB
/
llama2.mojo
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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
from algorithm import sum
from algorithm import vectorize, parallelize
from builtin import string
from math import round
from memory import memset_zero, memcpy, stack_allocation
from memory.unsafe import DTypePointer, bitcast
from tensor import rand
from sys.info import num_performance_cores
from sys import argv
from tensor import Tensor, TensorShape, TensorSpec
from collections import List, Dict
# The SIMD vector width.
from sys.info import simdwidthof
import math
import os
import random
import time
alias NUM_CONFIG_INT = 7
var workers = 0
alias nelts = (4 * simdwidthof[DType.float32]())
alias BufferPtrType = DTypePointer[DType.uint8]
alias BufferPtrFloat32 = DTypePointer[DType.float32]
alias PointerStrings = Pointer[String]
alias TensorF32 = Tensor[DType.float32]
@register_passable
struct Accumulator[T: DType, width: Int]:
# ideally this could be SIMD[T, width] but the width
# in accumulate() method is compared by identity
var data: DTypePointer[T]
@always_inline
fn __init__() -> Self:
# allocate a DTypePointer on stack that doesn't need to be freed.
var data = stack_allocation[width, T]()
memset_zero(data, width)
return Self {data: data}
@always_inline
fn accumulate[_width: Int](inout self, val: SIMD[T, _width]) -> None:
# This is a hack to make sure both SIMD have _width length.
# SIMD[T, width] += SIMD[T, _width] is always an error.
var newVal = self.data.load[width=_width]() + val
self.data.store[width=_width](newVal)
@always_inline
fn total(self) -> SIMD[T, 1]:
return self.data.load[width=width]().reduce_add()
@value
struct TensorSlice:
# Provides a view into a tensor representing a 1D slice on its first or first 2 dimensions.
# Same function signatures as Tensor but without owning the data.
var _data: BufferPtrFloat32
var _shape: TensorShape
fn __init__(inout self, t: TensorF32, layer: Int) raises:
var elements_per_layer = t.num_elements() // t.dim(0)
self._data = t.data().offset(layer * elements_per_layer)
if t.rank() == 2:
self._shape = TensorShape(t.dim(1))
elif t.rank() == 3:
self._shape = TensorShape(t.dim(1), t.dim(2))
else:
# Compiler complains if _shape not defined
self._shape = TensorShape(1)
raise Error("TensorSlice: rank greater than 3 not implemented.")
fn __init__(inout self, t: TensorF32, layer: Int, row: Int) raises:
var elements_per_layer = t.num_elements() // t.dim(0)
var elements_per_row = elements_per_layer // t.dim(1)
self._data = t.data().offset(
layer * elements_per_layer + row * elements_per_row
)
if t.rank() == 3:
self._shape = TensorShape(t.dim(2))
elif t.rank() == 1:
# Compiler complains if _shape not defined
self._shape = TensorShape(1)
raise Error(
"Trying to slice a 1D Tensor by layer and row. This requires a"
" 3D Tensor."
)
else:
# Compiler complains if _shape not defined
self._shape = TensorShape(1)
raise Error("TensorSlice: rank greater than 3 not implemented.")
fn data(self) -> BufferPtrFloat32:
return self._data
fn shape(self) -> TensorShape:
return self._shape
fn num_elements(self) -> Int:
return self._shape.num_elements()
fn dim(self, idx: Int) -> Int:
return self._shape[idx]
fn rank(self) -> Int:
return self._shape.rank()
fn load[width: Int](self, idx: Int) -> SIMD[DType.float32, nelts]:
return self._data.load[width=nelts](idx)
fn load[width: Int](self, *indices: Int) -> SIMD[DType.float32, nelts]:
if len(VariadicList(indices)) > 2:
print(
"Warning: TensorSlice only supports 1D and 2D indexing. "
" Results are unlikely to be correct."
)
return self.load[width=nelts](indices[0] * self._shape[1] + indices[1])
fn load[
width: Int
](self, indices: StaticIntTuple[2]) -> SIMD[DType.float32, nelts]:
return self._data.load[width=nelts](
indices[0] * self._shape[1] + indices[1]
)
fn __getitem__(self, idx: Int) -> SIMD[DType.float32, 1]:
return self._data.load[width=1](idx)
fn store[nelts: Int](self, idx: Int, val: SIMD[DType.float32, nelts]):
return self._data.store[width=nelts](idx, val)
fn __setitem__(self, idx: Int, val: SIMD[DType.float32, 1]):
return self.store[1](idx, val)
# not optimal concat
fn str_concat(s1: String, s2: String) -> String:
var l1 = len(s1)
var l2 = len(s2)
var str = List[Int8](capacity=l1 + l2 + 1)
memcpy(str.data, s1._buffer.data, l1)
memcpy(str.data + l1, s2._buffer.data, l2)
str[l1 + l2] = 0
str.size = l1 + l2 + 1
return str^
fn string_compare(a: String, b: String) -> Int:
var index = 0
while a._buffer[index] != 0 and b._buffer[index] != 0:
if a._buffer[index] < b._buffer[index]:
return -1
if a._buffer[index] > b._buffer[index]:
return 1
index += 1
if a._buffer[index] != 0 and b._buffer[index] == 0:
return 1
if a._buffer[index] == 0 and b._buffer[index] != 0:
return -1
_ = (a, b)
return 0
fn wrap(token: String) -> String:
alias a = String("\\n")
alias b = String("\\t")
alias c = String("'")
alias d = String('"')
if token == a:
return String(List[Int8](0x0A, 0))
if token == b:
return String(List[Int8](0x09, 0))
if token == c:
return String(List[Int8](0x27, 0))
if token == d:
return String(List[Int8](0x22, 0))
return token
fn string_from_bytes(owned bytes: List[Int8]) -> String:
bytes.append(0)
return bytes^
@value
struct Tokenizer:
var vocab: List[String]
var vocab_scores: List[Float32]
var max_token_length: Int
var vocab_size: Int
var map_vocab_to_index: Dict[String, Int]
fn __init__(inout self, vocab_size: Int, filename: String) raises:
with open(filename, "rb") as f:
@parameter
fn read_bytes_as[dtype: DType](size: Int) raises -> SIMD[dtype, 1]:
# a List that keeps ownership of the pointer
var bytes = f.read_bytes(size)
# copy one element of new type after casting pointer
var result = bytes.data.bitcast[SIMD[dtype, 1]]()[0]
# orginal List and data can be destroyed
_ = bytes
return result
self.vocab_size = vocab_size
self.vocab_scores = List[Float32](capacity=self.vocab_size)
self.vocab = List[String](capacity=self.vocab_size)
self.map_vocab_to_index = Dict[String, Int]()
self.max_token_length = int(read_bytes_as[DType.int32](4))
# read vocab_scores & vocab values (tokens)
for i in range(self.vocab_size):
var score = read_bytes_as[DType.float32](4)
var slen = int(read_bytes_as[DType.int32](4))
var token = string_from_bytes(f.read_bytes(slen))
self.vocab.append(token^)
self.vocab_scores.append(score)
self.map_vocab_to_index[self.vocab[i]] = i
fn find(self, token_o: String) -> Int:
var token = wrap(token_o)
var index = self.map_vocab_to_index.find(token)
if index:
return index.value()[]
return -1
@value
struct Config:
var dim: Int
var kv_dim: Int
var hidden_dim: Int
var n_layers: Int
var n_heads: Int
var n_kv_heads: Int
var kv_mul: Int
var vocab_size: Int
var seq_len: Int
var head_size: Int
var shared_weights: Bool
fn __init__(inout self, fileName: String, print_config: Bool) raises:
var f = open(fileName, "r")
# reading 7 vars of type DType.int32 from the file
var bytes_of_config_params = NUM_CONFIG_INT * sizeof[DType.int32]()
# config_data_raw id Tensor[DType.int8] with bytes_of_config_params elements
var config_data_raw = f.read_bytes(bytes_of_config_params)
f.close()
# correct Tensor type and shape for easy reading, without copying data
var int32_ptr = config_data_raw.steal_data().bitcast[Int32]()
var config_data = Tensor(TensorShape(NUM_CONFIG_INT), int32_ptr)
self.dim = int(config_data[0])
self.hidden_dim = int(config_data[1])
self.n_layers = int(config_data[2])
self.n_heads = int(config_data[3])
self.n_kv_heads = int(config_data[4])
self.vocab_size = int(config_data[5])
self.seq_len = int(config_data[6])
self.head_size = self.dim // self.n_heads
self.kv_dim = (self.n_kv_heads * self.dim) // self.n_heads
self.kv_mul = self.n_heads // self.n_kv_heads
# negative vocab size is hacky way of signaling unshared weights. bit yikes.
self.shared_weights = self.vocab_size > 0
if not self.shared_weights:
self.vocab_size = -self.vocab_size
if print_config:
print("config: dim, hidden_dim", self.dim, self.hidden_dim)
print("config: n_layers, n_heads", self.n_layers, self.n_heads)
print("config: vocab_size, seq_len", self.vocab_size, self.seq_len)
print("config: head_size", self.head_size)
print("config: kv_dim, kv_mul", self.kv_dim, self.kv_mul)
@value
struct RunState:
var x: TensorF32 # activation at current time stamp (dim,)
var xb: TensorF32 # same, but inside a residual branch (dim,)
var xb2: TensorF32 # an additional buffer just for convenience (dim,)
var hb: TensorF32 # buffer for hidden dimension in the ffn (hidden_dim,)
var hb2: TensorF32 # buffer for hidden dimension in the ffn (hidden_dim,)
var q: TensorF32 # query (dim,)
var k: TensorSlice # key (kv_dim,)
var v: TensorSlice # value (kv_dim,)
var att: TensorF32 # buffer for scores/attention values (n_heads, seq_len)
var logits: TensorF32 # output logits
var key_cache: TensorF32 # (layer, seq_len, dim)
var value_cache: TensorF32 # (layer, seq_len, dim)
fn __init__(inout self, config: Config) raises:
self.x = TensorF32(config.dim)
self.xb = TensorF32(config.dim)
self.xb2 = TensorF32(config.dim)
self.hb = TensorF32(config.hidden_dim)
self.hb2 = TensorF32(config.hidden_dim)
self.q = TensorF32(config.dim)
self.att = TensorF32(config.n_heads, config.seq_len)
self.logits = TensorF32(config.vocab_size)
self.key_cache = TensorF32(
config.n_layers, config.seq_len, config.kv_dim
)
self.value_cache = TensorF32(
config.n_layers, config.seq_len, config.kv_dim
)
# So their updates flow to the caches, k and v are slices with shared memory.
# Initialize with placeholders. The real tensors reference layer and position during forward pass.
self.k = TensorSlice(TensorF32(TensorShape(1, config.kv_dim)), 1)
self.v = TensorSlice(TensorF32(TensorShape(1, config.kv_dim)), 1)
@value
struct TransformerWeights:
var token_embedding_table: TensorF32
var freq_cis_real: TensorF32
var freq_cis_imag: TensorF32
var rms_att_weight: TensorF32
var wq: TensorF32
var wk: TensorF32
var wv: TensorF32
var wo: TensorF32
var rms_ffn_weight: TensorF32
var w1: TensorF32
var w3: TensorF32
var w2: TensorF32
var rms_final_weight: TensorF32
var wcls: TensorF32
fn __init__(inout self, file_name: String, config: Config) raises:
var bytes_read = 0
var f = open(file_name, "r")
# throw away config data
_ = f.read_bytes(NUM_CONFIG_INT * sizeof[DType.int32]())
bytes_read += NUM_CONFIG_INT * sizeof[DType.int32]()
@parameter
fn read_weights(*dims: Int) raises -> TensorF32:
var shape = TensorShape(dims)
# The created tensor takes a 1D shape equal to bytes read
# So we can't reshape to target shape because dims don't match
var tmp = f.read_bytes(
shape.num_elements() * sizeof[DType.float32]()
)
bytes_read += shape.num_elements() * sizeof[DType.float32]()
var data = tmp.steal_data().bitcast[Float32]()
return TensorF32(shape, data)
self.token_embedding_table = read_weights(config.vocab_size, config.dim)
self.rms_att_weight = read_weights(config.n_layers, config.dim)
self.wq = read_weights(config.n_layers, config.dim, config.dim)
self.wk = read_weights(config.n_layers, config.kv_dim, config.dim)
self.wv = read_weights(config.n_layers, config.kv_dim, config.dim)
self.wo = read_weights(config.n_layers, config.dim, config.dim)
self.rms_ffn_weight = read_weights(config.n_layers, config.dim)
self.w1 = read_weights(config.n_layers, config.hidden_dim, config.dim)
self.w2 = read_weights(config.n_layers, config.dim, config.hidden_dim)
self.w3 = read_weights(config.n_layers, config.hidden_dim, config.dim)
self.rms_final_weight = read_weights(config.dim)
# maybe need modifying for different model
# config.head_size // 2 for stories and tinyllama-1.1
self.freq_cis_real = read_weights(config.seq_len, config.head_size // 2)
self.freq_cis_imag = read_weights(config.seq_len, config.head_size // 2)
if config.shared_weights:
self.wcls = self.token_embedding_table
else:
self.wcls = read_weights(config.vocab_size, config.dim)
f.close()
print(
"Total bytes read:",
bytes_read,
"Estimated checkpoint size: ",
bytes_read // 1024 // 1024,
"MB",
)
@always_inline
fn rmsnorm(
inout o: BufferPtrFloat32,
x: BufferPtrFloat32,
weight: BufferPtrFloat32,
size: Int,
) -> None:
# Calculate sum of squares
var tmp = Accumulator[DType.float32, nelts]()
@parameter
fn _sum2[_nelts: Int](j: Int):
tmp.accumulate(x.offset(j).load[width=_nelts](0) ** 2)
vectorize[_sum2, nelts](size)
var ss: Float32 = tmp.total()
ss = ss / size + 1e-5
ss = 1.0 / math.sqrt(ss)
# Normalize and scale
@parameter
fn _norm[_nelts: Int](j: Int):
var val = weight.load[width=_nelts](j) * ss * x.load[width=_nelts](j)
o.offset(j).store[width=_nelts](0, val)
vectorize[_norm, nelts](size)
@always_inline
fn rmsnorm(inout o: TensorF32, x: TensorF32, weight: TensorF32):
rmsnorm(o._ptr, x.data(), weight.data(), weight.dim(weight.rank() - 1))
@always_inline
fn rmsnorm(inout o: TensorF32, x: TensorF32, weight: TensorSlice):
rmsnorm(o._ptr, x.data(), weight.data(), weight.dim(weight.rank() - 1))
@always_inline
fn softmax(inout x: TensorF32) -> None:
softmax(x, 0, x.dim(0))
@always_inline
fn softmax(inout x: TensorF32, start: Int, end: Int):
var max_val: Float32 = -1e9
@parameter
fn _max[_nelts: Int](ii: Int):
var val = x.load[width=_nelts](start + ii).reduce_max()
if val > max_val:
max_val = val
vectorize[_max, nelts](end - start)
var acc = Accumulator[DType.float32, nelts]()
@parameter
fn _exp[_nelts: Int](ii: Int):
var val = math.exp(x.load[width=_nelts](start + ii) - max_val)
x.store[width=_nelts](start + ii, val)
acc.accumulate(val)
vectorize[_exp, nelts](end - start)
var ssum = acc.total()
@parameter
fn _norm[_nelts: Int](ii: Int):
x.store[width=_nelts](
start + ii, x.load[width=_nelts](start + ii) / ssum
)
vectorize[_norm, nelts](end - start)
@always_inline
fn batch_matmul[
n: Int
](
C: StaticTuple[BufferPtrFloat32, n],
A: BufferPtrFloat32,
B: StaticTuple[BufferPtrFloat32, n],
rows: Int,
cols: Int,
):
@parameter
fn compute_row(i: Int):
var tmp = StaticTuple[Accumulator[DType.float32, nelts], n]()
@unroll
for k in range(n):
tmp[k] = Accumulator[DType.float32, nelts]()
var row_offset = i * cols
@parameter
fn dot[_nelts: Int](j: Int):
var a = A.load[width=_nelts](j)
@unroll
for k in range(n):
tmp[k].accumulate(a * B[k].load[width=_nelts](row_offset + j))
vectorize[dot, nelts](cols)
@unroll
for k in range(n):
C[k].store(i, tmp[k].total())
parallelize[compute_row](rows, workers)
@always_inline
fn matmul(C: TensorF32, A: TensorF32, B: TensorF32) raises:
# B (d,n) @ A (n,) -> C (d,)
matmul_dimension_checks(A.shape(), B.shape())
batch_matmul[1](
StaticTuple[BufferPtrFloat32, 1](C.data()),
A.data(),
StaticTuple[BufferPtrFloat32, 1](B.data()),
B.dim(0),
B.dim(1),
)
@always_inline
fn matmul(C: TensorF32, A: TensorF32, B: TensorSlice) raises:
# B (d,n) @ A (n,) -> C (d,)
matmul_dimension_checks(A.shape(), B.shape())
batch_matmul[1](
StaticTuple[BufferPtrFloat32, 1](C.data()),
A.data(),
StaticTuple[BufferPtrFloat32, 1](B.data()),
B.dim(0),
B.dim(1),
)
@always_inline
fn matmul(C: TensorSlice, A: TensorF32, B: TensorSlice) raises:
# B (d,n) @ A (n,) -> C (d,)
matmul_dimension_checks(A.shape(), B.shape())
batch_matmul[1](
StaticTuple[BufferPtrFloat32, 1](
C.data(),
),
A.data(),
StaticTuple[BufferPtrFloat32, 1](B.data()),
B.dim(0),
B.dim(1),
)
fn matmul_dimension_checks(a: TensorShape, b: TensorShape) raises:
if a[0] != b[1]:
raise Error(
"matmul dimension mismatch. A rows (dim 0) not equal to B columns"
" (dim 1)"
)
if b.rank() != 2:
raise Error("matmul expects B to be a 2D matrix")
# Apply RoPE rotation to the q and k vectors for each head
# rotate odd and even dim
@always_inline
fn rope_rotation_llama(
inout state: RunState,
freq_cis_real_row: TensorSlice,
freq_cis_imag_row: TensorSlice,
config: Config,
) -> None:
# stories model, llama2
var head_size = config.head_size
@parameter
fn head_loop(i: Int):
# Simple vectorization with (head_size // 2) steps gave junk transformer output.
# Maybe because the nelt ranges end up overlapping between the steps.
for j in range(0, config.head_size, 2):
var fcr = freq_cis_real_row[j // 2]
var fci = freq_cis_imag_row[j // 2]
var q0 = state.q[i * head_size + j]
var q1 = state.q[i * head_size + j + 1]
state.q[i * head_size + j] = q0 * fcr - q1 * fci
state.q[i * head_size + j + 1] = q0 * fci + q1 * fcr
if i < config.n_kv_heads:
var k0 = state.k[i * head_size + j]
var k1 = state.k[i * head_size + j + 1]
state.k[i * head_size + j] = k0 * fcr - k1 * fci
state.k[i * head_size + j + 1] = k0 * fci + k1 * fcr
parallelize[head_loop](config.n_heads, workers)
@always_inline
fn transformer(
token: Int,
pos: Int,
config: Config,
inout state: RunState,
weights: TransformerWeights,
) raises -> None:
# A few convenience variables
var dim = config.dim
var hidden_dim = config.hidden_dim
var head_size = config.head_size
var kv_dim = config.kv_dim
var kv_mul = config.kv_mul
# Copy the token embedding into x
var content_row = weights.token_embedding_table.data().offset(token * dim)
memcpy(state.x.data(), content_row, dim)
# Pluck out the "pos" row of freq_cis_real and freq_cis_imag
var freq_cis_real_row = TensorSlice(weights.freq_cis_real, pos)
var freq_cis_imag_row = TensorSlice(weights.freq_cis_imag, pos)
# Forward all the layers
for l in range(config.n_layers):
# Attention rmsnorm
rmsnorm(state.xb, state.x, TensorSlice(weights.rms_att_weight, l))
# QKV matmuls for this position
var loff = l * config.seq_len * config.kv_dim
state.k = TensorSlice(state.key_cache, l, pos)
state.v = TensorSlice(state.value_cache, l, pos)
if kv_dim == dim:
batch_matmul[3](
StaticTuple[BufferPtrFloat32, 3](
state.q.data(), state.k.data(), state.v.data()
),
state.xb.data(),
StaticTuple[BufferPtrFloat32, 3](
TensorSlice(weights.wq, l).data(),
TensorSlice(weights.wk, l).data(),
TensorSlice(weights.wv, l).data(),
),
dim,
dim,
)
else:
matmul(state.q, state.xb, TensorSlice(weights.wq, l))
batch_matmul[2](
StaticTuple[BufferPtrFloat32, 2](
state.k.data(), state.v.data()
),
state.xb.data(),
StaticTuple[BufferPtrFloat32, 2](
TensorSlice(weights.wk, l).data(),
TensorSlice(weights.wv, l).data(),
),
kv_dim,
dim,
)
# Apply RoPE rotation to the q and k vectors for each head
rope_rotation_llama(state, freq_cis_real_row, freq_cis_imag_row, config)
memset_zero(state.xb.data(), state.xb.num_elements())
# Multihead attention. Iterate over all heads in parallel.
@parameter
fn loop_over_heads(h: Int):
# Get the query vector for this head
var q_offset = h * head_size
# Index of attention scores for this head
var att_offset = h * config.seq_len
# Iterate over all timesteps, including the current one
for t in range(pos + 1):
# Starting index of the key vector for this head and at this timestep
var k_offset = loff + t * kv_dim + (h // kv_mul) * head_size
# Calculate the attention score as the dot product of q and k
var score: Float32 = 0.0
@parameter
fn score_fn[_nelts: Int](i: Int):
score += (
state.q.load[width=_nelts](q_offset + i)
* state.key_cache.load[width=_nelts](k_offset + i)
).reduce_add()
vectorize[score_fn, nelts](head_size)
score /= math.sqrt[DType.float32, 1](head_size)
# Save the score to the attention buffer
state.att[att_offset + t] = score
# Softmax the scores to get attention weights, from 0..pos inclusively
softmax(state.att, att_offset, att_offset + pos + 1)
# Weighted sum of the values, store back into xb
var xb_offset = h * head_size
for t in range(pos + 1):
# Starting index of the value vector for this head and at this timestep
var v_offset = loff + t * kv_dim + (h // kv_mul) * head_size
# Get the attention weight for this timestep
var a = state.att[att_offset + t]
# Accumulate the weighted value into xb
@parameter
fn xb_accumulate[_nelts: Int](i: Int):
var xbi = state.xb.load[width=_nelts](
xb_offset + i
) + a * state.value_cache.load[width=_nelts](v_offset + i)
state.xb.store[width=_nelts](xb_offset + i, xbi)
vectorize[xb_accumulate, nelts](head_size)
parallelize[loop_over_heads](config.n_heads, workers)
# Final matrix multiplication to get the output of the attention
matmul(state.xb2, state.xb, TensorSlice(weights.wo, l))
# Residual connection back into x
state.x = state.x + state.xb2
# FFN rmsnorm
rmsnorm(state.xb, state.x, TensorSlice(weights.rms_ffn_weight, l))
# Calculate self.w1(x) and self.w3(x) for FFN
batch_matmul[2](
StaticTuple[BufferPtrFloat32, 2](state.hb.data(), state.hb2.data()),
state.xb.data(),
StaticTuple[BufferPtrFloat32, 2](
TensorSlice(weights.w1, l).data(),
TensorSlice(weights.w3, l).data(),
),
hidden_dim,
dim,
)
@parameter
fn silu[_nelts: Int](i: Int):
var initial_hb = state.hb.load[width=_nelts](i)
# Apply SiLU activation function (silu(x) = x * sigmoid(x))
var hbi = initial_hb * (1.0 / (1.0 + math.exp(-initial_hb)))
# Elementwise multiply with w3(x)
state.hb.store[width=_nelts](
i, hbi * state.hb2.load[width=_nelts](i)
)
vectorize[silu, nelts](hidden_dim)
# Final matrix multiplication to get the output of the FFN
matmul(state.xb, state.hb, TensorSlice(weights.w2, l))
# Residual connection
state.x = state.x + state.xb
# Final rmsnorm
rmsnorm(state.x, state.x, weights.rms_final_weight)
# Classifier into logits
matmul(state.logits, state.x, weights.wcls)
fn sample(probabilities: TensorF32) -> Int:
var n = probabilities.dim(0)
# Sample index from probabilities, they must sum to 1
# get random value within (min, max) float32 range
var r = rand[DType.float32](1)
var cdf: Float32 = 0.0
for i in range(n):
cdf += probabilities[i]
if r[0] < cdf:
return i
return n - 1 # In case of rounding errors
fn bpe_encode(inout tokens: List[Int], text: String, tok: Tokenizer):
for pos in range(len(text)):
var char = text[pos]
var id = tok.find(char)
if id == -1:
print("Not a good prompt token at pos ", pos)
return
tokens.append(id)
while True:
var best_score = Float32(-1e10)
var best_id = -1
var best_idx = -1
for i in range(len(tokens) - 1):
# Check if we can merge the pair (tokens[i], tokens[i+1])
var str = str_concat(tok.vocab[tokens[i]], tok.vocab[tokens[i + 1]])
var id = tok.find(str)
if id != -1 and tok.vocab_scores[id] > best_score:
best_score = tok.vocab_scores[id]
best_id = id
best_idx = i
if best_idx == -1:
# We couldn't find any more pairs to merge, so we're done
break
# Merge the consecutive pair (best_idx, best_idx+1) into new token best_id
tokens[best_idx] = best_id
# Delete token at position best_idx+1, shift the entire sequence back 1
var _tokens = List[Int]()
for i in range(0, best_idx + 1):
_tokens.append(tokens[i])
for i in range(best_idx + 2, len(tokens)):
_tokens.append(tokens[i])
tokens = _tokens^
fn time_in_ms() -> Int:
# Returns time in milliseconds for benchmarking the model speed
return time.now() // 1_000_000
fn print_usage():
print("Usage: mojo llama2.mojo <checkpoint> [options]")
print(
'Example: mojo llama2.mojo stories15M.bin -s 99 -n 256 -t 0.5 -i "Llama'
' is an animal"'
)
print("Options:")
print(" -s <int> random seed, default time.now()")
print(" -t <float> temperature in [0,1.0], default 1.0")
print(
" -n <int> number of steps to run for, default 256. 0 = max_seq_len"
)
print(" -i <string> input prompt")
print(" -z tokenizer path")
print(" -j number of workers to use, default num_cores()")
fn main() raises:
workers = num_performance_cores()
var tokenizer = StringRef("tokenizer.bin")
var checkpoint = StringRef("stories15M.bin")
var temperature = 0.9
var steps = 256
var prompt = String("")
var rng_seed: Int = time.now()
var print_config = 0
@parameter
fn argparse() raises -> Int:
var args = argv()
if len(args) < 2:
return 0
checkpoint = args[1]
for i in range(2, len(args), 2):
if args[i] == "-p":
print("Option not supported: ", args[i])
if args[i] == "-n":
steps = atol(args[i + 1])
if args[i] == "-z":
tokenizer = args[i + 1]
if args[i] == "-s":
rng_seed = atol(args[i + 1])
if args[i] == "-i":
prompt = args[i + 1]
if args[i] == "-j":
workers = atol(args[i + 1])
if args[i] == "-pc":
print_config = atol(args[i + 1])
if args[i] == "-t":
var val = args[i + 1]
temperature = 0.0
# hacky parse float, keep only 1 digit
for c in range(0, len(val)):
if val[c] == ".":
temperature += atol(val[c + 1]) * 0.1
break
else:
temperature = atol(val[c])
if temperature < -1e9 or temperature > (1 + 1e9):
print("Wrong temperature value", temperature)
return 0
return 1
var res = argparse()
if res == 0:
print_usage()
return
print("num parallel workers:", workers, " SIMD width:", nelts)
random.seed(rng_seed)
var config = Config(checkpoint, print_config == 1)
var weights = TransformerWeights(checkpoint, config)
if steps <= 0 or steps > config.seq_len:
steps = config.seq_len
var tok = Tokenizer(config.vocab_size, tokenizer)
print(
"n layers:",
config.n_layers,
"| vocab size:",
tok.vocab_size,
)
# Create and initialize the application RunState
var state = RunState(config)
# Process the prompt, if any
var prompt_tokens = List[Int]()
if prompt:
bpe_encode(prompt_tokens, prompt, tok)
# Start the main loop
var start = 0 # Used to time our code, only initialized after the first iteration
var next_token = 0 # Will store the next token in the sequence
# Initialize with token 1 (=BOS), as done in Llama-2 sentencepiece tokenizer
var token = 1
# Position in the sequence
var pos = 0
while pos < steps:
# Forward the transformer to get logits for the next token
transformer(token, pos, config, state, weights)
if pos < len(prompt_tokens):
next_token = prompt_tokens[pos]
else:
# Sample the next token
if temperature == 0.0:
# Greedy argmax sampling: take the token with the highest probability
next_token = int(state.logits.argmax()[0])
else:
# Apply the temperature to the logits
for q in range(config.vocab_size):
state.logits[q] = state.logits[q] / temperature
# Apply softmax to the logits to get the probabilities for the next token
softmax(state.logits)
# Sample from this distribution to get the next token
next_token = sample(state.logits)
# Finish generating when EOS, BOS appear
if next_token == 1 or next_token == 2:
break
var token_str: String = tok.vocab[next_token]
if token == 1 and token_str._buffer[0] == ord(" "):
token_str = token_str[1:]
print(token_str, end="")
# Advance forward
token = next_token
pos += 1
if start == 0:
start = time_in_ms()
var end = time_in_ms()
print("\nachieved tok/s: ", (pos - 1) / (end - start) * 1000)