Skip to main content

lib.trainer.torch.utils

Memory-footprint estimators for PyTorch / Transformers models.

get_total_training_memory_transformers​

def get_total_training_memory_transformers(model: PreTrainedModel,
                                           batch_size: int,
                                           sequence_length: int) -> float

Estimate the total training memory (in MB) for a Transformers model.

Uses the formula from the EleutherAI Transformer Math reference.

Arguments:

  • model - A Hugging Face PreTrainedModel with config.hidden_size / num_hidden_layers / num_attention_heads and torch_dtype.
  • batch_size - Training batch size.
  • sequence_length - Input sequence length per sample.

Returns:

Estimated total training memory in MB, including a 20% buffer for fragmentation overhead.

Reference: https://blog.eleuther.ai/transformer-math/

estimate_activation_memory_non_transformer​

def estimate_activation_memory_non_transformer(layer_output_dims: dict,
                                               batch_size: int,
                                               bytes_per_value: int) -> float

Estimate activation memory (MB) given captured per-layer output shapes.

Arguments:

  • layer_output_dims - Mapping of nn.Module -> tensor shape captured via a forward hook.
  • batch_size - Training batch size.
  • bytes_per_value - Bytes per value in the activation tensor.

Returns:

Total activation memory in MB.

get_total_training_memory_nn_module​

def get_total_training_memory_nn_module(model: torch.nn.Module,
                                        batch_size: int,
                                        input_size: int) -> float

Estimate the total training memory (in MB) for a generic nn.Module.

Registers forward hooks on Linear / Conv* / Norm* / RNN* layers to capture activation shapes, then sums parameter + gradient + optimizer + activation memory.

Arguments:

  • model - The model to size.
  • batch_size - Training batch size.
  • input_size - Flat input size used to generate a sample input tensor.

Returns:

Estimated total training memory in MB, including a 20% buffer for fragmentation overhead.