Training: Accumulate Gradients with Equal Work

Lesson

Compute the weight gradient of mean squared error over all 256 examples. Split into microbatches, scale each loss by its share of the full dataset and accumulate gradients before returning. No optimizer step is performed; input weights are immutable. Validate the whole gradient against the full-batch FP32 reference on fresh inputs. Compare peak allocated/reserved memory, samples/sec and gradient error. Checkpointing would trade recomputation for saved activations; it does not save arithmetic. This is a separate training exercise, so inference output equivalence cannot substitute for gradient correctness.

Measurement contract: the real read-only driver defines WORKLOAD. Two warmups precede five paired trials with alternating baseline/candidate order. Each trial uses fresh deterministic inputs. Synchronized host time includes solve and completion, excludes input construction and reference verification. METRICS reports every sample, median and IQR; memory includes persistent buffers and allocator state. One additional solve is isolated in judged for Trace. No speedup is promised. Run the unchanged starter first.

What you will learn

• Identify the bottleneck from the appropriate evidence.
• Preserve the output contract on fresh inputs.
• Compare paired measurements and numerical error.

Starter code

    return baseline(case)