Memory, Batch Size and Precision

Start with the smallest possible question: how much memory does one tensor take?

A tensor holds numbers of one type. The commonest type is float32, which uses 4 bytes per number. So the memory is simply "how many numbers" times "bytes per number". PyTorch will tell you both.

That is the whole arithmetic. Everything else in this lesson is that multiplication applied to different things.

Now, what does a training run hold? Four kinds of tensor. Three of them scale with the size of the model and one scales with the size of the batch:

• parameters — the weights themselves, the thing you are training;
• gradients — one per parameter, saying how much that parameter contributed to the error. So, a second copy of the model;
• optimizer state — extra numbers the optimizer keeps between steps. Plain SGD keeps none. SGD with momentum keeps one per parameter. Adam and AdamW keep two;
• activations — the intermediate results of the forward pass, kept because the backward pass needs them.

Rather than take that on trust, measure it on the Module 3 network.

Gradients cost exactly what the parameters cost, as promised. AdamW costs twice the parameters. Look at what it is keeping:

exp_avg is a running average of recent gradients and exp_avg_sq is a running average of their squares. Both are the same shape as the weights, so both cost the same as the whole model.

That is the rule to remember: training with AdamW needs about four times the memory of the model itself, before a single activation is stored. One copy of weights, one of gradients, two of optimizer state.

That single factor of four explains a lot of Module 8. It is why fully fine-tuning a large model is out of reach on ordinary hardware, and why LoRA — which freezes almost all the weights, so they need no gradient and no optimizer state — changes what is possible rather than merely what is convenient.

The arithmetic above works for any model, so you can estimate a model you have never run, from its parameter count alone.

The function below does that. Its arguments:

• n_parameters — how many weights the model has;
• precision — how many bytes each number takes (Section 3 explains the choices);
• optimizer — how many extra copies it keeps per parameter;
• trainable — the share of parameters you are actually training. 1.0 means all of them; LoRA makes this tiny.

Activations are not included; Section 4 adds them.

The estimate is a few bytes under the measurement, because AdamW also keeps a tiny step counter per tensor. Close enough for planning.

Now use it on models you would not want to run by accident. A 7B model means seven billion parameters, the size of a small open-weights language model like the ones Module 8 fine-tunes.

Read those 7B rows against real hardware. A free Colab T4 has about 15 GB of GPU memory. A good consumer GPU has 24 GB. An A100 or H100 in a data centre has 80 GB.

• Fully fine-tuning a 7B model in fp32 with AdamW: 104 GB, before any activations. That does not fit on any single GPU you are likely to touch.
• Same, with fp16 weights: 52 GB. Still no.
• Loading it just to make predictions, in fp16: 13 GB. That fits on a free Colab GPU, which is why running a 7B model is routine while training one is not.
• LoRA, training 0.2% of the parameters: 13 GB plus 0.08 GB. The gradients and optimizer state almost vanish, and what is left is the frozen model itself. Quantise those frozen weights to 4 bits — that is QLoRA, in Module 8 — and even that shrinks.

This table is also how you answer "will this fit?" in a code review, without booking a GPU to find out.

Halving the bytes per number halves three of the four memory terms at once, so the number format is a memory decision before it is a speed decision.

Every one of these formats spends its bits on two different jobs:

• range — how big and how small a number it can hold at all;
• precision — how finely it can tell nearby numbers apart.

torch.finfo reports both.

In words:

• fp32 (float32) — 4 bytes. Big range, fine precision. The default everywhere.
• fp16 (half) — 2 bytes. Precision is decent, but the range is narrow: nothing above 65,504, and nothing usefully small below about 0.00006.
• bf16 (bfloat16) — 2 bytes. The same range as fp32, bought by giving up precision.

Watch both failures happen.

inf means the number was too big to represent, so it became infinity. 0.0 means it was too small and vanished. And bf16 cannot tell 1.0001 from 1 at all.

Now the consequence that matters for training. Gradients are small numbers. In fp16 many of them fall below that floor and become exactly zero, and a weight whose gradient is zero stops learning. Nothing crashes; the model just quietly trains worse.

The fix is loss scaling: multiply the loss by a big constant before the backward pass, so every gradient is multiplied by that constant too and lands inside fp16's range; then divide the gradients by the same constant before the optimizer uses them. torch.amp.GradScaler does this and adjusts the constant automatically.

Survived. bf16 needs no scaler, because its range is fp32's range; that is exactly what it bought with its precision. This is why bf16 is preferred on hardware that supports it.

Which you get depends on the chip. A free Colab T4 has fast fp16 but no hardware bf16. A100, H100 and recent Apple Silicon support bf16. Section 8 checks at runtime instead of assuming.

Parameters, gradients and optimizer state do not care about batch size. Activations do, and that is what makes batch size a memory control.

An activation is any intermediate result the forward pass produces and the backward pass will need. Every layer's output is one.

To measure them we use a forward hook: a small function PyTorch calls every time a layer produces an output. Ours just adds up the bytes.

Exactly 256 times as much. Activation memory is linear in batch size: each row costs the same, and you choose how many rows.

On the tiny Module 3 network nobody notices. Use a wider one, where each row costs more, and the picture changes.

At batch 32 the activations are a rounding error next to the weights. At batch 16,384 they are 61 times the weights.

That is the trade in one table. Model memory is fixed once you choose an architecture. Activation memory is yours to set, one batch at a time, which is why "halve the batch size" is the first thing to try when a run will not fit.

Two things this simple counting leaves out:

• Real architectures store more per row. Attention in a Transformer keeps intermediate results that grow with the square of the sequence length, which is why Module 6's talk of context windows is also a conversation about memory.
• Gradient checkpointing trades this memory for time: throw activations away during the forward pass and recompute them during the backward pass. It typically saves a large share of activation memory for roughly 30% more compute.

If small batches save memory, why not always use them? Because every step carries fixed overhead — the Python loop, the optimizer update, the cost of launching work on the hardware — and a tiny batch pays that overhead for very little work.

Throughput means rows processed per second. The cell measures it at several batch sizes by running as many steps as fit in one second.

The three warm-up steps matter: the first steps of any timing include one-off costs, so they are run and discarded.

These numbers belong to the machine that runs the cell. The shape of the curve is what transfers, not the values.

The curve climbs steeply and then flattens. Read your own output: on the machine that built this notebook, moving from batch 16 to batch 256 multiplied throughput several times over, moving from 256 to 1,024 added only a little more, and the largest batch was no faster than the one before it.

Past the flat part, a bigger batch buys no speed and costs memory in a straight line. That bend is the useful operating point, and it belongs to the machine rather than the model, so it is worth measuring once on the hardware you really use.

One warning. Batch size is not only an engineering knob. Module 3 showed it changes the training itself: bigger batches give a steadier gradient at each step and usually want a larger learning rate. Shrinking the batch to fit memory and leaving the learning rate alone is a common way to make a model worse for reasons that look mysterious afterwards.

Sometimes the batch size you want will not fit at all. Gradient accumulation gets it anyway: run several small batches, add their gradients together, and take one optimizer step. Memory follows the small batch. The weight update follows the big one.

There is one subtlety, and it is the usual bug. The loss is an average over the rows in the batch, so each small batch's gradient has to be divided by the number of small batches before it is added, or the total comes out too large.

Prove the two are the same.

The gradients agree to about 7 in the ninth decimal place, which is ordinary float32 rounding, not a difference in behaviour. And the accumulated version held a quarter of the activations.

Where the equivalence breaks:

• BatchNorm. It computes statistics over whatever is in the current small batch, so four batches of 64 normalise differently from one batch of 256. LayerNorm, which Module 6's Transformer block uses, works one row at a time and is unaffected.
• Forgetting to divide. Gradients come out four times too big, which behaves exactly like a four-times learning rate. Lesson 3.4 Case 1 shows what that looks like.
• Anything else you average inside the loop, such as a logged metric.

in Module 8 this is a routine setting. Fine-tuning recipes commonly pair a batch of 1 or 2 with 8 to 16 accumulation steps, to get an effective batch of 16 or 32 on one consumer GPU.

Mixed precision means running some operations in a smaller number format while keeping the weights themselves in fp32. PyTorch does this with torch.autocast, which picks a sensible format per operation: matrix multiplications in bf16 or fp16, and precision-sensitive things like losses in fp32.

On CPU, autocast supports bf16. Time it — remembering that the speed is a property of this machine.

Note what autocast does not do. The weights stay fp32. Only the operations run in the smaller format, and the results come back smaller. So the memory saving is in the activations, not in the model.

On the machine that built this notebook, bf16 ran roughly three times as many steps per second. On yours it may be slower, especially on a CPU without bf16 instructions. Speed from mixed precision is a hardware property, not a free gift. Measure it on the device you will really train on. The memory saving, by contrast, is arithmetic, and is always there.

Check against a result you already know: the Bank Marketing model from Lessons 4.1 and 4.2. First the data, in three compressed cells.

Now one training function with three settings, so we can put autocast in different places and compare:

• "fp32" — no autocast at all;
• "forward" — autocast around the forward pass only, which is how it is meant to be used;
• "whole" — autocast wrapped around the entire training loop, which looks tidier and is the mistake.

Read those four lines carefully, because the third one is the lesson.

The first two are identical to four decimal places: 0.7847. Running the forward pass in bf16 cost this model nothing.

The third dropped to 0.7584. The fourth, differing only by cache_enabled=False, came back to 0.7848.

So this is not a precision problem at all. It is a caching problem. Inside an autocast block, PyTorch caches the reduced-precision copies it makes of your fp32 weights, so that a weight used by several operations is converted once rather than repeatedly. That cache is designed for a block covering one forward pass. Stretch the block across many optimizer steps and the model keeps reusing copies made from older weights, so part of the network is training on stale values. Nothing errors. The loss curve looks plausible.

Put autocast around the forward pass, not around the loop. In an earlier draft of this lesson the block wrapped an entire 30-epoch run including the scoring, and validation AUC fell to 0.4079 — worse than guessing, a collapse that looks exactly like Lesson 3.4 Case 1 but has a completely different cause.

One more consequence, following from Lesson 4.2: mixed precision costs you bit-for-bit repeatability, because the arithmetic itself has changed. The run is still reproducible at Level 2, as long as the record says which precision was used. Make it a setting in your config rather than a switch you flip by hand.

Everything so far was arithmetic that holds anywhere. Three things are specific to an NVIDIA GPU, and these cells run only if one is present: reading how much GPU memory is actually in use, meeting a real out-of-memory error, and training in fp16 with a loss scaler.

VRAM is the GPU's own memory, separate from your computer's RAM. It is usually much smaller: 15 GB on a free Colab T4, against 16 GB or more of ordinary RAM on a laptop.

If you are on a laptop, run this section on a free Colab or Kaggle GPU. It is the same code; the cells below detect the machine and skip. Module 8's QLoRA work and Module 9's serving work need an NVIDIA GPU too, so it is worth doing once now, while the code is short enough to read.

Start with the recovery ladder, which is not GPU-specific logic at all and so can be tested anywhere: try the batch size you want; if it fails with an out-of-memory error, clean up and halve it.

On a GPU you pass the real training step and the real error type instead of the pretend ones. That is what the next cells do.

Three details in that recovery are worth remembering.

del on its own is not enough. Python releases the tensor, but PyTorch keeps the memory in its own pool to reuse later, so a tool like nvidia-smi still shows it as used. torch.cuda.empty_cache() gives it back to the driver. That matters when another process needs it; inside your own process the pool is a feature, not a leak.

A caught error can keep your memory alive. An exception carries a traceback, and the traceback holds every local variable in every function it passed through — including the huge tensors that caused the failure. In a notebook the variables from the failed cell survive too. Restarting the kernel after a bad out-of-memory error is not superstition.

nvidia-smi and torch.cuda.memory_allocated() measure different things. The first shows what the process has reserved from the driver; the second shows what your tensors are actually using. A large gap between them is the pool, not a leak.

Finally, fp16 training with the loss scaler from Section 3.

scaler.step() sometimes skips the optimizer step entirely. That happens when a gradient overflowed to infinity, in which case the step would be meaningless, so it is thrown away and the multiplier is lowered. Skipping a step occasionally is the correct behaviour, not a bug.

MPS is PyTorch's backend for the graphics processor built into Apple Silicon chips (M1 through M5). Two differences from an NVIDIA GPU matter here.

Memory is unified. The graphics processor shares the machine's RAM rather than having its own separate VRAM. So there is no separate memory figure to read, and no clean wall to hit. An oversized job tends to push the whole machine into swapping and slow everything down, rather than raising a tidy error. That is why this lesson does not include a deliberate out-of-memory demonstration for MPS. Raise batch sizes gradually and watch memory pressure in Activity Monitor.

Coverage is good but not complete. An operation with no MPS implementation raises an error, unless you set PYTORCH_ENABLE_MPS_FALLBACK=1, which quietly runs it on the CPU instead. That is correct but sometimes much slower than staying on the CPU for the whole job.

Worth planning around: bitsandbytes, vLLM and flash-attention are NVIDIA-only. Module 8's QLoRA work and Module 9's serving work go to Colab or Kaggle. Module 8 has a LoRA-only path that does run on a Mac.

Lesson 3.4 ordered its checks by cost, cheapest first. The same discipline applies here, because these fixes differ enormously in what they cost you.

Tier 0 — free, and changes nothing about the model

1. Restart the kernel or the process. Clears dead tensors held alive by tracebacks and by old notebook state.
2. Check you are not keeping tensors alive by accident. Appending loss to a list instead of loss.item() keeps the entire computation graph for that batch, epoch after epoch. It looks like a slow leak and it is the commonest cause.
3. Wrap scoring in torch.no_grad(). Without it, evaluation stores activations for a backward pass that never comes.
4. Free what you no longer need: del, then gc.collect(), then empty_cache().

Tier 1 — cheap, changes the run but not the model

5. Halve the batch size, and use the ladder from Section 8 to find the largest that fits.
6. Add gradient accumulation to get the effective batch back, as in Section 6.
7. Turn on mixed precision. Halves the activations, and may help speed if the hardware supports it.
8. Shorten the sequence length if you are working with text. The cost grows faster than linearly.

Tier 2 — changes the computation or the model

9. Gradient checkpointing: recompute activations instead of storing them, at roughly 30% more time.
10. Freeze most of the model, training only a head or small adapters. This is where LoRA belongs, and it removes gradients and optimizer state for everything frozen (Module 8).
11. Quantise the frozen weights, which is QLoRA in Module 8, or quantise for serving in Module 9.
12. A smaller model, or a bigger GPU. Last, not first, because both cost money or accuracy.

Work down the list. Most people start at 12, or at 5 with a guess. The first four cost nothing and fix a surprising share of real failures.

Training with AdamW needs about four times the model's own memory, before a single activation: weights, gradients, and two running averages. That factor is the reason Module 8 exists.

You can predict memory before running anything. Parameter count, times bytes per number, times the optimizer's copies. The 7B rows in Section 2 show why running such a model and training it sit on opposite sides of a hardware boundary.

Precision is a trade between range and fineness. bf16 keeps fp32's range and needs no loss scaler; fp16 keeps more fineness but its narrow range means gradients vanish without one.

Activations are the batch-dependent term, linear in batch size and constant per row. So batch size is the fastest memory control you have, and gradient accumulation gives back the effective batch you gave up.

Speed from mixed precision is a hardware property; the memory saving is arithmetic. Measure the first, count on the second.

Put autocast around the forward pass. Wrapping the whole loop cost 0.026 AUC here, silently, because of a weight cache that was never meant to live that long.

An out-of-memory error has a playbook. Restart, look for retained graphs, use no_grad, then batch size and accumulation, then checkpointing and freezing. Hardware last.

You now have three of the four pieces of a defensible experiment: data you can trust (4.1), a run that records itself and can be repeated (4.2), and a clear idea of what a run costs in memory (4.3).

Lesson 4.4 spends them. A gradient-boosting model and the Module 3 neural network compete on Bank Marketing, on identical splits, with the tuning budget written down in advance and every claim checked against the noise you measured in 4.2. The blueprint's instruction for that contest is explicit: the conclusion is not allowed to be "neural networks win". It has to be a decision you could defend to somebody who disagrees with you.