stackformer 0.0.1__tar.gz → 0.1.9__tar.gz

stackformer-0.1.9/PKG-INFO

@@ -0,0 +1,185 @@
+Metadata-Version: 2.4
+Name: stackformer
+Version: 0.1.9
+Summary: Modular transformer blocks built in PyTorch
+Author-email: Gurumurthy <gurumurthy.00300@gmail.com>
+License: MIT
+Project-URL: Repository, https://github.com/Gurumurthy30/Stackformer
+Project-URL: Releases, https://github.com/Gurumurthy30/Stackformer/releases
+Project-URL: Issue-Tracker, https://github.com/Gurumurthy30/Stackformer/issues
+Project-URL: Discussions, https://github.com/Gurumurthy30/Stackformer/discussions
+Project-URL: Documentation, https://github.com/Gurumurthy30/Stackformer/tree/main/docs
+Keywords: transformer,pytorch,deep-learning,attention,llm,machine-learning
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Typing :: Typed
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch<2.7,>=2.0
+Requires-Dist: numpy<3.0,>=1.23
+Requires-Dist: tqdm<5.0,>=4.5
+Requires-Dist: tensorboard<3.0,>=2.14
+Requires-Dist: wandb<1.0,>=0.17
+Provides-Extra: dev
+Requires-Dist: pytest<9.0,>=8.0; extra == "dev"
+Requires-Dist: pytest-cov<6.0,>=5.0; extra == "dev"
+Requires-Dist: build<2.0,>=1.2; extra == "dev"
+Dynamic: license-file
+
+<p align="center">
+  <img src="assets/logo.png" alt="StackFormer logo" width="560" />
+</p>
+
+<p align="center">
+  <a href="https://pypi.org/project/stackformer/"><img src="https://img.shields.io/pypi/v/stackformer.svg" alt="PyPI version" /></a>
+  <a href="https://pypi.org/project/stackformer/"><img src="https://img.shields.io/pypi/pyversions/stackformer.svg" alt="Python versions" /></a>
+  <a href="LICENSE"><img src="https://img.shields.io/badge/license-MIT-blue.svg" alt="License" /></a>
+  <a href="https://github.com/Gurumurthy30/Stackformer/actions">
+    <img src="https://img.shields.io/github/actions/workflow/status/Gurumurthy30/Stackformer/core-tests.yml?branch=main&label=CI" alt="CI status" />
+  </a>
+</p>
+
+# StackFormer
+
+StackFormer is a modular PyTorch framework for building, training, and experimenting with Transformer architectures.
+
+## Overview
+
+StackFormer is designed for fast experimentation with reusable Transformer building blocks and model implementations. It supports both language and vision workflows in a single modular codebase. The framework is built for research, prototyping, and iterative model development with practical training infrastructure.
+
+## Key Features
+
+- Modular transformer components
+- GPT / LLaMA / Gemma-style model implementations
+- Vision models (ViT, SegFormer)
+- Trainer infrastructure with AMP mixed precision and DDP support
+- Logging and metrics utilities
+- Checkpointing and resume training
+- CI-tested training infrastructure
+
+## Project Structure
+
+```text
+Stackformer/
+├── assets/
+├── docs/
+│   ├── user_docs/
+│   └── developer_docs/
+├── examples/
+├── stackformer/
+│   ├── __init__.py
+│   ├── generate.py
+│   ├── metrics.py
+│   ├── py.typed
+│   ├── amp/
+│   │   └── scaler.py
+│   ├── distributed/
+│   │   └── ddp.py
+│   ├── engine/
+│   │   ├── checkpoint.py
+│   │   ├── engine.py
+│   │   ├── state.py
+│   │   └── trainer.py
+│   ├── logging/
+│   │   ├── csv_logger.py
+│   │   ├── logger.py
+│   │   ├── metrics.py
+│   │   ├── tensorboard_logger.py
+│   │   ├── wandb_logger.py
+│   │   └── wb_logger.py
+│   ├── models/
+│   │   ├── OpenAI.py
+│   │   ├── Meta.py
+│   │   ├── Google.py
+│   │   └── Transformer.py
+│   ├── modules/
+│   │   ├── Attention.py
+│   │   ├── Feed_forward.py
+│   │   ├── Masking.py
+│   │   ├── Normalization.py
+│   │   ├── position_embedding.py
+│   ├── optim/
+│   │   ├── factories.py
+│   │   └── loss_fn.py
+│   ├── training/
+│   │   └── loops.py
+│   ├── utils/
+│   │   ├── device.py
+│   │   └── utils.py
+│   └── vision/
+│       ├── vit.py
+│       └── segformer.py
+├── tests/
+│   ├── integration/
+│   ├── models/
+│   ├── modules/
+│   ├── trainer/
+│   ├── utils/
+│   ├── conftest.py
+│   ├── test_distributed.py
+│   └── test_vision.py
+├── LICENSE
+├── pyproject.toml
+└── README.md
+```
+
+## Installation
+
+Python >= 3.10
+
+### Install from PyPI
+
+```bash
+pip install stackformer
+```
+
+### Install from source
+
+```bash
+git clone https://github.com/Gurumurthy30/Stackformer.git
+cd Stackformer
+pip install -e .
+```
+
+## Quick Start
+
+```python
+from stackformer.engine import Trainer
+import torch.nn as nn
+
+model = nn.Linear(10, 1)
+
+trainer = Trainer(model=model)
+trainer.fit(dataset)
+```
+
+## Examples
+
+More runnable examples are available in:
+
+```text
+examples/
+```
+
+```text
+examples/simple_train.py
+examples/simple_trainer_v2.py
+examples/train_ddp.py
+```
+
+## Documentation
+
+- User documentation: [docs/user_docs/installation.md](docs/user_docs/installation.md)
+- Developer documentation: [docs/developer_docs/architecture.md](docs/developer_docs/architecture.md)
+
+## Community
+
+- Issues: https://github.com/Gurumurthy30/Stackformer/issues
+- Discussions: https://github.com/Gurumurthy30/Stackformer/discussions
+- Releases: https://github.com/Gurumurthy30/Stackformer/releases

stackformer-0.1.9/README.md

@@ -0,0 +1,151 @@
+<p align="center">
+  <img src="assets/logo.png" alt="StackFormer logo" width="560" />
+</p>
+
+<p align="center">
+  <a href="https://pypi.org/project/stackformer/"><img src="https://img.shields.io/pypi/v/stackformer.svg" alt="PyPI version" /></a>
+  <a href="https://pypi.org/project/stackformer/"><img src="https://img.shields.io/pypi/pyversions/stackformer.svg" alt="Python versions" /></a>
+  <a href="LICENSE"><img src="https://img.shields.io/badge/license-MIT-blue.svg" alt="License" /></a>
+  <a href="https://github.com/Gurumurthy30/Stackformer/actions">
+    <img src="https://img.shields.io/github/actions/workflow/status/Gurumurthy30/Stackformer/core-tests.yml?branch=main&label=CI" alt="CI status" />
+  </a>
+</p>
+
+# StackFormer
+
+StackFormer is a modular PyTorch framework for building, training, and experimenting with Transformer architectures.
+
+## Overview
+
+StackFormer is designed for fast experimentation with reusable Transformer building blocks and model implementations. It supports both language and vision workflows in a single modular codebase. The framework is built for research, prototyping, and iterative model development with practical training infrastructure.
+
+## Key Features
+
+- Modular transformer components
+- GPT / LLaMA / Gemma-style model implementations
+- Vision models (ViT, SegFormer)
+- Trainer infrastructure with AMP mixed precision and DDP support
+- Logging and metrics utilities
+- Checkpointing and resume training
+- CI-tested training infrastructure
+
+## Project Structure
+
+```text
+Stackformer/
+├── assets/
+├── docs/
+│   ├── user_docs/
+│   └── developer_docs/
+├── examples/
+├── stackformer/
+│   ├── __init__.py
+│   ├── generate.py
+│   ├── metrics.py
+│   ├── py.typed
+│   ├── amp/
+│   │   └── scaler.py
+│   ├── distributed/
+│   │   └── ddp.py
+│   ├── engine/
+│   │   ├── checkpoint.py
+│   │   ├── engine.py
+│   │   ├── state.py
+│   │   └── trainer.py
+│   ├── logging/
+│   │   ├── csv_logger.py
+│   │   ├── logger.py
+│   │   ├── metrics.py
+│   │   ├── tensorboard_logger.py
+│   │   ├── wandb_logger.py
+│   │   └── wb_logger.py
+│   ├── models/
+│   │   ├── OpenAI.py
+│   │   ├── Meta.py
+│   │   ├── Google.py
+│   │   └── Transformer.py
+│   ├── modules/
+│   │   ├── Attention.py
+│   │   ├── Feed_forward.py
+│   │   ├── Masking.py
+│   │   ├── Normalization.py
+│   │   ├── position_embedding.py
+│   ├── optim/
+│   │   ├── factories.py
+│   │   └── loss_fn.py
+│   ├── training/
+│   │   └── loops.py
+│   ├── utils/
+│   │   ├── device.py
+│   │   └── utils.py
+│   └── vision/
+│       ├── vit.py
+│       └── segformer.py
+├── tests/
+│   ├── integration/
+│   ├── models/
+│   ├── modules/
+│   ├── trainer/
+│   ├── utils/
+│   ├── conftest.py
+│   ├── test_distributed.py
+│   └── test_vision.py
+├── LICENSE
+├── pyproject.toml
+└── README.md
+```
+
+## Installation
+
+Python >= 3.10
+
+### Install from PyPI
+
+```bash
+pip install stackformer
+```
+
+### Install from source
+
+```bash
+git clone https://github.com/Gurumurthy30/Stackformer.git
+cd Stackformer
+pip install -e .
+```
+
+## Quick Start
+
+```python
+from stackformer.engine import Trainer
+import torch.nn as nn
+
+model = nn.Linear(10, 1)
+
+trainer = Trainer(model=model)
+trainer.fit(dataset)
+```
+
+## Examples
+
+More runnable examples are available in:
+
+```text
+examples/
+```
+
+```text
+examples/simple_train.py
+examples/simple_trainer_v2.py
+examples/train_ddp.py
+```
+
+## Documentation
+
+- User documentation: [docs/user_docs/installation.md](docs/user_docs/installation.md)
+- Developer documentation: [docs/developer_docs/architecture.md](docs/developer_docs/architecture.md)
+
+## Community
+
+- Issues: https://github.com/Gurumurthy30/Stackformer/issues
+- Discussions: https://github.com/Gurumurthy30/Stackformer/discussions
+- Releases: https://github.com/Gurumurthy30/Stackformer/releases

stackformer-0.1.9/pyproject.toml

@@ -0,0 +1,61 @@
+[project]
+name = "stackformer"
+version = "0.1.9"
+description = "Modular transformer blocks built in PyTorch"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT" }
+
+authors = [
+  { name = "Gurumurthy", email = "gurumurthy.00300@gmail.com" }
+]
+
+keywords = ["transformer", "pytorch", "deep-learning", "attention", "llm", "machine-learning"]
+
+dependencies = [
+  "torch>=2.0,<2.7",
+  "numpy>=1.23,<3.0",
+  "tqdm>=4.5,<5.0",
+  "tensorboard>=2.14,<3.0",
+  "wandb>=0.17,<1.0"
+]
+
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Typing :: Typed",
+  "Operating System :: OS Independent",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+  "Topic :: Software Development :: Libraries :: Python Modules"
+]
+
+[project.optional-dependencies]
+dev = [
+  "pytest>=8.0,<9.0",
+  "pytest-cov>=5.0,<6.0",
+  "build>=1.2,<2.0"
+]
+
+[project.urls]
+Repository = "https://github.com/Gurumurthy30/Stackformer"
+Releases = "https://github.com/Gurumurthy30/Stackformer/releases"
+Issue-Tracker = "https://github.com/Gurumurthy30/Stackformer/issues"
+Discussions = "https://github.com/Gurumurthy30/Stackformer/discussions"
+Documentation = "https://github.com/Gurumurthy30/Stackformer/tree/main/docs"
+
+[tool.setuptools.packages.find]
+include = ["stackformer*"]
+exclude = ["assets*"]
+
+[tool.setuptools.package-data]
+stackformer = ["py.typed"]
+
+[build-system]
+requires = ["setuptools>=61", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-q"

stackformer-0.1.9/stackformer/__init__.py

@@ -0,0 +1,74 @@
+"""Top-level public API for Stackformer."""
+
+from .generate import text_generate
+from .config import GenerationConfig, ModelConfig, TrainingConfig
+from .models import GPT_1, GPT_2, gemma_1_2b, gemma_1_7b, llama_1, llama_2, transformer
+from .modules import (
+    AbsolutePositionEmbedding,
+    Cross_MultiHead_Attention,
+    FF_GELU,
+    FF_GeGLU,
+    FF_LeakyReLU,
+    FF_ReLU,
+    FF_Sigmoid,
+    FF_SiLU,
+    FF_SwiGLU,
+    Group_query_Attention,
+    Group_query_Attention_With_RoPE,
+    LayerNormalization,
+    Multi_Head_Attention,
+    Multi_Head_Attention_With_RoPE,
+    Multi_query_Attention,
+    Multi_query_Attention_With_RoPE,
+    RMSNormalization,
+    RoPE,
+    Self_Attention,
+    SinusoidalPositionalEmbedding,
+    kv_cache_group_query,
+    kv_cache_multihead,
+    make_mask,
+)
+from .vision import SegFormerB0, ViT
+
+from .engine import Trainer
+
+__all__ = [
+    "AbsolutePositionEmbedding",
+    "Cross_MultiHead_Attention",
+    "FF_GELU",
+    "FF_GeGLU",
+    "FF_LeakyReLU",
+    "FF_ReLU",
+    "FF_Sigmoid",
+    "FF_SiLU",
+    "FF_SwiGLU",
+    "GPT_1",
+    "GPT_2",
+    "Group_query_Attention",
+    "Group_query_Attention_With_RoPE",
+    "LayerNormalization",
+    "Multi_Head_Attention",
+    "Multi_Head_Attention_With_RoPE",
+    "Multi_query_Attention",
+    "Multi_query_Attention_With_RoPE",
+    "RMSNormalization",
+    "RoPE",
+    "SegFormerB0",
+    "Self_Attention",
+    "SinusoidalPositionalEmbedding",
+    "ViT",
+    "gemma_1_2b",
+    "gemma_1_7b",
+    "kv_cache_group_query",
+    "kv_cache_multihead",
+    "llama_1",
+    "llama_2",
+    "make_mask",
+    "text_generate",
+    "transformer",
+    "GenerationConfig",
+    "ModelConfig",
+    "TrainingConfig",
+]
+
+__all__.append("Trainer")

stackformer-0.1.9/stackformer/amp/__init__.py

@@ -0,0 +1,5 @@
+"""Automatic mixed precision utilities."""
+
+from .scaler import AMPScaler, initialize_scaler, scale_loss, step_optimizer, update_scaler
+
+__all__ = ["AMPScaler", "initialize_scaler", "scale_loss", "step_optimizer", "update_scaler"]

stackformer-0.1.9/stackformer/amp/scaler.py

@@ -0,0 +1,84 @@
+"""Utilities for automatic mixed precision (AMP)."""
+
+from __future__ import annotations
+
+from contextlib import nullcontext
+
+import torch
+
+
+class AMPScaler:
+    """Small wrapper around PyTorch AMP scaler/autocast.
+
+    AMP is automatically disabled on CPU-only runs.
+    """
+
+    def __init__(self, enabled: bool = True):
+        self.enabled = bool(enabled and torch.cuda.is_available())
+
+        # Prefer modern torch.amp API, with fallback for older versions.
+        if hasattr(torch, "amp") and hasattr(torch.amp, "GradScaler"):
+            self.scaler = torch.amp.GradScaler("cuda", enabled=self.enabled)
+            self._autocast = lambda: torch.amp.autocast(device_type="cuda", enabled=self.enabled)
+        else:
+            self.scaler = torch.cuda.amp.GradScaler(enabled=self.enabled)
+            self._autocast = lambda: torch.cuda.amp.autocast(enabled=self.enabled)
+
+    def autocast(self):
+        if not self.enabled:
+            return nullcontext()
+        return self._autocast()
+
+    def scale(self, loss: torch.Tensor) -> torch.Tensor:
+        if not self.enabled:
+            return loss
+        return self.scaler.scale(loss)
+
+    def step(self, optimizer: torch.optim.Optimizer) -> None:
+        if not self.enabled:
+            optimizer.step()
+            return
+        self.scaler.step(optimizer)
+
+    def update(self) -> None:
+        if self.enabled:
+            self.scaler.update()
+
+    def unscale_(self, optimizer: torch.optim.Optimizer) -> None:
+        if self.enabled:
+            self.scaler.unscale_(optimizer)
+
+    def state_dict(self):
+        if not self.enabled:
+            return {}
+        return self.scaler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        if self.enabled:
+            self.scaler.load_state_dict(state_dict)
+
+    @property
+    def is_enabled(self) -> bool:
+        return self.enabled
+
+
+def initialize_scaler(enabled: bool = True) -> AMPScaler:
+    return AMPScaler(enabled=enabled)
+
+
+def scale_loss(loss: torch.Tensor, scaler: AMPScaler | None) -> torch.Tensor:
+    if scaler is None:
+        return loss
+    return scaler.scale(loss)
+
+
+def step_optimizer(optimizer: torch.optim.Optimizer, scaler: AMPScaler | None) -> None:
+    if scaler is None:
+        optimizer.step()
+        return
+    scaler.step(optimizer)
+
+
+def update_scaler(scaler: AMPScaler | None) -> None:
+    if scaler is not None:
+        scaler.update()

stackformer-0.1.9/stackformer/config.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class ModelConfig:
+    vocab_size: int
+    embed_dim: int
+    num_layers: int
+    num_heads: int
+    seq_len: int
+    hidden_dim: int
+    dropout: float = 0.0
+
+
+@dataclass(slots=True)
+class TrainingConfig:
+    max_epochs: int = 1
+    max_train_steps: int | None = None
+    max_eval_steps: int | None = None
+    eval_every_n_epochs: int = 1
+    save_every_n_epochs: int = 1
+    grad_accumulation_step: int = 1
+    max_grad_norm: float | None = None
+    lr: float = 3e-4
+    weight_decay: float = 0.01
+    optimizer_name: str = "adamw"
+    scheduler_name: str = "none"
+    warmup_steps: int = 0
+    total_steps: int | None = None
+
+
+@dataclass(slots=True)
+class GenerationConfig:
+    max_context_len: int = 128
+    max_new_tokens: int = 50
+    temperature: float = 1.0
+    top_k: int | None = None
+    top_p: float = 1.0
+    eos_token_id: int | None = None

stackformer-0.1.9/stackformer/distributed/__init__.py

@@ -0,0 +1,31 @@
+"""Distributed training helpers for StackFormer."""
+
+from .ddp import (
+    barrier,
+    cleanup,
+    cleanup_distributed,
+    distributed_sampler,
+    get_local_rank,
+    get_rank,
+    get_world_size,
+    init_distributed,
+    is_distributed,
+    is_main_process,
+    setup_ddp,
+    wrap_model_ddp,
+)
+
+__all__ = [
+    "init_distributed",
+    "setup_ddp",
+    "wrap_model_ddp",
+    "distributed_sampler",
+    "get_rank",
+    "get_world_size",
+    "get_local_rank",
+    "is_distributed",
+    "is_main_process",
+    "barrier",
+    "cleanup",
+    "cleanup_distributed",
+]