[OOM] Huge memory usage in MoE

Hi, I encountered Out of Memory Error when I train a 19B Llama MoE model (~190B total parameters and 19B activated parameters) on v5p-256, which uses 400G memory, significantly exceeding 95G memory limit.

Here is the key model config modification on the [base.yaml ](https://github.com/hilbertmeng/maxtext/blob/main/MaxText/configs/base.yml).
```
# batch_size and sequence length
max_target_length = 4096
per_device_batch_size = 8.0

# Dense config
base_emb_dim = 5120
base_num_query_heads = 40
base_num_kv_heads = 40
base_mlp_dim = 12288
base_num_decoder_layers = 60
head_dim = 128
vocab_size = 151936

# MoE config
base_mlp_dim = 6144 # base_mlp_dim refers to hidden_dim of each expert
num_experts_per_tok = 2
num_experts = 32
megablox = True
sparse_matmul = True
shared_experts = 0
```
To analyze the memory usage, I only reduced `base_num_decoder_layers` config from 60 to 6, so that it can run without OOM and I profiled the memory usage below.

I found the memory usage constantly increases during backpropagation (Program Order 1500-4000). A further analysis reveals that [tgmm](https://github.com/hilbertmeng/maxtext/blob/main/MaxText/kernels/megablox/gmm.py#L619) kernel saves each layer's expert weight of shape (32, 6144, 5120), consuming 1.9G memory, so three `linear` (w0, w1 and wo) count ~6G. Thus for a 60-layer model, `tgmm` consumes 60x6=~360G, close to the observed 400G memory usage. Here I provide a typical piece of expert memory usage in `tgmm`.

```
Size:
1920.00 MiB
Unpadded Size:
1920.00 MiB
Shape (and minor-to-major order):
bf16[32,6144,5120]{2,1,0:T(8,128)(2,1)}
Tf Op Name::
jit(train_step)/jit(main)/transpose(jvp(Transformer))/decoder/checkpoint/rematted_computation/layers_5/sub_0/moe/moe.sparse_matmul/jit(shmap_body)/jit(tgmm)/pallas_call
Allocation Type:
Temporary
```
Is this a memory leakage of `tgmm` or is there something wrong in my training config? How can I fix it?