Harish241412/llama-3.2-1b-from-scratch

Llama 3.2 1B — From Scratch

A PyTorch implementation of Llama 3.2 1B built from scratch for educational and research purposes.

The project reimplements the model architecture and inference pipeline from the ground up and loads the original Llama 3.2 1B weights into the custom implementation.

Educational implementation only.

Architechture

Model Details

Property	Value
Base model	Llama 3.2 1B
Parameters	~1.24B
Architecture	Decoder-only Transformer
Hidden size	2048
Layers	16
Attention heads	32
KV heads	8
Head dimension	64
Intermediate size	8192
Vocabulary size	128,256
Maximum context	131,072 tokens
RoPE base	500,000
Weight dtype	bfloat16

What's Implemented

• Token embeddings
• RMSNorm
• Grouped Query Attention (GQA)
• Rotary Positional Embeddings (RoPE)
• SwiGLU feed-forward network
• Transformer decoder blocks
• Weight tying
• Causal language modeling
• KV-cache inference
• Greedy decoding
• Temperature sampling
• Top-k sampling
• Top-p sampling
• Stop-token handling
• Llama 3 tokenizer
• Chat-template encoding and decoding
• BF16 checkpoint loading

KV Cache

The implementation includes a dynamic KV cache for autoregressive decoding.

Per-token decode latency was benchmarked on an RTX 3050:

Context	No KV	KV	Speedup
16	28.84 ms	20.90 ms	1.38×
32	24.29 ms	26.97 ms	0.90×
64	27.47 ms	21.52 ms	1.28×
128	32.22 ms	21.50 ms	1.50×
256	60.62 ms	21.39 ms	2.83×
512	119.99 ms	26.41 ms	4.54×
1024	268.96 ms	19.44 ms	13.83×

The benefit of KV caching becomes increasingly significant as context length grows.

Verification

The scratch implementation was verified against the Hugging Face implementation using the original Llama 3.2 1B weights.

Hugging Face vs Scratch

FP32 verification across five prompts produced:

• Identical Top-1 predictions
• Identical Top-5 predictions
• Maximum absolute logit difference below 2e-5

Example:

Prompt:
The capital of France is

Next token:
Paris

Generation

Example greedy generation:

The capital of France is Paris, and the capital of the United States is

KV-cache and non-KV generation were also verified for consistency.

Setup

The model checkpoint is designed to be used with the custom PyTorch implementation available in the GitHub repository.

Clone the repository and install the required dependencies using uv:

git clone https://github.com/N-Harish/llama3.2-1b-scratch.git
cd llama3.2-1b-scratch

uv sync

The project uses uv for environment and dependency management.

Note: This is a from-scratch implementation of Llama 3.2 1B and is not intended to be loaded directly with the standard Hugging Face Transformers API. Use the implementation provided in the GitHub repository to run the model.

Tokenizer

The project includes a custom Llama 3 tokenizer implementation using the original tokenizer vocabulary and configuration.

from llama3_2_1b_scratch.tokenizer import Llama3Tokenizer

tokenizer = Llama3Tokenizer(
    "tokenizer.model",
    "tokenizer_config.json",
)

ids = tokenizer.encode("The capital of France is")
text = tokenizer.decode(ids)

print(ids)
print(text)

Download and Run

The model weights and tokenizer files can be downloaded from this repository using the provided utility.

import torch

from llama3_2_1b_scratch.utils import download_model
from llama3_2_1b_scratch.model import LlamaForCausalLM
from llama3_2_1b_scratch.tokenizer import Llama3Tokenizer


local_dir = download_model(
    output_dir="llama3-2-1b",
    repo_id="Harish241412/llama-3.2-1b-from-scratch",
)

model = LlamaForCausalLM.from_pretrained(
    local_dir / "llama3_2_1b_scratch_bf16.safetensors",
    device="cuda" if torch.cuda.is_available() else "cpu",
    dtype=torch.bfloat16,
)

tokenizer = Llama3Tokenizer(
    local_dir / "tokenizer.model",
    local_dir / "tokenizer_config.json",
)

prompt = "The capital of France is"

input_ids = torch.tensor(
    [tokenizer.encode(prompt)],
    dtype=torch.long,
    device=next(model.parameters()).device,
)

generated_ids = model.generate(
    input_ids,
    max_new_tokens=50,
    do_sample=False,
)

print(tokenizer.decode(generated_ids[0].tolist()))

Example output

The capital of France is Paris, and the capital of the United States is

Downloaded files

The download utility saves the following files to the specified directory:

• llama3_2_1b_scratch_bf16.safetensors
• tokenizer.model
• tokenizer_config.json

from_pretrained() loads the checkpoint from the local filesystem and does not perform model downloads itself.

Sampling

The implementation supports greedy and stochastic generation.

generated_ids = model.generate(
    input_ids,
    max_new_tokens=50,
    do_sample=True,
    temperature=0.7,
    top_p=0.9,
    top_k=50,
)

For greedy decoding:

generated_ids = model.generate(
    input_ids,
    max_new_tokens=50,
    do_sample=False,
)

Checkpoint

The repository provides a BF16 checkpoint containing the Llama 3.2 1B weights mapped to the scratch implementation.

The checkpoint contains approximately 1.24B parameters.

Project Goals

This project is intended as a hands-on implementation for understanding:

• Transformer architecture
• Llama 3.2 internals
• Grouped Query Attention
• Rotary embeddings
• KV caching
• Autoregressive generation
• Sampling
• Tokenization
• LLM inference

Educational Use

This project is an educational implementation only and is intended for learning and experimentation with LLM architecture and inference.

It should not be considered an official implementation of Llama 3.2.