PyTorch发布TinyTorch：从零手写ML框架的开源课程

A framework you write yourself, tensors through transformers

TL;DR

Every mature systems project eventually needs a teaching version. TinyTorch is a free, open-source curriculum where you build a working ML framework from scratch, tensors through transformers, in pure Python, using PyTorch’s own API. Twenty modules. Runs on a laptop with 4 GB of RAM and no GPU. It is at tinytorch.ai.

The rest of this post is why we think it needed to exist, and what six years of running it taught us that might be useful to anyone else doing open curriculum work.

Every Systems Field Eventually Builds Its Teaching Version

Unix got too big to hold in your head, so Andrew Tanenbaum wrote MINIX. Small enough that a student could actually finish it. It went on to shape a generation of systems engineers and famously inspired Linux.

Compilers went the same way. LLVM and GCC are decades of excellent engineering and close to unreadable as a first text, so courses teach the Tiger compiler instead. MIT rewrote xv6 from x86 to RISC-V for the same reason, stripping out historical complexity to expose clean abstractions. Before all of them, Nachos and Pintos and SICP.

None of these were trying to replace the production system. They taught what production systems have to hide.

PyTorch is at that point now, which is a compliment. There is genuinely good writing on its internals, most of it by the people who wrote them, and nearly all of it assumes you already think like a framework engineer. What has been missing is the rung below that. Something you build yourself, with grading infrastructure a university can actually adopt for credit, that a self-learner can get through without a GPU budget.

Why This Matters to PyTorch, Not Only to Students

There is a version of this argument that only educators care about. This is not that version.

Every framework runs on a small population of people who can reason about it from the inside. The ones who spot a memory leak in tensor caching, who know when gradient checkpointing is worth the recompute, who can look at a slow training run and name the bottleneck before they open a profiler.

That population does not grow by itself. Right now most people get to PyTorch’s internals by accident, because something broke badly enough to force the trip. They learn the codebase under deadline pressure, from the outside in, in whatever order the bug happened to demand. Anyone who has onboarded a systems engineer knows the result, which is competence with holes in it.

Building the thing yourself changes the arrival path, and it changes it permanently. Once you have implemented autograd you cannot unsee the computational graph. Once you have profiled your own memory allocation you cannot unknow the cost.

Figure 1: The same linear layer, two ways. On the left, the framework call that works right up until it does not. On the right, the version you wrote and can open when something breaks. The difference is not syntax, it is whether the abstraction is a wall or a door.

A student who has written backward() themselves, who allocated the momentum and variance buffers and watched Adam’s memory footprint triple, shows up to PyTorch’s real autograd with the mental model already loaded. They read the production code as a more sophisticated version of something they understand. That is a cheaper engineer to onboard.

The part we did not anticipate is that companies want this too. Teams have used TinyTorch for new-hire onboarding as a two to three week intensive, as internal training spread across a quarter, and as targeted debugging workshops where somebody works through Module 06 on autograd or Module 12 on attention because that is the subsystem they keep losing to. You do not have to be enrolled anywhere.

What We Built

Twenty modules in four tiers, driven by a CLI called tito, delivered as Jupyter notebooks with the hard parts cut out for you to fill in. You need Python and to be comfortable with NumPy. You do not need a GPU, a cloud account, or any prior ML systems background.

Figure 2: The four tiers. Each one depends on the tier below it, so you cannot skip ahead to optimization without having built the training loop you are optimizing. Foundation fits a half-semester module, all twenty fit a four-credit course, and self-paced learners take anywhere from a few intense weeks to several unhurried months.

Three design decisions do most of the pedagogical work.

Systems from day one. Module 01 ships a memory_footprint() method before it ships matrix multiplication. You learn that one batch of 32 ImageNet images costs 19 MB by computing it, not by reading it somewhere. Later you find out Adam needs roughly 3× the optimizer memory of SGD, because you allocated those buffers yourself and then measured them.

Progressive disclosure. The Tensor class stays clean through Module 05, with no gradient machinery cluttering up data layout and arithmetic. Then in Module 06 you implement enable_autograd(), which bolts requires_grad, .grad, and .backward() onto the class you already understand. Your old code keeps working.

We went back and forth on this one. The tasteful way to do it is inheritance. What we shipped is runtime monkey-patching, which is going to offend somebody reading this. It won because it keeps one Tensor class across all twenty modules instead of two, and because the moment your tensors visibly grow new powers turns out to be the thing students remember. It also happens to mirror PyTorch 0.4, which merged Variable into Tensor for roughly the same reason.

Figure 3: The gradient of matrix multiplication, as a learner meets it in Module 06. The docstring gives you the mathematical rule and the numbered approach breaks it into steps, but the implementation is yours. All twenty modules look like this.

Build to validate. Six historical milestones prove your implementation works. Rosenblatt’s Perceptron in 1958, the XOR crisis in 1969, the backpropagation revival in 1986, the CNN breakthrough in 1998 where your network has to clear 75% on CIFAR-10, the transformer in 2017, and finally MLPerf-style benchmarking. The history is there for motivation. The task performance is there because it is a much harder thing to fake than a unit test. Your autograd is correct because your network learns.

Figure 4: The milestone ladder. A milestone unlocks only when the modules beneath it produce a working implementation, so the timeline doubles as a progress tracker and a correctness proof. Students recreate 67 years of ML history running nothing but code they wrote.

Throughout, TinyTorch mirrors PyTorch’s API deliberately. The API is the transfer mechanism and it carries both ways. Somebody who builds loss.backward() here can open PyTorch’s version afterward and recognize the shape of it. A PyTorch developer can read our attention module in an afternoon and see how the graph gets constructed.

Figure 5: A TinyTorch training loop next to the equivalent PyTorch loop. The imports differ and almost nothing else does, which is the entire point. Nothing you learn here has to be unlearned later.

Two things it is not. The resemblance stops at the API surface, so there is no dispatcher, no C++ or CUDA layer, no JIT, nothing distributed. And it is slow. Pure Python runs somewhere between 100 and 10,000 times slower than PyTorch, which we will come back to, because it turned out to matter.

Where It Came From

TinyTorch did not start as a framework. It started as a course with a problem.

CS 249r launched at Harvard in 2020, a graduate seminar on TinyML, and there was no textbook to assign. So the course notes became one. We put the book up as an open repository, students and educators started fixing examples and proposing chapters, and by 2024 it had outgrown its TinyML origins because people kept asking about training at scale. It split into two volumes.