[RFC] Universal checkpoint: non-semantic per-parameter geometric shard map

This is a design proposal (RFC) that came out of reviewing #8185 (uneven sharding) and #8168 (ZeRO-3 checkpoint consolidation).

The universal-checkpoint conversion/restore schema is currently semantic (regex-classified categories or semantic driven during restore), which limits it to a closed set of layouts — anything outside that set gets marked unsupported and refused, even though the model can shard and train those layouts fine at runtime. I'd like to propose moving to a non-semantic, per-parameter geometric shard map, so any layout a model can shard becomes convertible, and the converter stops being model-specific. Full proposal below; feedback very welcome before any implementation.

---

• Status: Draft
• Author: delock
• Related: PR #8185 (uneven sharding), PR #8168 (ZeRO-3 checkpoint consolidation)

1. Summary

Replace the current semantic, regex-category-based conversion/restore schema in the
universal checkpoint (UC) pipeline with a non-semantic, per-parameter geometric shard
map. Each parameter records the geometric correspondence between its tensor-parallel
shards and its full logical tensor. The converter and restorer execute this map; they no
longer classify parameters by architectural role. This makes conversion universal (no
unsupported set), model-agnostic, and extensible without converter code changes.

2. Motivation

The universal checkpoint is DeepSpeed's topology-agnostic interchange format: it stores
each parameter as a full logical tensor so a model trained with one TP/PP/DP grid can be
restored with another.

To produce and consume UC, the converter (deepspeed/checkpoint/ds_to_universal.py) and
the restorer (deepspeed/checkpoint/universal_checkpoint.py) must know how each parameter
was sharded across TP ranks. Today this knowledge is encoded semantically, in
universal_checkpoint_info (uc_info), as a set of regex-classified categories:

• TP_REPLICATED_PARAMETER_PATTERNS
• PARAMETER_WITH_ROW_PARALLELISM_PATTERNS
• VOCABULARY_PARAMETER_PATTERNS
• PARAMETER_WITH_SUB_PARAMS (+ SUB_PARAM_SHARD_WIDTHS)
• AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS （Adding from PR 8185）

The merge/slice logic branches on which category a parameter name matches (see
merge_tp_slices and _resolve_autotp_partition).

2.1 Problems with the semantic scheme

1. It is closed-world. Only layouts the schema can name can be converted. Layouts
   outside the schema are marked unsupported and refused — currently BLOOM/CodeGen
   interleaved fused-QKV, Yuan non-contiguous shared-QK, and generic fused layouts with no
   per-sub-parameter description. These models can be sharded and trained at runtime,
   yet cannot round-trip through UC.

2. Each new layout type costs converter code. Supporting a new architecture means
   adding a new category, a new regex set, and a new merge branch — a code change in the
   converter, not just new data. The converter's capability is bounded by the schema's
   vocabulary, not by what models can actually shard.

The root cause is that the scheme describes what a parameter is (its role) rather than
how it is laid out (its geometry).

2.2. The opportunity from RL/OPD

New training schema such as reinforcement learning and on-policy distillation requires trained weight to be available for both training and rollout. This requires either co-location (inference and training share same activated weights), or trained weights needs to be reshaped/transferred between training/inference phase. If training and inference use different parallelism schema, reshape weight sharding during phase change will be a big problem. Naive method (convert to topology-agnostic tensor and convert back) could bring additional overhead.

A non-semantic UC info can help this situation. Each non-semantic UC info would describe a mapping. The source mapping maps source sharded weights into UCP, the target mapping maps UCP into target sharded weights. No parallelism specific knowledge is needed to execute these mapping. More over, a properly designed algorithm can take the mapping and 'solve' it into a data movement plan, so data can be moved from source directly to target without moving into UCP and move out. This give us a chance to optimally execute the phase change, and it is fully decoupled with parallelism implemented or yet to be implemented in DeepSpeed. Makes this RFC a future proof proposal.

3. Goals

• Non-semantic. Each parameter carries a geometric shard map; the converter/restorer
  execute geometry, not categories.
• Universal. Any layout a model can shard at runtime is convertible. Eliminate (or
  shrink to truly degenerate cases) the unsupported set.
• Model-agnostic converter. No architecture-specific code or regex patterns in the
  converter.
• Backward compatible. No change to UCP format.
• Scope Cover optimizer states, also cover different sharding schema such as Zero1, Zero3, TP, EP, PP

4. Design

4.1 Core idea: a per-parameter shard map

For each parameter we record a shard map M: for every element of the full logical tensor F, where it is read from and where it is written to. M is data, not code, and agnostic to why an element lives where it does — TP, DP, EP are indistinguishable at element granularity; they only produce different location assignments.

• Source provides M_s → converter assembles F from per-rank shards.
• Target provides M_t → restorer slices F into per-rank shards.

Location is hierarchical: a rank, optionally refined by memory tier (GPU / host / NVMe / disk). Phase 1 uses rank granularity; Phase 2 refines to the tier. "Materializing F" is simply routing elements through a sub-location (e.g. disk); Phase 2 routes around it — same language, no separate placement machinery.

Edge cases (intra-shard placement, gather/scatter for replicated/reduced layouts) are deferred to the subtask.

4.2 Cross-topology restore

The full tensor F is the topology-agnostic bridge:

• The source records M_s (how it sharded). Convert = M_s^{-1}: {s_r} -> F.
• The target records M_t (how it wants to shard). Restore = M_t: F -> {s'_r}.

M_s and M_t are independent; both reference the same logical F. Changing topology
works as long as each side can describe its own geometry. This generalizes today's split
(source uc_info for convert; target parameter metadata for restore) — but replaces
semantic categories with geometric maps at both ends.

Again, this is an over simplified idea, further discussion on the shard map is needed.

4.3 Geometry language

Leave to further discussion in sub task.

4.4 Where the map lives

• At save time: AutoTP, which already computes the split, writes each parameter's M_s
  into the DeepSpeed checkpoint alongside PARAM_SHAPES.
• At convert time: ds_to_universal reads M_s, executes M_s^{-1}, and writes F
  to UC. The UC output may carry a normalized M for tooling, but the full tensor remains
  the canonical, topology-agnostic artifact.
• At restore time: the target parameter carries M_t (set when the target is
  AutoTP-sharded). The restorer executes M_t on F.

4.5 Converter simplification

merge_tp_slices collapses from seven category lookups plus a branched merge into one
operation: read the parameter's map, execute it. The get_matched_pattern /
get_matched_sub_params_pattern regex machinery and the unsupported refusal both
disappear — the latter becomes "the map is Explicit; execute it."

5. Costs and Trade-offs

• Schema design. The geometry language must be specified once. It is arguably simpler
  than the current category set (fewer, composable concepts).
• Engineering. Converter/restorer become a single map-execution engine — a net
  reduction in code, but a one-time rewrite.
• Optimizer states. ZeRO partitions optimizer states; their map is derived from the
  parameter's map plus the DP partition. Needs specifying but follows the same principle.

6. Open Questions

1. Final geometry vocabulary: minimal structured set vs. richer; the exact rule for when to
   escalate from structured to explicit indices.
2. Co-location of the map (DeepSpeed checkpoint vs. UC vs. both) and versioning.
3. Optimizer-state and expert-parallel (EP) interaction.
4. Whether to retain a residual unsupported for genuinely degenerate (non-invertible)
   layouts, if any exist.
5. Migration timeline; whether to gate the new format behind a version flag.

7. Summary

The semantic, regex-category scheme is a compact historical approximation that works for
common models but cannot scale to arbitrary layouts without per-architecture converter
changes. A non-semantic, per-parameter geometric shard map makes conversion universal
and model-agnostic, eliminates the unsupported set, and simplifies the converter — while
remaining backward compatible via one-time translation of legacy metadata.

8. Tasks

First step: on disk conversion

• Define sharding mapping language
• Convert sharding map to execution plan (pe rank sharding to UCP, UCP to per rank sharding)
• Implement executor
• Legacy uc_info --> sharding map translator
• Executor integration
  • AutoTP write checkpoint along with sharding map. Converter convert the checkpoint on disk to UCP, loader load the UCP and shard it according to sharding map.
  • Zero 1 integration
  • Zero 3 integration
  • AutoEP integration
    Second step: direct convertion
• Cross topology executor -- Define how shards of weight are transfered from one device to another device