Neural networks, deep learning, NLP · insight-0469

Run the first training step on the largest batch and longest sequence you will see

Situation
A text model trains with dynamic padding, and batch memory depends on the longest document in each batch.
Action
Before the real run, build a synthetic batch at the maximum batch size and maximum sequence length, run forward and backward, and record torch.cuda.max_memory_allocated(). Set batch size, length cap or token-based batching from that peak plus a margin, and checkpoint regularly.
Reason
Activation memory grows with batch size and sequence length, and standard attention memory grows quadratically with length, so the worst case is set by the longest batch, not the average one.
Boundary
Training and batch inference with variable-length inputs or variable image sizes.

Symptom or risk

Early batches fit, then hours into training a batch with a rare very long document triggers an out-of-memory error and the run is lost.

Fixed-shape inputs, where every batch has the same footprint.

A hard length cap may truncate useful text; token-based batching keeps memory flat but complicates the loader.

Testing the peak in evaluation mode without gradients, which needs far less memory than training.

PyTorch exposes max_memory_allocated and related counters to monitor peak memory occupied by tensors.
CUDA semantics (PyTorch documentation)
Dao et al. note that self-attention time and memory are quadratic in sequence length in standard implementations.
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness