LLMs in Java - part 1

Interpretability of LLMs is important in many real world situations. For instance, there is an EU law that says if a bank's model refuses a loan then the client has the right to know why. But neural nets are notoriously hard to interpret.

There are a few techniques:

Activation Patching: this is basically A/B testing a neural net where the activations from a clean prompt are transplanted to that of a corrupted prompt and we check if that fixes the corruption.

Logit lens: take the residual stream vector (that is, the vector that represents the working memory) and apply it to the vocabulary logits at each layer. This way, we can see where an output word (probably) starts to form.

Sparse Auto Encoders: the auto encoder is trained on the model's input and output, and deliberately made sparse for low values so prominent features become clear.

The neural net itself

Most of the Java code that iterates over the tokens lives in InferenceEngine.generateTokensGPULlama. In turn, this delegates to the Tornado graphs that are executed in TornadoVMMasterPlan.tornadoVMForwardExecuteLayered. Note that the positionHolder is an IntArray that has a single element.

However, the real heavy lifting is done in LlamaFP16FFNLayers.setupSingleFFNLayer and all the functions in TransformerComputeKernelsLayered it uses. Note, that it doesn't call them. Instead, it builds them up into a graph of lazily evaluated functions that are executed in the code mentioned above.

Note that values whose name ends with Cache are a distillation of all that has come before.

The functions that are called are (in this order):

reductionOneBlockWithLayer

reductionFinalNormalization

fusedQKVMatmulX

ropeRotationWithCacheCopy

processHeadsFlashAttention

processHeadsParallel

matrixVectorGenericWithResidual

reductionOneBlockWithLayer

reductionFinalNormalization

fusedRmsNormFFNGateUp

matrixVectorGenericWithResidual

The italicised functions configure the attention.

Let's go through them:

reductionOneBlockWithLayer calculates the root mean square of the input and takes its inverse. This is calculated per work group.

reductionFinalNormalization reduces these calculations of all the workgroups to a single value.

 

fusedQKVMatmulX multiplies the vector x by the vectors wq, wk, and wv, putting the results in q, k, and v respectively. 

ropeRotationWithCacheCopy here we see a standard, two-dimensional rotation on pairs of elements in q and k. The angle we rotate by is a function of both the word's position and which attention head this is for (which is just the global thread ID mod the head size). Note we're imposing a geometry on vectors q and k that doesn't exist before this point. 

processHeadsFlashAttention this calculates softmax(Q K^T) V. It uses tiling [previous post] as this is a massive matrix operation.

This is done for all the FFN layers plus similar calls in LogitsFP16Layer.setupLogitsTaskGraph for the logit layer. Here, the functions are (in this order):

reductionOneBlockWithLayer

reductionFinalNormalization

mapContextWithQuantizeLogits

matrixVectorGeneric

Note that these are all called for each of the layers. There appear to be 16 for Llama.

There are more layers but this is enough for my first post on the subject. More in Part 2.

Rope

Upon loading the model, RoPE.precomputeFreqsCis is called when just before the weights are loaded.

We do a 2-d rotation on pair elements of the vector, imposing a geometry that was not there before. The choice of a 2-d rotation is arbitrary but it's computationally the cheapest dimension to rotate as the number of elements in a rotation matrix is d² (where d is the number of dimensions we're working in).

What happens in an LLM? (part 1)

A nice overview that's detailed but not too intricate is here [blog of SteelPh0enix AKA Wojciech Olech]

Note that when using a fully trained LLM, things are conceptually much simpler because it is more or less just a feedforward network. That is, the weights are immutable. State lives outside of the ANN and is updated by the output after each token runs through the feedforward network.

Attention!

The attention mechanism of an RNN takes the encodings and for each token, augments the input vector both forwards and backwards. "The rationale behind this is to capture additional information since current inputs may have a dependence on sequence elements that came either before or after it in a sentence, or both." [1] The two vectors are then concatenated to make one long vector. "We can consider this concatenated hidden state as the annotation of the source word since it contains the information of the jth word in both directions." [1]

The self-attention mechanism for each word has three vectors: query, key and and value. It compares the query of each word to the key of the others. This process is done in parallel ("multi-head attention") using different Q/K/V weights and the results combined.

The transformer architecture has superceded RNNs.

The cat sat on the mat

Conceptually, q, represents the query (eg, the word "sat" is looking for something to in the sentence to do the sitting); k is the key where the word "cat" is saying I am a noun that can sit; and v links the verb looking for a noun and the noun looking for a verb.

This is similar to when we apply singular value decomposition (SVD) to a document/term space and create a concept space.

Basic self attention

Imagine T input vectors x⁽ⁱ⁾. 

Tokens are embedded in a space of size d.

The T output vectors of self-attention are vectors z⁽ⁱ⁾.

These vectors are calculated thus:

z^{(i) =} Σ_j=1^T α_ij x⁽ⁱ⁾

where α is a matrix of the dot products of all the x vectors with softmax applied to it (remember, softmax does not change the relative sizes of the logits but does convert them into probabilities).

Multi-headed attention is the same algorithm calculated for multiple heads that is, mutliple swimlanes of datat that represent nuances in language structure. Note that these nuances (grammar, sentiment etc) are not deliberately targeted. They are an emergent property like concept space in SVD/NLP.

Self attention with learnable parameters

Here we project each x vector onto U_k, U_v and U_q. Note that Q=xU_q etc and U_q etc are fixed for a feed forward network. That is, somebody has done the hard work of calculating them during training.

The projections onto q and v are then multiplied together and the result is put into a matrix ω_ij where i is a token and j any other token.

The matrix ω is divided by (typically) √d and softmaxed. 

That is:

Attention(Q, V, K) = softmax ( Q K^T / √d_k ) V

Tiling

Tiling is a technique when performing matrix operations on data that won't fit into memory.

"With naïve algorithm, to compute each element of the result, we gonna need to fetch S elements from both matrices. The output matrix has S² elements, therefore the total count of loaded elements is 2S³." [Penny Xu's blog]

This breaks down as 2 [vectors - one from each matrix] * S [the size of those vectors] * S² [the number of elements that are the result of all these dot products - that is, the size of the matrix].

"With 32×32 tiling, to compute each 32×32 tile of the result, we gonna need to fetch S/32 tiles from both matrices. The output size in tiles is (S/32)², the total count of loaded [elements] is 2*(S/32)³. Each 32×32 tile contains 32² elements, the total count of loaded elements is therefore (32²)*2*(S/32)³ = (2/32)*S³. Therefore, the tiling reduced global memory bandwidth by the factor of 32, which is a huge performance win." [ibid]

In other words, having broken the matrix down into a grid of size 32×32, and each block in the grid involving 2*(S/32)³ operations, the total number of operations is this first number times the second - that is, (2/32)*S³.

Note that the resulting matrix is tiled also. So, if we doing the matrix multiplication of C=AB, the total memory needed is one tile of each of A, B and C.

There's a further optimization. If we're looking for the maximum value (which is very common in neural nets where we typically employ the softmax function), we only need to store one value per tile per column/row.

[1] Machine Learning with PyTorch and SciKitLearn.

GPU Programming pt 1.

This is a nice Java implementation of various LLMs that can also defer to TornadoVM to process the maths on a GPU. Looking at the code gives a good idea of the impedance mismatch between CPU and GPU programming as the Java code covers both. Here are some miscellaneous notes.

GPU Terminology

In TornadoVM's KernelContext you'll see various references to Global, Group and Local IDs.

The global ID is the overall ID of the thread. Note that it can be virtualized - that is, it can be larger than the number of threads your hardware supports.

The group ID refers to a logical grouping of threads. Note that a warp is a physical (ie, hardware dependent) grouping of threads. Work groups are made of integer multiples of warps. Warps always have 32 threads in NVidia hardware and process in lockstep. Work groups however can execute their warps asynchronously.

If a warp hits an if/else statement, the branches are executed sequentially and you lose parallelism!

Threads in a work group can share memory. The local ID is the ID of a thread within that group.

GPU algorithms

Writing algorithms is different in the world of GPUs. This Java/TornadoVM code that is turned into GPU code:

context.localBarrier();

for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) {

    if (localId < stride) {

        localSum[localId] += localSum[localId + stride];

    }

    context.localBarrier();

}

is actually reducing an array to an int using the GPUs threads. It first with half the warf of 32 threads. Starting with 16 threads, then 8, then 4 etc each thread takes 2 elements in the array and adds them. Only half the number of threads are needed on the next iteration and so on. All the other threads are "masked", that is, not used.

Mapping TornadoVM to GPU concepts

There is a line in fusedQKVMatmulX that basically says:

if (context.globalId < input_array_size) ...

Yeah, but what if the maximum globalId (the actual thread ID) is much lower than the array size? Do we ignore the rest of the array? 

The answer is no because globalId is a virtual ID and does not represent the physical limits of your hardware. As it happens, my RTX 4070 (Laptop) has 4608 CUDA cores whereas the model I am running (Llama-3.2-1B-Instruct-F16) has a hidden size of 4096 so it seems like it all fits into memory without resorting to virtualization tricks.

The functions above don't generally have loops in them. The reason is that the loop is implicit. Each GPU thread calls the function.

Graal- and TornadoVM

Note that TornadoVM heavily relies on GraalVM. If you look at the stack, you'll see the code in PTXCompiler.emitFrontEnd appears to exploit GraalVM's ability to examine the bytecode of the functions mentioned above. It does this so it can convert them into CUDA code.

Consequently, you'll never see any breakpoints hit in these TransformerComputeKernelsLayered functions.

BERT

I've been applying HuggingFace's ModernBert to a sequence of medical codes in the hope we can identify people who may have type 2 diabetes before even their doctors. The results have been pretty good and met expectations (the same as an American study; this is UK data). 

But what I'm surprised about is that a very simple approach performed as well as the more complicated one recommended by our American-based colleagues. Their recommendation was to use transfer learning on their model. This simply didn't work as their model used a medical encoding system called ICD10 when the British more commonly use the richer SNOMED. The theory was with fine tuning, the model would learn. It didn't.

As an aside, mapping one encoding to another is a non-trivial task. Avoid it if you can.

Instead, we pretrained BERT ourselves exclusively on our data using a Masked Language Model. This is where random words are omitted from a sequence and the model is asked to guess them. Having built the model, we then finetune it with the medical code sequence for those with and without T1D seeing if BERT can classify them.

Note that in this case, there is no pretrained BERT. We build it entirely from just our data. A sequence of medical codes is a very simple 'language' so you don't need massive compute resources to train a model. A single Tesla V100 with 16gb of VRAM was sufficient.

Although this achieved good results, it took literally days. The surprising things was that when we (perhaps naively) treated the sequence of codes as a sentence and used a plain BertForSequenceClassification object to train the model in just one step (no pretraining and finetuning) we got the same results.

Anyway, here are some notes I made along the way:

The neural net

Each of the heads defined by num_attention_heads focusses on different aspects of the text. More can better represent the data. Too many can cause overfitting. Its default is 12.

num_hidden_layers impacts accuracy and the model's ability to generalize. Its default is 12. Recommended values depend on use case. Less than 6 is good for resource-constrained environments; 6-12 for fine-tuning; 12-24 for large pretraining.

ChatGPT suggests for BERT classification, final loss should be 0.2-0.4. "BERT can overfit fast, especially with small datasets."

per_device_train_batch_size is purely a runtime optimization.

call model.eval() and put all access to the model in a with torch.no_grad() block.

Note that BERT does not evaluate model improvement by being told the labels. It just uses the labels to calculate the validation loss.

Pretraining

It's important to note that when pretraining BERT for classification, it's not told what the answers are. Instead, it's fed token sequences where some are masked and it tries to guess what the missing tokens can be. This is the accuracy metric it spits out every so often.

Finetuning

When BERT is training, you'll see Training and Validation loss. Of the two, validation is the most important. Training loss is telling you how well the model deals with data it has already seen. The model is not trained on the evaluation set so Validation indicates how well it will fare with data it has not been trained on.  Consequently, a falling training loss and a constant validation loss indicates that the model is overfitting. This can be addressed with dropout and/or weight_decay.

The conventional way of training BERT is to use BertForMaskedLM then use BertForSequenceClassification for fine tuning. The first takes sentences from the training data, randomly blanks out some words and guesses what they might be. However, we found that for small data sets that are not natural language (sequences of medical codes), we could get good performance by just using BertForSequenceClassification which does all the back propagation in a fraction of the time.

When calling BertForSequenceClassification.from_pretrained you might see:

Some weights of BertForSequenceClassification were not initialized from the model checkpoint at ... and are newly initialized: ['bert.pooler.dense.bias', 'bert.pooler.dense.weight', 'classifier.bias', 'classifier.weight']

You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

Apparently, this is OK. If you've been given a model that you wish to then use for classification. The classifier.* is the last layer of the neural net that takes a model that's been trained (possibly by masking but not classification) to which we will add a classification layer.

The special CLS token represents a summary of the string that's being encoded. The bert.pooler.* is a single, final, dense layer applied to this token that's used in classification (is not used at all in masking).

The Tokenizer(s)

Typically, you'll see tokenizers from both the Tokenizers and Tranformers libraries. They're both from HuggingFace so it's not surprising that you can tokenize with one and save it and then load the data with a tokenizer from the other.

Note that the tokenizers do not save the words as vectors. This embedding is done in the neural net. Instead, they offer different approaches  - for instance, balancing a vocabulary that can cover words it's not seen yet with having a vocab that's too big and makes the training inefficient. 

In the Transfomers library, you'll see some tokenizer names appended with Fast. These really are much faster to run as they're Rust implementations.

Almost as an aside, we tried to build our own tokenizer vocabulary by hand as we wanted to stuck as closely to the methodology used by our American colleagues as we could. This is perfectly possible (just point tokenizers.models.WordPiece.from_file at your vocab.txt file) but we found when we prepended all the codes with their classification system in both vocab.txt and the actual data we were tokenizing, accuracy was an unrealistically high 0.999 from the get-go. Why this was the case remains a mystery. Needless to say, we backed out that change. [Addendum: it's still a mystery why but using just the 30k most common whole words as tokens rather than >100k previously resulted in good accuracy].

PyTorch

Backing the Transformer's implementation of BERT if PyTorch. Browsing through the code you'll see references to unsqueeze. This is an interesting fella that can best be described by this StackOverflow answer. Basically, it can take a 2-d matrix and turn it into a 3-d tensor with the argument dictating whether the elements are in the x-, y- or z-plane. From a programming point of view, you'll see that the matrice's elements become nested one layer deeper in more square brackets, the exact configuration of those brackets depend upon which plane the elements have been unsqueezed into.

This diagram from the SO answer is great:

.

What is a (local) LLM?

The short answer: just a bunch of matrices and config files.

The long answer: if you download an LLM from, say, HuggingFace, you'll see a directory structure a little like this:

$ tree ~/.cache/huggingface/hub/models--mlabonne--TwinLlama-3.1-8B/

├── refs

│   └── main

└── snapshots

    └── 5b8c30c51e0641d36a2f9007c88425c93db0cb07

        ├── config.json

        ├── generation_config.json

        ├── model-00001-of-00004.safetensors

        ├── model-00002-of-00004.safetensors

        ├── model-00003-of-00004.safetensors

        ├── model-00004-of-00004.safetensors

        ├── model.safetensors.index.json

        ├── special_tokens_map.json

        ├── tokenizer_config.json

        └── tokenizer.json

[slightly modified for clarity].

You can see the actual matrices with something like:

from transformers import AutoModel, AutoTokenizer

import torch

model = AutoModel.from_pretrained("mlabonne/TwinLlama-3.1-8B")

embedding_layer = model.get_input_embeddings()  # This contains token ID → vector mappings

embedding_weights = embedding_layer.weight  # Shape: (vocab_size, embedding_dim)

print(embedding_weights.size())

Output:

torch.Size([128256, 4096])

But if you want to get even more raw, you can just read the files directly:

from safetensors.torch import load_file

state_dict = load_file("model-00001-of-00004.safetensors")

embeddings = state_dict['model.embed_tokens.weight']

print(embeddings.size())

Output:

torch.Size([128256, 4096])

That is the size of the vocabulary and we can demonstrate this with:

$ grep -rl 128256 ~/.cache/huggingface/hub/models--mlabonne--TwinLlama-3.1-8B/

...

config.json:  "vocab_size": 128256 

You can see what the actual tokens (note: not words) are in tokenizer.json.

File formats

Safe Tensors are a HuggingFace format for storing tensors. The "safe" properties come from the possible security implications of PyTorch's .pt and .bin formats that allow the pickling of code in the files.

GGUF (Giant GGML Unified Format) is a format used by Llama (and compatible) models. The format stores everything (both weights and metadata) in one file, is good for CPU processing and allows quantization.

You can convert from Safe Tensors to GGUF by running something like this from the llama.cpp repo:

python  convert_hf_to_gguf.py --outtype f16  --outfile /home/henryp/Temp/models--mlabonne--TwinLlama-3.1-8B-DPO.gguf /home/henryp/.cache/huggingface/hub/models--mlabonne--TwinLlama-3.1-8B-DPO/snapshots/4a76d14414118c00fbbaed96cf1dd313a10f1409/

That directory being where the model's config.json lives.

Tensor files store their weights in various formats. FP16 and BF16 are both 16-bit encodings with the difference being which bits they allocate to the mantissa and exponent. BF16 has a larger range (same as FP32) but lower precision. FP16 is the traditional format for GPUs with more modern GPUs now also supporting BF16. 

You can even just use integers like INT8 and Q4_0, the latter being a 4-bit integer with the _0 refering to which technique was used to quantize it (eg, scaling). These different strategies trade off accuracy, speed, complexity etc.

The model

There's a nice Java implementation of Llama.cpp here. It's very nicely written and even allows you to compile the Java code to a native binary (use the latest GraalVM). This class represents the weights that are loaded from disk. In the case of a LLama model, the GGUF file is read from disk, parsed and just becomes an object this class.

Hands on LLM training

Here are a collection of tips on tuning LLMs. A lot of these comments are taken from the Unsloth discord server.

What is Unsloth

Unsloth is a Python library that allows you to fine tune models without using silly amounts of compute. It does this by using smaller numeric precision.

"Almost everyone here is doing load_in_4bit. You load the model as a (dynamically) quantized 4bit version, then create an (uninitialized) bf16 QLoRA train said QLoRA merge the QLoRA onto the quantized model, generating an fp16 or bf16 model (up-quantized). [Optionally,] quantize it back down to your desired level" [Discord]

"LORAs are always fp16" [Discord]

AWQ is Activation-aware Weight Quantization "preserves a small fraction of the weights that are important for LLM performance to compress a model to 4-bits" HuggingFace.

Models

What size model can I fit into my GPU's VRAM?

A "safe bet would be to divide by 1.5 - 2x your total available memory (ram + vram) to get the B parameters vram / 1.5 or 2 to get the active params, so for example, I have 8gb vram i can load a 8b 4bit quantized model directly onto my gpu because of its relatively small approx 4-5gb vram footprint." [Discord]

This sentiment seems popular at the moment:

Mistral Small 2501

Hyperparameters

You want a rank "that is large enough to imitate the writing style and copy the knowledge from our instruction samples. You can increase this value to 64 or 128 if your results are underwhelming.

"The importance of the reference model is directly controlled via a beta parameter between 0 and 1. The reference model is ignored when beta is equal to 0, which means that the trained model can be very different from the SFT one. In practice, a value of 0.1 is the most popular one, but this can be tweaked" 

LLM Engineering Handbook

In the event that you're not "able to achieve consistently accurate answers to my questions... the model's responses are sometimes correct and sometimes incorrect, even for questions directly from the fine-tuning dataset", the advice is:

"You may want to up the rank and the alpha

Rank = how many weights it effects. Alpha is how strong they are effected. PEFT only effects the last few layers" [Discord]

(PEFT is Parameter Efficient Fine Tuning. LORA is just one of these methodologies.)

"Alpha should at least equal to the rank number, and rank should be bigger for smaller models/more complex datasets; it usually is between 4 and 64." [Discord]

Batchsize takes up a lot of VRAM so if you are having OOMs, choose smaller batches. This will mean training takes longer. To counter this, you can increase the of the gradient accumulation. This in effect batches the batches and writes back deltas to the matrix less often.

Data

Data makes or breaks an LLM. An old hand who goes by the name MrDragonFox on Discord says: 

"Most problems are really in the data it self has to be prepped. 80% of the work is there" [Discord]

"You will need good data. ML 101: garbage in, garbage out. That is a science it it self." [Discord]

"Over 80% of the moat is in the data - once you have that - you can start with experimenting" [Discord]

It's a view echoed in LLM Engineer's Handbook: "In most use cases, creating an instruction dataset is the most difficult part of the fine-tuning process."

Cleaning the data is essential:

"Avoid abbreviations if you can. As a general rule, you shouldn't expect [a] model to understand and be able to do a task a human with the same context would be unable to accomplish" - etrotta, Discord

But quantity is not quality. In general, "if I have 30k samples (user-assistant), will finetuning a base or instruct model result in a strictly better model?"

"No. Fine tuning can worsen the model's performance if your data is bad or poorly formatted. Even if the data is good, it could cause the model to 'forget' things outside of the dataset you used to fine tune it (although that is a bit less likely with LoRA). Fine tuning will make the outputs closer to whatever you trained it on. For it to make a better model, "whatever you trained it on" must be better than what it's currently responding and follow the same structure you will use later." [Discord]

Overfitting and underfitting

What causes loss to make a staircase pattern when hitting a new epoch?"

"It's very normal, it's overfitting" [Discord]

"Underfitting is probably seen more often as a common phenomenon where a low rank model fails to generalize due to a lack of learnable params

"Generally speaking you are looking for a smooth descent of loss and a gradual convergence of around 0.5. Making data more diverse and novel compared to what the model has seen already is good to combat overfitting and reducing LR/epochs can help" [Discord]

Increasing rank and alpha should be included here.

Apparently this does not apply to RSLoRA (Rank Stabilisation LoRA) which addresses training instability and performance degradation issues that can arise in low-rank adaptation methods when fine-tuning on complex datasets..

"I think for RSLoRA, alpha is sqrt(hidden_size), not sqrt(r) as claimed in the blog post. You can have whatever LoRA size (rank) you want for GRPO, with or without RS. But note that for GRPO it's usually a good idea to have max_grad_norm=0.1 (or some other low value) because GRPO models tend to go insane if you overcook them. It's always either a=r or a=sqrt(hidden_size) if use_rslora=True"

If a fine tuned model doesn't zero shot a relatively simple question it was trained it on, you "might need more examples or deeper rank" [Discord]