CHICKEN Scheme API

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

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

[procedure] (make-args checkpoint tokenizer temperature steps prompt seed)

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? x)(args-checkpoint args)(args-tokenizer args)(args-temperature args)(args-steps args)(args-prompt args)(args-seed args)

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)

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? x)(config-dim config)(config-hidden-dim config)(config-n-layers config)(config-n-heads config)(config-n-kv-heads config)(config-vocab-size config)(config-seq-len config)(config-shared-weights config)

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)

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? x)(transformer-weights-token-embedding-table w)(transformer-weights-rms-att-weight w)(transformer-weights-wq w)(transformer-weights-wk w)(transformer-weights-wv w)(transformer-weights-wo w)(transformer-weights-rms-ffn-weight w)(transformer-weights-w1 w)(transformer-weights-w2 w)(transformer-weights-w3 w)(transformer-weights-rms-final-weight w)(transformer-weights-freq-cis-real w)(transformer-weights-freq-cis-imag w)(transformer-weights-wcls w)

Predicate and field accessors.

Checkpoints and tokenizer

[procedure] (read-config port checkpoint)

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

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)

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)

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.

Prints the summary used by verify-checkpoint-data for flat-layout weights.

Tokens and sampling

[procedure] (bpe-encode text vocab vocab-scores)

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)

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

[procedure] (sample probabilities state)

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

[procedure] (softmax-f32! out x size)

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 . vs)(f64vector-fold f init v . vs)

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)

Returns s without leading whitespace.

Inference backends

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

[procedure] (inference-backend? x)

True for inference backends.

[procedure] (backend-name backend)

Returns a symbol naming the backend.

[procedure] (backend-prepare! backend config port file-size)

Reads the weights from the checkpoint port, positioned just after the configuration header, and allocates the run state.

[procedure] (backend-forward! backend token pos)

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)

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.

[procedure] (run args)

Generates text with the blas backend.

[procedure] (make-blas-inference-backend)

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)

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 state)(run-state-xb state)(run-state-xb2 state)(run-state-q state)(run-state-k state)(run-state-v state)(run-state-hb state)(run-state-hb2 state)(run-state-key-cache state)(run-state-value-cache state)(run-state-logits state)

Field accessors.

Forward pass

[procedure] (transformer token pos config state weights)

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)

Copies the embedding of token into x.

[procedure] (get-rope-frequencies weights pos head-size)

Returns two values, the cosine and sine rows of the rotary embedding tables for position pos.

[procedure] (attention-rmsnorm state weights layer-idx config)

xb := RMS-normalized x, scaled by the attention norm weights.

[procedure] (compute-qkv state weights layer-idx config)

Computes q, k and v from xb.

[procedure] (apply-rope state config freq-real freq-imag)

Rotates q and k with the rotary embedding rows from get-rope-frequencies.

[procedure] (cache-kv state layer-idx pos config)

Stores k and v in the key/value cache at position pos.

[procedure] (compute-attention state layer-idx pos config)

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)

xb2 := Wo xb.

[procedure] (ffn-rmsnorm state weights layer-idx config)

xb := RMS-normalized x, scaled by the feed-forward norm weights.

[procedure] (compute-ffn-w1w3 state weights layer-idx config)

hb := W1 xb and hb2 := W3 xb.

[procedure] (apply-swiglu state config)

hb := SiLU(hb) * hb2, element-wise.

[procedure] (ffn-output state weights layer-idx config)

xb := W2 hb.

[procedure] (process-transformer-layer state weights layer-idx pos config freq-real freq-imag)

Runs one layer: the attention block and the feed-forward block, each followed by its residual addition to x.

[procedure] (final-rmsnorm state weights)

RMS-normalizes x in place with the final norm weights.

[procedure] (compute-logits state weights config)

logits := Wcls x.

Vector operations

[procedure] (rmsnorm out x weight)

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

[procedure] (matmul out x w n d)

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

[procedure] (softmax out x size)

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

[procedure] (accum a b)

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.

[procedure] (run args #!key threads)

Generates text with the crunch backend.

[procedure] (make-am-inference-backend #!key threads)

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)

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)

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)

Allocates the run state and compiles the graphs for am-weights.

[procedure] (am-transformer token pos config state am-weights)

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)

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

Author

Ivan Raikov

Repository

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