File size: 1,660 Bytes
805d830 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
from torch import nn
from torchtitan.tools.logging import logger
def get_nparams_and_flops(model: nn.Module, model_config, seq_len: int) -> tuple[int, int]:
nparams = sum(p.numel() for p in model.parameters())
nparams_embedding = sum(
sum(p.numel() for p in m.parameters())
for m in model.children()
if isinstance(m, nn.Embedding)
)
if hasattr(model_config, "num_heads"):
num_heads = model_config.num_heads
elif hasattr(model_config, "num_attention_heads"):
num_heads = model_config.num_attention_heads
else:
num_heads = 1
logger.warning("num_heads not found in model_config, defaulting to 1. ")
l, h, q, t = (
model_config.num_hidden_layers,
num_heads,
model_config.hidden_size // num_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
num_flops_per_token = 6 * (nparams - nparams_embedding) + 12 * l * h * q * t
return nparams, num_flops_per_token
|