[ENHANCEMENT] Reduce torch_dist checkpoint preparation overhead for prepended-axis tensors

**Problem**

We observed substantial CPU overhead during distributed checkpoint saving in a v0.16.0-based fork. A large part of the training pause was spent in `mcore_to_pyt_state_dict`, converting shard metadata before asynchronous checkpoint writing could begin.

The affected case is a regular, non-flattened `ShardedTensor` with `prepend_axis_num > 0`, commonly used to represent layer/expert axes. These tensors are routed through the legacy PyTorch `ShardedTensor` conversion even when the PyTorch version supports the checkpointable path.

The legacy conversion enumerates the entire global shard grid, including remote shards, for each converted key. Consequently, preparation cost grows with the number of fragments across the logical axes, rather than only the shards held locally. This becomes expensive for MoE checkpoints with many layer/expert shards.

Because this conversion happens synchronously before the async write request is returned, enabling asynchronous checkpointing does not hide this overhead. The same routing condition remains in the [inspected upstream revision](https://github.com/NVIDIA/Megatron-LM/blob/f6c33bde4cdc10805ee7157db746d47b758f5f02/megatron/core/dist_checkpointing/strategies/torch.py#L293-L315).

**Observed impact**

Our workload was a 30B MoE model with 48 layers, 128 experts and EP=8, running on eight GPUs with PyTorch 2.10.0. The profiling record identified 14,832 of 15,447 sharded-tensor records on the affected path.

With the downstream workaround, recorded `mcore_to_pyt_state_dict` time dropped from **32.1 seconds to 0.13 seconds**. This measures conversion only, not total checkpoint completion time. The timings are from our downstream fork; upstream was checked by source inspection rather than a fresh benchmark.

**Downstream workaround**

Our MR allows eligible prepended-axis tensors to use the existing checkpointable path. A prepended axis has extent one within each local chunk, so we expose those singleton dimensions consistently in the local chunk metadata and tensor view. The adapter can then describe local shards directly, avoiding enumeration of the global grid.

This preserves the logical tensor shapes and offsets and uses views rather than copying dense tensor payloads. Existing version and flattened-shard restrictions are retained. Quantized tensors need representation-preserving views; GPU validation of the final FP8 changes is still pending.

Would maintainers consider extending the checkpointable path to this case? We would appreciate feedback on this approach and any compatibility requirements that should be addressed before upstreaming it.

Start in megatron/core/dist_checkpointing/strategies/torch.py around the checkpointable-path condition at lines 293-315, and trace mcore_to_pyt_state_dict for regular prepended-axis ShardedTensor values. Compare the existing checkpointable adapter with the downstream workaround described in the issue. Done means eligible tensors avoid global shard-grid enumeration while preserving logical shapes, offsets, views, version restrictions, flattened-shard behavior, and quantized representations; GPU validation remains part of the work.

Written by the indexing model from the issue text.