Incorrect MFU computation

[thepowerfuldeez]

Hi! In your code you mention https://arxiv.org/abs/2204.02311 (PaLM) paper for MFU computation which use total number of parameters with embedding
Line of code:

torchtitan/train.py

Line 225 in 7afe902

get_num_params(whole_model, exclude_embedding=True),

You instead compute based on number of parameters with embedding excluded (older OpenAI style). This is wrong and leads to larger than usual numbers.

[awgu]

You instead compute based on number of parameters with embedding excluded (older OpenAI style). This is wrong and leads to larger than usual numbers.

I might be missing something, but why does excluding the embedding led to higher MFU?

Stepping through the logic:

torchtitan/train.py

Line 434 in 7afe902

mfu = 100 * num_flop_per_token * wps / gpu_peak_flops

torchtitan/train.py

Lines 224 to 228 in 7afe902

	num_flop_per_token = get_num_flop_per_token(
	get_num_params(whole_model, exclude_embedding=True),
	model_config,
	job_config.training.seq_len,
	)

torchtitan/torchtitan/utils.py

Lines 123 to 138 in 7afe902

	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

• mfu is directly proportional to num_flop_per_token.
• num_flop_per_token is directly proportional to num_params.

Excluding the embedding leads to lower num_params --> leads to lower mfu?

---

For context, we changed the MFU calculation to exclude embedding parameters in #280 with suggestion in #274. Could you also point me to where in the PaLM paper it says that the MFU calculation includes the embedding parameters?

I was reading from Appendix B:

Excluding self-attention, an N-parameter decoder-only model requires 6N matmul FLOPs per token seen (2N for forward and 4N for backward), because each matmul performs one multiplication and one addition per pair of input values, and the backward pass includes two matmuls for each one in the forward pass

For us, num_params affects the matmul FLOPs term of 6N, and embedding does not include matmul.

[thepowerfuldeez]

@awgu PaLM paper, Appendix B, page 66:
We discuss the calculation of MFU numbers reported in Table 3. The MFU calculation for Megatron–Turing NLG
530B is as follows: it was trained on A100 GPUs with 312 peak matmul TFLOP/s, and achieved a training throughput
of 65.43K tokens/second (1920 × 2048/60.1) on 2240 GPUs (Smith et al., 2022), resulting in MFU of 29.7% without
attention ((65.43 × 6 × 530)/(312 × 2240)) or 30.2% with it. I

So for 530B parameters with embedding, they use number 530 in the computation

[awgu]

@thepowerfuldeez Thanks for pointing that out!

I did not see the Megatron-Turing NLG paper mention their vocabulary size or tokenizer. Were you able to find them by any chance? They mention that their hidden dimension is 20480. If their vocab size is 32k, then their embedding weight is <0.7B numel. Even if their vocab size is 128k (which I think is less likely), then their embedding weight is still <2.7B.

Changing their 530B to 527B changes their MFU without attention from 29.7% to 29.6% (and similarly from 530B to 529B is still 29.7%), which I think is not a major difference. In other words, for such a large dense transformer, I think the embedding weight being included or not is less important (compared to say Llama3-8B).

Also, I still did not follow why excluding embedding leads to higher MFU? It seems to lead to lower MFU.

[thepowerfuldeez]

@awgu You are correct that for such big models it doesn't matter. Hmm, you are also correct that with excluding embedding size you get lower MFU, since it's proportional to the number of parameters (6ND in original formula).
I think I misunderstood. Then it is really good MFU numbers in torchtitan repo! You can close this issue

[awgu]

@thepowerfuldeez We are using a factor of 12 for the self-attention FLOPs, which could be considered unfair since flash attention technically only uses 7. This would decrease our MFU. From previous discussion, we kept it as using factor of 12 since this seems to be the more common practice outside of the flash attention repo.