llama

Llama 2 inference in CHICKEN Scheme, based on Andrej Karpathy's llama2.c and its OCaml port 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 blas egg and a system BLAS library.
crunch: the llama-am module, built on 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
srfi-1
srfi-69
vector-lib
blas
random-mtzig
endian-sequence
endian-port
getopt-long
yasos
datatype
array-morphisms
array-morphisms-crunch

The tests also need test.

Installation

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

Command-Line Interface

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

-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

(import llama-common)

Generation arguments

make-args checkpoint tokenizer temperature steps prompt seedprocedure

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

args? xprocedure
args-checkpoint argsprocedure
args-tokenizer argsprocedure
args-temperature argsprocedure
args-steps argsprocedure
args-prompt argsprocedure
args-seed argsprocedure: Predicate and field accessors.

Model configuration

make-config dim hidden-dim n-layers n-heads n-kv-heads vocab-size seq-len shared-weightsprocedure

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

config? xprocedure
config-dim configprocedure
config-hidden-dim configprocedure
config-n-layers configprocedure
config-n-heads configprocedure
config-n-kv-heads configprocedure
config-vocab-size configprocedure
config-seq-len configprocedure
config-shared-weights configprocedure: Predicate and field accessors.

Model weights

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 wclsprocedure: 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).

transformer-weights? xprocedure
transformer-weights-token-embedding-table wprocedure
transformer-weights-rms-att-weight wprocedure
transformer-weights-wq wprocedure
transformer-weights-wk wprocedure
transformer-weights-wv wprocedure
transformer-weights-wo wprocedure
transformer-weights-rms-ffn-weight wprocedure
transformer-weights-w1 wprocedure
transformer-weights-w2 wprocedure
transformer-weights-w3 wprocedure
transformer-weights-rms-final-weight wprocedure
transformer-weights-freq-cis-real wprocedure
transformer-weights-freq-cis-imag wprocedure
transformer-weights-wcls wprocedure: Predicate and field accessors.

Checkpoints and tokenizer

read-config port checkpointprocedure: 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.

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

tokenizer-init config portprocedure: 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.

verify-checkpoint-data checkpoint-file #!optional detailedprocedure: 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.

print-transformer-weights weights config #!optional detailedprocedure: Prints the summary used by verify-checkpoint-data for flat-layout weights.

Tokens and sampling

bpe-encode text vocab vocab-scoresprocedure: 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.

argmax vprocedure: Returns the index of the largest element of the f32vector v; the first one when several are equal.

sample probabilities stateprocedure: Draws an index from the distribution in the f32vector probabilities, using the random-mtzig state state.

softmax-f32! out x sizeprocedure: Stores the softmax of the first size elements of x in out and copies the remaining elements unchanged. out may be x. Returns out.

f32vector-fold f init v #!rest vsprocedure
f64vector-fold f init v #!rest vsprocedure: Fold over the elements of one or more vectors of equal length. f receives one element of each vector followed by the accumulated value.

trim-left sprocedure: Returns s without leading whitespace.

Inference backends

An inference backend is a yasos object that implements the following operations.

inference-backend? xprocedure: True for inference backends.

backend-name backendprocedure: Returns a symbol naming the backend.

backend-prepare! backend config port file-sizeprocedure: Reads the weights from the checkpoint port, positioned just after the configuration header, and allocates the run state.

backend-forward! backend token posprocedure: 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.

run-with-backend args backendprocedure: 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

(import llama)

The forward pass on the 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.

run argsprocedure: Generates text with the blas backend.

make-blas-inference-backendprocedure: Returns the blas inference backend, for run-with-backend.

Run state

make-run-state x xb q k v att key-cache value-cache xb2 hb hb2 logitsprocedure: 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.

run-state-x stateprocedure
run-state-xb stateprocedure
run-state-xb2 stateprocedure
run-state-q stateprocedure
run-state-k stateprocedure
run-state-v stateprocedure
run-state-hb stateprocedure
run-state-hb2 stateprocedure
run-state-key-cache stateprocedure
run-state-value-cache stateprocedure
run-state-logits stateprocedure: Field accessors.

Forward pass

transformer token pos config state weightsprocedure

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.

token-embedding-lookup state weights tokenprocedure: Copies the embedding of token into x.

get-rope-frequencies weights pos head-sizeprocedure: Returns two values, the cosine and sine rows of the rotary embedding tables for position pos.

attention-rmsnorm state weights layer-idx configprocedure: xb := RMS-normalized x, scaled by the attention norm weights.

compute-qkv state weights layer-idx configprocedure: Computes q, k and v from xb.

apply-rope state config freq-real freq-imagprocedure: Rotates q and k with the rotary embedding rows from get-rope-frequencies.

cache-kv state layer-idx pos configprocedure: Stores k and v in the key/value cache at position pos.

compute-attention state layer-idx pos configprocedure: Multi-head attention of q over the cached positions 0..pos; the result is stored in xb.

attention-output state weights layer-idx configprocedure: xb2 := Wo xb.

ffn-rmsnorm state weights layer-idx configprocedure: xb := RMS-normalized x, scaled by the feed-forward norm weights.

compute-ffn-w1w3 state weights layer-idx configprocedure: hb := W1 xb and hb2 := W3 xb.

apply-swiglu state configprocedure: hb := SiLU(hb) * hb2, element-wise.

ffn-output state weights layer-idx configprocedure: xb := W2 hb.

process-transformer-layer state weights layer-idx pos config freq-real freq-imagprocedure: Runs one layer: the attention block and the feed-forward block, each followed by its residual addition to x.

final-rmsnorm state weightsprocedure: RMS-normalizes x in place with the final norm weights.

compute-logits state weights configprocedure: logits := Wcls x.

Vector operations

rmsnorm out x weightprocedure: out := weight * x / sqrt(mean(x^2) + 1e-5), element-wise. Returns out.

matmul out x w n dprocedure: out := w x, where w is a row-major d x n matrix. Returns out.

softmax out x sizeprocedure: Softmax of the first size elements of x into out; the remaining elements are copied. Returns out.

accum a bprocedure

a := a + b. Returns a.

scopy is re-exported from blas.

Module llama-am

(import llama-am)

The forward pass on 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.

run args #!key threadsprocedure: Generates text with the crunch backend.

make-am-inference-backend #!key threadsprocedure: Returns the crunch inference backend, for run-with-backend. threads, when given, is the number of threads used by the crunch kernels.

register-crunch-backend! #!key threadsprocedure: 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.

transformer-weights->am-weights config weightsprocedure: 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.

make-am-state config am-weightsprocedure: Allocates the run state and compiles the graphs for am-weights.

am-transformer token pos config state am-weightsprocedure: Runs the model on token id token at position pos and returns the logits f32vector, which is the one am-state-logits returns.

am-state-logits stateprocedure: Returns the logits vector of state.

Examples

Generate text with each backend:

(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))

Greedy decoding of a few tokens with the llama components:

(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)))))

A backend is any yasos object that implements the backend operations. This one wraps another backend and counts forward passes:

(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)))

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

chickadee » llama

Identifier search

TOC »