AleGPT

AleGPT is a custom PyTorch-based transformer language model designed for efficient autoregressive text generation. It implements modern transformer architectures with rotary positional embeddings, SwiGLU feed-forward layers, RMS normalization, and dynamic caching for inference acceleration. AleGPT includes a custom tokenizer based on byte pair encoding (BPE) and supports top-k and top-p (nucleus) sampling strategies for generation.

Features

• Transformer Language Model

  • Multi-layer transformer blocks with multi-head self-attention.
  • Support for rotary positional embeddings (RoPE) for better long-context handling.
  • SwiGLU activation in feed-forward layers for improved expressivity.
  • RMS normalization for stable training.

• Efficient Inference

  • Dynamic caching mechanism to store key-value pairs for each attention layer, enabling fast autoregressive generation.
  • Optional cache offloading to CPU to save GPU memory.

• Custom Byte Pair Encoding (BPE) Tokenizer

  • Supports pre-tokenization, byte-pair merges, and special tokens.
  • Generates vocab.json, index2token.json, and merges.json for reproducibility.
  • Handles both UTF-8 and arbitrary byte sequences.

• Sampling Methods

  • Top-k sampling: Limits next token selection to the top-k highest probability tokens.
  • Top-p (nucleus) sampling: Selects smallest set of tokens whose cumulative probability exceeds a threshold p.
  • Temperature control for adjusting randomness in generation.

• Flexible Configuration

  • Customizable model parameters: d_model, num_layers, num_heads, d_ff, vocab_size, context_length.
  • Supports mixed precision via configurable dtype.
  • Optimized or standard RoPE computation.

• Training Utilities

  • Cross-entropy loss implementation compatible with autoregressive training.
  • Supports sequence caching for long-context batch training.
  • Configurable learning rate schedule, weight decay, and optimizer parameters.

---

Installation

git clone https://github.com/AlejandroParedesLT/alegpt.git
cd alegpt
pip install -r requirements.txt

Dependencies include: torch, einops, regex, and standard Python libraries.

---

Usage

Initialize Model

from alegpt import TransformerLM
import torch

model = TransformerLM(
    vocab_size=31744,
    context_length=1024,
    d_model=512,
    num_layers=12,
    num_heads=8,
    d_ff=2048,
    rope_theta=1e5,
    device=torch.device('cuda'),
    dtype=torch.float32,
)

Forward Pass

input_ids = torch.randint(0, 31744, (1, 128), device='cuda')
logits = model(input_ids)

Sampling

from alegpt.core.inference import neural_operation

next_token = neural_operation.top_p_sampling(logits[:, -1, :], p=0.9, temperature=0.8)

Training

from alegpt.train import training_together

training_together(
    model=model,
    train_path="data/train.txt",
    valid_path="data/valid.txt",
    vocab_size=31744,
    context_length=1024,
    batch_size=4,
    lr=1e-4,
    max_iters=100000,
)

---

File Structure

alegpt/
├─ core/
│  ├─ inference.py          # Dynamic caching, attention, sampling
├─ tokenization/
│  ├─ bpe.py                # Byte Pair Encoding tokenizer
├─ models/
│  ├─ transformer.py        # TransformerLM and submodules
├─ data/
│  ├─ TinyStoriesV2-GPT4-valid.txt
├─ assets/
│  ├─ vocab.json
│  ├─ index2token.json
│  ├─ merges.json
├─ train.py                 # Training script
├─ README.md

---

Model Architecture

• Input Embedding: Custom learned embeddings.

• Transformer Block:

  • RMSNorm → Multihead Self-Attention → Residual → RMSNorm → SwiGLU Feed-Forward → Residual.

• Output Projection: Linear layer mapping back to vocabulary logits.

• Positional Encoding: Rotary embeddings (RoPE) with optional optimization.

• Cache Support: Dynamic key-value caching for faster autoregressive decoding.

---

License

MIT License.