Memory Management › Key Parameters

The key configuration parameters that control memory usage during transformer training are:

• batch_size: Controls activation memory linearly. Halving batch size halves activation memory but may require more gradient accumulation steps to maintain the effective batch size.
• seq_len / max_position_embeddings: Controls both the $T^2$ attention term and the linear activation terms. Doubling sequence length roughly quadruples attention activation memory and doubles all other activation memory, for an overall factor slightly less than 4x depending on the model's dimension.
• num_layers: Controls how many layers' activations must be stored simultaneously. Gradient checkpointing amortizes this to one layer regardless of depth.
• hidden_dim: Controls the linear activation memory per token and the parameter count per layer. Scaling hidden dimension has a quadratic effect on parameter count (through the $d^2$ attention projection terms) and a linear effect on per-token activation memory.
• optimizer: Adam requires 8 bytes per parameter for states; alternatives like Lion (4 bytes), SGD (4 bytes), or 8-bit Adam (2 bytes) can significantly reduce this. AdaFactor's factored representation provides sub-linear scaling but requires careful learning rate tuning.
• param_dtype: BF16 parameters halve parameter memory compared to FP32. Mixed precision typically uses FP16/BF16 parameters with FP32 gradients and optimizer states.
• gradient_checkpointing: Reduces activation memory from $O(L \times B \times T \times d)$ to $O(B \times T \times d)$ at the cost of approximately 33% more forward pass compute.