Micro-optimizations and instrumentation of your training loop

Deliverableartifact

An updated training module with timing instrumentation, configurable precision/accumulation options, and a brief text report comparing at least two configurations by tokens/sec and loss curves.

Goals

Instrument your code to separately time data loading and forward/backward passes and report their percentages.
Add a configuration flag for mixed precision (even if you only prototype it) and explain its impact in logs.
Run two training configurations and compare tokens/sec and convergence behavior.
Summarize in a short note which changes gave the largest throughput improvement in your setup.

Resources

📄stanford-cs336/assignment2-systems 📄bcefghj/learn-cs336

Overview: what you’re building

You’re going to upgrade your training loop so you can:

Time data loading vs model compute separately
Flip precision (fp32, mixed precision) and gradient accumulation via config
Run small, repeated experiments with consistent logging/checkpointing
Produce a tiny report comparing configs by tokens/sec and loss curves

We’ll walk through this by gradually shaping a small but real PyTorch-style training module. You can adapt it to other frameworks; the structure is what matters.

1 / 6