Skip to content

README.md

Latest commit

 

History

History
148 lines (107 loc) · 4 KB

README.md

File metadata and controls

148 lines (107 loc) · 4 KB
Raw

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.