kleincode / llama2.jl

Llama 2 Inference in Julia. Developed for the JuliaML course at TU Berlin.

Llama2.jl

This is a port of Andrew Karpathy's llama2.c to Julia.

Important

This is part of the JuliaML course at Technical University Berlin. Therefore, the repository will not be maintained for a long time.

Features

• Read tokenizer and model weights from .bin files specified in llama2.c
• Inference, generation loop and chat loop
• argmax, multinomial and top-p sampling
• Tokenizer for encoding text to LLM input and decoding LLM output to text
• Multi-threading in transformer forward function (#37)
• Compatibility tested with all of Andrew Karpathy's models

Getting started

Add the package to your local environment via Pkg by running

add https://github.com/kleincode/Llama2.jl

To get started, check out the docs.

llama2.jl's Issues

Read the Llama2.c file format

lines 111-162 in llama2.c
Depends on #3

Use in-place operations in forward function

• Replace matrix multiplications with mul!(out, in1, in2)
• Add a @view where applicable, maybe even an @views in front of the whole forward function? But make sure to copy token_embedding_table into x in line 144!
• Replace rmsnorm, softmax, swiglu with in-place versions
• Remove mutable from mutable struct RunState

Feedback before Code Review 2

Hi all,

here's some quick feedback before the second code review session next week.
I'm opening this as a single issue to not spam your repository, which already looks really good!

Documentation

• Please describe in the README what the repo is about and that it is part of course work at TU Berlin
• Your docstrings still contain TODOs and notes to yourselves (e.g. https://kleincode.github.io/Llama2.jl/dev/#Llama2.open_file-Tuple{String}). Try to write your documentation for potential users, not for yourselves!
• Please add a second doc page with a small "Getting started" guide for the second code review session. Even just demonstrating your tokenizer is enough.

Code

• It's a bit odd to keep the Tokenizer struct and Tokenizer constructor in separate places. Use a single docstring for both and document the way you expect users to use your struct.
  

  Llama2.jl/src/tokenizer.jl

  Line 10 in a4c3241

  struct Tokenizer

  
  

  Llama2.jl/src/tokenizer.jl

  Line 48 in a4c3241

  function Tokenizer(tokenizer_path::String, vocab_size::Int)

• You could use DocStringExtensions.jl to generate parts of your documentation (see example here)
• Most of your structs hardcode element types to be Float32. I understand that this makes sense for a direct port of the C code base. However, in general, a parametric type Sampler{T} would make sense here.
  

  Llama2.jl/src/sampler.jl

  Lines 20 to 23 in a4c3241

  	struct Sampler
  	temperature::Float32
  	topp::Float32
  	rng_state::MersenneTwister

• Your type annotations are a bit too strong in general:
  

  Llama2.jl/src/math_llama.jl

  Line 32 in a4c3241

  function softmax(x::Vector{Float32})::Vector{Float32}

  • The code looks perfectly valid for any AbstractArray, not just Vector{Float32}
  • Return types generally don't need to be annotated, since they force Julia to cast outputs, which can hide bugs and slow down code
  • During development, I sometimes also like to use very restrictive types to catch initial bugs, but types have to be relaxed after a while. To give you an example why: your softmax function isn't compatible with ForwardDiff.jl's forward mode AD, since that uses a custom dual number type, which you prohibit.

Tests

• Your test suite looks great!

Minimum Code Example/ Small Examples of Components

Could you provide more code examples for individual components? Especially with ML libraries they get really complex.

Build tokenizer & read vocab

lines 367-409 in llama2.c

Build neural net blocks

• lines 182-215 in llama2.c: rmsnorm, softmax
• lines 338-345 in llama2.c: SwiGLU

Documentation for functions

A parameter list for the individual functions would be helpful.
F.E.
rmsnorm
x [point in space]
weight [weight for multiplication]

Real dummy example, but with a more complex function that does not do some obvious math it might be helpful.

Add clarifications to function documentations

Something similar to doxygen @params (which describes the function input arguments) and @return (which describes the return value of a function) could be added to functions to clarify the usage.

Tokenizer decode

lines 418-429 in llama2.c
Depends on #4

weights download

document that it's necessary to download the weights or better yet do that automatically when installing the package.

Generation loop

lines 729-783
Depends on #7 and #8 and #9 and #11

Sampler structs, argmax, mult

lines 577-614

Math Methods - Testing

Maybe test small valued vectors on math functions. Especially with multiplication.
Maybe even try to brute force bad behaviour.

Setup structs

• lines 19-75 in llama2.c
• lines 77-96, 111-140 should be implemented as external constructors

Automatic weight download

I saw that you have weight download functionalities in your test.
https://github.com/kleincode/Llama2.jl/blob/main/test/utils.jl
Would be nice to have this feature also as part of your main package. Would be much easier to get started here:
https://kleincode.github.io/Llama2.jl/dev/example/#Getting-Started

documentation

add documentation how set up and use the package
also a example would be nice

Tokenizer encode

lines 445-571
Depends on #4

Transformer forward function

lines 231-362 in llama2.c

Depends on #3 and #6

Keyword/Optional Arguments for Predict

I noticed that you do not use keyword arguments (https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments) for e.g. the predict function. It would be nice if you could run e.g.:

generate(transformer, tokenizer, sampler, testprompt, max_steps=10)
instead of how it is currently:
generate(transformer, tokenizer, sampler, testprompt, false, true, false, 10)

PS: Also as discussed, max_steps is not working currently...

Put sampler together

lines 667-714
Depends on #2 and #10

Chat loop

lines 802-884
Depends on #7 and #8 and #9 and #11

Sampler top-p

lines 616-665
Depends on #2