Skip to content

central documentation

Jump to bottom
GURUMURTHY edited this page Feb 24, 2026 · 1 revision

Stackformer Wiki

--  

Note

This page is the central documentation hub for the Stackformer library.  

Table of contents

 

What is Stackformer?

  Stackformer is a modular PyTorch library for Transformer-based models across NLP and vision.   It is designed for:  

  • quick experimentation,
  • reusable building blocks,
  • clean model prototyping,
  • production-friendly project structure.  

Repository structure

 

Stackformer/
├── stackformer/
│   ├── modules/
│   │   ├── Attention.py
│   │   ├── Feed_forward.py
│   │   ├── Normalization.py
│   │   └── position_embedding.py
│   ├── models/
│   │   ├── OpenAI.py
│   │   ├── Meta.py
│   │   ├── Google.py
│   │   └── Transformer.py
│   ├── vision/
│   │   ├── vit.py
│   │   └── segformer.py
│   ├── generate.py
│   └── trainer.py
├── tests/
├── docs/
└── README.md

 

Installation

  Install the library:  

pip install stackformer

  Install with optional extras:  

pip install "stackformer[train]"
pip install "stackformer[dev]"

 

Tip

For local development, use an editable install: pip install -e ".[dev]".  

Quick start

 

import torch
from stackformer.models.OpenAI import GPT_2
 
model = GPT_2(
vocab_size=128,
num_layers=2,
embed_dim=32,
num_heads=4,
seq_len=16,
dropout=0.0,
hidden_dim=64,
)
 
x = torch.randint(0, 128, (2, 16))
logits = model(x)
print(logits.shape)  # [2, 16, 128]

 

Core modules

 

stackformer.modules.Attention

 

  • Multi_Head_Attention
  • Multi_query_Attention
  • Group_query_Attention
  • RoPE and KV-cache helper utilities  

stackformer.modules.Feed_forward

 

  • FF_ReLU
  • FF_GELU
  • FF_SwiGLU  

stackformer.modules.Normalization

 

  • LayerNormalization
  • RMSNormalization  

stackformer.modules.position_embedding

 

  • absolute/learned embeddings
  • sinusoidal embeddings
  • RoPE positional encoding  

Model families

 

OpenAI-style (stackformer.models.OpenAI)

 

  • GPT_1
  • GPT_2  

Meta-style (stackformer.models.Meta)

 

  • llama_1
  • llama_2  

Google-style (stackformer.models.Google)

 

  • gemma_1_2b
  • gemma_1_7b  

Generic baseline (stackformer.models.Transformer)

 

  • transformer  

Vision models

 

  • stackformer.vision.vit.ViT
  • stackformer.vision.segformer.SegFormerB0  

Training workflow

  The standard training loop with stackformer.trainer.Trainer:  

  1. Prepare tokenized tensors/datasets.
  2. Create a model instance.
  3. Initialize Trainer with datasets and hyperparameters.
  4. Run trainer.train().
  5. Evaluate and use saved checkpoints.  

Text generation

  Use stackformer.generate for autoregressive text generation with language models.  

Development workflow

  Run tests:  

pytest -q

  Install in editable mode:  

pip install -e ".[dev]"

 

Contribution guidelines

 

  • Keep modules composable and independent.
  • Prefer clear, typed APIs.
  • Add or update tests in tests/ for behavior changes.
  • Keep default configurations practical for smoke tests.  

Project links

 

  If you are new to the project, start with:  

  1. README.md for a high-level overview.
  2. stackformer/modules/ for reusable core blocks.
  3. tests/ for usage examples and expected behavior.

Clone this wiki locally