[Feature] Support DCP save/load for InterleavedShard (tpep) MoE weights

[HAOCHENYE]

Stacked on top of InternLM#1837 (ep_tp). Makes per-expert column-parallel fused MoE weights
(InterleavedShard DTensors, produced under expert TP) participate in DCP snapshots instead of
being dropped and refilled from HF safetensors on resume.

Root cause

An InterleavedShard local tensor is several interleaved runs of the global tensor, not one
contiguous slice. DCP's default planner models each DTensor as a single chunk via
compute_local_shape_and_global_offset — which cannot describe this placement: it raises on
torch >= 2.9 (_StridedShard split_factor != aggregate mesh size) and silently returns a wrong
(size, offset) on 2.8. So these params were dropped from DCP entirely.

Fix

• xtuner/v1/patch/dcp_interleaved_planner.py: InterleavedShardSavePlanner /
  InterleavedShardLoadPlanner emit one WriteItem / ReadItem per contiguous run (from
  compute_runs) at the run's global offset, so the checkpoint is stored in global coordinates
  and reshards across tp/ep topologies. Both planners iterate the state_dict directly and never
  hand an interleaved DTensor to torch's default builder.
• train_engine.save_dcp / load_dcp pass the planners explicitly and drop the
  _drop_interleaved_* filtering path.
• The cache-storage patch injects the interleaved-aware save planner so the torch 2.7.x
  incremental-save path stays correct.

Test plan

• tests/patch/test_dcp_interleaved_planner.py (world=4): save->load round-trip and reshard
  (ep2/tp2 -> ep1/tp4), bit-equal. Passes on torch 2.8 and 2.9.
• Real training resume (ep2/tp4, deepep, compile) on torch 2.9: DCP model + optimizer save at
  step 5/10 and auto-resume continues training with no error.

Not yet covered: fsdp_size > 1 post-FSDP end-to-end (validated so far at fsdp_size == 1).

[gemini-code-assist]

Code Review

This pull request introduces InterleavedShardSavePlanner and InterleavedShardLoadPlanner to properly support InterleavedShard (tpep) DTensors during Distributed Checkpoint (DCP) save and load operations, replacing the previous workaround of dropping these parameters. It also adds comprehensive regression tests for round-trip validation and topology resharding. The reviewer identified potential crashes in both planners when handling non-tensor objects (such as dicts, lists, or primitives) in the state dict, and provided code suggestions to safely fall back to byte array handling.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The else branch in InterleavedShardSavePlanner.create_local_plan assumes that any non-DTensor object is a torch.Tensor and calls _create_write_items(fqn, obj). However, state dicts (especially optimizer state dicts) can contain non-tensor objects (such as dicts, lists, or primitives). Passing these to _create_write_items will cause a crash (e.g., AttributeError). We should explicitly check for torch.Tensor and handle other types as WriteItemType.BYTE_ARRAY, matching the behavior of DefaultSavePlanner.

Suggested change

	elif isinstance(obj, DTensor):
	if obj.device_mesh.get_coordinate() is not None:
	requests.extend(_create_write_items(fqn, obj))
	else:
	requests.extend(_create_write_items(fqn, obj))
	elif isinstance(obj, DTensor):
	if obj.device_mesh.get_coordinate() is not None:
	requests.extend(_create_write_items(fqn, obj))
	elif isinstance(obj, torch.Tensor):
	requests.extend(_create_write_items(fqn, obj))
	else:
	requests.append(
	WriteItem(
	index=MetadataIndex(fqn),
	type=WriteItemType.BYTE_ARRAY,
	)
	)

[gemini-code-assist]

Similarly, InterleavedShardLoadPlanner.create_local_plan assumes any non-DTensor object is a torch.Tensor and calls _create_read_items(fqn, md, obj). Non-tensor objects in the state dict will cause a crash. We should explicitly check for torch.Tensor and handle other types as WriteItemType.BYTE_ARRAY, matching DefaultLoadPlanner.

Suggested change

	elif isinstance(obj, DTensor):
	if obj.device_mesh.get_coordinate() is not None:
	requests.extend(_create_read_items(fqn, md, obj))
	else:
	requests.extend(_create_read_items(fqn, md, obj))
	elif isinstance(obj, DTensor):
	if obj.device_mesh.get_coordinate() is not None:
	requests.extend(_create_read_items(fqn, md, obj))
	elif isinstance(obj, torch.Tensor):
	requests.extend(_create_read_items(fqn, md, obj))
	else:
	requests.append(
	ReadItem(
	type=WriteItemType.BYTE_ARRAY,
	dest_index=MetadataIndex(fqn),
	dest_offsets=torch.Size(),
	storage_index=MetadataIndex(fqn),
	storage_offsets=torch.Size(),
	lengths=torch.Size(),
	)
	)