Megatron style TFLOPs Calculation

[abhinavgoel95]

@rwitten this is a draft.

This type of change would be specific to a few transformer models (e.g., Gemma, LLama, GPT, etc.). It wouldn't work with MoE, or some new architectures.

I was thinking that walking through the train-step and calculating the FLOPs layer-by-layer would be a very intrusive change.

What do you think?

[abhinavgoel95]

Made the changes as requested in the meeting @rwitten

[abhinavgoel95]

cc @rwitten following up on this