Machine FLOPs Utilization (MFU)

def openai_flops_per_token(n_layers, n_heads, d_model, n_ctx, n_vocab, ff_ratio=4):
    """
    Open AI method for forward pass FLOPs counting of decoder-only Transformer
    """
    d_attn = d_model // n_heads
    d_ff = d_model * ff_ratio
 
    embeddings = 4 * d_model
    attn_qkv = 2 * n_layers * d_model * 3 * (d_attn * n_heads)
    attn_mask = 2 * n_layers * n_ctx * (d_attn * n_heads)
    attn_project = 2 * n_layers * (d_attn * n_heads) * d_model
    ff = 2 * n_layers * 2 * d_model * d_ff
    logits = 2 * d_model * n_vocab
 
    return embeddings + attn_qkv + attn_mask + attn_project + ff + logits

def deepmind_flops_per_sequence(n_layers, n_heads, d_model, n_ctx, n_vocab, ff_ratio=4):
    """
    DeepMind method for forwad pass FLOPs counting of decoder-only Transformer
    """
    d_attn = d_model // n_heads
    d_ff = d_model * ff_ratio
 
    embeddings = 2 * n_ctx * n_vocab * d_model
 
    attn_qkv = 2 * n_ctx * 3 * d_model * (d_attn * n_heads)
    attn_logits = 2 * n_ctx * n_ctx * (d_attn * n_heads)
    attn_softmax = 3 * n_heads * n_ctx * n_ctx
    attn_reduce = 2 * n_ctx * n_ctx * (d_attn * n_heads)
    attn_project = 2 * n_ctx * (d_attn * n_heads) * d_model
    total_attn = attn_qkv + attn_logits + attn_softmax + attn_reduce + attn_project
 
    ff = 2 * n_ctx * (d_model * d_ff + d_model * d_ff)
 
    logits = 2 * n_ctx * d_model * n_vocab
 
    return embeddings + n_layers * (total_attn + ff) + logits

def karpathy_flops(n_layer, n_embd, n_head, block_size):
    # we only count Weight FLOPs, all other layers (LayerNorm, Softmax, etc) are effectively irrelevant
    # we count actual FLOPs, not MACs. Hence 2* all over the place
    # basically for any matrix multiply A (BxC) @ B (CxD) -> (BxD) flops are 2*B*C*D

    out = OrderedDict()
    head_size = n_embd // n_head

    # attention blocks
    # 1) the projection to key, query, values
    out['attention/kqv'] = 2 * block_size * (n_embd * 3*n_embd)
    # 2) calculating the attention scores
    out['attention/scores'] = 2 * block_size * block_size * n_embd
    # 3) the reduction of the values (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
    out['attention/reduce'] = 2 * n_head * (block_size * block_size * head_size)
    # 4) the final linear projection
    out['attention/proj'] = 2 * block_size * (n_embd * n_embd)
    out['attention'] = sum(out['attention/'+k] for k in ['kqv', 'scores', 'reduce', 'proj'])

    # MLP blocks
    ffw_size = 4*n_embd # feed forward size
    out['mlp/ffw1'] = 2 * block_size * (n_embd * ffw_size)
    out['mlp/ffw2'] = 2 * block_size * (ffw_size * n_embd)
    out['mlp'] = out['mlp/ffw1'] + out['mlp/ffw2']

    # the transformer and the rest of it
    out['block'] = out['attention'] + out['mlp']
    out['transformer'] = n_layer * out['block']
    out['dense'] = 2 * block_size * (n_embd * vocab_size)

    # forward,backward,total
    out['forward_total'] = out['transformer'] + out['dense']
    out['backward_total'] = 2 * out['forward_total'] # use common estimate of bwd = 2*fwd
    out['total'] = out['forward_total'] + out['backward_total']

    return out
    
# compare our param count to that reported by PyTorch
f = karpathy_flops()
flops_total = f['forward_total']
print(f"{'name':20s} {'flops':14s} {'ratio (%)':10s}")
for k,v in f.items():
    print(f"{k:20s} {v:14d} {v/flops_total*100:10.4f}")

def palm_flops():
    """estimate of the model flops following PaLM paper formula"""
    # non-embedding model parameters. note that we do not subtract the
    # embedding/token params because those are tied and get used in the last layer.
    N = params()['total'] - params()['emebedding/position']
    L, H, Q, T = n_layer, n_head, n_embd//n_head, block_size
    mf_per_token = 6*N + 12*L*H*Q*T
    mf = mf_per_token * block_size
    return mf

print(f"palm_flops: {palm_flops():d}, flops: {flops()['total']:d}, ratio: {palm_flops()/flops()['total']:.4f}")

def simplified_mfu(
    num_params,
    TP_degree,
    MP_degree,
    batch_size,
    block_size,
    fwdbwd_per_iter,
    grad_accum,
):
    # a100_promised = 300e12*MP_degree*TP_degree
    a100_promised = 312e12*MP_degree*TP_degree # bfloat16
    achieved = (6*num_params/fwdbwd_per_iter)*(batch_size*block_size)*grad_accum # 6ND is approximated FLOPs per token for forward + backward 
    mfu = achieved/a100_promised

    print(f"num_params: {num_params * 1e-9}B")
    print(f"fwdbwd_per_iter: {fwdbwd_per_iter}")
    print(f"MFU: {mfu*100:.2f}%")
    print()
    return mfu

def estimate_mfu(
   num_params,
   n_layer,
   n_head,
   n_embd,
   block_size,
   fwdbwd_per_iter, 
   dt
):
   """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
   # first estimate the number of flops we do per iteration.
   # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
   N = num_params
   L, H, Q, T = n_layer, n_head, n_embd//n_head, block_size

   flops_per_token = 6*N + 12*L*H*Q*T
   flops_per_fwdbwd = flops_per_token * T
   flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter

   # express our flops throughput as ratio of A100 bfloat16 peak flops
   flops_achieved = flops_per_iter * (1.0/dt) # per second
   flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
   mfu = flops_achieved / flops_promised

   print(f"MFU: {mfu}")
   return mfu

# https://github.com/pytorch/torchtitan/blob/main/docs/performance.md

non_embedding_params = 8e9
# non_embedding_params = 8e9 - 1*4096*128256 # torch titan exclude embedding 1 times
# non_embedding_params = 8e9 - 2*4096*128256 # if embedding is tied and exclude them all

TP_degree, MP_degree = 1, 1
micro_batch_size, seq_len, accum, num_nodes = 1, 8192, 1, 8
num_steps = 3000
num_tokens = num_steps * micro_batch_size * accum * num_nodes * 8 * seq_len
wps = 2904
fwdbwd_per_iter = 1/wps*micro_batch_size*seq_len # 2.82 sec
# fwdbwd_per_iter = 2.3*3600/num_steps
# wps = micro_batch_size*seq_len/fwdbwd_per_iter

args = {
   'num_params': non_embedding_params,
   'TP_degree': TP_degree,
   'MP_degree': MP_degree,
   'batch_size': micro_batch_size,
   'block_size': seq_len,
   'fwdbwd_per_iter': fwdbwd_per_iter, # sec
   'grad_accum': accum,
}
print(f'non_embedding_params: {non_embedding_params/1e9:.2f}B')
print(f'num_tokens: {num_tokens/1e9:.2f}B')
print(f'wps: {wps:.2f}')
mfu = simplified_mfu(**args)

additional_mfu = (12*32*32*128*seq_len)*wps/312e12
print(f'for torchtitan, approx {100*additional_mfu:.2f} should be added')
print(f'so, PaLM MFU is {100*(mfu+additional_mfu):.2f}')

# https://github.com/pytorch/torchtitan/blob/81c555f8bd6791ccd575a52bdcaaf816f4de7ee6/train.py#L345C1-L354C70
def simplified_train_loop_with_mfu(dataloader, model, model_config, optimizer, seq_len):
    
    # https://github.com/pytorch/torchtitan/blob/81c555f8bd6791ccd575a52bdcaaf816f4de7ee6/torchtitan/utils.py#L118C1-L133C26
    def get_num_flop_per_token(num_params: int, model_config, seq_len) -> int:
        l, h, q, t = (
            model_config.n_layers,
            model_config.n_heads,
            model_config.dim // model_config.n_heads,
            seq_len,
        )
        # Reasoning behind the factor of 12 for the self-attention part of the formula:
        # 1. each self-attention has 2 matmul in the forward and 4 in the backward (6)
        # 2. the flash attention does 1 more matmul recomputation in the backward
        #    but recomputation should not be counted in calculating MFU           (+0)
        # 3. each matmul performs 1 multiplication and 1 addition                 (*2)
        # 4. we follow the convention and do not account for sparsity in causal attention
        flop_per_token = 6 * num_params + 12 * l * h * q * t
        return flop_per_token

    # get num_params, flops per token
    num_params = get_num_params(model)
    num_flop_per_token = get_num_flop_per_token(num_params, model_config, seq_len)
    gpu_peak_flops = utils.get_peak_flops(gpu_memory_monitor.device_name) # 312e12
    
    # training loop
    time_last_log = time.perf_counter()
    for inputs, labels in dataloader():

        # num tokens
        ntokens_since_last_log += labels.numel()
        
        # impl fwd+bwd and count time
        optimizer.zero_grad()
        loss = model(inputs)
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        time_delta = time.perf_counter() - time_last_log

        # tokens per second, abbr. as wps by convention (model parallel (tensor parallel) degree should be count)
        # https://github.com/pytorch/torchtitan/blob/81c555f8bd6791ccd575a52bdcaaf816f4de7ee6/torchtitan/parallelisms/parallel_dims.py#L69-L70
        wps = ntokens_since_last_log / (time_delta * parallel_dims.model_parallel_size)
        
        # model FLOPS utilization
        # For its definition and calculation, please refer to the PaLM paper:
        # https://arxiv.org/abs/2204.02311
        mfu = 100 * num_flop_per_token * wps / gpu_peak_flops # \approx 100 * 6 * N * D / fwd_bwd_iter_time / (312e12 * MP_degree * TP_degree)
        time_last_log = time.perf_counter()

def cal_with_mfu(N, D, nodes, mfu):
    C=6*N*D
    a100_promised=312e12
    gpus=nodes*8
    training_hours=C/(a100_promised*mfu*gpus)/3600
    training_days=training_hours/24
    print(f'training_hours: {training_hours:.2f} hours (computed by MFU)')
    print(f'training_days: {training_days:.2f} days (computed by MFU)')
    
def cal_with_iter_time(D, bsz, fwd_bwd_time):
    num_steps=D//bsz
    training_hours=num_steps*fwd_bwd_time/3600
    training_days=training_hours/24 # should be sec
    print(f'training_hours: {training_hours:.2f} hours (computed by fwd bwd iter time)')
    print(f'training_days: {training_days:.2f} days (computed by fwd bwd iter time)')

cal_with_mfu(non_embedding_params, num_tokens, num_nodes, mfu)
cal_with_iter_time(num_tokens, micro_batch_size*accum*num_nodes/TP_degree/MP_degree*8*seq_len, fwdbwd_per_iter)

# https://github.com/NVIDIA/Megatron-LM/blob/54f1f78529cbc2b9cddad313e7f9d96ac0420a27/megatron/training/training.py#L90
def num_floating_point_operations_megatron_lm(
    padded_vocab_size,
    num_layers,
    hidden_size,
    num_attention_heads,
    group_query_attention,
    num_query_groups,
    kv_channels,
    ffn_hidden_size,
    swiglu,
    num_experts,
    moe_router_topk,
    
    world_size,
    global_batch_size,
    seq_length,

    elapsed_time_per_iteration,
):
    # Attention projection size.
    query_projection_size = kv_channels * num_attention_heads
    query_projection_to_hidden_size_ratio = query_projection_size / hidden_size
    # Group Query Attention.
    if not group_query_attention:
        num_query_groups = num_attention_heads
    # MoE.
    num_experts_routed_to = 1 if num_experts is None else moe_router_topk
    gated_linear_multiplier = 3 / 2 if swiglu else 1

    # The 12x term below comes from the following factors; for more details, see
    # "APPENDIX: FLOATING-POINT OPERATIONS" in https://arxiv.org/abs/2104.04473.
    # - 3x: Each GEMM in the model needs to be performed 3 times (forward pass,
    #       backward wgrad [weight gradient], backward dgrad [data gradient]).
    # - 2x: GEMMs of a particular size are stacked twice in the standard Transformer model
    #       architectures implemented in this codebase (e.g., h->ffn_h GEMM and ffn_h->h GEMM
    #       in MLP layer).
    # - 2x: A GEMM of a m*n tensor with a n*k tensor requires 2mnk floating-point operations.
    expansion_factor = 3 * 2 * 2

    FLOPs = (
        expansion_factor
        * global_batch_size
        * seq_length
        * num_layers
        * hidden_size
        * hidden_size
        * (
            # Attention.
            (
                (
                    1
                    + (num_query_groups / num_attention_heads)
                    + (seq_length / hidden_size)
                ) * query_projection_to_hidden_size_ratio
            )
            # MLP.
            + (
                (ffn_hidden_size / hidden_size)
                * num_experts_routed_to
                * gated_linear_multiplier
            )
            # Logit.
            + (padded_vocab_size / (2 * num_layers * hidden_size))
        )
    )
    TFLOPs_per_gpu = FLOPs / (elapsed_time_per_iteration * 10**12 * world_size)
    MFU = TFLOPs_per_gpu/312
    print(f'throughput per GPU (TFLOP/s/GPU): {TFLOPs_per_gpu:.2f}')
    print(f'MFU using megatron-LM: {100*(MFU):.2f}%')
    return MFU

args = {
    'padded_vocab_size': 128256,
    'num_layers': 32,
    'hidden_size': 4096,
    'num_attention_heads': 32,
    'group_query_attention': True,
    'num_query_groups': 8,
    'kv_channels': 128,
    'ffn_hidden_size': 14336,
    'swiglu': True,
    'num_experts': None,
    'moe_router_topk': None,
    
    'world_size': 64,
    'global_batch_size': 64*1,
    'seq_length': 8192,
    
    'elapsed_time_per_iteration': fwdbwd_per_iter, # sec
}
megatron_lm_mfu = num_floating_point_operations_megatron_lm(**args)

params: 8.00B (not exclude embedding)
num_tokens: 1.57B
wps: 2904.00
num_params: 8.00B
fwdbwd_per_iter: 2.82
MFU: 44.68%

==================================================
training_hours: 2.35 hours (computed by MFU, num_tokens: 1.5729B, nodes: 8,  MFU: 44.68)
training_days: 0.10 days (computed by MFU), num_tokens: 1.5729B, nodes: 8, MFU: 44.68
training_hours: 2.35 hours, num_tokens: 1.5729B, fwd_bwd_time: 2.82 (computed by fwd bwd iter time)
training_days: 0.10 days, num_tokens: 1.5729B, fwd_bwd_time: 2.82 (computed by fwd bwd iter time)

==================================================
for torchtitan, approx 11.99 should be added
so, MFU using PaLM is 56.67%

==================================================
throughput per GPU (TFLOP/s/GPU): 168.18
MFU using megatron-LM: 53.90%