Forge-LM is an end-to-end, hardware-optimized Generative AI ecosystem. It bridges the gap between low-level neural network training and high-level full-stack containerized deployment.
Building upon the foundations of NanoTransformer, this repository upgrades the PyTorch architecture to handle * multi-hour training runs* on the TinyStories dataset, featuring production-grade mechanics like Train/Val loss splits, Gradient Accumulation, and Checkpoint Callbacks.
The resulting model is fully serialized and deployed via a FastAPI microservice and consumed by a **Streamlit Web UI **, all orchestrated via Docker Compose, with robust ONNX quantization support for low-memory cloud environments.
Training a small model on Wikipedia results in confused, garbled text due to the immense semantic complexity of the data. Instead, this repository trains on TinyStories—a dataset of short narratives generated entirely with the vocabulary of a 4-year-old child. Because the grammatical structure is highly consistent and the vocabulary is constrained, this lightweight Transformer learns to output perfectly coherent, beginning-to-end narratives in just a few hours of training on a single RTX 3060.
Training Hardware: NVIDIA GeForce RTX 3060 (140W Mobile)
Dataset: TinyStories (100,000 story subset)
Micro-Batch (B) x Context (T) x Accumulation: 8 x 512 x 8 (Effective Batch Size: 64)
Parameters: ~25M (8 Layers, 8 Heads, 512 Embedding)
- Throughput: ~74,000 Tokens/sec
- Validation Loss (5000 steps): ~1.69
Prompt: "Once upon a time, there was a"
Output:
Once upon a time, there was a boy named Timmy. Timmy had a smooth motorcycle with a shiny toy pedals. He loved to ride it around his house, but now he had to manage to sit still while slowing.
One day, Timmy's mom took him to the park. Timmy was so happy because he could ride his smooth motorcycle. But then, he saw a girl who was getting ready to ride away. She wanted to help but didn't know how.
Timmy decided to manage
- Gradient Accumulation: Simulates a massive global batch size without causing Out-of-Memory (OOM) crashes on 6GB VRAM GPUs.
- Validation Loss Tracking: Prevents Overfitting by dynamically tracking loss on an unseen subset of the data.
- Dynamic Checkpointing: Automatically monitors the validation loss curve and saves
storyweaver.ptstrictly when a new historical best is achieved. - Full-Stack Docker Deployment: Decouples inference from the training loop, packaging the
.ptweights into an NVIDIA CUDA base image, allowing stateless web clients to generate stories via REST API. - ONNX Quantization & Lightweight Inference: Export models to INT8 ONNX format (
storyweaver_quant.onnx) to deploy inference APIs on strictly resource-constrained PaaS environments (<512MB RAM), bypassing massive PyTorch dependencies entirely.
For deep dives into the mathematics, scaling, and deployment, see our documentation:
If you already have storyweaver.pt and lexibyte_vocab.json generated from a training run, launch the entire
ecosystem (API + Web UI).
Requires NVIDIA Container Toolkit.
- Setup Environment Variables:
cp .env.example .env
- Launch Ecosystem:
docker-compose up --build -d
Web UI: http://localhost:8501 | API Docs: http://localhost:8000/docs
If you want to deploy just the FastAPI microservice and build your own custom frontend elsewhere, target the API service directly:
- Setup Environment Variables:
cp .env.example .env
- Launch API:
docker-compose up --build -d api
If you want to train the model from scratch and run it locally without Docker:
- Install Requirements: (Python 3.12+ recommended)
pip install -r requirements.min.txt
- Download Dataset & Train:
python src/prepare_data.py python src/train.py
- Run Local FastAPI (Headless Mode):
uvicorn src.api:app --reload
- Run Streamlit UI (Full-Stack Mode): (In a second terminal)
streamlit run src/app.py