Running Gradian

Why does gradient indexing run out of memory?

Short answer

The peak allocation is the fp32 upcast of the logits, sized batch_size x supervised_tokens x vocab x 4 bytes. Lower grads.batch_size first, then grads.max_seq_length. Those are the only two settings that touch the tensor that actually fails. The settings that shrink the stored index on disk, storage_dtype, module_filter and projection_dim, will not rescue this error.

The tensor that fails#

One allocation dominates, and it is the fp32 upcast of the logits:

peak_bytes = batch_size * supervised_tokens * vocab_size * 4

With a 128k-vocab Llama-3 tokenizer at batch 8 and 512 supervised tokens, that single tensor is about 2.1 GB, on top of the fp16 logits it was upcast from.

Note that the term is supervised_tokens, not sequence length. Gradian gathers the supervised positions before the upcast, so prompt tokens and padding cost nothing here. On instruction data, where the prompt is often most of the sequence, that is a large reduction. It also means two runs at the same max_seq_length can have very different peak memory when one has longer answers than the other.

What to lower#

grads:
  batch_size: 2        # first lever, linear in the peak
  max_seq_length: 512  # second lever, bounds supervised_tokens

Batch size first, because it is linear in the allocation and costs you nothing but wall time. Sequence length second, because lowering it far enough to matter risks truncating targets, which creates a different and worse problem. The truncated-completions run shows what that looks like when it goes unnoticed.

Measured on real hardware: a 16GB V100 runs out of memory at a gradient batch of 8 on a 1B model, and runs clean at 2.

What will not help#

Three settings look like memory knobs and are not, at least not for this error:

setting	what it shrinks
`grads.storage_dtype`	the index on disk
`grads.module_filter`	the index on disk
`grads.projection_dim`	the index on disk

All three reduce the stored index, not the live tensor that fails during indexing. Reach for them when the index does not fit on disk, which is a separate problem with separate trade-offs. module_filter at down_proj is about 4.3x smaller, storage_dtype: float16 roughly halves it at almost no accuracy cost, and projection_dim is independent of model size but does cost ranking fidelity, so tune it rather than assuming a width.

Before you tune anything#

Run gradian doctor. It reports installed versions against the supported window, visible devices, determinism settings, and which backends are available along with the reason any are not. Its output is also the right thing to paste into an issue.

The configuration reference documents every setting above, and troubleshooting covers the other errors you are most likely to hit.