You can edit this page using wiki syntax for markup.

Article contents:[[tags: egg ai machine-learning nlp llm transformer blas]] [[toc:]] == llama Llama 2 inference in CHICKEN Scheme, based on Andrej Karpathy's [[https://github.com/karpathy/llama2.c|llama2.c]] and its OCaml port [[https://github.com/jackpeck/llama2.ml|llama2.ml]]. === Description The library reads a llama2.c checkpoint and tokenizer and generates text one token at a time. The transformer forward pass has two implementations, which compute the same result: ; {{blas}} : the {{llama}} module, built on the [[/eggref/6/blas|blas]] egg and a system BLAS library. ; {{crunch}} : the {{llama-am}} module, built on [[/eggref/6/array-morphisms|array-morphisms]] with the crunch kernels of {{array-morphisms-crunch}}. The parts that do not depend on the numerical backend (records, checkpoint and tokenizer readers, BPE encoding, sampling and the generation loop) are in the {{llama-common}} module. === Requirements * A BLAS library (OpenBLAS, Intel MKL or the system BLAS), for the {{blas}} backend * [[/eggref/6/srfi-1|srfi-1]] * [[/eggref/6/srfi-69|srfi-69]] * [[/eggref/6/vector-lib|vector-lib]] * [[/eggref/6/blas|blas]] * [[/eggref/6/random-mtzig|random-mtzig]] * {{endian-sequence}} * {{endian-port}} * [[/eggref/6/getopt-long|getopt-long]] * [[/eggref/6/yasos|yasos]] * [[/eggref/6/datatype|datatype]] * [[/eggref/6/array-morphisms|array-morphisms]] * {{array-morphisms-crunch}} The tests also need [[/eggref/6/test|test]]. === Installation <enscript highlight="shell"> chicken-install llama # Example model (15M parameters, about 60 MB) and its tokenizer wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin wget https://github.com/iraikov/llama-chicken/raw/refs/heads/main/tokenizer.bin </enscript> === Command-Line Interface <enscript highlight="shell"> llama-cli -c stories15M.bin -p "Once upon a time" llama-cli -c stories15M.bin -t 0.8 -s 100 -p "The meaning of life is" llama-cli -c stories15M.bin --backend=crunch --threads=4 -p "Once upon a time" llama-cli -c stories15M.bin --verify-checkpoint </enscript> ; {{-c}}, {{--checkpoint FILE}} : model checkpoint (required) ; {{-k}}, {{--tokenizer FILE}} : tokenizer file (default {{tokenizer.bin}}) ; {{-t}}, {{--temperature NUM}} : sampling temperature; 0 selects the most likely token (default 0) ; {{-s}}, {{--steps NUM}} : number of tokens to generate, at most the model's sequence length (default 256) ; {{-p}}, {{--prompt TEXT}} : prompt text (default empty) ; {{-r}}, {{--seed NUM}} : random seed for sampling (default: the current time) ; {{-b}}, {{--backend NAME}} : {{blas}} (default) or {{crunch}} ; {{-j}}, {{--threads NUM}} : number of threads used by the {{crunch}} backend ; {{--verify-checkpoint}} : print a summary of the checkpoint weights and exit ; {{-h}}, {{--help}} : print usage and exit Long options take their value after {{=}}, as in {{--steps=100}}. Without {{--threads}}, the {{crunch}} backend uses the {{AM_CRUNCH_THREADS}} environment variable, or the number of processors up to 8. === Module llama-common <enscript highlight="scheme"> (import llama-common) </enscript> ==== Generation arguments <procedure>(make-args checkpoint tokenizer temperature steps prompt seed)</procedure> Returns the arguments of a generation run. ; {{checkpoint}} : path of the model checkpoint ; {{tokenizer}} : path of the tokenizer file ; {{temperature}} : sampling temperature; {{0.0}} selects the most likely token ; {{steps}} : number of tokens to generate; values below 1 or above the sequence length mean the sequence length ; {{prompt}} : prompt text, possibly empty ; {{seed}} : random seed, or {{#f}} to use the current time <procedure>(args? x)</procedure> <procedure>(args-checkpoint args)</procedure> <procedure>(args-tokenizer args)</procedure> <procedure>(args-temperature args)</procedure> <procedure>(args-steps args)</procedure> <procedure>(args-prompt args)</procedure> <procedure>(args-seed args)</procedure> Predicate and field accessors. ==== Model configuration <procedure>(make-config dim hidden-dim n-layers n-heads n-kv-heads vocab-size seq-len shared-weights)</procedure> Returns a model configuration. ; {{dim}} : embedding dimension ; {{hidden-dim}} : hidden dimension of the feed-forward network ; {{n-layers}} : number of transformer layers ; {{n-heads}} : number of attention heads ; {{n-kv-heads}} : number of key/value heads ; {{vocab-size}} : vocabulary size ; {{seq-len}} : maximum sequence length ; {{shared-weights}} : true when the classifier shares the token embedding table <procedure>(config? x)</procedure> <procedure>(config-dim config)</procedure> <procedure>(config-hidden-dim config)</procedure> <procedure>(config-n-layers config)</procedure> <procedure>(config-n-heads config)</procedure> <procedure>(config-n-kv-heads config)</procedure> <procedure>(config-vocab-size config)</procedure> <procedure>(config-seq-len config)</procedure> <procedure>(config-shared-weights config)</procedure> Predicate and field accessors. ==== Model weights <procedure>(make-transformer-weights token-embedding-table rms-att-weight wq wk wv wo rms-ffn-weight w1 w2 w3 rms-final-weight freq-cis-real freq-cis-imag wcls)</procedure> Returns the model weights. The fields hold f32vectors. Matrices are stored row-major, with one row per output element. A layered tensor (the attention and feed-forward weights and their norm weights) holds either one f32vector with all layers one after another (the flat layout) or a vector of one f32vector per layer (the per-layer layout). <procedure>(transformer-weights? x)</procedure> <procedure>(transformer-weights-token-embedding-table w)</procedure> <procedure>(transformer-weights-rms-att-weight w)</procedure> <procedure>(transformer-weights-wq w)</procedure> <procedure>(transformer-weights-wk w)</procedure> <procedure>(transformer-weights-wv w)</procedure> <procedure>(transformer-weights-wo w)</procedure> <procedure>(transformer-weights-rms-ffn-weight w)</procedure> <procedure>(transformer-weights-w1 w)</procedure> <procedure>(transformer-weights-w2 w)</procedure> <procedure>(transformer-weights-w3 w)</procedure> <procedure>(transformer-weights-rms-final-weight w)</procedure> <procedure>(transformer-weights-freq-cis-real w)</procedure> <procedure>(transformer-weights-freq-cis-imag w)</procedure> <procedure>(transformer-weights-wcls w)</procedure> Predicate and field accessors. ==== Checkpoints and tokenizer <procedure>(read-config port checkpoint)</procedure> Reads the configuration header from the little-endian {{endian-port}} {{port}}. {{checkpoint}} is the file name, used in error messages. A negative vocabulary size in the header marks a checkpoint whose classifier weights are stored separately; the returned configuration then has a positive vocabulary size and {{shared-weights}} false. <procedure>(checkpoint-init-weights config port shared-weights file-size #!key (layout 'flat))</procedure> Reads the weights that follow the configuration header and returns a {{transformer-weights}} record. ; {{shared-weights}} : true when the classifier is the token embedding table; otherwise the classifier is read from the rest of the file ; {{file-size}} : size of the checkpoint file in bytes ; {{layout}} : {{'flat}} or {{'per-layer}}, as described under Model weights <procedure>(tokenizer-init config port)</procedure> Reads the tokenizer file from {{port}} and returns three values: the list of token strings (decoded as UTF-8), the list of their merge scores, and the maximum token length. <procedure>(verify-checkpoint-data checkpoint-file #!optional detailed)</procedure> Loads a checkpoint, prints a summary of its weights and checks that they are finite. With {{detailed}} true it also prints sample values. Returns the weights. <procedure>(print-transformer-weights weights config #!optional detailed)</procedure> Prints the summary used by {{verify-checkpoint-data}} for flat-layout weights. ==== Tokens and sampling <procedure>(bpe-encode text vocab vocab-scores)</procedure> Encodes {{text}} as a list of token ids. Each character is first mapped to its token; adjacent tokens are then merged, highest merge score first, while the merged string is in the vocabulary. Exits with an error message when a character has no token. <procedure>(argmax v)</procedure> Returns the index of the largest element of the f32vector {{v}}; the first one when several are equal. <procedure>(sample probabilities state)</procedure> Draws an index from the distribution in the f32vector {{probabilities}}, using the [[/eggref/6/random-mtzig|random-mtzig]] state {{state}}. <procedure>(softmax-f32! out x size)</procedure> Stores the softmax of the first {{size}} elements of {{x}} in {{out}} and copies the remaining elements unchanged. {{out}} may be {{x}}. Returns {{out}}. <procedure>(f32vector-fold f init v . vs)</procedure> <procedure>(f64vector-fold f init v . vs)</procedure> Fold over the elements of one or more vectors of equal length. {{f}} receives one element of each vector followed by the accumulated value. <procedure>(trim-left s)</procedure> Returns {{s}} without leading whitespace. ==== Inference backends An inference backend is a [[/eggref/6/yasos|yasos]] object that implements the following operations. <procedure>(inference-backend? x)</procedure> True for inference backends. <procedure>(backend-name backend)</procedure> Returns a symbol naming the backend. <procedure>(backend-prepare! backend config port file-size)</procedure> Reads the weights from the checkpoint {{port}}, positioned just after the configuration header, and allocates the run state. <procedure>(backend-forward! backend token pos)</procedure> Runs the model on token id {{token}} at position {{pos}} and returns the f32vector of logits, one per vocabulary entry. Positions are given in order, starting at 0. <procedure>(run-with-backend args backend)</procedure> Generates text as described by {{args}} with {{backend}}. Starting from the beginning-of-sequence token, the prompt tokens are fed in first; after them each next token is sampled from the logits. Tokens are printed as they are produced, followed by the generation rate. === Module llama <enscript highlight="scheme"> (import llama) </enscript> The forward pass on the [[/eggref/6/blas|blas]] egg, with weights in the flat layout. The module also exports {{make-args}}, the {{args}} accessors, {{make-config}}, {{config-dim}}, {{config-n-layers}}, {{make-transformer-weights}}, {{bpe-encode}}, {{argmax}} and {{verify-checkpoint-data}} from {{llama-common}}. <procedure>(run args)</procedure> Generates text with the {{blas}} backend. <procedure>(make-blas-inference-backend)</procedure> Returns the {{blas}} inference backend, for {{run-with-backend}}. ==== Run state <procedure>(make-run-state x xb q k v att key-cache value-cache xb2 hb hb2 logits)</procedure> Returns the activations of the forward pass, all f32vectors: {{x}}, {{xb}}, {{xb2}}, {{q}}, {{k}} and {{v}} of length {{dim}}; {{att}} of length {{n-heads * seq-len}}; {{key-cache}} and {{value-cache}} of length {{n-layers * seq-len * dim}}; {{hb}} and {{hb2}} of length {{hidden-dim}}; {{logits}} of length {{vocab-size}}. <procedure>(run-state-x state)</procedure> <procedure>(run-state-xb state)</procedure> <procedure>(run-state-xb2 state)</procedure> <procedure>(run-state-q state)</procedure> <procedure>(run-state-k state)</procedure> <procedure>(run-state-v state)</procedure> <procedure>(run-state-hb state)</procedure> <procedure>(run-state-hb2 state)</procedure> <procedure>(run-state-key-cache state)</procedure> <procedure>(run-state-value-cache state)</procedure> <procedure>(run-state-logits state)</procedure> Field accessors. ==== Forward pass <procedure>(transformer token pos config state weights)</procedure> Runs the model on token id {{token}} at position {{pos}}, updating {{state}} in place; the logits are left in {{(run-state-logits state)}}. Returns {{state}}. The steps of the forward pass are exported individually. Each one reads and writes the vectors of {{state}}; {{layer-idx}} selects the layer. <procedure>(token-embedding-lookup state weights token)</procedure> Copies the embedding of {{token}} into {{x}}. <procedure>(get-rope-frequencies weights pos head-size)</procedure> Returns two values, the cosine and sine rows of the rotary embedding tables for position {{pos}}. <procedure>(attention-rmsnorm state weights layer-idx config)</procedure> {{xb}} := RMS-normalized {{x}}, scaled by the attention norm weights. <procedure>(compute-qkv state weights layer-idx config)</procedure> Computes {{q}}, {{k}} and {{v}} from {{xb}}. <procedure>(apply-rope state config freq-real freq-imag)</procedure> Rotates {{q}} and {{k}} with the rotary embedding rows from {{get-rope-frequencies}}. <procedure>(cache-kv state layer-idx pos config)</procedure> Stores {{k}} and {{v}} in the key/value cache at position {{pos}}. <procedure>(compute-attention state layer-idx pos config)</procedure> Multi-head attention of {{q}} over the cached positions {{0..pos}}; the result is stored in {{xb}}. <procedure>(attention-output state weights layer-idx config)</procedure> {{xb2}} := Wo {{xb}}. <procedure>(ffn-rmsnorm state weights layer-idx config)</procedure> {{xb}} := RMS-normalized {{x}}, scaled by the feed-forward norm weights. <procedure>(compute-ffn-w1w3 state weights layer-idx config)</procedure> {{hb}} := W1 {{xb}} and {{hb2}} := W3 {{xb}}. <procedure>(apply-swiglu state config)</procedure> {{hb}} := SiLU({{hb}}) * {{hb2}}, element-wise. <procedure>(ffn-output state weights layer-idx config)</procedure> {{xb}} := W2 {{hb}}. <procedure>(process-transformer-layer state weights layer-idx pos config freq-real freq-imag)</procedure> Runs one layer: the attention block and the feed-forward block, each followed by its residual addition to {{x}}. <procedure>(final-rmsnorm state weights)</procedure> RMS-normalizes {{x}} in place with the final norm weights. <procedure>(compute-logits state weights config)</procedure> {{logits}} := Wcls {{x}}. ==== Vector operations <procedure>(rmsnorm out x weight)</procedure> {{out}} := {{weight}} * {{x}} / sqrt(mean({{x}}^2) + 1e-5), element-wise. Returns {{out}}. <procedure>(matmul out x w n d)</procedure> {{out}} := {{w}} {{x}}, where {{w}} is a row-major {{d}} x {{n}} matrix. Returns {{out}}. <procedure>(softmax out x size)</procedure> Softmax of the first {{size}} elements of {{x}} into {{out}}; the remaining elements are copied. Returns {{out}}. <procedure>(accum a b)</procedure> {{a}} := {{a}} + {{b}}. Returns {{a}}. {{scopy}} is re-exported from [[/eggref/6/blas|blas]]. === Module llama-am <enscript highlight="scheme"> (import llama-am) </enscript> The forward pass on [[/eggref/6/array-morphisms|array-morphisms]] with the crunch kernels of {{array-morphisms-crunch}}. Each layer's position-independent work is written as two array morphism graphs, one for the query/key/value projection and one for the attention output and feed-forward network, plus one graph for the classifier. Each graph is compiled once into an SSA program and replayed for every token into buffers planned by a morphism context, so a token allocates little memory. Attention over the key/value cache, whose length grows with the position, calls the crunch softmax, rotary embedding and BLAS kernels directly. Weights are read in the per-layer layout. <procedure>(run args #!key threads)</procedure> Generates text with the {{crunch}} backend. <procedure>(make-am-inference-backend #!key threads)</procedure> Returns the {{crunch}} inference backend, for {{run-with-backend}}. {{threads}}, when given, is the number of threads used by the crunch kernels. <procedure>(register-crunch-backend! #!key threads)</procedure> Makes crunch the active array-morphisms BLAS and element-wise backend, and sets the number of crunch threads when {{threads}} is given. Must be called before {{make-am-state}}. <procedure>(transformer-weights->am-weights config weights)</procedure> Converts weights in either layout to the layout of this module: one array morphism per matrix and layer, with Wq, Wk and Wv stacked into one matrix per layer. <procedure>(make-am-state config am-weights)</procedure> Allocates the run state and compiles the graphs for {{am-weights}}. <procedure>(am-transformer token pos config state am-weights)</procedure> Runs the model on token id {{token}} at position {{pos}} and returns the logits f32vector, which is the one {{am-state-logits}} returns. <procedure>(am-state-logits state)</procedure> Returns the logits vector of {{state}}. === Examples Generate text with each backend: <enscript highlight="scheme"> (import llama-common (only llama make-blas-inference-backend) (only llama-am make-am-inference-backend)) (define args (make-args "stories15M.bin" "tokenizer.bin" 0.8 100 "Once upon a time" #f)) (run-with-backend args (make-blas-inference-backend)) (run-with-backend args (make-am-inference-backend threads: 4)) </enscript> Greedy decoding of a few tokens with the {{llama}} components: <enscript highlight="scheme"> (import scheme (chicken base) (chicken file posix) (srfi 4) endian-port llama-common (only llama make-run-state run-state-logits transformer)) (define port (open-endian-port 'read "stories15M.bin")) (set-littlendian! port) (define config (read-config port "stories15M.bin")) (define weights (checkpoint-init-weights config port (config-shared-weights config) (file-size "stories15M.bin"))) (close-endian-port port) (define dim (config-dim config)) (define hidden (config-hidden-dim config)) (define cache (* (config-n-layers config) (config-seq-len config) dim)) (define state (make-run-state (make-f32vector dim 0.0) (make-f32vector dim 0.0) (make-f32vector dim 0.0) (make-f32vector dim 0.0) (make-f32vector dim 0.0) (make-f32vector (* (config-n-heads config) (config-seq-len config)) 0.0) (make-f32vector cache 0.0) (make-f32vector cache 0.0) (make-f32vector dim 0.0) (make-f32vector hidden 0.0) (make-f32vector hidden 0.0) (make-f32vector (config-vocab-size config) 0.0))) ;; Token ids of the first five greedy steps after the start token 1 (let loop ((token 1) (pos 0) (out '())) (if (= pos 5) (reverse out) (let ((next (argmax (run-state-logits (transformer token pos config state weights))))) (loop next (+ pos 1) (cons next out))))) </enscript> A backend is any yasos object that implements the backend operations. This one wraps another backend and counts forward passes: <enscript highlight="scheme"> (import scheme (chicken base) llama-common (only llama make-blas-inference-backend) (except yasos object object-with-ancestors operate-as)) (define (make-counting-backend inner) (let ((count 0)) (operations () ((inference-backend? self) #t) ((backend-name self) 'counting) ((backend-prepare! self config port file-size) (backend-prepare! inner config port file-size)) ((backend-forward! self token pos) (set! count (+ count 1)) (backend-forward! inner token pos))))) (run-with-backend (make-args "stories15M.bin" "tokenizer.bin" 0.0 20 "" #f) (make-counting-backend (make-blas-inference-backend))) </enscript> === Author [[/users/ivan-raikov|Ivan Raikov]] === Repository [[https://github.com/iraikov/llama-chicken|https://github.com/iraikov/llama-chicken]] === License MIT === Version History ; 2.0 : Port to CHICKEN 6; {{llama-common}} module with the inference backend protocol; {{llama-am}} module with the array-morphisms/crunch backend; {{--backend}} and {{--threads}} options and {{-r}} short option for the seed in {{llama-cli}}; UTF-8 tokenizer strings ; 1.2 : Added {{getopt-long}} dependency ; 1.1 : Added {{endian-port}} dependency ; 1.0 : Initial release === See Also * [[https://github.com/karpathy/llama2.c|llama2.c]], the original C implementation * [[https://github.com/jackpeck/llama2.ml|llama2.ml]], the OCaml port * [[https://arxiv.org/abs/2307.09288|Llama 2: Open Foundation and Fine-Tuned Chat Models]] Description of your changes:

I would like to authenticate

Spam control

What do you get when you multiply 9 by 8?