Add a remat checkpoint after the MoE expert-parallel combine

[NuojCheng]

Description

In the sparse-matmul MoE path (_moe_body in src/maxtext/layers/moe.py), the expert-parallel combine collective sits inside the remat region, so the backward pass recomputes it:

• unsort_output_and_ra2a(...) — the ragged all-to-all that returns tokens to their original shards, and
• the psum_scatter over the expert axis on the use_ring_of_experts path.

Both are pure communication, so recomputing them buys nothing but latency. This PR names their outputs moe_combine via adc.checkpoint_name, so a custom remat policy can save the combine output on device (or offload it) and skip re-running the RS / A2A in the backward pass:

remat_policy=custom moe_combine=device

The new knob is registered alongside the existing moe_mlpwo-style tensors (base.yml, configs/types.py) and defaults to 'remat', so behavior is unchanged unless it is set.

Tests

Unit: python -m pytest tests/unit/moe_test.py -q → 11 passed, 8 skipped, 1 failed (test_gmm_grad_equivalence_tokamax_v2_fp8_static_ep1, which fails identically on the base commit without this change).

Small-scale e2e on a v5p-8, EP=4 (DeepSeek-v3 shape, base_emb_dim=2048, base_moe_mlp_dim=1408, 4 layers / 1 dense, 16 experts, top-8, per_device_batch_size=2, max_target_length=2048, synthetic data, remat_policy=custom):

python -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
  model_name=deepseek3-tiny override_model_config=true \
  base_emb_dim=2048 base_moe_mlp_dim=1408 base_num_decoder_layers=4 first_num_dense_layers=1 \
  num_experts=16 num_experts_per_tok=8 per_device_batch_size=2 max_target_length=2048 \
  steps=12 enable_checkpointing=false dataset_type=synthetic \
  ici_expert_parallelism=4 ici_fsdp_parallelism=1 \
  remat_policy=custom moe_combine=<remat|device> run_name=...

	moe_combine=remat (today)	moe_combine=device
steady-state step time	0.128 s	0.120 s (-6%)
TFLOP/s/device	98.1	104.6
compiler HBM estimate	6.2 GB (3.5 GB temp)	6.9 GB (4.2 GB temp)
AOT temp_size_in_bytes (train_compile.py compile_topology=v5p-8)	3.72 GB	4.53 GB

Loss is bit-identical between the two runs at every step (e.g. step 11: 11.486 both), as expected for a remat-policy-only change.

The optimized HLO confirms the intended effect — counting ragged-all-to-all in jit_train_step, the combine A2A recompute in the backward pass disappears:

• moe_combine=remat: 6 (fwd dispatch + fwd combine, plus recomputed dispatch and recomputed combine in bwd)
• moe_combine=device: 5 (the recomputed combine is gone)

Checklist

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

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• 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, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces a new rematerialization policy key, moe_combine, representing the output of the expert-parallel combine collective (reduce-scatter or ragged all-to-all). This key is added to the configuration files base.yml, pyconfig_deprecated.py, and types.py. Additionally, checkpoints are registered in moe.py after the expert-parallel combine collective to allow saving this tensor on device, which avoids recomputing the collective during the backward pass. I have no feedback to provide as there are no review comments.