[refinedheart]

Background

EPMoELayer keeps only the current EP rank's expert shards. MoECheckpointIO gathers expert parameters marked ep_fused when saving, but the layer did not mark its fused expert parameters or slice a full checkpoint tensor back to the local shard when loading.

Changes

• Mark gate_up_proj and down_proj as ep_fused so the existing checkpoint IO gathers complete expert tensors.
• Add a load-state-dict pre-hook that slices full [num_experts, ...] tensors using the layer's local_expert_range.
• Preserve already-local tensors and leave malformed shapes to the standard load_state_dict size-mismatch error.
• Add CPU/gloo tests for parameter tagging, full/local loading, malformed shapes, and a save/load round trip through MoECheckpointIO.

No public Python API or checkpoint format is introduced. This PR does not change expert computation or gradient synchronization.

Test environment

• Image: wangzhijian/colossalai-ci:torch2.13-cu130-tf5.17.0-peft0.21.2-20261022-091529
• OS: Ubuntu 22.04; Python 3.10
• PyTorch: 2.13.0+cu130; Transformers: 5.17.0
• CUDA Toolkit: 13.0.3; host NVIDIA driver: 595.71.05
• The following tests ran on CPU/gloo; no GPU result is claimed.

Tests

python -m pytest -q tests/test_moe/test_ep_moe_layer_cpu.py
# 5 passed

.github/scripts/ci/run_cpu_unit_tests.sh
# 133 passed, 3 skipped; exit code 0

The three skips are pre-existing CUDA/TensorNVMe-gated optimizer async-save tests: test_real_async_writer and test_general_repeated_async_save[False/True]. They are unrelated to this change.

Known limits

Public Qwen3-MoE Booster/checkpoint/HF-readback coverage is being developed separately with the Qwen3-MoE policy. This PR does not claim that model-level GPU checkpoint acceptance or full L2 validation has completed.