Built GPT-2 on Custom Deep Learning Framework I built from scratch in C++

Developer Akshay Muchaklavya built Forge, a custom deep learning framework written from scratch in C++, and used it to implement GPT-2 that matches Hugging Face's transformers token-for-token under greedy decoding. The framework includes core primitives like Linear, LayerNorm, SelfAttention, optimizers, and AVX2 kernels, and supports loading real pretrained weights via safetensors. Forge is available on GitHub and requires CMake 3.20+, a C++20 compiler, OpenBLAS, and a CPU with AVX2 support.

Overview overview Installation installation Tensor tensor-documentation Linear Layer linear-layer-documentation Activations activations Loss functions loss-functions-documentation Optimizers optimizers-documentation LayerNorm layernorm Self-Attention self-attention Embeddings embeddings Tokenizer tokenizer Parameters extraction & Model Load/Save safetensors reflection-based-parameters--model-loadsave-safetensors Forge implements core deep learning abstractions and primitives, enabling you to build, train, and optimize neural network models efficiently. The framework is designed for simplicity, making it suitable for learning use cases. Most people learn deep learning by calling model.fit and trusting PyTorch got the internals right. I wanted to know why it's right - how a tensor actually sits in memory, why GEMM dominates a forward pass, why attention needs a causal mask and what breaks silently if you get a transpose wrong. So Forge is a deep learning framework built from scratch. Every primitive here - Linear, LayerNorm, SelfAttention, Optimizers, and some of the the AVX2 kernels underneath - exists because I wrote it myself and verified it against a real reference. The proof: Forge's GPT-2, loaded with real pretrained weights, matches Hugging Face's transformers token-for-token under greedy decoding. Not close - exact. That's what convinces me this actually taught me how modern AI works, not just how to produce something that looks like it does. - CMake 3.20+ - A C++20 compiler - Windows: MinGW-w64 tested with the toolchain bundled in CLion - Linux: GCC or Clang - OpenBLAS - A CPU with AVX2 support any x86-64 CPU from roughly the last 10 years Eigen, reflect-cpp and ctti are fetched automatically via CMake's FetchContent - no manual setup needed for either. - Install OpenBLAS. If you don't already have it, grab a prebuilt release from the OpenBLAS releases page https://github.com/OpenMathLib/OpenBLAS/releases and extract it somewhere, e.g. C:/Libs/OpenBLAS . - Clone the repo: git clone https://github.com/muchlakshay/Forge cd Forge - Configure and build: cmake -B build -G "MinGW Makefiles" -DCMAKE BUILD TYPE=Release -DOPENBLAS ROOT="C:/Libs/OpenBLAS" cmake --build build If OpenBLAS is somewhere other than C:/Libs/OpenBLAS , point -DOPENBLAS ROOT at wherever you extracted it. Forge builds as a static library at build/libForge.a , with headers under core/ and primitives/ in the repo itself. - Install OpenBLAS and a compiler toolchain: sudo apt install build-essential cmake libopenblas-dev - Clone the repo: git clone https://github.com/muchlakshay/Forge cd Forge - Configure and build: cmake -B build -DCMAKE BUILD TYPE=Release cmake --build build OPENBLAS ROOT defaults to /usr , which matches where apt installs it - no extra flag needed unless you built OpenBLAS from source somewhere custom. Forge builds as a static library at build/libForge.a . Forge includes gpt2 and mnist test executables that demonstrate the framework in action. They're off by default so a plain build only produces the library. To build them too: cmake -B build -DFORGE BUILD TESTS=ON -DCMAKE BUILD TYPE=Release cmake --build build This adds gpt2 / mnist Linux or gpt2.exe / mnist.exe Windows to build/ . Prebuilt version of these for windows are also available on the Releases page https://github.com/muchlakshay/Forge/releases/tag/0.1 if you'd rather skip building them yourself. Since Forge links OpenBLAS as a shared library rather than statically, any executable you link against Forge needs the OpenBLAS runtime library available at runtime, not just at link time: Windows: copy libopenblas.dll found under OPENBLAS ROOT/bin into the same folder as your built executable. Without it, the executable will fail to launch with a missing-DLL error. Linux: either install OpenBLAS system-wide sudo apt install libopenblas0 , or ensure libopenblas.so / libopenblas.so.0 is somewhere on your LD LIBRARY PATH . The tests/CMakeLists.txt in this repo already handles this automatically for the gpt2 / mnist executables via a post-build copy step on Windows - if you're linking your own executable against Forge outside that setup, you'll need to do this step yourself. | Platform | Arch | Tested OS Version | Most Thoroughly Tested | |---|---|---|---| | Windows | x64 | Windows 11 Pro version 25H2 | yes | | Linux | x64 | Ubuntu 26.04 via WSL2 | yes | Currently, Forge uses Eigen as its underlying math library for tensor operations. However, Eigen's abstraction and lazy evaluation expression templates introduce significant performance bottlenecks , especially during backpropagation. Performance Optimization Plan: - Replace Eigen with OpenBLAS for optimized linear algebra operations - Implement custom vectorized element-wise kernels for backward pass computations - These optimizations will dramatically improve training speed and throughput Multi-dimensional tensors supporting up to 4 dimensions for handling batch data, spatial dimensions, and feature channels- Optimized for deep learning use cases not a general-purpose tensor library - Core abstraction layer for all neural network computations CPU radix tree-based memory allocator for efficient memory allocation and deallocation- Optimized for the allocation patterns typical in deep learning workflows Linear Layer : Fully connected layer with learnable weights and biases Layer Normalization : Normalization across features Embedding Layer : Token/index embedding lookups Sinusoidal Encoding : Positional encodings for sequence models Multi-Head Self Attention : Core component for transformer-based architectures- BPE Tokenizer - ReLU - GELU - Sigmoid - Tanh - Softmax - Cross-Entropy Loss, with fused softmax for multi-class classification - Binary Cross-Entropy Loss, with fused softmax - Mean Squared Error MSE SGD : Standard stochastic gradient descent SGD with Momentum : Accelerated gradient-based optimization Adam : Adaptive learning rate optimizer AdamW : Adam with weight decay regularization CPU Backend with Eigen-based math, performance optimizations in progress CUDA Backend GPU acceleration for NVIDIA devices Optimized kernels OpenBLAS + custom vectorized kernels Forge is ideal for: - Learning deep learning systems implementation from first principles - Experimenting with neural network architectures The Tensor class is the core abstraction in Forge for multi-dimensional data. It encapsulates: Data storage via CPU memory allocator radix tree-based Shape and stride information for multi-dimensional indexing Automatic differentiation support for backpropagation Data type dtype and device information class Tensor { Device m device; // CPU or GPU Dtype m dtype; // float32, float64, int32, int16, etc. std::shared ptr<StorageAbstract m storage; // Actual data buffer std::vector<std::size t m shape; // Dimensions: batch, height, width, channels std::vector<std::size t m strides; // Row-major strides for efficient indexing std::size t m size; // Total element count bool m need grads; // Requires gradients during backward? DispatchKey m dispatch key; // CPU/CUDA dispatch key mutable std::shared ptr<NodeAbstract m node; // Autograd computation graph node std::shared ptr<Tensor m grads; // Accumulated gradients tensor }; include "Forge.h" using namespace Forge; // Create empty tensor Tensor t; // Create tensor with shape batch=32, seq len=128, features=64 std::vector<std::size t shape = {32, 128, 64}; Tensor embeddings shape, Dtype::float32, true, Device::CPU ; // Create batch of 10 samples with 5 features, initialized to zero Tensor zeros = Tensor::Zeros {10, 5}, true ; // Create ones matrix for bias initialization Tensor ones = Tensor::Ones {1, 64}, true ; // Create constant tensor with value 0.5 Tensor half tensor = Tensor::Constant {8, 8}, 0.5f, true ; // Create range tensor: 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 Tensor indices = Tensor::Range 0, 10, 1, true ; // Convert numpy/C array to Forge tensor float data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; Tensor from data = Tensor::FromHostPtr data, {2, 3}, true ; // Direct initializer list initialization Tensor matrix {2, 3} ; matrix = {{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}}; // Create flat tensor Tensor flat = Tensor::Range 0, 24, 1, true ; // 0, 1, 2, ..., 23 // Reshape to 2, 3, 4 - batch=2, height=3, width=4 Tensor reshaped = flat.reshape 2, 3, 4 ; // Elements remain the same, just different view // Another reshape: 2, 3, 4 - 6, 4 Tensor flattened = reshaped.reshape 6, 4 ; Tensor original = Tensor::Ones {3, 3}, true ; // Deep clone - independent copy Tensor independent copy = original.clone ; // Copy data from one tensor to another Tensor destination {3, 3} ; destination.copy original ; Tensor batch = Tensor::Zeros {32, 128, 64} ; // batch size=32 // Get first sample from batch Tensor first sample = batch 0 ; // Shape: 128, 64 // Get tenth sample Tensor tenth sample = batch 9 ; // Shape: 128, 64 Forge integrates with Eigen for efficient element-wise operations: // Create a weight matrix for linear layer Tensor weights {64, 32}, true ; // 64 input features, 32 output features weights = 0.1f; // Get Eigen map for direct element access auto w map = weights.as eigen<float, 2 ; // Set specific elements w map 0, 0 = 0.5f; // weights 0 0 w map 10, 5 = -0.3f; // weights 10 5 // Element-wise operations via Eigen w map = w map 2.0f; // Scale all weights by 2 w map = w map + 0.01f; // Add small constant regularization Tensor t = Tensor::Range 0, 60, 1 ; t = t.reshape 3, 4, 5 ; // Get shape information const auto& shape = t.shape ; // {3, 4, 5} std::cout << "Shape: " << shape 0 << ", " << shape 1 << ", " << shape 2 << " \n"; // Get memory layout const auto& strides = t.strides ; // {20, 5, 1} for row-major layout // Data type and device Dtype dtype = t.dtype ; // Dtype::float32 Device device = t.device ; // Device::CPU // Memory information std::size t total elements = t.size ; // 60 void data ptr = t.data ; // Raw pointer to buffer // Gradient requirement bool requires grad = t.need grads ; // true or false // Input features Tensor x {Tensor::Ones {32, 10}, true }; // 32 samples, 10 features Linear linear {10, 32}; // Forward pass simplified Tensor output = linear x ; //trigger a backward pass output.backward //access gradients of x std::cout<<x.gradients ; Tensor output = Linear x ; // Backward with graph retention for multiple backward passes output.backward true ; // keep graph=true // Use gradients... auto grad1 = x.gradients .clone ; // Clear and do another backward x.clear grads ; output.backward true ; // Second backward pass // Now gradients accumulated again The Linear layer is a fully connected dense layer that applies an affine transformation to the input. It's one of the fundamental building blocks in neural networks, implementing the operation: output = input @ weights^T + bias Where: input has shape batch size, ..., input features supports up to 4D tensors weights has shape input features, output features bias has shape 1, 1, 1, output features broadcasted output has shape batch size, ..., output features class Linear { Tensor m weights; // Shape: input size, output size Tensor m bias; // Shape: 1, 1, 1, output size DispatchKey m dispatch key; // CPU or GPU dispatch key std::size t m input size; // Number of input features std::size t m output size; // Number of output features bool m using bias; // Use bias or not Dtype m dtype; // Data type float32, float64, etc. Device m device; // Device CPU or GPU }; include "Forge.h" using namespace Forge; // Create a linear layer: 128 input features - 64 output features // Uses Xavier Normal initialization by default Linear linear 1 128, 64 ; // With custom data type Linear linear 2 256, 128, Initializers::he normal, Dtype::float32 ; // Without bias term Linear linear no bias 100, 50, Initializers::xavier normal, Dtype::float32, Device::CPU, false ; Linear std::size t input size, // Number of input features std::size t output size, // Number of output features Initializers initializer = Initializers::xavier normal, // Weight initialization Dtype dtype = Dtype::float32, // Data type Device device = Device::CPU, // Device CPU/GPU bool bias = true // Use bias term? ; // Xavier Normal: Good for networks with sigmoid/tanh activations Linear layer xavier 128, 64, Initializers::xavier normal ; // He Normal: Recommended for ReLU-based networks Linear layer he 128, 64, Initializers::he normal ; // Xavier Uniform: Uniform distribution variant Linear layer xu 128, 64, Initializers::xavier uniform ; // He Uniform: He initialization with uniform distribution Linear layer hu 128, 64, Initializers::he uniform ; // Create layer and input Linear fc 784, 128 ; Tensor input {32, 784} ; // batch size=32, features=784 // Forward pass using operator Tensor output = fc input ; // Returns shape: 32, 128 std::cout << "Input shape: "; for auto s : input.shape std::cout << s << " "; std::cout << "\nOutput shape: "; for auto s : output.shape std::cout << s << " "; std::cout << "\n"; js Linear layer 64, 32 ; // Get read-only references const auto& weights = layer.weights ; // Shape: 64, 32 const auto& bias = layer.bias ; // Shape: 1, 1, 1, 32 // Get mutable references for optimization auto& weights mut = layer.weights ; auto& bias mut = layer.bias ; // Inspect dimensions std::size t in size = layer.input size ; // 64 std::size t out size = layer.output size ; // 32 // Print shapes std::cout << "Weights shape: " << weights.shape 0 << " x " << weights.shape 1 << "\n"; std::cout << "Bias shape: " << bias.size << " elements\n"; Linear layer 512, 256 ; // Query layer metadata std::cout << "Layer configuration:\n"; std::cout << " Input size: " << layer.input size << "\n"; std::cout << " Output size: " << layer.output size << "\n"; std::cout << " Data type: " << static cast<int layer.dtype << "\n"; std::cout << " Device: " << static cast<int layer.device << "\n"; // Weights storage const auto& w = layer.weights ; std::cout << " Weights total elements: " << w.size << "\n"; std::cout << " Bias total elements: " << layer.bias .size << "\n"; This module provides the activation functions used inside models. Each one is a small stateless functor - construct it or use a temporary and call it like a function on a Tensor . Gradients are wired up automatically whenever the input requires them, so you never call a grads function directly; it runs as part of the normal .backward pass. js Tensor operator const Tensor& input const; Returns a new tensor of the same shape as input , with ReLU applied element-wise. Relu x = max x, 0 Forge::Relu relu; Tensor x = ...; // any shape Tensor y = relu x ; // same shape as x y.backward ; js Tensor operator const Tensor& input const; Returns a new tensor of the same shape as input , with the sigmoid logistic function applied element-wise. Sigmoid x = 1 / 1 + exp -x Forge::Sigmoid sigmoid; Tensor x = ...; // any shape Tensor y = sigmoid x ; // same shape as x, values in 0, 1 y.backward ; js Tensor operator const Tensor& input const; Returns a new tensor of the same shape as input , with the hyperbolic tangent applied element-wise. Tanh x = tanh x Forge::Tanh tanh ; Tensor x = ...; // any shape Tensor y = tanh x ; // same shape as x, values in -1, 1 y.backward ; js Tensor operator const Tensor& input const; Returns a new tensor of the same shape as input , with Leaky ReLU applied element-wise. LeakyRelu x = x if x 0 = a x otherwise a = 0.01, fixed internally Forge::LeakyRelu leaky relu; Tensor x = ...; // any shape Tensor y = leaky relu x ; // same shape as x y.backward ; The negative slope a is hardcoded at 0.01 - there's currently no constructor parameter to change it. js Tensor operator const Tensor& input const; Returns a new tensor of the same shape as input , with GELU applied element-wise, using the tanh-based approximation the same one used in GPT-style implementations . Gelu x = 0.5 x 1 + tanh k x + c x^3 k = sqrt 2 / pi = 0.7978845608028654 c = 0.044715 Forge::Gelu gelu; Tensor x = ...; // any shape Tensor y = gelu x ; // same shape as x y.backward ; This is the tanh approximation, not the exact erf-based formula 0.5 x 1 + erf x / sqrt 2 . Results are very close but not bit-identical to an exact GELU implementation. js Tensor operator const Tensor& input const; Computes softmax along the last axis of a 4D tensor, returning a tensor of the same shape. This is the same shape convention used by attention scores, e.g. batch, heads, seq len, seq len . For each row x along the last axis: Softmax x i = exp x i - max x / sum j exp x j - max x The - max x subtraction is purely for numerical stability and doesn't change the result. Forge::Softmax softmax; Tensor scores = ...; // shape batch, heads, seq len, seq len Tensor probs = softmax scores ; // same shape, normalized over the last axis probs.backward ; Softmax specifically reduces over axis index 3 the last of four dims - passing a 2D or 3D tensor won't normalize correctly. If you need softmax over a different rank/axis, reshape first. All six are stateless - there's nothing to configure or store, so Forge::Relu{} x works just as well as keeping a named instance around.and work on any tensor shape/rank, always preserving it. Relu , Sigmoid , Tanh , LeakyRelu , and Gelu are purely element-wise Softmax is the one exception, requiring a 4D input as noted above. Loss functions quantify the difference between predictions and ground truth. MSE mse; Tensor loss = mse predictions, targets ; Definition: L = 1/N Σ predictions i - targets i ² Use Cases: Regression, reconstruction tasks API: js Tensor operator const Tensor& predictions, const Tensor& targets ; CrossEntropy ce; Tensor loss = ce logits, targets ; // targets: one-hot encoded Definition: L = - 1/N Σ targets i log softmax logits i Use Cases: Multi-class classification Note: Softmax is fused internally for numerical stability. Pass logits, not probabilities. API: js Tensor operator const Tensor& predictions, const Tensor& ground truth ; BinaryCrossEntropy bce; Tensor loss = bce logits, targets ; // targets: 0 or 1 Definition: L = - 1/N Σ targets i log sigmoid logits i + 1 - targets i log 1 - sigmoid logits i Use Cases: Binary classification, multi-label classification Note: Sigmoid is fused internally for numerical stability. Pass logits, not probabilities. API: js Tensor operator const Tensor& predictions, const Tensor& ground truth ; This module provides the public-facing optimizers used to train models. Each optimizer owns the parameters it was constructed with, along with any internal state momentum/moment buffers , and updates those parameters in place each time update is called. Internally, update dispatches to a device-specific backend currently only CPU supported . Two optimizers are currently available: - stochastic gradient descent, with optional momentum Forge::SGD - adaptive moment estimation, with optional weight decay Forge::Adam explicit SGD std::vector<Parameter parameters, float lr = 0.01f, float momentum coef = 0.0f ; | Argument | Meaning | |---|---| parameters | The parameters this optimizer will update. Must all share the same device and dtype; duplicates are silently removed. | lr | Learning rate step size . Default 0.01 . | momentum coef | Momentum coefficient, must be in 0, 1 . 0.0 default = plain SGD. 0.0 = momentum SGD. | If momentum coef 0 , a zero-initialized velocity buffer is allocated internally for each parameter - no manual buffer setup needed. | Method | Meaning | |---|---| update | Applies one optimization step to all owned parameters, using their currently accumulated gradients. | clear grads | Zeroes the gradients of all owned parameters. Call after update , before the next backward pass. | setLearningRate lr | Updates the learning rate. | setMomentumCoef momentum coef | Updates the momentum coefficient must be in 0, 1 . | parameters | Read-only access to the owned parameters. | learningRate | Returns the current learning rate. | momentumCoef | Returns the current momentum coefficient. | For each parameter p with gradient g : - Without momentum momentum coef == 0 : p = p - lr g - With momentum momentum coef 0 , using internal velocity buffer V : V = momentum coef V + g p = p - lr V V accumulates an exponentially-weighted sum of past gradients, smoothing the descent direction and accelerating convergence along consistent gradient directions. Forge::SGD optimizer model.parameters , / lr= /0.01f, / momentum coef= /0.9f ; Forge::MSE loss fn; for auto& batch : batches { optimizer.clear grads ; pred = model batch auto loss = loss fn ground truth, pred ; loss.backward ; optimizer.update ; } js explicit Adam const std::vector<Parameter & parameters, float lr = 0.01f, float beta 1 = 0.9f, float beta 2 = 0.999f, float decay factor = 0.01f ; | Argument | Meaning | |---|---| parameters | The parameters this optimizer will update. Must all share the same device and dtype; duplicates are silently removed. | lr | Learning rate. Default 0.01 . | beta 1 | Exponential decay rate for the first moment estimate, must be in 0, 1 . Default 0.9 . | beta 2 | Exponential decay rate for the second moment estimate, must be in 0, 1 . Default 0.999 . | decay factor | Weight-decay coefficient, applied only to parameters whose need decay flag is set. Default 0.01 . | Zero-initialized first- and second-moment buffers are allocated internally for each parameter. The step counter used for bias correction epoch is managed internally, starting at 1 and incrementing automatically on every update call - no need to track it yourself. | Method | Meaning | |---|---| update | Applies one Adam step to all owned parameters, then advances the internal step counter. | clear grads | Zeroes the gradients of all owned parameters. Call after update , before the next backward pass. | reset | Resets the internal step counter back to 1 e.g. when restarting training . Does not reset the moment buffers. | setLearningRate lr | Updates the learning rate. | setBeta 1 beta 1 | Updates beta 1 must be in 0, 1 . | setBeta 2 beta 2 | Updates beta 2 must be in 0, 1 . | setDecayFactor decay rate | Updates the weight-decay coefficient must be in 0, 1 . | parameters | Read-only access to the owned parameters. | firstMoment | Read-only access to the internal first-moment buffers M . | secondMoment | Read-only access to the internal second-moment buffers V . | beta 1 / beta 2 | Current beta values. | learningRate | Current learning rate. | epoch | Current internal step counter. | decayFactor | Current weight-decay coefficient. | For each parameter p with gradient g , at step t = epoch : M = beta 1 M + 1 - beta 1 g first moment V = beta 2 V + 1 - beta 2 g^2 second moment M hat = M / 1 - beta 1^t bias-corrected first moment V hat = V / 1 - beta 2^t bias-corrected second moment p = p - lr M hat / sqrt V hat + e e = 1e-8, fixed internally if need decay: p = p - lr decay factor p weight decay, applied after the Adam step Forge::Adam optimizer model.parameters , / lr= /0.001f, / beta 1= /0.9f, / beta 2= /0.999f, / decay factor= /0.01f ; Forge::MSE loss fn; for auto& batch : dataloader { optimizer.clear grads ; pred = model batch auto loss = loss fn ground truth, pred ; loss.backward ; optimizer.update ; } - Construction validates inputs. Constructing either optimizer with an empty parameter list, or with parameters spread across mixed devices/dtypes, throws std::invalid argument . Duplicate parameters are detected and removed automatically, no need to de-duplicate yourself. - to momentum coef , beta 1 , beta 2 , and decay factor are all range-checked 0, 1 , both at construction and whenever set via their setters. - Per-parameter weight decay: decay is opt-in via each Parameter 's need decay flag, not a global switch - set this when building your model if you want decay to skip certain parameters e.g. biases/LayerNorm scales . - - it starts at Adam::epoch is managed for you 1 and increments on every update call. Use reset if you need to restart bias correction from scratch e.g. after loading a fresh set of parameters into an existing optimizer . - - remember to call it once per step typically right before clear grads is separate from update loss.backward , or gradients will keep accumulating across steps. Forge::LayerNorm implements layer normalization over the last dimension of a 3D input tensor batch, seq len, d model , with learnable per-feature scale gamma and shift beta parameters. Like the optimizers, the actual math runs through a device-specific backend CPU, GPU, ... resolved internally - callers just construct a LayerNorm and call it like a function. js LayerNorm std::size t d model, Dtype dtype, const Device& device, bool need grads = true ; | Argument | Meaning | |---|---| d model | Size of the last dimension to normalize over the feature dimension . | dtype | Data type for gamma , beta , and the computation. | device | Device the parameters live on e.g. CPU, GPU . | need grads | Whether gamma / beta are trainable accumulate gradients . Default true . | On construction, gamma is initialized to ones and beta to zeros, both with shape d model . js Tensor operator const Tensor& input ; Applies layer normalization to input and returns a new tensor of the same shape. input must be a 3D tensor shaped batch, seq len, d model , where the last dimension matches the d model this LayerNorm was constructed with - otherwise it throws std::invalid argument . input must reside on the same device as the LayerNorm instance - otherwise it throws std::invalid argument .- If input , gamma , or beta require gradients, the necessary autograd node is attached automatically - no manual backward wiring needed. | Method | Meaning | |---|---| gamma | Mutable reference to the learnable scale parameter, shape d model . | beta | Mutable reference to the learnable shift parameter, shape d model . | d model | Returns the configured feature dimension size. | d device | Returns the device this LayerNorm operates on. | dtype | Returns the configured data type. | For each row x along the last dimension i.e. each batch, seq len position, a vector of length d model : mean = mean x average over d model var = mean x - mean ^2 variance over d model x norm = x - mean / sqrt var + eps eps = 1e-5, fixed internally y = gamma x norm + beta gamma, beta broadcast over batch, seq len gamma and beta are applied element-wise per feature broadcast across batch and seq len , so every position in the sequence is rescaled/shifted the same way along the feature axis. Forge::LayerNorm ln / d model= /512, Dtype::Float32, Device::CPU ; Tensor x = ...; // shape batch, seq len, 512 Tensor y = ln x ; // normalized output, same shape batch, seq len, 512 y.backward ; // gradients flow back into x, ln.gamma , and ln.beta automatically // gamma/beta are ordinary learnable parameters - hand them to an optimizer like any other: Forge::Adam optimizer ln.parameters , / lr= /0.001f ; Input must be 3D. LayerNorm expects batch, seq len, d model - reshape lower- or higher-rank tensors to this layout before calling it. Normalization is over the last axis only. Each batch, seq len row is normalized independently across its d model features; there is no cross-sequence or cross-batch normalization.and is not currently configurable from the constructor. eps is fixed at 1e-5 Even if you pass need grads only controls gamma / beta . need grads = false , the output tensor will still require gradients if the input you call it with requires gradients - gradient tracking is need grads OR input.need grads . Device/shape mismatches throw immediately std::invalid argument rather than silently broadcasting or casting - construct one LayerNorm per device/ d model combination you need. Forge::SelfAttention implements multi-head scaled dot-product self-attention with an optional causal mask, followed by a linear output projection. As with the other modules, the actual math runs through a device-specific backend resolved internally - you only ever interact with the SelfAttention object itself. SelfAttention std::size t d model, std::size t Q K dims, std::size t V dims, std::size t heads, bool mask, Device device, Dtype dtype = Dtype::float32, Initializers initializer = Initializers::he normal ; | Argument | Meaning | |---|---| d model | Input/output feature dimension. | Q K dims | Dimension of the query/key projection, per head. | V dims | Dimension of the value projection, per head. | heads | Number of attention heads. | mask | Whether this instance uses causal masking. If true , you must call createMask seq len before the first forward pass. | device | Device the weights live on. | dtype | Data type for weights and computation. Default float32 . | initializer | Initialization scheme used for the query/key/value projection weights. Default he normal . | On construction, the query/key projection weights are allocated with shape heads, d model, Q K dims , the value projection weights with shape heads, d model, V dims , and all three are initialized using initializer . An internal Linear layer no bias projects the concatenated multi-head output, heads V dims , back down to d model . js Tensor operator const Tensor& input const; Runs the forward pass and returns a tensor of shape batch, seq len, d model . input is expected as batch, seq len, d model rank 3 throws std::invalid argument .- If this instance was constructed with mask = true , you must have called createMask seq len first - otherwise it throws std::runtime error . The mask's size must also match the input's sequence length, or it throws std::invalid argument . - Gradients are wired up automatically for input and the query/key/value weights whenever any of them require gradients - no manual backward bookkeeping needed. void createMask std::size t seq len ; Builds a seq len, seq len causal mask and stores it internally. Only valid for instances constructed with mask = true otherwise throws std::runtime error . Call this once before the first forward pass, and again whenever seq len changes. | Method | Meaning | |---|---| query | Read-only access to the query projection weights, shape heads, d model, Q K dims . | key | Read-only access to the key projection weights, shape heads, d model, Q K dims . | value | Read-only access to the value projection weights, shape heads, d model, V dims . | linear | Mutable access to the internal output-projection Linear layer its own weights are also learnable - see the Linear module's docs . | mask | Read-only access to the stored causal mask tensor. | using mask / useMask | Whether this instance was constructed with masking enabled both return the same flag . | heads | Number of attention heads. | d model | Configured model dimension. | device | Device the weights live on. | dtype | Configured data type. | dispatch key | Internal device-routing key - not generally needed by callers. | Shapes: input is batch, seq len, d model ; query W / key W are heads, d model, Q K dims ; value W is heads, d model, V dims . @ denotes matrix multiplication, batched over batch and heads . For each head h : Q = input @ query W h batch, seq len, Q K dims K = input @ key W h batch, seq len, Q K dims V = input @ value W h batch, seq len, V dims scores = Q @ K^T / sqrt Q K dims batch, seq len, seq len if using mask: scores = scores + mask mask has -10000 at future positions j i attn = softmax scores, axis=-1 batch, seq len, seq len head out = attn @ V batch, seq len, V dims The per-head outputs are concatenated along the last axis and projected back to d model : concat = concat heads head out batch, seq len, heads V dims output = concat @ W out^T batch, seq len, d model , W out = linear weights, no bias // Encoder-style self-attention, no mask Forge::SelfAttention attn / d model= /512, / Q K dims= /64, / V dims= /64, / heads= /8, / mask= /false, Device::CPU ; Tensor x = ...; // shape batch, seq len, 512 Tensor y = attn x ; // shape batch, seq len, 512 y.backward ; // Decoder-style self-attention, causal mask Forge::SelfAttention attn / d model= /512, / Q K dims= /64, / V dims= /64, / heads= /8, / mask= /true, Device::CPU ; attn.createMask seq len ; // required before the first forward pass Tensor y = attn x ; y.backward ; // query/key/value projection weights, plus the internal Linear's own weights, // are all ordinary learnable parameters - register them with your optimizer Forge::Adam optimizer attn.parameters , / lr= /0.0001f ; Forgetting it throws createMask must be called before forward when mask = true . std::runtime error "create a mask prior to forward pass" . The mask is fixed at the sequence length you build it for. If your sequence length changes between batches, call createMask new seq len again. Heads up: the input-vs-mask length check in operator compares input.shape 0 against the mask's sequence length. For a batch, seq len, d model input, shape 0 is the batch size, not seq len - it looks like it should be comparing against shape 1 instead. As written, this check only behaves correctly when batch size == seq len , so don't rely on it to catch a real mismatch. The output projection has no bias - linear is constructed with bias disabled. Weight initialization applies only to Q/K/V. The internal Linear layer initializes itself independently see the Linear module's own docs for its default scheme . Forge::Embedding implements a learned token embedding lookup table - mapping integer token IDs to dense d model -dimensional vectors, the first layer in a typical transformer. Forge::Embedding embd d model, // embedding dimension vocab size, // number of distinct tokens Dtype::float32, // must be float32 or float64 - int32 is rejected at construction / need grads= /true, Initializers::xavier normal, Device::CPU ; Tensor out = embd seq ids ; // seq ids: 1-D int32 Tensor of token IDs - seq len, d model seq ids must be a 1-D int32 tensor - any other dtype or rank throws std::invalid argument .- The embedding table itself m embeddings must be float32 for now - constructing with Dtype::int32 throws immediately. - Dispatches through primitive dispatcher to the correct backend currently only CPU implementation EmbeddingsImplAbstract based on the output tensor's dispatch key, keeping the op backend-agnostic at the call site. - When need grads is true, a gradient node EmbeddingsGradsAbstract is automatically attached to the autograd graph after the forward pass. all embeddings exposes the underlying weight tensor directly e.g. for weight tying with an output projection . parameters returns the embedding weight wrapped as a Parameter , named embd.<n .w , for use with an optimizer. SimpleTokenizer implements a from-scratch Byte Pair Encoding BPE tokenizer with GPT-2-style pre-tokenization - the same word-splitting regex GPT-2 uses handling contractions, letter runs, digit runs, and whitespace separately before BPE merges are applied within each chunk. SimpleTokenizer tok SimpleTokenizer::Type::BPE ; tok.on file "corpus.txt", / max vocab= /8000 ; - Starts from a base vocabulary of 3 special tokens <pad , <bos , <eos plus all 256 raw byte values, then iteratively merges the most frequent adjacent token pair in the corpus until max vocab is reached or no pair occurs more than once. max vocab must be at least 300 MIN VOCAB SIZE - lower throws std::invalid argument . tok.load "vocab.bin" ; Vocabularies are stored in a simple binary format save / load - useful for loading a pretrained vocabulary e.g. reconstructed GPT-2 merge rules without retraining from a corpus. php Forge::Tensor ids = tok.encode "he went to the" ; // - 1-D int32 Tensor of token IDs StringVec tokens = tok.decode ids ; // - token strings, one per ID encode returns a Forge::Tensor int32 , ready to feed directly into Embedding::operator . decode throws std::invalid argument if it encounters a token ID outside the loaded vocabulary.- An optional transformation callback can be passed to encode to preprocess each pre-tokenized chunk e.g. lowercasing before BPE merges are applied. Forge uses reflect-cpp https://github.com/getml/reflect-cpp to automatically discover a model's trainable parameters from its struct layout, with no manual registration required - any member exposing a parameters method is picked up automatically, including nested submodules. Any aggregate struct whose members expose parameters - std::vector<Parameter e.g. Linear , Embedding , LayerNorm is automatically reflectable: struct GPT2Block { Forge::SelfAttention attn; Forge::LayerNorm ln1; Forge::Linear mlp fc; Forge::LayerNorm ln2; }; No base class, no manual parameter registration - rfl::to view introspects the struct's members directly. auto params = Forge::extract parameters model ; // flat std::vector<Parameter Forge::Adam optimizer params, / lr= /... ; Walks every member satisfying the HasParameters concept and collects their Parameter s into one flat list - this is what feeds the optimizer without you having to hand-list every layer's weights. Forge::save model, "weights.safetensors" ; Writes an actual safetensors https://github.com/huggingface/safetensors -format file: an 8-byte header length, a JSON header describing each tensor's dtype, shape, and byte offset padded to an 8-byte boundary , followed by the raw tensor data. Parameter names are built as <member name .<param name recursing into nested submodules via get state dict , matching the naming convention each layer's own parameters already establishes. Forge::load model, "weights.safetensors" ; Reads the file, matches tensor names against the model's own state dict, and copies loaded data directly into each parameter in place. Throws if the number of tensors doesn't match the model's parameter count - a basic architecture-mismatch guard. This is the mechanism used to load real GPT-2 weights directly into Forge's GPT-2 implementation. Forge::load safetensors filename is also available standalone, returning a raw std::map<std::string, Tensor without requiring a matching model struct - useful for inspecting a checkpoint's contents directly.