stackformer 0.0.1__tar.gz → 0.1.10__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {stackformer-0.0.1 → stackformer-0.1.10}/LICENSE +21 -21
- stackformer-0.1.10/PKG-INFO +347 -0
- stackformer-0.1.10/README.md +312 -0
- stackformer-0.1.10/pyproject.toml +66 -0
- {stackformer-0.0.1 → stackformer-0.1.10}/setup.cfg +4 -4
- stackformer-0.1.10/stackformer/__init__.py +167 -0
- stackformer-0.1.10/stackformer/amp/__init__.py +14 -0
- stackformer-0.1.10/stackformer/amp/scaler.py +170 -0
- stackformer-0.1.10/stackformer/cache/__init__.py +11 -0
- stackformer-0.1.10/stackformer/cache/paged.py +245 -0
- stackformer-0.1.10/stackformer/cache/static.py +144 -0
- stackformer-0.1.10/stackformer/config.py +113 -0
- stackformer-0.1.10/stackformer/distributed/__init__.py +42 -0
- stackformer-0.1.10/stackformer/distributed/ddp.py +126 -0
- stackformer-0.1.10/stackformer/engine/__init__.py +16 -0
- stackformer-0.1.10/stackformer/engine/checkpoint.py +496 -0
- stackformer-0.1.10/stackformer/engine/engine.py +317 -0
- stackformer-0.1.10/stackformer/engine/state.py +150 -0
- stackformer-0.1.10/stackformer/engine/trainer.py +473 -0
- stackformer-0.1.10/stackformer/generate.py +264 -0
- stackformer-0.1.10/stackformer/language/__init__.py +21 -0
- stackformer-0.1.10/stackformer/language/decoder.py +57 -0
- stackformer-0.1.10/stackformer/language/encoder_decoder.py +65 -0
- stackformer-0.1.10/stackformer/logging/__init__.py +33 -0
- stackformer-0.1.10/stackformer/logging/csv_logger.py +79 -0
- stackformer-0.1.10/stackformer/logging/logger.py +111 -0
- stackformer-0.1.10/stackformer/logging/metrics.py +208 -0
- stackformer-0.1.10/stackformer/logging/tensorboard_logger.py +83 -0
- stackformer-0.1.10/stackformer/logging/wandb_logger.py +103 -0
- stackformer-0.1.10/stackformer/logging/wb_logger.py +10 -0
- stackformer-0.1.10/stackformer/metrics.py +29 -0
- stackformer-0.1.10/stackformer/models/__init__.py +38 -0
- stackformer-0.1.10/stackformer/models/bert.py +128 -0
- stackformer-0.1.10/stackformer/models/gemma.py +270 -0
- stackformer-0.1.10/stackformer/models/gpt.py +280 -0
- stackformer-0.1.10/stackformer/models/llama.py +389 -0
- stackformer-0.1.10/stackformer/models/roberta.py +134 -0
- stackformer-0.1.10/stackformer/models/transformer.py +166 -0
- stackformer-0.1.10/stackformer/modules/Attention.py +1320 -0
- stackformer-0.1.10/stackformer/modules/Feed_forward.py +370 -0
- stackformer-0.1.10/stackformer/modules/Normalization.py +122 -0
- stackformer-0.1.10/stackformer/modules/__init__.py +104 -0
- stackformer-0.1.10/stackformer/modules/attention_engine.py +169 -0
- stackformer-0.1.10/stackformer/modules/layer.py +528 -0
- stackformer-0.1.10/stackformer/modules/masks.py +495 -0
- stackformer-0.1.10/stackformer/modules/position_embedding.py +230 -0
- stackformer-0.1.10/stackformer/optim/__init__.py +33 -0
- stackformer-0.1.10/stackformer/optim/factories.py +270 -0
- stackformer-0.1.10/stackformer/optim/loss_fn.py +140 -0
- stackformer-0.1.10/stackformer/training/__init__.py +14 -0
- stackformer-0.1.10/stackformer/training/loops.py +104 -0
- stackformer-0.1.10/stackformer/utils/__init__.py +45 -0
- stackformer-0.1.10/stackformer/utils/attn_utils.py +49 -0
- stackformer-0.1.10/stackformer/utils/cache.py +33 -0
- stackformer-0.1.10/stackformer/utils/device.py +100 -0
- stackformer-0.1.10/stackformer/utils/ff_utils.py +40 -0
- stackformer-0.1.10/stackformer/utils/utils.py +99 -0
- stackformer-0.1.10/stackformer/vision/__init__.py +15 -0
- stackformer-0.1.10/stackformer/vision/segformer.py +359 -0
- stackformer-0.1.10/stackformer/vision/vit.py +313 -0
- stackformer-0.1.10/stackformer.egg-info/PKG-INFO +347 -0
- stackformer-0.1.10/stackformer.egg-info/SOURCES.txt +65 -0
- stackformer-0.1.10/stackformer.egg-info/requires.txt +11 -0
- stackformer-0.1.10/stackformer.egg-info/top_level.txt +1 -0
- stackformer-0.1.10/tests/test_distributed.py +30 -0
- stackformer-0.1.10/tests/test_vision.py +133 -0
- stackformer-0.0.1/PKG-INFO +0 -75
- stackformer-0.0.1/README.md +0 -54
- stackformer-0.0.1/models/GPT_2.py +0 -181
- stackformer-0.0.1/models/__init__.py +0 -0
- stackformer-0.0.1/modules/Attention.py +0 -533
- stackformer-0.0.1/modules/Feed_forward.py +0 -59
- stackformer-0.0.1/modules/Normalization.py +0 -41
- stackformer-0.0.1/modules/__init__.py +0 -0
- stackformer-0.0.1/modules/mask.py +0 -36
- stackformer-0.0.1/modules/position_embedding.py +0 -43
- stackformer-0.0.1/modules/tokenizer.py +0 -25
- stackformer-0.0.1/pyproject.toml +0 -25
- stackformer-0.0.1/setup.py +0 -31
- stackformer-0.0.1/stackformer.egg-info/PKG-INFO +0 -75
- stackformer-0.0.1/stackformer.egg-info/SOURCES.txt +0 -18
- stackformer-0.0.1/stackformer.egg-info/requires.txt +0 -2
- stackformer-0.0.1/stackformer.egg-info/top_level.txt +0 -2
- {stackformer-0.0.1 → stackformer-0.1.10}/stackformer.egg-info/dependency_links.txt +0 -0
|
@@ -1,21 +1,21 @@
|
|
|
1
|
-
MIT License
|
|
2
|
-
|
|
3
|
-
Copyright (c) 2025
|
|
4
|
-
|
|
5
|
-
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
-
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
-
in the Software without restriction, including without limitation the rights
|
|
8
|
-
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
-
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
-
furnished to do so, subject to the following conditions:
|
|
11
|
-
|
|
12
|
-
The above copyright notice and this permission notice shall be included in all
|
|
13
|
-
copies or substantial portions of the Software.
|
|
14
|
-
|
|
15
|
-
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
-
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
-
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
-
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
-
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
-
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
-
SOFTWARE.
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Stackformer Labs
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,347 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: stackformer
|
|
3
|
+
Version: 0.1.10
|
|
4
|
+
Summary: Modular transformer blocks built in PyTorch
|
|
5
|
+
Author-email: Stackformer Labs <stackformer.dev@gmail.com>
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Repository, https://github.com/stackformer-labs/Stackformer
|
|
8
|
+
Project-URL: Releases, https://github.com/stackformer-labs/Stackformer/releases
|
|
9
|
+
Project-URL: Issue-Tracker, https://github.com/stackformer-labs/Stackformer/issues
|
|
10
|
+
Project-URL: Discussions, https://github.com/stackformer-labs/Stackformer/discussions
|
|
11
|
+
Project-URL: Documentation, https://github.com/stackformer-labs/Stackformer/tree/main/docs
|
|
12
|
+
Keywords: transformer,pytorch,deep-learning,attention,llm,machine-learning
|
|
13
|
+
Classifier: Programming Language :: Python :: 3
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
17
|
+
Classifier: Typing :: Typed
|
|
18
|
+
Classifier: Operating System :: OS Independent
|
|
19
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
20
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
21
|
+
Requires-Python: >=3.10
|
|
22
|
+
Description-Content-Type: text/markdown
|
|
23
|
+
License-File: LICENSE
|
|
24
|
+
Requires-Dist: torch<3.0,>=2.3
|
|
25
|
+
Requires-Dist: numpy<3.0,>=1.23
|
|
26
|
+
Requires-Dist: tqdm<5.0,>=4.5
|
|
27
|
+
Requires-Dist: tensorboard<3.0,>=2.14
|
|
28
|
+
Requires-Dist: wandb<1.0,>=0.17
|
|
29
|
+
Requires-Dist: safetensors<1.0,>=0.4
|
|
30
|
+
Provides-Extra: dev
|
|
31
|
+
Requires-Dist: pytest<9.0,>=8.0; extra == "dev"
|
|
32
|
+
Requires-Dist: pytest-cov<6.0,>=5.0; extra == "dev"
|
|
33
|
+
Requires-Dist: build<2.0,>=1.2; extra == "dev"
|
|
34
|
+
Dynamic: license-file
|
|
35
|
+
|
|
36
|
+
<p align="center">
|
|
37
|
+
<img src="assets/logo.png" alt="StackFormer logo" width="560" />
|
|
38
|
+
</p>
|
|
39
|
+
|
|
40
|
+
<p align="center">
|
|
41
|
+
<a href="https://pypi.org/project/stackformer/"><img src="https://img.shields.io/pypi/v/stackformer.svg" alt="PyPI version" /></a>
|
|
42
|
+
<a href="https://pypi.org/project/stackformer/"><img src="https://img.shields.io/pypi/pyversions/stackformer.svg" alt="Python versions" /></a>
|
|
43
|
+
<a href="LICENSE"><img src="https://img.shields.io/badge/license-MIT-blue.svg" alt="License" /></a>
|
|
44
|
+
<a href="https://github.com/stackformer-labs/Stackformer/actions">
|
|
45
|
+
<img src="https://img.shields.io/github/actions/workflow/status/stackformer-labs/Stackformer/core-tests.yml?branch=main&label=CI" alt="CI status" />
|
|
46
|
+
</a>
|
|
47
|
+
</p>
|
|
48
|
+
|
|
49
|
+
# StackFormer
|
|
50
|
+
|
|
51
|
+
**Composable PyTorch Transformer building blocks, modular architecture zoo, and lightweight training engine in a single clean library.**
|
|
52
|
+
|
|
53
|
+
---
|
|
54
|
+
|
|
55
|
+
## Why StackFormer?
|
|
56
|
+
|
|
57
|
+
Building custom Transformer architectures today usually forces a frustrating choice: either fork thousands of lines of monolithic, copy-pasted model code from single-file implementations, or navigate bloated multi-layer abstraction frameworks designed exclusively for pretrained weight conversion.
|
|
58
|
+
|
|
59
|
+
**StackFormer** provides a third way: modular PyTorch building blocks that let researchers and engineers assemble custom vision and language Transformers declaratively via `BlockConfig`. Instead of rewriting multi-head attention, rotary positional embeddings, SwiGLU feed-forward networks, or pre-norm connections for every experiment, StackFormer exposes orthogonal layer primitives that compose cleanly and execute on pure PyTorch runtime.
|
|
60
|
+
|
|
61
|
+
---
|
|
62
|
+
|
|
63
|
+
## Key Features
|
|
64
|
+
|
|
65
|
+
- **Declarative `BlockConfig` Layer Composability:** Arbitrarily combine 10 attention implementations (`mha`, `gqa`, `mqa`, `cross_mha`, and RoPE variants), 7 feed-forward activation layers (`gelu`, `swiglu`, `geglu`, `leaky_relu`, `relu`, `sigmoid`, `silu`), 2 normalization types (`layernorm`, `rmsnorm`), and 3 positional embedding strategies (`absolute`, `sinusoidal`, `rope`).
|
|
66
|
+
- **Comprehensive Architecture Zoo (Language & Vision):** Production-faithful PyTorch implementations of **GPT-1, GPT-2, LLaMA-1, LLaMA-2, Gemma-1 (2B & 7B), BERT, RoBERTa, Vaswani Encoder-Decoder Transformer, ViT, and SegFormer-B0**.
|
|
67
|
+
- **Lightweight Engine & Trainer:** Built-in `Trainer` featuring Automatic Mixed Precision (`AMPScaler` for FP16/BF16), Distributed Data Parallel (`DDP`), gradient accumulation, max gradient clipping, warmup/decay learning rate schedulers, and zero-dependency SafeTensors checkpointing (`CheckpointManager`).
|
|
68
|
+
- **Stateful KV-Cache Autoregressive Generation:** Standardized `prefill()` and `decode()` model contract supporting fast KV-cache accelerated text generation (`text_generate()`) with temperature, top-k, and top-p (nucleus) sampling strategies.
|
|
69
|
+
- **Zero-Bloat Foundation:** Typed codebase (`py.typed`), 136-test suite across unit and integration targets, with zero heavy required dependencies beyond PyTorch, NumPy, SafeTensors, and tqdm.
|
|
70
|
+
|
|
71
|
+
---
|
|
72
|
+
|
|
73
|
+
## Framework Comparison
|
|
74
|
+
|
|
75
|
+
| Dimension | StackFormer | Hugging Face Transformers | nanoGPT | x-transformers |
|
|
76
|
+
| :--- | :--- | :--- | :--- | :--- |
|
|
77
|
+
| **Primary Focus** | Modular custom block composition & native training engine | Pretrained weight distribution & pipeline inference | Minimal teaching codebase for GPT-2 | Experimental attention module collection |
|
|
78
|
+
| **Architecture Scope** | Language + Vision (GPT, LLaMA, Gemma, BERT, RoBERTa, ViT, SegFormer) | Comprehensive pretrained repository | Decoder-only GPT models | Modular Transformer blocks |
|
|
79
|
+
| **Block Composability** | Declarative `BlockConfig` (swap attention, FFN, norm, pos-emb seamlessly) | Monolithic per-model source files | Single-file script | PyTorch module blocks |
|
|
80
|
+
| **Training Infrastructure** | Built-in `Trainer` (AMP, DDP, SafeTensors, gradient accum, schedulers) | Requires `Trainer` or Accelerate | Custom training loop in `train.py` | External (user writes loop) |
|
|
81
|
+
| **Checkpoint Format** | SafeTensors (`.safetensors`) + JSON metadata | SafeTensors / PyTorch bin | PyTorch `.pt` dictionary | External |
|
|
82
|
+
|
|
83
|
+
---
|
|
84
|
+
|
|
85
|
+
## Installation
|
|
86
|
+
|
|
87
|
+
### Prerequisites
|
|
88
|
+
- Python `>= 3.10`
|
|
89
|
+
- PyTorch `>= 2.0`
|
|
90
|
+
|
|
91
|
+
### Install via PyPI
|
|
92
|
+
|
|
93
|
+
```bash
|
|
94
|
+
pip install stackformer
|
|
95
|
+
```
|
|
96
|
+
|
|
97
|
+
### Install from Source
|
|
98
|
+
|
|
99
|
+
```bash
|
|
100
|
+
git clone https://github.com/stackformer-labs/Stackformer.git
|
|
101
|
+
cd Stackformer
|
|
102
|
+
pip install -e .
|
|
103
|
+
```
|
|
104
|
+
|
|
105
|
+
### Optional Dependencies
|
|
106
|
+
|
|
107
|
+
For logging integrations (TensorBoard / Weights & Biases):
|
|
108
|
+
|
|
109
|
+
```bash
|
|
110
|
+
pip install tensorboard wandb
|
|
111
|
+
```
|
|
112
|
+
|
|
113
|
+
---
|
|
114
|
+
|
|
115
|
+
## Quick Start
|
|
116
|
+
|
|
117
|
+
### 1. Build a Custom Architecture with `BlockConfig`
|
|
118
|
+
|
|
119
|
+
```python
|
|
120
|
+
import torch
|
|
121
|
+
from stackformer.modules import BlockConfig, TransformerEncoder
|
|
122
|
+
|
|
123
|
+
# Define a custom Transformer block configuration
|
|
124
|
+
config = BlockConfig(
|
|
125
|
+
embed_dim=512,
|
|
126
|
+
num_heads=8,
|
|
127
|
+
hidden_dim=2048,
|
|
128
|
+
attention="gqa_rope", # Grouped-Query Attention with Rotary Embeddings
|
|
129
|
+
num_kv_heads=2,
|
|
130
|
+
ffn="swiglu", # SwiGLU Feed-Forward Network
|
|
131
|
+
norm="rmsnorm", # RMSNorm normalization
|
|
132
|
+
pre_norm=True,
|
|
133
|
+
dropout=0.1,
|
|
134
|
+
)
|
|
135
|
+
|
|
136
|
+
# Instantiate a 6-layer Encoder backbone
|
|
137
|
+
encoder = TransformerEncoder(config, num_layers=6)
|
|
138
|
+
|
|
139
|
+
x = torch.randn(2, 64, 512)
|
|
140
|
+
output = encoder(x, mask=True)
|
|
141
|
+
print("Encoder output shape:", output.shape) # torch.Size([2, 64, 512])
|
|
142
|
+
```
|
|
143
|
+
|
|
144
|
+
### 2. Instantiate a Model from the Architecture Zoo
|
|
145
|
+
|
|
146
|
+
```python
|
|
147
|
+
import torch
|
|
148
|
+
from stackformer.models import Llama2
|
|
149
|
+
|
|
150
|
+
# Initialize LLaMA-2 with GQA and stateful KV-cache support
|
|
151
|
+
model = Llama2(
|
|
152
|
+
vocab_size=32000,
|
|
153
|
+
num_layers=4,
|
|
154
|
+
embed_dim=512,
|
|
155
|
+
num_query_heads=8,
|
|
156
|
+
num_kv_heads=2,
|
|
157
|
+
batch_size=1,
|
|
158
|
+
kv_seq_len=128,
|
|
159
|
+
)
|
|
160
|
+
|
|
161
|
+
input_ids = torch.randint(0, 32000, (1, 16))
|
|
162
|
+
logits = model(input_ids, start_pos=0)
|
|
163
|
+
print("Logits shape:", logits.shape) # torch.Size([1, 16, 32000])
|
|
164
|
+
```
|
|
165
|
+
|
|
166
|
+
### 3. Train with the High-Level `Trainer` Engine
|
|
167
|
+
|
|
168
|
+
```python
|
|
169
|
+
import torch
|
|
170
|
+
from torch.utils.data import DataLoader, TensorDataset
|
|
171
|
+
from stackformer.engine import Trainer
|
|
172
|
+
from stackformer.models import GPT2
|
|
173
|
+
|
|
174
|
+
# Synthetic dataset
|
|
175
|
+
x = torch.randint(0, 50257, (64, 16))
|
|
176
|
+
y = torch.randint(0, 50257, (64, 16))
|
|
177
|
+
train_loader = DataLoader(TensorDataset(x, y), batch_size=8)
|
|
178
|
+
|
|
179
|
+
model = GPT2(vocab_size=50257, num_layers=2, embed_dim=256, num_heads=4, seq_len=64)
|
|
180
|
+
|
|
181
|
+
trainer = Trainer(
|
|
182
|
+
model=model,
|
|
183
|
+
train_dataloader=train_loader,
|
|
184
|
+
val_dataloader=train_loader,
|
|
185
|
+
device="cpu",
|
|
186
|
+
use_amp=False,
|
|
187
|
+
use_ddp=False,
|
|
188
|
+
max_epochs=1,
|
|
189
|
+
max_train_steps=5,
|
|
190
|
+
lr=3e-4,
|
|
191
|
+
checkpoint_dir="checkpoints",
|
|
192
|
+
)
|
|
193
|
+
|
|
194
|
+
trainer.fit()
|
|
195
|
+
```
|
|
196
|
+
|
|
197
|
+
### 4. Autoregressive Text Generation
|
|
198
|
+
|
|
199
|
+
```python
|
|
200
|
+
import torch
|
|
201
|
+
from stackformer import text_generate
|
|
202
|
+
from stackformer.models import Llama2
|
|
203
|
+
|
|
204
|
+
model = Llama2(
|
|
205
|
+
vocab_size=32000,
|
|
206
|
+
num_layers=2,
|
|
207
|
+
embed_dim=256,
|
|
208
|
+
num_query_heads=4,
|
|
209
|
+
num_kv_heads=2,
|
|
210
|
+
batch_size=1,
|
|
211
|
+
kv_seq_len=128,
|
|
212
|
+
)
|
|
213
|
+
|
|
214
|
+
prompt = torch.randint(0, 32000, (1, 8))
|
|
215
|
+
generated_ids = text_generate(
|
|
216
|
+
model=model,
|
|
217
|
+
prompt_ids=prompt,
|
|
218
|
+
max_new_tokens=20,
|
|
219
|
+
temperature=0.8,
|
|
220
|
+
top_p=0.9,
|
|
221
|
+
)
|
|
222
|
+
print("Generated shape:", generated_ids.shape) # torch.Size([1, 28])
|
|
223
|
+
```
|
|
224
|
+
|
|
225
|
+
---
|
|
226
|
+
|
|
227
|
+
## Project Structure
|
|
228
|
+
|
|
229
|
+
```text
|
|
230
|
+
Stackformer/
|
|
231
|
+
├── assets/ # Branding logos and documentation images
|
|
232
|
+
├── ci/ # CI runner scripts (Kaggle GPU integration)
|
|
233
|
+
├── docs/ # User & developer documentation
|
|
234
|
+
│ ├── roadmap.md # Detailed technical roadmap (Phases 0–5)
|
|
235
|
+
│ ├── user_docs/ # Installation, quickstart, API reference
|
|
236
|
+
│ └── developer_docs/ # Architecture deep-dive and scope/design non-goals
|
|
237
|
+
├── examples/ # Runnable usage examples
|
|
238
|
+
│ ├── simple_trainer.py # Trainer engine execution demo
|
|
239
|
+
│ └── train_ddp.py # Multi-GPU DistributedDataParallel (DDP) demo
|
|
240
|
+
├── reviews/ # Internal engineering reviews and roadmap specifications
|
|
241
|
+
├── stackformer/ # Core library package
|
|
242
|
+
│ ├── __init__.py # Top-level API exports
|
|
243
|
+
│ ├── config.py # ModelConfig, TrainingConfig, GenerationConfig dataclasses
|
|
244
|
+
│ ├── generate.py # Autoregressive decoding engine and KV-cache dispatcher
|
|
245
|
+
│ ├── metrics.py # Public metric utilities
|
|
246
|
+
│ ├── py.typed # PEP 561 inline typing indicator
|
|
247
|
+
│ ├── amp/ # Automatic Mixed Precision (AMPScaler)
|
|
248
|
+
│ ├── cache/ # KV-cache strategies (StaticKVCache, PagedKVCache scaffold)
|
|
249
|
+
│ ├── distributed/ # DistributedDataParallel (DDP) wrappers and process helpers
|
|
250
|
+
│ ├── engine/ # High-level Trainer, Engine, State, and CheckpointManager
|
|
251
|
+
│ ├── language/ # Abstract decoder and encoder-decoder bases
|
|
252
|
+
│ ├── logging/ # Metrics tracking, CSV, TensorBoard, WandB loggers
|
|
253
|
+
│ ├── models/ # GPT-1/2, LLaMA-1/2, Gemma-1, BERT, RoBERTa, Transformer
|
|
254
|
+
│ ├── modules/ # Attention, FFN, Norm, Positional Embedding, BlockConfig, Layer
|
|
255
|
+
│ ├── optim/ # Optimizer and scheduler factory constructors and loss functions
|
|
256
|
+
│ ├── training/ # Engine loop helper routines
|
|
257
|
+
│ ├── utils/ # Device helpers, seed utility, shape formatting
|
|
258
|
+
│ └── vision/ # ViT and SegFormer-B0 architectures
|
|
259
|
+
└── tests/ # 136 passing tests across unit, integration, and model suites
|
|
260
|
+
```
|
|
261
|
+
|
|
262
|
+
---
|
|
263
|
+
|
|
264
|
+
## Architecture & Model Zoo
|
|
265
|
+
|
|
266
|
+
StackFormer ships verified implementations of canonical language and vision models built directly on `stackformer.modules`:
|
|
267
|
+
|
|
268
|
+
| Model Architecture | Category | Attention Mechanism | Positional Embedding | Normalization | FFN Activation |
|
|
269
|
+
| :--- | :--- | :--- | :--- | :--- | :--- |
|
|
270
|
+
| **GPT-1** | Causal Language Model | MHA (Causal) | Absolute Learned | Post-LayerNorm | GELU |
|
|
271
|
+
| **GPT-2** | Causal Language Model | MHA (Causal) | Absolute Learned | Pre-LayerNorm | GELU |
|
|
272
|
+
| **LLaMA-1** | Causal Language Model | MHA + RoPE | RoPE | Pre-RMSNorm | SwiGLU |
|
|
273
|
+
| **LLaMA-2** | Causal Language Model | GQA + Stateful KV Cache | RoPE | Pre-RMSNorm | SwiGLU |
|
|
274
|
+
| **Gemma-1 (2B/7B)** | Causal Language Model | MQA / GQA + RoPE | RoPE | Pre-RMSNorm | GeGLU |
|
|
275
|
+
| **BERT** | Bidirectional Language Model | MHA (Bidirectional) | Absolute Learned + Segment | Post-LayerNorm | GELU |
|
|
276
|
+
| **RoBERTa** | Bidirectional Language Model | MHA (Bidirectional) | Absolute Learned Offset | Post-LayerNorm | GELU |
|
|
277
|
+
| **Vaswani Transformer** | Seq2Seq Encoder-Decoder | Causal MHA + Cross-MHA | Sinusoidal Fixed | Post-LayerNorm | ReLU |
|
|
278
|
+
| **ViT** | Vision Classification | MHA (Bidirectional) | Absolute Learned | Pre-LayerNorm | GELU |
|
|
279
|
+
| **SegFormer-B0** | Semantic Segmentation | Spatial Reduction Attention | Efficient Mix-FFN | Pre-LayerNorm | GELU |
|
|
280
|
+
|
|
281
|
+
---
|
|
282
|
+
|
|
283
|
+
## Roadmap
|
|
284
|
+
|
|
285
|
+
Development is structured into linear engineering phases per [`reviews/00_FUTURE_PLAN.md`](docs/roadmap.md):
|
|
286
|
+
|
|
287
|
+
- **Phase 0 — Unblock Core Generation (Current):** Standardize `prefill()` and `decode()` KV-cache contracts across decoder architectures, update `text_generate()`, and enforce cache parity testing. *(Completed for LLaMA-2).*
|
|
288
|
+
- **Phase 1 — Attention Engine & Compiler Integration (Planned):** Unified `AttentionEngine` kernel dispatcher (`stackformer/modules/attention_engine.py`) routing between SDPA, FlexAttention, `flash-attn` v2/v4, and custom Triton fallbacks with zero `torch.compile` graph breaks.
|
|
289
|
+
- **Phase 2 — Budget Planner & PyTorch Native FSDP2 Scaling (Planned):** Analytical and empirical execution planner (`plan_training()`), native PyTorch FSDP2 sharding (`torch.distributed.fsdp.fully_shard`) + `DTensor` + Distributed Checkpoint (`DCP`).
|
|
290
|
+
- **Phase 3 — `torchao` Quantization & Native Adapters (Planned):** Native low-precision quantization via `torchao` (INT8, INT4, FP8, NF4) and zero-bloat `LoRALinear` / `DoRALinear` modules on `BlockConfig`.
|
|
291
|
+
- **Phase 4 — Multimodal Vision-Language & Document AI (Planned):** Unified vision-text backbone combining SegFormer/ViT patch encoders with LLaMA decoder stacks via cross-attention and projection adapters.
|
|
292
|
+
- **Phase 5 — Paged KV Cache & Serving Engine (Planned):** `KVCacheManager` featuring PagedAttention virtual block tables and high-concurrency continuous-batching server (`FastAPI` + Model Context Protocol server).
|
|
293
|
+
|
|
294
|
+
*For the complete deep technical specification, see [docs/roadmap.md](docs/roadmap.md).*
|
|
295
|
+
|
|
296
|
+
---
|
|
297
|
+
|
|
298
|
+
## Documentation
|
|
299
|
+
|
|
300
|
+
- **User Documentation:**
|
|
301
|
+
- [Installation Guide](docs/user_docs/installation.md)
|
|
302
|
+
- [Quickstart Guide](docs/user_docs/quickstart.md)
|
|
303
|
+
- [API Reference](docs/user_docs/api_reference.md)
|
|
304
|
+
- **Developer & Technical Documentation:**
|
|
305
|
+
- [Architecture Specification](docs/developer_docs/architecture.md)
|
|
306
|
+
- [Library Scope & Design Non-Goals](docs/developer_docs/scope.md)
|
|
307
|
+
- [Technical Roadmap](docs/roadmap.md)
|
|
308
|
+
- **Code Examples:**
|
|
309
|
+
- [Runnable Example Scripts](examples/)
|
|
310
|
+
|
|
311
|
+
---
|
|
312
|
+
|
|
313
|
+
## Community & Resources
|
|
314
|
+
|
|
315
|
+
- **GitHub Repository:** [https://github.com/stackformer-labs/Stackformer](https://github.com/stackformer-labs/Stackformer)
|
|
316
|
+
- **Issue Tracker:** [https://github.com/stackformer-labs/Stackformer/issues](https://github.com/stackformer-labs/Stackformer/issues)
|
|
317
|
+
- **Discussions:** [https://github.com/stackformer-labs/Stackformer/discussions](https://github.com/stackformer-labs/Stackformer/discussions)
|
|
318
|
+
- **Releases:** [https://github.com/stackformer-labs/Stackformer/releases](https://github.com/stackformer-labs/Stackformer/releases)
|
|
319
|
+
|
|
320
|
+
---
|
|
321
|
+
|
|
322
|
+
## Contributing
|
|
323
|
+
|
|
324
|
+
We welcome community contributions to StackFormer! Please review our [Library Scope & Non-Goals](docs/developer_docs/scope.md) before opening pull requests to ensure alignment with project design principles.
|
|
325
|
+
|
|
326
|
+
---
|
|
327
|
+
|
|
328
|
+
## Citation
|
|
329
|
+
|
|
330
|
+
If you use StackFormer in your research or project, please consider citing:
|
|
331
|
+
|
|
332
|
+
```bibtex
|
|
333
|
+
@software{stackformer2026,
|
|
334
|
+
author = {Stackformer Labs},
|
|
335
|
+
title = {StackFormer: A Modular PyTorch Framework for Transformer Architecture Composability},
|
|
336
|
+
year = {2026},
|
|
337
|
+
publisher = {GitHub},
|
|
338
|
+
journal = {GitHub repository},
|
|
339
|
+
howpublished = {\url{https://github.com/stackformer-labs/Stackformer}}
|
|
340
|
+
}
|
|
341
|
+
```
|
|
342
|
+
|
|
343
|
+
---
|
|
344
|
+
|
|
345
|
+
## License
|
|
346
|
+
|
|
347
|
+
This project is licensed under the [MIT License](LICENSE).
|