reshard checkpoint

[shuningjin]

Description

FIXES: b/504714612

Tests

Full test details in b/504714612 (comment2, comment4, comment8), which shows re-sharding checkpoint effectively reduces loading time.

CKPT1=gs://ml-auto-solutions/output/unowned/maxtext_nightly_deepseek2-16b-v5p-8-2026-04-20-06-52-15/scanned/0/items \    
CKPT2_DIR=gs://shuningjin-multipod-dev/conversion/ds2-fsdp-$(date +%Y-%m-%d-%H-%M)                                       
CKPT2=$CKPT2_DIR/0/items                                    
echo $CKPT2        

python3 -m maxtext.checkpoint_conversion.reshard_checkpoint \
model_name=deepseek2-16b attention=dot_product mla_naive_kvcache=false \
scan_layers=True load_parameters_path=$CKPT1 \
base_output_directory=$CKPT2_DIR \
weight_dtype=bfloat16 \
checkpoint_storage_concurrent_gb=1024 checkpoint_storage_use_ocdbt=True checkpoint_storage_use_zarr3=True \
skip_jax_distributed_system=True ici_fsdp_parallelism=16 \
--simulated_cpu_devices_count=16

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.