LM Head FLOPs

[Muennighoff]

Why are we not multiplying the LM Head flops per iteration with the checkpoint_activations_factor?

Megatron-LM/megatron/utils.py

Line 253 in bd0aaba

flops_per_iteration += (6 * batch_size * seq_len * num_layers * (hidden_size**2)) * (vocab_size / (num_layers * hidden_size))

Afaik the factor of 4 means 1 forward, 2 backward & 1 forward, where the last forward is needed for ckpt acts. Don't we also need all 4 for the LM Head? cc @RaymondLi0 @NouamaneTazi

[NouamaneTazi]

In selective recomputation, we only checkpoint attention, so it shouldnt affect LM Head
In full recomputation, we assume that the last hidden_states are the LM Head activs, so no need for recomputing them

Megatron-LM/megatron/model/transformer.py

Line 1905 in bd0aaba

hidden_states = self._checkpointed_forward(hidden_states,

[Muennighoff]

In selective recomputation, we only checkpoint attention, so it shouldnt affect LM Head In full recomputation, we assume that the last hidden_states are the LM Head activs, so no need for recomputing them

Megatron-LM/megatron/model/transformer.py

Line 1905 in bd0aaba

hidden_states = self._checkpointed_forward(hidden_states,

I see but then don't we at least need to multiply it by 3 to account for the backward pass?