Added support of training with NANOO fp8 GEMM on AMD MI300/MI325 GPUs.

[wenchenvincent]

Description

This PR added support of training with NANOO fp8 GEMM on AMD MI300/MI325 GPUs.

There are several different genres of fp8 formats used by different HW vendors. Two popular genres include

• OCP fp8, which is used natively on NVIDIA H100.
• NANOO fp8, which is used natively on AMD MI300/MI325.

These two genres of fp8 formats work very similarly. The support of training with NANOO fp8 GEMM in Maxtext is based on this PR in Flax: google/flax#3993

References:

OCP fp8 paper: https://arxiv.org/abs/2209.05433
NANOO fp8 paper: https://arxiv.org/abs/2206.02915
JAX PR: jax-ml/jax#21376
XLA PR: openxla/xla#9531
Flax PR: google/flax#3993

Tests

I had run llama2-7b with NANOO fp8 from this PR and verified it was functional and the loss went down quickly with the synthetic dataset. I was not able to run the full unit tests locally due to the requirement of Google Cloud API credentials.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[anfals]

LGTM; also using FP8Ops but just with different types; should work the same as everything else we have with _overwrite_with_gradient, so no additional lift needed