Motivation

mcore_to_pyt_state_dict routes every ShardedTensor with a non-zero prepend_axis_num
through the legacy PyTorch ShardedTensor conversion
(sharded_tensor_to_torch_sharded_tensor). That conversion enumerates the entire global
shard grid — remote shards included — for each converted key:

for fragment_offsets in product(*map(range, some_sh_ten.axis_fragmentations)):

so its cost grows with the number of fragments across all logical axes, not with the number of
locally held shards.

Prepended axes are not an exotic corner case (closes #7275):

• TransformerBlock.sharded_state_dict passes (0, global_layer_offset, num_layers) as the
  sharded_offsets of every layer whenever the stack is homogeneous, and
  sharded_state_dict_default turns those offsets into a prepended layer axis of every
  parameter of the block.
• SequentialMLP.sharded_state_dict adds a prepended global-expert axis on top of it for
  each routed expert.

For a 48-layer MoE model with 128 experts and EP=8 the per-rank local shards for one weight key
therefore describe a 48 x 128 grid — 4.7M grid entries per rank in total — and all of that work
happens synchronously before the async write request is returned, so asynchronous
checkpointing cannot hide it from the training loop.

What changed

A prepended axis has extent one within a local chunk (that is what prepend_axis_num means:
the local data covers a single index along the prepended axes). The DCP checkpointable path can
describe that directly, so prepended-axis shards no longer need the legacy conversion:

• CheckpointableShardedTensor now reports prepended axes as singleton dimensions of the local
  chunk (sizes = (1,) * prepend_axis_num + local_shape) and exposes the local data through a
  matching view (data.view((1,) * prepend_axis_num + local_shape)) — a view, no payload copy.
• mcore_to_pyt_state_dict routes prepended-axis shards to the checkpointable path, keeping the
  existing restrictions (PyTorch >= 2.6, regular grid, no flattened shards).

The resulting DCP metadata is identical to the metadata produced by the legacy conversion
(same chunk offsets, chunk sizes, global size and dtype — asserted by the new tests), so:

• checkpoints written before this change still load, and
• checkpoints written after it are indistinguishable from the old format.

Measurements (8xH200, PyTorch 2.10, torch_dist)

Local (non-TransformerEngine) MoE model, 48 layers, 128 experts, EP=8, 1,109 local shards per
rank, 1,646 unique keys, 4,734,725 grid entries per rank:

A synthetic prepended-axis workload with 3,072 keys/rank x 128 grid entries (393,216 entries)
lands in the same place: 3.69 s -> 0.027 s.

Testing

• New: tests/unit_tests/dist_checkpointing/test_checkpointable.py::TestPrependedAxisCheckpointable
  — write-item metadata parity with the legacy conversion, save/load round trip, and loading a
  checkpoint written by the legacy conversion.
• New: tests/unit_tests/dist_checkpointing/models/test_moe_expert_sharding.py — MoE model
  round trips for EP=1/2/4 and EP=4 -> EP=2 resharding (asserts the prepended layer/expert axes
  are present). Runs without TransformerEngine.
• pytest tests/unit_tests/dist_checkpointing/{test_checkpointable,test_torch_dist, test_serialization,test_validation,test_local,test_fully_parallel,test_replication}.py
  on 8 GPUs: 65 passed, 2 skipped.
• black, isort, ruff, pylint clean on the changed files.

Follow-ups (not in this PR)

The legacy branch of mcore_to_pyt_state_dict (and sharded_tensor_to_torch_sharded_tensor)
is now reachable only for torch < 2.6; removing it once the minimum PyTorch version allows
would be a separate cleanup.