Lesson 4: Where a training step spends its time and memory

Module 5 · Lesson 4 · About 55 minutes · Needs: Module 4's lab (tinyformer at v0.4.0); Lesson 2 of this module (the fused attention path)

The situation

Sol 409 at Aurora Station. Hana Aoki, the Compute Officer, has a question for the lab: "base trains in half a minute on your laptop. Where does that time go, and what would you change before we train anything bigger on the station's CPUs?"

You do not guess. You profile one training step of base and follow the time and the memory. The running example is that one step, on the first batch base ever trains on: 32 windows of 128 characters, whose row 0 begins "ra.Assignment sheet for sol 288".