TinyTorch: Building Machine Learning Systems from First Principles

TL;DR

TinyTorch tackles the fragmentation between ML research and production systems education by teaching ML systems engineering on CPU-only hardware. It presents a 20-module, pure-Python curriculum that builds PyTorch internals (tensors, autograd, optimizers, networks) with memory profiling and performance analysis embedded from day one. The framework combines a Build→Use→Reflect pedagogy, a 3-tier learning journey, and milestone validation that traces history from the Perceptron to Transformers, culminating in MLPerf-style benchmarking. The result is a reproducible, accessible path for students and institutions to become engineers who understand why ML systems work and how to scale them, with practical impact on debugging and deployment.

Abstract

Machine learning systems engineering requires a deep understanding of framework internals. Yet most current education separates algorithms from systems. Students learn gradient descent without measuring memory usage, and attention mechanisms without profiling computational cost. This split leaves graduates unprepared to debug real production failures and widens the gap between machine learning research and reliable deployment. We present TinyTorch, a 20 module curriculum in which students implement the core components of PyTorch, including tensors, autograd, optimizers, and neural networks, entirely in pure Python. The curriculum is built around three pedagogical principles. Progressive disclosure gradually introduces complexity as students build confidence. Systems first integration embeds memory and performance awareness from the very beginning. Historical milestone validation guides students to recreate key breakthroughs, from the Perceptron in 1958 to modern Transformers, using only code they have written themselves. TinyTorch requires only a laptop with 4GB of RAM and no GPU, making machine learning systems education accessible worldwide. Its goal is to prepare the next generation of AI engineers, practitioners who understand not only what machine learning systems do, but why they work and how to make them scale. The curriculum is available as open source at mlsysbook.ai slash tinytorch.

Figures (3)

• Figure 1: From User to Engineer. (a) PyTorch's high-level APIs hide framework internals. (b) TinyTorch students implement components like Adam, learning memory costs and update rules firsthand. (c) By Module 13, every import is student-built code, and transformers train on infrastructure they fully understand.
• Figure 2: Module Dependency Graph. TinyTorch's 20 modules form a directed acyclic graph with two architectural paths. Foundation modules (blue, M01--08) build core infrastructure sequentially, culminating in Training (M08). From Training, two paths branch: the Vision path (M09) builds CNNs for spatial processing; the Language path (M10--13) builds tokenization through transformers. Both paths converge at Profiling (M14), then branch into parallel optimization tracks---Model-level (quantization, compression) and Runtime (acceleration, memoization)---before final convergence at Benchmarking (M19) and Capstone (M20).
• Figure 3: Progressive Disclosure Timeline. Runtime capability expansion manages cognitive load. Modules 01--05 use a clean Tensor class focused on fundamentals (blue): data storage, arithmetic, matrix operations, and memory profiling. In Module 06, enable_autograd() enhances Tensor with gradient tracking (orange): requires_grad, .grad, and .backward() are dynamically added via monkey-patching. Three learning benefits: (1) students master fundamentals without distraction from unused features; (2) Module 01--05 code continues working unchanged after enhancement (backward compatibility); (3) the enhancement moment in Module 06 becomes a concrete "aha" experience as familiar tensors gain new capabilities.