bridge.training.utils.flop_utils#

Module Contents#

Functions#

num_floating_point_operations

Calculate the theoretical number of floating-point operations for a training step.

API#

bridge.training.utils.flop_utils.num_floating_point_operations(
cfg: megatron.bridge.training.config.ConfigContainer,
batch_size: int,
) → float#

Calculate the theoretical number of floating-point operations for a training step.

Computes the FLOPs based on model configuration (hidden size, layers, vocab size, attention specifics, MLP size, MoE config) and batch size.

Parameters:

  • cfg (ConfigContainer) – The main configuration container.

  • batch_size (int) – The global batch size for the training step.

Returns:

The estimated number of floating-point operations for one training step (forward + backward pass).

Return type:

float