link#

• llm.c | C언어로 구현한 chatgpt2 오리지날 코드 | github.com/karpathy/llm.c

• rust로 구현한 예시 코드

코드 구성#

• https://github.com/YoungHaKim7/rust_gpt-2_impl/blob/main/docs/lib.rs.md

├── dev
│  └── verify_vs_c.sh
├── llm.c
├── Makefile
├── src
│  ├── bin
│  │  ├── gen_synth.rs
│  │  ├── test_gpt2.rs
│  │  └── train_gpt2.rs
│  ├── lib.rs
│  └── llmc
│     ├── dataloader.rs
│     ├── mod.rs
│     ├── rand.rs
│     ├── tokenizer.rs
│     └── utils.rs
└── tests
   └── unit.rs

lib.rs#

Excellent. This is essentially a complete GPT-2 training engine written in safe Rust. It implements the major pieces needed to train GPT-2 from scratch:

Tokens → Embeddings → Transformer blocks → Logits → Softmax → Loss → Backpropagation → AdamW update

Because this file is very large, explaining literally every line in one response would become difficult to follow. So I’ll explain it in the order the program works, while still examining the important Rust syntax and individual operations closely. We can then zoom into any function in later messages.

1. What is lib.rs?#

lib.rs is the conventional root module of a Rust library crate.

pub mod llmc;

This declares a public module named llmc. Code outside this crate can access its public items.

use rayon::prelude::*;

This imports Rayon traits needed for parallel iteration, such as:

.par_chunks_exact_mut(...)
.for_each(...)

Rayon is the Rust equivalent of much of what OpenMP provides in the original C implementation.

2. The design philosophy of this implementation#

The opening comment is extremely important:

Where the C code carves all tensors out of one big allocation with raw pointers,
this port keeps the exact same layout but tracks tensors as (start, len) views
into a single Vec

The original C implementation probably looks conceptually like this:

float *memory = malloc(total_size * sizeof(float));

float *wte = memory + offset1;
float *wpe = memory + offset2;
float *ln1w = memory + offset3;

This Rust implementation keeps the same memory layout, but avoids raw pointers.

Instead:

Vec<f32>

owns one large allocation, and each tensor is represented by:

TensorView {
    start: ...,
    len: ...,
}

So conceptually:

┌──────────────────────────────────────────────┐
│              Vec<f32>                       │
├──────────┬──────────┬──────────┬────────────┤
│   wte    │   wpe    │   ln1w   │    ...     │
└──────────┴──────────┴──────────┴────────────┘
     ▲          ▲          ▲
 TensorView  TensorView  TensorView

This is one of the most interesting Rust design decisions in the entire program.

3. Compiler attributes#

#![allow(non_snake_case)]
#![allow(clippy::too_many_arguments)]
#![allow(clippy::needless_range_loop)]

The #![...] syntax is an inner attribute, applying to the entire crate.

let batch_size = 32;

B, T, C, V

out, inp, weight, bias, B, T, C, OC

4. Tensor dimensions: (B, T, C)#

The comment says:

// B = batch_size, T = sequence_length, C = channels, V = vocab_size

These symbols appear everywhere:

A tensor with shape (B, T, C) is stored in a flat &[f32].

The element (b, t, c) is located at:

b * T * C + t * C + c

This is called row-major flattening.

5. encoder_forward: converting tokens into vectors#

pub fn encoder_forward(
    out: &mut [f32],
    inp: &[i32],
    wte: &[f32],
    wpe: &[f32],
    B: usize,
    T: usize,
    C: usize,
)

Let’s look at the types.

out: &mut [f32]

inp: &[i32]

"Hello world"
    ↓ tokenizer
[15496, 995]

wte: &[f32]

(V, C)

wpe: &[f32]

(maxT, C)

for b in 0..B {
    for t in 0..T {

let out_bt = &mut out[b * T * C + t * C..][..C];

B = 2
T = 4
C = 3

b * T * C + t * C

[..C]

out[b][t][0..C]

let ix = inp[b * T + t] as usize;

let wte_ix = &wte[ix * C..][..C];

let wpe_t = &wpe[t * C..][..C];

for i in 0..C {
    out_bt[i] = wte_ix[i] + wpe_t[i];
}

6. encoder_backward: backpropagation through addition#

The forward operation is:

out = token_embedding + position_embedding

The derivative of addition is simple.

If:

then:

and:

Therefore:

dwte_ix[i] += d;
dwpe_t[i] += d;

Notice the +=, not =.

Multiple positions can contain the same token:

"The cat and the dog"

Both occurrences of "the" contribute gradients to the same token embedding. Therefore the gradients must be accumulated.

7. Layer Normalization#

The function:

pub fn layernorm_forward(...)

implements LayerNorm for each (b, t) vector independently.

For a vector:

let mut m = 0.0f32;
for i in 0..C {
    m += x[i];
}
m /= C as f32;

let mut v = 0.0f32;
for i in 0..C {
    let xshift = x[i] - m;
    v += xshift * xshift;
}
v /= C as f32;

let eps = 1e-5f32;
let s = 1.0f32 / (v + eps).sqrt();

let n = s * (x[i] - m);
let o = n * weight[i] + bias[i];

8. Matrix multiplication and Rayon parallelism#

The naive implementation begins with:

out.par_chunks_exact_mut(OC)
    .enumerate()
    .for_each(|(bt, out_bt)| {

This deserves careful attention.

Normally, you might write:

for bt in 0..B * T {

But Rayon allows:

.par_chunks_exact_mut(OC)

to divide the output buffer into chunks of size OC.

Conceptually:

out:

┌──────────── OC ────────────┐
│ output for token position 0│
├────────────────────────────┤
│ output for token position 1│
├────────────────────────────┤
│ output for token position 2│
└────────────────────────────┘

Each chunk is a separate mutable slice. Since they cannot overlap, Rust can safely let Rayon process them concurrently.

This is a major advantage of Rust’s ownership model:

Parallel mutation is safe because each thread receives a distinct &mut [f32].

Inside:

let inp_bt = &inp[bt * C..(bt + 1) * C];

gets one input vector.

Then:

for o in 0..OC {
    let mut val = bias.map_or(0.0f32, |bias| bias[o]);

bias has the type:

Option<&[f32]>

So it may be either:

Some(bias)

or:

None

map_or means:

If there is a bias, use bias[o]; otherwise use 0.0.

Then the actual matrix multiplication is:

for i in 0..C {
    val += inp_bt[i] * weight[o * C + i];
}

Mathematically:

9. Optimized matrix multiplication#

const LOOP_UNROLL: usize = 8;

Instead of calculating one (B,T) position at a time, this implementation calculates eight positions simultaneously.

let mut result = [0.0f32; LOOP_UNROLL];

This creates a fixed-size Rust array:

[result0, result1, ..., result7]

The compiler may keep these values in CPU registers.

The important optimization is:

let w = weight[i + o * C];

for ibt in 0..LOOP_UNROLL {
    let bt = obt + ibt;
    result[ibt] += inp[bt * C + i] * w;
}

The same weight value w is loaded once and reused eight times.

This improves cache and register usage without requiring CPU-specific SIMD intrinsics.

10. Attention: the heart of GPT#

The attention input has shape:

(B, T, 3C)

The 3C consists of:

Query | Key | Value
  C       C      C

let C3 = C * 3;
let hs = C / NH;

For GPT-2 Small:

C  = 768
NH = 12
hs = 64

Each attention head therefore works with a vector of 64 numbers.

The scale:

let scale = 1.0f32 / (hs as f32).sqrt();

implements:

from the Transformer attention formula.

for t2 in 0..=t {

0, 1, 2, ..., t

let mut val = 0.0f32;
for i in 0..hs {
    val += query_t[i] * key_t2[i];
}
val *= scale;

let expv = (preatt_bth[t2] - maxval).exp();

att_bth[t2] *= expsum_inv;

att_bth[t2] = 0.0f32;

out_bth[i] += att_btht2 * value_t2[i];

11. GELU activation#

pub const GELU_SCALING_FACTOR: f32 = 0.797_884_560_802_865_4;

The comment says this is approximately:

The forward implementation:

let x = inp[i];
let cube = 0.044715f32 * x * x * x;
out[i] = 0.5f32 * x
    * (1.0f32 + (GELU_SCALING_FACTOR * (x + cube)).tanh());

This is the approximate GELU function used by GPT-2:

It is applied element-by-element.

12. Residual connections#

The entire function is:

pub fn residual_forward(out: &mut [f32], inp1: &[f32], inp2: &[f32], N: usize) {
    for i in 0..N {
        out[i] = inp1[i] + inp2[i];
    }
}

Mathematically:

This is the residual connection used throughout Transformers.

Backward is:

dinp1[i] += dout[i];
dinp2[i] += dout[i];

The gradient flows through both branches.

13. Softmax and cross-entropy#

GPT produces raw scores called logits.

For example:

"logits"

cat    2.1
dog    1.5
bird  -0.3

Softmax converts them into probabilities.

probs_bt[i] = (logits_bt[i] - maxval).exp();

Then:

probs_bt[i] /= sum;

Now:

The padded vocabulary is interesting:

V  = actual vocabulary size
Vp = padded vocabulary size

GPT-2 might have:

V  = 50257
Vp = 50304

The extra values exist for computational efficiency but don’t represent real tokens.

losses[b * T + t] = -probs_bt[ix].ln();

14. The beautiful Softmax + Cross-Entropy derivative#

This function combines both operations:

crossentropy_softmax_backward(...)

The core line is:

dlogits_bt[i] += (p - indicator) * dloss;

where:

let indicator = if i == ix { 1.0f32 } else { 0.0f32 };

This represents the famous simplification:

where y is the one-hot target vector.

This is one of the most important equations in neural network training.

15. TensorView: safe tensor pointers#

#[derive(Clone, Copy, Debug)]
pub struct TensorView {
    pub start: usize,
    pub len: usize,
}

This struct doesn’t own any memory.

It simply says:

“My tensor begins at index start and contains len floating-point numbers.”

For example:

TensorView {
    start: 1000,
    len: 768,
}

means:

memory[1000..1768]

let a = self.acts.unwrap();
let p = self.params;

16. split_disjoint: one of the key Rust techniques#

This function solves a fundamental Rust borrowing problem.

Suppose one Vec<f32> contains multiple tensors. You want:

let tensor_a = &mut memory[0..100];
let tensor_b = &mut memory[100..200];

Rust can sometimes understand this directly, but when ranges are dynamically calculated, borrowing becomes harder.

split_disjoint safely creates multiple mutable slices after verifying they don’t overlap.

Its signature uses const generics:

pub(crate) fn split_disjoint<'a, const N: usize>(
    buf: &'a mut [f32],
    ranges: [(usize, usize); N],
) -> [&'a mut [f32]; N]

N is known at compile time.

If you call it with three ranges:

split_disjoint(buffer, [range1, range2, range3])

then Rust infers:

N = 3

and returns:

[&mut [f32]; 3]

Before borrowing, it checks:

assert!(start + len <= buf.len(), "tensor view out of bounds");

Then:

assert!(start >= prev_end, "tensor views must be pairwise disjoint");

This guarantees that no two returned mutable slices overlap.

That is the crucial safety property:

One big allocation
       │
       ▼
┌─────────────────────────────┐
│           Vec<f32>          │
└─────────────────────────────┘
       │
       ▼ split_disjoint
 ┌─────┼─────┐
 ▼     ▼     ▼
&mut A &mut B &mut C

All regions are proven disjoint.

This lets the implementation retain C-like memory efficiency while remaining 100% safe Rust.

17. Parameter tensors#

pub struct ParameterTensors {
    pub wte: TensorView,
    pub wpe: TensorView,
    ...
}

This struct contains the locations of every trainable parameter.

For example:

pub qkvw: TensorView, // (L, 3*C, C)

means the Query/Key/Value projection weights for all L Transformer layers.

Importantly, ParameterTensors does not contain the weights themselves.

The actual data is here:

pub params_memory: Vec<f32>,

So:

GPT2
 ├── params_memory ───────► [all floating-point parameters]
 │
 └── params
       ├── wte ──► start + length
       ├── wpe ──► start + length
       ├── qkvw ─► start + length
       └── ...

18. Activation tensors#

The activations are intermediate results produced during the forward pass.

For example:

pub qkv: TensorView,
pub att: TensorView,
pub fch_gelu: TensorView,
pub logits: TensorView,
pub probs: TensorView,
pub losses: TensorView,

Why store all these?

Because backpropagation needs many forward-pass values.

For example, the GELU backward pass needs the original input:

gelu_backward(dl_fch, l_fch, dl_fch_gelu, ...)

Training therefore requires substantially more memory than inference.

19. The GPT2 struct: the complete model state#

pub struct GPT2 {

This represents the entire trainable model.

The most important groups are:

pub config: GPT2Config,

pub params_memory: Vec<f32>,

pub grads_memory: Option<Vec<f32>>,

pub m_memory: Option<Vec<f32>>,
pub v_memory: Option<Vec<f32>>,

pub acts_memory: Option<Vec<f32>>,

20. Loading the checkpoint#

pub fn build_from_checkpoint(checkpoint_path: &str) -> GPT2

The argument:

&str

is a borrowed string slice. The function doesn’t need to own the path.

let mut model_file = fopen_check(checkpoint_path, "rb");

Opens the checkpoint.

let model_header = read_i32s(&mut model_file, 256);

Reads 256 i32 values from the file header.

Then:

if model_header[0] != 20240326 {

checks the magic number. This verifies that the file is in the expected format.

if model_header[1] != 3 {

checks the file format version.

21. The forward pass: the complete GPT-2 pipeline#

The most important function is:

pub fn gpt2_forward(&mut self, ...)

The &mut self means this function can modify the model’s internal state, including:

• activation buffers
• cached inputs
• cached targets
• mean loss

The pipeline is:

Input Tokens
     │
     ▼
Token + Position Embeddings
     │
     ▼
┌─────────────────────────────┐
│ Transformer Layer 0         │
│ LayerNorm                   │
│ QKV Projection              │
│ Causal Multi-Head Attention │
│ Residual Connection         │
│ LayerNorm                   │
│ MLP → GELU → MLP            │
│ Residual Connection         │
└──────────────┬──────────────┘
               │
              ...
               │
     ▼
Final LayerNorm
     │
     ▼
Vocabulary Projection
     │
     ▼
Softmax
     │
     ▼
Cross-Entropy Loss

The loop:

for l in 0..L {

runs every Transformer layer.

Inside, these lines are essentially the architecture of GPT-2 written directly as code:

layernorm_forward(...)
matmul_forward(...)      // create Q, K, V
attention_forward(...)
matmul_forward(...)      // attention output projection
residual_forward(...)
layernorm_forward(...)
matmul_forward(...)      // MLP expansion: C → 4C
gelu_forward(...)
matmul_forward(...)      // MLP projection: 4C → C
residual_forward(...)

That’s the GPT-2 Transformer block.

22. Why is split_disjoint used in the forward pass?#

This code:

let [l_ln1, l_ln1_mean, ..., residual] =
    split_disjoint(acts_memory, [...]);

destructures an array of mutable slices.

It is similar in spirit to:

let [a, b, c] = array;

But here every variable becomes a separate &mut [f32].

Without a mechanism like this, Rust would reject multiple mutable references into the same acts_memory buffer because overlapping mutable references could cause undefined behavior.

split_disjoint proves that they don’t overlap.

This is an excellent example of adapting a low-level C memory layout to Rust’s ownership rules.

23. Backpropagation goes in reverse#

The backward function contains:

for l in (0..L).rev() {

This means:

L - 1
L - 2
...
1
0

Backpropagation must run in the reverse order of the forward computation because of the chain rule.

For example, the forward MLP is:

LayerNorm
   ↓
Matrix Multiplication
   ↓
GELU
   ↓
Matrix Multiplication

The backward pass is:

Matrix Multiplication Backward
   ↑
GELU Backward
   ↑
Matrix Multiplication Backward
   ↑
LayerNorm Backward

The code mirrors this exactly.

24. Starting backpropagation#

let dloss_mean = 1.0f32 / (B * T) as f32;

The final loss is the mean:

Therefore:

The code initializes every loss gradient with this value:

dlosses[i] = dloss_mean;

This is the starting point for the entire chain rule.

25. AdamW parameter update#

The final major function is:

pub fn gpt2_update(...)

For every parameter:

for i in 0..*num_parameters {

the optimizer updates two running averages.

let m = beta1 * m_memory[i]
    + (1.0f32 - beta1) * grad;

let v = beta2 * v_memory[i]
    + (1.0f32 - beta2) * grad * grad;

let m_hat = m / (1.0f32 - beta1.powf(t as f32));
let v_hat = v / (1.0f32 - beta2.powf(t as f32));

params_memory[i] -= learning_rate
    * (m_hat / (v_hat.sqrt() + eps)
       + weight_decay * param);

26. Rust ownership and memory management in this entire program#

One of the nicest things about this implementation is how naturally Rust handles cleanup.

The original C implementation likely has something conceptually like:

gpt2_free(model);

But the Rust code ends with:

// gpt2_free() has no Rust equivalent: all model memory is freed when the GPT2 is dropped

This is RAII, which connects directly to the C++ and Rust ownership discussions we’ve had.

When:

GPT2

goes out of scope, Rust automatically drops its fields:

GPT2 dropped
   │
   ├── Vec<f32> parameters freed
   ├── Vec<f32> gradients freed
   ├── Vec<f32> activations freed
   ├── AdamW buffers freed
   └── other fields dropped

No manual free() is necessary.

The most important overall idea#

I would summarize this lib.rs as having three layers of design:

LayerNorm
Matrix Multiplication
Attention
GELU
Softmax
Cross-Entropy
Backpropagation
AdamW

Vec<f32>

TensorView { start, len }
split_disjoint(...)
&[f32]
&mut [f32]
Option<Vec<f32>>