---

license: mit language: - en pipeline_tag: text-generation tags: - text-generation-inference - training - from_scratch_model - cupy - numpy

Model_1

This is the first working model I made. I built it from scratch in Python using only CuPy/NumPy. I did this project to learn about transformer architecture, optimizers, training loops, and tokenizers.

1. Architecture

Total params = 6,912,000 dmodel = 256, heads = 4, layers = 6

For the model architecture, I chose something similar to LLaMA but with small changes:

input -> embedding -> [ RMSNorm -> RoPE (on K and Q) -> MHAttention -> +embedding -> RMSNorm -> FFN (ReLU based) -> +prenorm tensor -> ] x layers -> RMSNorm -> linear (projection) -> softmax -> probabilities (result) -> decoder

1.1 Embeddings

This model has a vocab length of 4256, created with a custom-made tokenizer (similar to tiktoken) with a merge number of 4000.

1.2 Attention

I used normal multi-head attention with future token masking [Vaswani et al. (2017)] for the training loop, and used it with KV-caching for inference.

1.3 FFN

For the feed-forward network, I used ReLU activation:

FFN(x) = ReLU(x @ Wu + bu) @ Wd + bd

where Wu is the upscaling weight and Wd is the downscaling weight (same as in [Vaswani et al. (2017)]).

2. Training

I trained the model on 2 T4 GPUs on the Kaggle platform, on 5M tokens of the WikiText-2 dataset.

• Total training batches = 4992
• Training batches per step = 16
• Epoch = 312 steps
• Eval batches = 32
• seq_len = 1024

I used one T4 GPU for the training loop and the other for evaluation each epoch.

I stopped training when the eval loss plateaued, with a minimum eval loss of 2.4635 and training loss of 1.9883 at step 5304, epoch 17.

Training speed was constant at ~23 minutes per epoch.

I used the Adam optimizer with beta1 = 0.9, beta2 = 0.999, eps = 1e-5, and total norm clipping on the gradients with max_norm = 1.0, along with a dynamic learning rate [Vaswani et al. (2017)]:

warm_up_steps = 2000
lr = (dmodel ** -0.5) * min((step ** -0.5), step * (warm_up_steps ** -1.5))

This yielded a smooth downward loss curve with no oscillations or gradient exploding.

I initialized the weights using the He algorithm, since it's compatible with ReLU-based networks:

weight = np.random.randn(*shape).astype(np.float32) * np.sqrt(2 / shape[-1])

Biases were set to 0, and gamma (for RMSNorm) was set to 1.

3. Inference

For inference, the maximum context window is 1024. I modified the MHAttention to include a KV-caching mechanism, which yielded a 3x improvement in speed. I used top-p sampling to pick the result token.

4. Results

After training, I tested the model with several prompts. It can generate grammatically sound English sentences and can accurately follow text patterns like bullet points, spacing, and punctuation marks (open/close). However, it couldn't learn the actual meaning behind the sentences or memorize information, as the dataset is smaller than the model's parameter count, and the small size of the model isn't enough to learn meaning.

5. Challenges

While building the model, I faced multiple challenges. The most important ones were:

• Small computation power (using just Kaggle's free GPUs), which made me choose a small dataset and a small model scale.
• Fitting the model in VRAM: this was the hardest challenge. It required freeing the GPU's memory pool after each step, using a small batch size (16), using np.float32 precision, and dynamically allocating each new tensor while deleting the old intermediate ones after use each step. This directly made the training speed slower.

6. Files

This repo contains the model's main code and final weights (without the Adam weights) in the model/ directory.

The code is split into several files. The main model file contains the forward/backward passes, the inference code, the weight initialization, and the weight update logic.

The parent folder contains the model folder, the dataset folder, and the custom-made libraries used in the model code:

• transformer_block.py — contains the transformer's components, such as MHAttention and RMSNorm, along with their gradient-calculating code and the optimizer.
• BPE_training.py — contains the code for translating input to tokens (def encoder) and tokens back to text (def decoder), as well as the BPE algorithm training code (def train).
• embeding.py — contains the code for positional encoding (class RoPE), the lookup algorithm (class lookup, which extracts the sequence from the embedding table and contains its backward operation), and the code for tokenizing/importing the tokenized data (class embeding_data).