[content] jasmine/utils/train_utils.py (python)


def _calculate_ffn_tflops_per_device(args):
    """Helper function to calculate matmul TFLOP in ffn based on MLP dimension.
    Adapted from https://github.com/AI-Hypercomputer/maxtext/blob/7898576359bacde81be25cb3038e348aac1f943b/MaxText/max_utils.py#L473

    Applies to:
      - Dense FFN layers (mlp_dim = config.mlp_dim).
      - MoE FFN layers (mlp_dim = config.moe_mlp_dim),
        need to scale by shared_experts or num_experts_per_tok.
    """


def calculate_tflops_per_device_tokenizer(args: TokenizerArgs):
    """Calculate training TFLOP for the Tokenizer."""
    input_dim = args.image_channels * (args.patch_size**2)

    tflops_per_device_encoder = calculate_tflops_per_device_transformer(
        args,
        input_dim,
        args.model_dim,
        args.ffn_dim,
        args.latent_dim,
        fully_causal=False,
    )
    tflops_per_device_decoder = calculate_tflops_per_device_transformer(
        args,
        args.latent_dim,
        args.model_dim,
        args.ffn_dim,
        input_dim,
        fully_causal=False,
    )

    batch_per_device = args.batch_size // jax.device_count()
    hn = (args.image_height + args.patch_size - 1) // args.patch_size
    wn = (args.image_width + args.patch_size - 1) // args.patch_size
    N = int(hn * wn)
    L = float(args.latent_dim)
    K = float(args.num_latents)

    # VQ distance compute (forward only; gradients do not flow through distances with STE)
    num_tokens = batch_per_device * args.seq_len * N
    vq_flops = 2.0 * num_tokens * L * K

    forward_flops = tflops_per_device_encoder + tflops_per_device_decoder

    # Training FLOPs: scale forward by ~3 for fwd+bwd+update, add VQ once
    total_training_flops = 3.0 * forward_flops + vq_flops

    to_tflops = 1.0 / 1e12
    result = {
        "encoder_tflops": tflops_per_device_encoder * to_tflops * 3.0,
        "decoder_tflops": tflops_per_device_decoder * to_tflops * 3.0,
        "vq_tflops": vq_flops * to_tflops * 3.0,
        "total_tflops": total_training_flops * to_tflops,
    }

    return result


def calculate_tflops_per_device_transformer(
    args: TokenizerArgs | LAMArgs | DynamicsArgs,
    input_dim: int,
    model_dim: int,
    ffn_dim: int,
    out_dim: int,
    fully_causal: bool = False,
):
    """Calculate training TFLOP for the STTransformer."""
    num_devices = jax.device_count()
    batch_per_device = args.batch_size // num_devices

    T = int(args.seq_len)
    H = int(args.image_height)
    W = int(args.image_width)

    # Number of patches per frame
    hn = (H + args.patch_size - 1) // args.patch_size
    wn = (W + args.patch_size - 1) // args.patch_size
    N = int(hn * wn)

    # Token counts for projections
    num_tokens = batch_per_device * T * N

    # Per-block FLOPs
    # Spatial attention (non-causal)
    qkv_spatial = 3.0 * 2.0 * num_tokens * model_dim * model_dim
    causal_multiplier = 0.5 if fully_causal else 1.0
    attn_spatial = (
        causal_multiplier * 2.0 * 2.0 * (batch_per_device * T) * (N * N) * model_dim
    )
    proj_spatial = 2.0 * num_tokens * model_dim * model_dim

    # Temporal attention (causal)
    qkv_temporal = 3.0 * 2.0 * num_tokens * model_dim * model_dim
    attn_temporal = (
        2.0 * (batch_per_device * N) * (T * T) * model_dim
    )  # = 0.5 * (2 * 2 * B*N * T^2 * M)
    proj_temporal = 2.0 * num_tokens * model_dim * model_dim

    # FFN per block (2 matmuls)
    # FIXME (f.srambical):
    ffn_block = 4.0 * batch_per_device * T * N * model_dim * ffn_dim

    # NOTE (f.srambical): The calculation omits layer norm flops
    block_total = (
        qkv_spatial
        + attn_spatial
        + proj_spatial
        + qkv_temporal
        + attn_temporal
        + proj_temporal
        + ffn_block
    )

    # FIXME (f.srambical):
    input_output = 2.0 * num_tokens * (input_dim * model_dim + model_dim * out_dim)
    return block_total + input_output