Transformer

A polished PyTorch implementation of the current State-Of-The-Art(SOTA) Transformer. Designed for clarity, reproducibility, and interoperability with HuggingFace Transformers, this repository provides a robust baseline for Research and Engineering being Fully Configurable. The codebase emphasizes readable and well-documented components so you can iterate on Feed-Forward, Attention and Normalization blocks and other architectural variants with minimal friction.

Features

• Fully Configurable architecture (layers, heads, model dimensions, dropout, etc.)
• HuggingFace-compatible API alignment with past_key_values support for efficient generation
• KV-Cache support for fast incremental decoding
• Encoder-Decoder architecture support with cross-attention
• Vision Transformer (ViT) support for image classification and feature extraction
• Multiple Attention Variants: MHA, GQA (Grouped Query Attention), CrossAttention
• Flexible Position Encodings: RoPE, PartialRoPE, ALiBi
• LoRA Integration for parameter-efficient fine-tuning
• Flash Attention support for accelerated training and inference
• Compact and easily extensible design for rapid prototyping and research experiments
• Clear, well-documented modules to facilitate experimentation with attention, FFNs, etc.

Download the code

git clone --depth=1 https://github.com/lof310/transformer
cd transformer

Installation

# Install dependencies
pip install -r requirements.txt
# Install on developer mode (Recommended)
pip install -e .# Install Normally
pip install .

Quick Start

Decoder-Only Model

importtorchfromtransformerimportTransformer, TransformerConfig# Configure the modelconfig=TransformerConfig(
n_layers=12,
n_heads=32,
d_model=1536,
attn_qk_norm=True,
tied_weights=False,
seq_len=1024,
max_seq_len=4096,
)
# Initialize modelmodel=Transformer(config)
# Forward PassB, N=16, 1024input_ids=torch.randint(low=0, high=config.vocab_size, size=(B, N))
output=model(input_ids=input_ids, return_states=False)

Vision Transformer (ViT) for Image Processing

importtorchfromtransformerimportTransformer, TransformerConfig# Configure ViT for image classificationconfig=TransformerConfig(
n_layers=12,
n_heads=16,
d_model=1024,
vocab_size=1000, # Number of output classespatch_size=16, # Patch size (16x16 pixels per token)img_size=224, # Input image sizein_channels=3, # RGB imagesmax_seq_len=512, # Must accommodate (img_size/patch_size)^2 + 1 CLS tokenpos_encoding="RoPE",
)
model=Transformer(config)
model.eval()
# Process images: shape (batch_size, channels, height, width)images=torch.randn(4, 3, 224, 224)
withtorch.no_grad():
output=model(images=images)
logits=output.logits# Shape: (4, 197, 1000)# Use CLS token (first position) for classificationcls_logits=logits[:, 0, :] # Shape: (4, 1000)predictions=cls_logits.argmax(dim=-1)

Incremental Decoding with KV-Cache

importtorchfromtransformerimportTransformer, TransformerConfigconfig=TransformerConfig(n_layers=6, n_heads=8, d_model=512)
model=Transformer(config)
model.eval()
# Initial promptinput_ids=torch.randint(0, config.vocab_size, (1, 10))
# First forward pass (no cache)withtorch.no_grad():
output=model(input_ids, use_cache=True)
logits=output.logitspast_key_values=output.past_key_values# Cache for next step# Incremental decoding (one token at a time)next_token_id=logits[:, -1:].argmax(dim=-1)
withtorch.no_grad():
output=model(next_token_id, past_key_values=past_key_values, use_cache=True)
new_past_key_values=output.past_key_values# Updated cache

Encoder-Decoder Model

importtorchfromtransformerimportTransformer, TransformerConfig, EncoderDecoderModel# Encoder configencoder_config=TransformerConfig(
n_layers=6,
n_heads=8,
d_model=512,
pos_encoding="RoPE",
)
# Decoder config (with cross-attention)decoder_config=TransformerConfig(
n_layers=6,
n_heads=8,
d_model=512,
attn_class="GQA",
n_kv_heads=4,
pos_encoding="RoPE",
add_cross_attention=True, # Enable cross-attention
)
# Create encoder-decoder modelmodel=EncoderDecoderModel(encoder_config, decoder_config)
# Forward passencoder_input=torch.randint(0, encoder_config.vocab_size, (4, 20))
decoder_input=torch.randint(0, decoder_config.vocab_size, (4, 10))
output=model(
input_ids=decoder_input,
encoder_input_ids=encoder_input,
return_dict=True
)

Applying LoRA Adapters

fromtransformerimportapply_lora_to_model# After creating your modelmodel=Transformer(config)
# Apply LoRA to specific layers (e.g., query/key/value projections)apply_lora_to_model(model, target_modules=["qkv_proj"], lora_rank=8, lora_alpha=16)
# Now only LoRA parameters are trainableforparaminmodel.parameters():
param.requires_grad=Falseforname, paraminmodel.named_parameters():
if"lora_"inname:
param.requires_grad=True

Default Configuration

The default configuration implements the latest SOTA Transformer design.

fromtransformerimportTransformerConfigTransformerConfig(
n_layers=12,
d_model=1536,
n_heads=32,
n_kv_heads=None, # GQA Disabled (MHA by default)vocab_size=50000,
d_ff=None, # Chosen Automatically (ratio 8/3 ≈ 2.666)norm_design="pre_norm",
norm_class="rms_norm",
ffn_class="SwiGLU",
attn_class="MHA", # Options: "MHA", "GQA", "CrossAttention"block_class=None, # Uses default TransformerBlockattn_bias=False,
ffn_bias=True,
lm_head_bias=False,
attn_qk_norm=True,
attn_dropout=0.0,
tied_weights=False,
seq_len=1024,
pos_encoding="RoPE", # Options: "RoPE", "PartialRoPE", "ALiBi"rope_base=10000.0,
max_seq_len=4096,
add_cross_attention=False, # Enable for encoder-decoder
)

Architecture Overview

Attention Mechanisms

• MHA (Multi-Head Attention): Standard self-attention with equal query/key/value heads
• GQA (Grouped Query Attention): Efficient attention with fewer KV heads, sharing across query groups
• CrossAttention: Attention between decoder queries and encoder key/values for seq2seq tasks

Position Encodings

• RoPE (Rotary Position Embeddings): Rotates query/key vectors based on absolute positions
• PartialRoPE: Applies RoPE to only a subset of dimensions
• ALiBi (Attention with Linear Biases): Adds distance-based biases to attention scores

Normalization Designs

• pre_norm: Normalize before attention/FFN (recommended for deep models)
• post_norm: Normalize after attention/FFN (original Transformer)
• parallel: Apply normalization once, then both attention and FFN in parallel
• both: Normalize both before and after (not compatible with CrossAttention)

Documentation

Full documentation available at This Page

Contributing

Contributions are welcome!

License

Distributed under the Apache License 2.0. See LICENSE for more information.

Citation

If you use transformer in your research, please cite:

@software{transformer2026,
author = {Leinier Orama},
title = {transformer: PyTorch implementation of the current State-Of-The-Art(SOTA) Transformer},
year = {2026},
publisher = {GitHub},
url = {https://github.com/lof310/transformer}
}