p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] utils/train_utils.py (python)
  """Helper function to calculate matmul TFLOP in ffn based on MLP dimension.

  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.
  """
[content] utils/train_utils.py (python)
    to_tflops = 1.0 / 1e12
    result = {
        "encoder_attention_tflops": ((qkv_spatial + attn_spatial + proj_spatial + qkv_temporal + attn_temporal + proj_temporal) * num_blocks) * to_tflops * 3.0,
        "encoder_ffn_tflops": (ffn_block * num_blocks) * to_tflops * 3.0,
        "encoder_io_tflops": encoder_io * to_tflops * 3.0,
        "decoder_attention_tflops": ((qkv_spatial + attn_spatial + proj_spatial + qkv_temporal + attn_temporal + proj_temporal) * num_blocks) * to_tflops * 3.0,
        "decoder_ffn_tflops": (ffn_block * num_blocks) * to_tflops * 3.0,
        "decoder_io_tflops": decoder_io * to_tflops * 3.0,
        "vq_tflops": vq_flops * to_tflops,
        "total_tflops": total_training_flops * to_tflops,
        "per_device_batch_size": batch_per_device,
        "num_patches_per_frame": N,
    }
[content] utils/train_utils.py (python)
    # Per-device batch
    num_devices = max(1, jax.device_count())
    batch_per_device = args.batch_size // num_devices

    # Shapes
    T = int(args.seq_len)
    H = int(args.image_height)
    W = int(args.image_width)
    patch = int(args.patch_size)
    # Number of patches per frame (ceil division with padding as in patchify)
    hn = (H + patch - 1) // patch
    wn = (W + patch - 1) // patch
    N = int(hn * wn)

    # Model dims
    M = int(args.model_dim)
    D_ffn = int(args.ffn_dim)
    num_blocks = int(args.num_blocks)

    # Convenience aliases
    B = float(batch_per_device)
    T_f = float(T)
    N_f = float(N)
    M_f = float(M)
    D_ffn_f = float(D_ffn)

    # Token counts for projections
    tokens_spatial = B * T_f * N_f           # spatial attention runs per frame over N tokens
    tokens_temporal = B * N_f * T_f          # temporal attention runs per spatial position over T tokens

    # Per-block FLOPs
    # Spatial attention (non-causal)
    qkv_spatial = 3.0 * 2.0 * tokens_spatial * M_f * M_f
    attn_spatial = 4.0 * (B * T_f) * (N_f * N_f) * M_f
    proj_spatial = 2.0 * tokens_spatial * M_f * M_f

    # Temporal attention (causal mask halves the non-causal complexity)
    qkv_temporal = 3.0 * 2.0 * tokens_temporal * M_f * M_f
    attn_temporal = 2.0 * (B * N_f) * (T_f * T_f) * M_f  # = 0.5 * (4 * B*N * T^2 * M)
    proj_temporal = 2.0 * tokens_temporal * M_f * M_f

    # FFN per block (2 matmuls)
    ffn_block = 4.0 * B * T_f * N_f * M_f * D_ffn_f

    block_total = (
        qkv_spatial
        + attn_spatial
        + proj_spatial
        + qkv_temporal
        + attn_temporal
        + proj_temporal
        + ffn_block
    )

    # Encoder/Decoder IO (input/output linears executed once per forward)
    P = float(args.image_channels * (patch ** 2))      # patch token dim
    L = float(args.latent_dim)                          # latent dim

    # Encoder: input_dim=P -> M, output_dim=M -> L
    encoder_io = 2.0 * B * T_f * N_f * (P * M_f + M_f * L)
    # Decoder: input_dim=L -> M, output_dim=M -> P
    decoder_io = 2.0 * B * T_f * N_f * (L * M_f + M_f * P)

    # VQ distance compute (forward only; gradients do not flow through distances with STE)
    K = float(args.num_latents)
    vq_flops = 2.0 * (B * T_f * N_f) * L * K

    # Forward FLOPs
    encoder_forward = num_blocks * block_total + encoder_io
    decoder_forward = num_blocks * block_total + decoder_io
    forward_flops = encoder_forward + decoder_forward
[content] utils/train_utils.py (python)

    # Encoder/Decoder IO (input/output linears executed once per forward)
    P = float(args.image_channels * (args.patch_size ** 2))      # patch token dim
    L = float(args.latent_dim)                          # latent dim
[content] utils/train_utils.py (python)

    # Encoder: input_dim=P -> M, output_dim=M -> L
    encoder_io = 2.0 * num_tokens * (P * args.model_dim + args.model_dim * L)
    # Decoder: input_dim=L -> M, output_dim=M -> P
    decoder_io = 2.0 * num_tokens * (L * args.model_dim + args.model_dim * P)

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

    # Forward FLOPs
    encoder_forward = args.num_blocks * block_total + encoder_io
    decoder_forward = args.num_blocks * block_total + decoder_io
    forward_flops = encoder_forward + decoder_forward

    # 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_attention_tflops": ((qkv_spatial + attn_spatial + proj_spatial + qkv_temporal + attn_temporal + proj_temporal) * num_blocks) * to_tflops * 3.0,
        "encoder_ffn_tflops": (ffn_block * args.num_blocks) * to_tflops * 3.0,
        "encoder_io_tflops": encoder_io * to_tflops * 3.0,
        "decoder_attention_tflops": ((qkv_spatial + attn_spatial + proj_spatial + qkv_temporal + attn_temporal + proj_temporal) * num_blocks) * to_tflops * 3.0,
        "decoder_ffn_tflops": (ffn_block * args.num_blocks) * to_tflops * 3.0,
        "decoder_io_tflops": decoder_io * to_tflops * 3.0,
        "vq_tflops": vq_flops * to_tflops,
        "total_tflops": total_training_flops * to_tflops,
        "per_device_batch_size": batch_per_device,
        "num_patches_per_frame": N,
    }

    return result
[content] utils/train_utils.py (python)

    # Decoder: input_dim=L -> M, output_dim=M -> P
    decoder_io = 2.0 * num_tokens * (L * args.model_dim + args.model_dim * P)

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

    # Forward FLOPs
    encoder_forward = args.num_blocks * block_total + encoder_io
    decoder_forward = args.num_blocks * block_total + decoder_io
    forward_flops = encoder_forward + decoder_forward

    # 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_attention_tflops": ((qkv_spatial + attn_spatial + proj_spatial + qkv_temporal + attn_temporal + proj_temporal) * num_blocks) * to_tflops * 3.0,
        "encoder_ffn_tflops": (ffn_block * args.num_blocks) * to_tflops * 3.0,
        "encoder_io_tflops": encoder_io * to_tflops * 3.0,
        "decoder_attention_tflops": ((qkv_spatial + attn_spatial + proj_spatial + qkv_temporal + attn_temporal + proj_temporal) * num_blocks) * to_tflops * 3.0,
        "decoder_ffn_tflops": (ffn_block * args.num_blocks) * to_tflops * 3.0,
        "decoder_io_tflops": decoder_io * to_tflops * 3.0,
        "vq_tflops": vq_flops * to_tflops,
        "total_tflops": total_training_flops * to_tflops,
        "per_device_batch_size": batch_per_device,
        "num_patches_per_frame": N,
    }

    return result