Update tflops calculation

[RissyRan]

Description

Update tflops calculation to match Megatron style with sources:

• MaxText reference PR
• Megatron-LM link

Test

List models before/after the change: total_ffn_flops value

Model Name	TFLOPS/step before	TFLOPS/step after
llama2-7b	1033.2	1053.4
llama2-13b	1981.1	2018.7
mistral-7b	1107.4	1127.6
llama2-70b	10368.9	10528.0
mixtral-8x7b	6926.1	6946.5
gemma-7b	1293.6	1328.2
gemma-2b	380.7	391.8
gpt3-6b	423.1	430.4
gpt3-52k	0.01256	0.01253

Test code snippet

class FlopsTest(unittest.TestCase):

  def setUp(self):
    import os
    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
    pyconfig.initialize(
      [None, 'configs/base.yml'],
      run_name='test',
      model_name='mixtral-8x7b',
      flops_change=False,
    )
    self.cfg = pyconfig.config
    devices_array = max_utils.create_device_mesh(self.cfg)
    self.mesh = Mesh(devices_array, self.cfg.mesh_axes)
    quant = quantizations.configure_quantization(self.cfg)
    self.model = Transformer(self.cfg, mesh=self.mesh, quant=quant)

  def test_flops(self):
    rng = random.PRNGKey(0)

    tx = optax.adam(learning_rate=0.001)
    state, _, _ = max_utils.setup_training_state(self.model, None, tx, self.cfg, rng, self.mesh, None)
    num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
    maxtext_utils.calculate_tflops_training_per_device(num_model_parameters, self.cfg)

[rwitten]

Let's make sure Abhinav Goel at NV is aligned

[abhinavgoel95]

@rwitten @RissyRan this PR looks good to me

[rwitten]

Looks great! But please don't merge without Abhinav's approval.

Oh he already approved