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

[delock]

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), 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 [AutoTP] Complete uneven sharding and universal checkpoint support #8185 (uneven sharding), PR [AutoTP] Fix ZeRO-3 checkpoint consolidation to gather across TP and DP #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

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.

3. Patterns are model-coupled. Different model families name parameters differently;
   patterns must be maintained per family. AutoTP auto-detection mitigates this but does
   not remove the fragility.

4. The default fallback is a footgun. Parameters matching no pattern fall to a blind
   torch.cat(dim=0), correct only for plain column-parallel weights. Misclassification
   silently produces wrong weights.

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

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. Existing checkpoints continue to load; no forced re-save.

4. Non-Goals

• Changing the UC on-disk output (full logical tensors, topology-agnostic — stays as is).
• AutoTP runtime sharding (orthogonal; this RFC concerns only checkpoint
  conversion/restore).
• Converter throughput optimization.

5. Design

5.1 Core idea: a per-parameter shard map

For each sharded parameter, record a shard map M: a geometric description of the
correspondence between the full logical tensor F and the per-rank shards {s_r}.

• At convert time, the converter executes M_s^{-1} to assemble F from the source
  shards {s_r}.
• At restore time, the restorer executes a target map M_t to slice F into the
  target's shards {s'_r}.

M is recorded per-parameter as data, describing pure structure — which elements of
F go to which rank, in what order. The converter becomes a single map-execution engine
with no per-category branches.

5.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.

5.3 Geometry language

The map must express: contiguous slicing, concatenation of blocks, permutations
(interleaving), replication, and padding/trim. Two representation strategies, used in
tandem:

Compact structured form (default). A small set of composable descriptors sufficient
for the overwhelming majority of real layouts:

• Contiguous(dim, widths_per_rank) — even/uneven contiguous split. Covers plain
  column/row weights and GQA kv-head-boundary splits.
• SubParamConcat([Contiguous, ...]) — fused contiguous sub-parameters. Covers fused QKV
  and fused gate-up.
• Replicate() — all ranks hold F.
• Trim(logical_size) — strip padding.
• Interleave(head_grouping, stride) — BLOOM/CodeGen-style head interleaving.

These compose: a parameter's map is a tree. Note the primitive names describe
structure (how elements are arranged), not role (vocab vs attention). For example,
padding is expressed as Trim(N), not as "this is a vocabulary parameter."

Explicit index form (fallback). For layouts that resist compact description, allow a
rank-to-indices map: Explicit({rank: index_array}). This is fully general (any
permutation) at O(numel) cost, but it is used only for the rare exotic parameter, bounding
total overhead.

The converter's capability is now bounded by the union of (structured primitives ∪ explicit
indices) = any layout, not by a closed category set.

5.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.

5.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."

6. Backward Compatibility

Every current semantic category has a well-defined geometric equivalent (Table 1), so an
automatic translator can convert legacy uc_info into shard maps on read:

Table 1 — Semantic category -> geometric equivalent

Semantic category	Geometric equivalent
TP_REPLICATED	Replicate()
PARAMETER_WITH_ROW_PARALLELISM	Contiguous(dim=1, widths)
PARAMETER_WITH_SUB_PARAMS + widths	SubParamConcat([Contiguous, ...])
VOCABULARY (strip padding)	Contiguous(...) ∘ Trim(N)
AUTOTP_UNSUPPORTED (interleaved)	Interleave(...) or Explicit(...)

Thus:

• Old checkpoints (semantic uc_info): read via auto-translation; no re-save required.
• New checkpoints: write geometric maps directly.
• The legacy semantic writer can be deprecated on a release schedule.

7. Costs and Trade-offs

• Storage. Compact form is comparable to today's SUB_PARAM_SHARD_WIDTHS.
  Explicit-index form is O(numel) but used only for exotic parameters, so total overhead is
  bounded.
• 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.

8. 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.

9. Alternatives Considered

• Status quo (semantic). Current scheme; rejected for the reasons in §2.
• Explicit index maps everywhere. Fully general but O(numel) for every parameter — too
  expensive.
• Per-architecture plugins (each model family registers gather/scatter callbacks).
  Re-introduces model coupling in the converter; rejected.

10. 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.

[delock]

Additional motivation / future work: the schema also enables live, cross-topology weight transfer

Worth flagging a direction beyond offline conversion that this schema would unlock.

Today, moving weights between two instances with different parallelism (e.g. teacher tp=4 ↔ student tp=2 in online distillation, or resizing TP/DP mid-training) forces a disk round-trip: save → ds_to_universal convert → load. The converter is the only thing that knows how to translate, and it only runs on files.

With a non-semantic per-parameter shard map, that translation is just "execute M_s⁻¹ on the source side and M_t on the target side" — both are pure geometry, no model knowledge. So a model-agnostic component could stream weights directly from source shards to target shards over the network, parametrized by the two schemas. Concretely target_shard = M_t(M_s⁻¹(source_shard)) per parameter — i.e. "UC conversion without the disk."

Two things make this tractable specifically because the maps are geometric:

1. No full-tensor materialization. Since M_s and M_t describe element-level correspondences, you can compute each source element's (target rank, offset) directly and run a sparse rank-to-rank all-to-all — never reconstructing the full logical tensor on any single rank. A semantic/regex scheme can't do this; it has no per-element geometry.
2. Model-agnostic. The component doesn't need to know what attention, fused-QKV, or MoE are — it only executes two maps. With the current semantic schema it would need every model family's knowledge baked in.

Use cases this would unlock:

• Online distillation with teacher and student at different TP (the case that came up in [AutoTP] Complete uneven sharding and universal checkpoint support #8185 review).
• Elastic training: resize TP/DP without stopping or saving to disk.
• Inference rollout / hot-load: train at one topology, serve at another, stream weights in.
• Fault-tolerant migration across clusters/topologies.

Scope caveats. Live transfer concerns the full parallelism description, not just TP — it needs to reason about which source DP replica to read and how EP experts align, so the schema (or its consumer) would extend beyond the TP map this RFC focuses on. The distributed-systems parts (rendezvous, snapshot consistency while the source is training, flow control, transport) are real engineering and orthogonal to the schema itself. The point of this note is that the non-semantic schema is the prerequisite that makes such a component buildable as a generic tool — a strong additional motivation for the RFC, not something to implement in the same PR.

For reference, ZeRO-3's gather_params is already a special case (live gather across DP within one topology); the component above generalizes it to cross-topology source→target transfer.

[goransh-buh]

Clear repro steps, that'll make this much easier for a maintainer to pick up.