Define affine IR for checkpointing

[delock]

This sub-issue (feature request) tracks the task of define IR spec for non-semantic checkpointing. The idea is describe shard operations as affine operations, this will help use achieve goal of both phase 1 (UC_info for UCP conversion) and phase 2 (direct checkpointing transfer between topologies)

The outcome is defined as a spec document where we can implement UC_info description (both in memory structure and in file format), checkpoint conversion in phase 1, and checkpoint transfer component in phase 2.

[Achyuthan-S]

@delock Consolidating my part of the #8230 discussion here as requested, updated for your last two replies.

The boundary. The IR describes; the planner decides. Concretely the IR should be a relation, not a function — for each element of F it names the set of shard elements that legitimately hold it, and the planner resolves that relation into a function by choosing sources under a cost model. Anything the IR resolves, the planner can no longer optimize. Cost model, scheduling and replica choice stay outside.

On all-reduce / average — agreed, and I think dropping it is structurally necessary rather than just simpler. I had argued for an optional combine op on a piece. Your read is better, and the repo supports it: parameter_to_average_patterns has exactly three references in DeepSpeed — the constant, the import, and the .get in merge_tp_slices — with no producer in tree, and collect_autotp_universal_checkpoint_info never emits it.

There is also a structural reason to drop it. A reduce is not invertible: N source elements collapse into one element of F and M_s^-1 cannot recover them. If the IR admits arithmetic, then invertibility — which you named as an intentional requirement — holds only over whichever subset of the language happens to be view-only. Removing the combine op is what makes M_s^-1 well-defined for the entire language. So: pieces are pure views, no arithmetic. A legacy checkpoint carrying the pattern lowers to replicated-with-assert-equal in the task-4 translator, or is refused explicitly.

On subset + holding ranks. I think this collapses your extension and the geometry/placement split into a single concept: a piece is an affine view of F plus a location set — the ranks, and later tiers, that hold that view.

• TP sharded: one location per piece
• TP replicated: one piece, N locations — exactly the freedom case 2 needs preserved
• ZeRO-1/3: the DP partition subdivides the byte range; the location set is the ranks holding it
• AutoEP: expert pieces carry location sets over the EP group
• ZeRO offload: the same field refines to a memory tier rather than needing separate placement machinery, which matches section 4.1

Phase 1 then falls out as the degenerate case: the converter needs any one element of the location set, since every element holds identical data by construction. That is your "the only piece that is holding a tensor will be automatically filled into UCP." Phase 2 needs the whole set. One field, no phase-1 cost.

The one thing I would keep separate is what gets serialized. The geometry is a property of the parameter and is stable across jobs; the location set is a property of a particular grid and changes with it. Both are needed at transfer time, but only the geometry has to survive in UCP for it to stay portable.

Two open cases carried over from #8230.

1. Truncation. merge_tp_slices concatenates and then truncates vocabulary parameters to original_vocab_size, so where a source padded its vocab those rows have no destination in F. A pure view language cannot express "goes nowhere." Does F carry the padding with truncation as a separate post-step, or does the IR admit pieces that map to nothing, making M_s non-total over the shards?

2. Whether an escalation rule from structured affine to explicit indices is needed at all, or whether N cat'd affine pieces already cover every layout in section 2.1 including the Yuan non-contiguous shared-QK case.

What I'll draft, following your outcome definition: the affine representation as both in-memory structure and file format, a lowering table for each current uc_info key, the two cases above, and the phase-1 converter algorithm expressed over the IR. I'll post it here for review while #8185 lands. Happy to be assigned this for tracking.

[Achyuthan-S]

@delock

Coverage check: the layouts #8185 marks unsupported are affine-expressible

Before writing the spec I wanted to try to falsify its central assumption rather than assert it. #8185 (merged as aa3914df8) added AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS with three distinct reasons, and those are exactly the cases the affine IR has to cover if it is going to replace the semantic categories. So I ran each one through the real partition code.

Method. Per layout, per rank: run the actual DeepSpeed partition function on a marker tensor to recover which element of F each shard element came from; compress that index map into maximal runs to get affine (offset, shape, strides) pieces; then re-run the same partition function on independent random data and check that the pieces, materialised with torch.as_strided, reproduce the shard bit-exactly. That last step is the real test — if pieces derived from markers reproduce a random tensor's shard, the mapping is a pure view of F, not something data-dependent.

Result. All four pass.

Layout	Reason string in master	F	tp	Pieces/rank	Covers F
bigcodetype (fused_LinearLayer)	"interleaves or replicates blocks … cannot be reassembled from the shards"	48×8	4	2 (1 on the last rank)	yes
codegentype (fused_LinearLayer)	same	24×8	2	11	yes
Yuan value, dim 0 (Yuan_LinearLayer)	"selects noncontiguous head groups … cannot currently describe"	32×8	2	2	yes
Yuan o_proj, dim 1 (Yuan_LinearAllreduce)	same	8×32	2	2, strides (32, 1)	yes

Piece count is bounded by block structure, not model size. CodeGen holds at 11 pieces/rank while rows-per-rank grows 32×:

 hidden   F rows   tp  kv_heads  rows/rank  pieces/rank
      8       24    2         8         12           11
     16       48    2         8         24           11
     32       96    2         8         48           11
     64      192    2         8         96           11
    128      384    2         8        192           11
    256      768    2         8        384           11

Three things fall out of this that I think should shape the spec.

1. No information is lost in any of these cases. In all four the shards cover F exactly, and Yuan's are disjoint — it was always a clean partition. The reason string says the parameter "cannot be reassembled from the shards"; what is actually true is that the current description format cannot say how. That is a vocabulary limit rather than a data limit, which is what makes it fixable by an IR change alone, with no change to how anything is partitioned at runtime.

2. bigcodetype is a concrete case for the location set. Its kv block is byte-identical on every rank — partial replication within a single parameter. Today's schema carries one partition_dim per parameter and so cannot express "these rows sharded, those rows replicated", which is exactly why the layout is unsupported. Your framing in #8230 — a subset of a parameter plus the list of ranks holding that subset — expresses it directly, with no extra mechanism. GPTBigCode/StarCoder is an in-tree model this unblocks.

3. The o_proj case needs real strides, not ranges. Yuan's dim-1 shard comes out as strides=(32, 1) — a column selection, not a contiguous span. So the language can't be narrowed to "a list of contiguous ranges"; the full (offset, strides, shape) primitive you proposed is the right one.

I'd like to fold this in as the coverage section of the spec, with the harness cleaned up into a real test so the claim stays checked rather than asserted. Working on the rest now: the lowering table for each current uc_info key, the on-disk format, and the phase-1 converter.

One open case I don't yet have a clean answer for — vocabulary truncation. ds_to_universal.py:451-452 does param = param[:original_vocab_size, :], which drops rows, so M_s is not total over the shards. Two ways out: either F carries the padding and truncation becomes a separate post-step outside the IR, or the IR admits pieces that map to nothing, which costs invertibility. I lean to the first — it keeps every piece a view and keeps M_s^-1 well defined for the whole language — but it means the padded size has to be recorded somewhere. PADDED_VOCAB_SIZE is already declared in checkpoint/constants.py:68 but is never read or written anywhere in the repo. Was it intended for this?

[delock]

Hi @Achyuthan-S thank you for your detailed analysis and validation of affine IR for different layouts. It make this concept more concrete.

For vocab padding this is what I think for option A:

1. This issue only happens for legacy TP mode which requires even sharding. AutoTP already support uneven sharding and should not require padding in F.
2. I'm good with map to F then truncate to UC, this makes the mapping invertible, this is good for phase 1. But we need to aware that in phase 2 there needs to be different F for each instance. For example, when source has TP=4 and target has TP=2, it is possible that source F has 3 padding and target F has 1 padding. So the mapping is like:
   source shard --affine--> Fsrc --> truncate or pad to Ftgt --affine--> target shard
   Note in phase2 Fsrc and Ftgt we expect them to be virtual, so has F per intance is fine, but we need to check whether these mapping steps allows us to generate affine from source shards to target shard with truncation (or padding) If this is possible, then it proves your option A is compatible with phase 2. If it is not possible, then we may exclude legacy TP from phase 2.

For UC disk format, maybe its okay to allow padding enters UC in legacy mode. @sfc-gh-truwase for comments.

@jinyouzhi for necessarty to support uneven sharding on vocab. I see you have a related PR.

[Achyuthan-S]

@delock Thanks — I ran your phase-2 check, and option A does compose. Legacy TP does not need excluding.

I built your example literally: V_orig=13, source TP=4 → Fsrc 16 rows (3 padding), target TP=2 → Ftgt 14 rows (1 padding). The composed source-shard → target-shard plan, with neither Fsrc nor Ftgt ever materialised:

src[0][0:4] -> tgt[0][0:4]
src[1][0:3] -> tgt[0][4:7]
src[1][3:4] -> tgt[1][0:1]
src[2][0:4] -> tgt[1][1:5]
src[3][0:1] -> tgt[1][5:6]
tgt[1] row 6 has no source   (target-side padding)

Applying it reproduces the shards the TP=2 topology would have produced itself, exactly.

Why it composes. Fsrc and Ftgt agree on the real rows [0, V_orig) and differ only above them, so the truncate/pad step is the identity restricted to a range. Restricting an affine map to a range leaves it affine, so the composition is just overlap arithmetic between two runs — no new operator, and nothing that has to know the parameter is a vocabulary.

The one thing it needs. The composed map is partial over the target: target padding rows have no source. That is a different and much weaker property than M_s being non-total, which is what we were trying to avoid — padding is not data, and the target instance has to initialise it regardless of where its weights came from. So this costs nothing structurally.

Padding is never read. I filled the source padding with NaN rather than zeros, since legacy TP promises nothing about what lives above the real vocabulary. Every target still comes out exact, so the plan provably never touches those rows.

It holds generally, and stays cheap. All 64 source/target pairs over TP ∈ {1,2,4,8} and vocab ∈ {13, 32, 100, 50257} compose exactly, in both directions — truncating (Fsrc taller) and padding (Ftgt taller). Piece count is bounded by the two TP degrees rather than by the vocabulary: across vocab 13 → 262144 and TP up to 16, it never exceeds src_tp + tgt_tp. The exact count within that bound depends on how the two padded heights align, so it varies between vocabularies even though it does not grow with them.

 vocab  src_tp  tgt_tp   Fsrc   Ftgt  pieces  exact
    13       4       2     16     14       5   True
    13       8       4     16     16       7   True
 50257       4       2  50260  50258       5   True
 50257       4       8  50260  50264      11   True
 50257       8       4  50264  50260      11   True

So your reading is right that phase 2 needs a different F per instance, and that is exactly what makes it work: Fsrc and Ftgt are only coordinate systems, and V_orig is the single fact both ends must agree on.

On the UC disk format — this is evidence that padding does not need to enter UC even in legacy mode. The composition needs only V_orig; it never reads a padding row, and it reconstructs the target's padding count from the target topology. Writing rows whose contents nothing promises seems worth avoiding, since readers tend to start depending on them. Though this is @sfc-gh-truwase's call and there may be compatibility reasons I can't see from here.

One limit I should be explicit about. The above covers contiguous row ranges, which is what legacy TP vocab sharding is. I have not verified that composition stays affine when both sides carry non-unit strides — e.g. a Yuan-style strided shard on one end and a fused layout on the other. Intersecting two strided lattices is not obviously a single affine piece, and if it isn't, that case needs either a piece split or an escalation rule. I'll check it next and report, since it bears directly on §2.1's open question about whether explicit indices are ever needed.

Spec draft is otherwise in good shape — the lowering table for the current uc_info keys is done, and I have a working reference implementation of the piece/map structure with tests covering the four layouts from my last comment. Writing the truncation result into it now.

[Achyuthan-S]

@delock Checked the strided case I flagged. Composition stays affine — no escalation to explicit indices is needed. But there is a real subtlety in how the pieces are formed, and getting it wrong makes the affine language look insufficient when it isn't.

First, a correction to my own framing. I described the case as "a Yuan-style shard on one end and a fused layout on the other" — that isn't a real scenario, since you don't reshard a Yuan model into a CodeGen layout. The scenario that actually matters for phase 2 is the same layout at two different TP degrees, which still puts non-unit strides on both ends. That is what I tested.

Method. Mark F with a unique value per element, run the real partition function at both TP degrees to get shard_position -> F_offset for every rank, invert the source side, then compress each target shard into affine runs and check the result against random data with F never materialised.

Result — every layout, every pair, exact.

layout        src  tgt   elems    1-D   N-D grouped
codegen         2    4     768     44           20
codegen         4    2     768     44           20
codegen         2    8     768     96            8
codegen         8    2     768     96            8
codegen         4    8     768     88           40
codegen         8    4     768     88           40
bigcode         2    4    1536      7            7
bigcode         4    2    1024      6            6
bigcode         2    8    2560     15           15
bigcode         8    2    1024     10           10
bigcode         4    8    2560     15           15
bigcode         8    4    1536     12           12
yuan_value      2    4     512      8            4
yuan_value      4    2     512      8            4
yuan_oproj      2    4    2048    256            4
yuan_oproj      4    2    2048    256            4

The subtlety — two things the composer must do, or affine looks inadequate.

1. Compress in N dimensions, not 1-D. A column split emits one run per row: same length, same strides, starts advancing by a constant. That is an outer dimension, so the whole stack is one piece. Compress only in 1-D and yuan_oproj needs one piece per row — I measured it growing 32 → 1024 as rows went 4 → 128. With N-D pieces it is a constant 4.

2. Group by source rank before folding. A target shard interleaves blocks from several source ranks, so the runs that stack into one piece are usually not adjacent in shard order. For yuan_oproj 4→2 the runs alternate rank 0, rank 1, rank 0, ... and a merge that only fuses neighbours never fires — 256 pieces, again O(rows). Grouping by source rank first drops it to 4, flat from rows 8 to 128.

Either mistake alone produces a piece count that grows with model size, which would look like evidence that the language needs explicit indices. It doesn't; the compressor was just leaving structure on the floor.

Worth noting this is only safe because a piece carries its own destination offset — reordering pieces is permitted. I had originally assumed a shard was just its pieces concatenated in order, and a round-trip test on the column-split case falsified that: a dim-1 shard interleaves its pieces row by row rather than appending them. So each piece now records where it sits in F and where it sits in the shard, same shape on both sides. That fix is what makes grouping by rank legal here.

Everything is bounded by topology and block structure, not model size. Across an 8× hidden-size sweep (F rows 96 → 768): codegen 4→8 stays at 40 pieces, bigcode 2→4 stays at 7, yuan_oproj stays at 4.

So §2.1's open question closes: N cat'd affine pieces do cover every layout in tree, including Yuan. The spec will carry the two compression rules above as normative, since they are exactly where an implementation would silently produce a correct-but-unusable map.

On @jinyouzhi's #8309 — I looked at it for the vocab question you raised. VocabParallelLinear derives vocab_start_index = sum(self._partition_sizes[:self.tp_index]), i.e. it uses the uneven partition_sizes mechanism with no padding at all. That lowers to exactly one affine piece per rank with offset at the prefix sum, which is already the partition_sizes row of the lowering table — a new layer class needing no extension to the language. It also supports your point that padding is a legacy-TP-only concern.

[Achyuthan-S]

@delock Spec draft, covering the outcome you defined: the UC_info description as both in-memory structure and file format, the phase-1 converter, and what phase 2 needs from it. Written against master at 92843ad70.

Decisions it makes, if you want the short version:

Piece	(shape, source_offset, source_strides, dest_offset, dest_strides, locations) — both ends described, shape shared. The shard side is not optional; see §2.1.
Average op	dropped. No producer in tree, and a reduce would break invertibility for the whole language (§4.3).
Vocab padding	option A, resolved with you above. Truncation is a post-step outside the IR (§8.1).
Explicit indices	not needed. Affine covers every in-tree layout and composes across TP degrees (§8.2).
On disk	M_s only, never M_t. Additive key, UCP 0.4 → 0.5, plain scalars (§6).
Out of the IR	output_shape and target_partition_shape — they describe a destination, so by the boundary rule they are planner inputs (§4.2).

Two things are deliberately marked unfinished rather than smoothed over: AutoEP expert placement (§8.3), and the ZeRO-1/3 claim in that same section, which is reasoned rather than measured — unlike §5 and §8.2, which are both verified against the real partition functions.

There is a reference implementation of §2 with tests for §5 (AffinePiece / ParamAffineMap, 14 cases, no process group needed). Happy to open it as a PR once you have had a look at the design — or sooner if you would rather review them together.

---

Affine IR for Universal Checkpoint

Spec draft for #8252 (sub-issue of #8230).
Written against master at 92843ad70, which includes #8185 (aa3914df8).

---

1. Scope and boundary

Universal Checkpoint stores every parameter as one full logical tensor F. Getting into
and out of that form currently requires knowing what a parameter means — is it a
vocabulary embedding, is it row-parallel, does it hold fused sub-parameters — and that
knowledge is encoded as regex lists over parameter names. This spec replaces the meaning
with geometry.

The rule this document follows throughout:

The IR states what is true. The planner decides what to do.

Everything that is a fact about how bytes are laid out belongs in the IR. Everything that
is a choice — which replica to read from, what order to move things in, where to place a
tensor next — belongs to the caller. Section 4 applies this test to every key in the
current schema, and two of them (output_shape, target_partition_shape) fail it: they
describe a destination, not a fact about the source, and so they are not part of the IR.

Non-goals. This spec does not change how any layer partitions its weights at runtime.
It changes only how that partitioning is described. Every layout in tree today keeps the
exact bytes it has on every rank.

---

2. The representation

2.1 Piece

A piece is a block of elements, described by where it sits in the full tensor F and
where it sits in the shard:

Piece = (shape, source_offset, source_strides, dest_offset, dest_strides, locations)

• shape — the extent of the block, shared by both sides: a piece holds the same
  elements in the same arrangement wherever it lives
• source_offset, source_strides — where the block sits in F
• dest_offset, dest_strides — where it sits in the rank's shard
• locations — the set of ranks holding this piece

locations is not the same thing as which rank's shard a piece belongs to. A map lists
pieces per rank because that is what reading or writing one shard needs; locations says
which other ranks hold identical data for that same block. For an ordinary sharded
piece the two coincide. For bigcodetype's replicated kv block they do not: the block
appears in every rank's list, each entry naming all four ranks, and a converter reads it
once from whichever rank is cheapest.

Each side is exactly torch.as_strided's argument list, so a piece is directly executable
with no interpretation step, in either direction.

The shard side has to be described explicitly. It is tempting to say a shard is simply its
pieces concatenated in order, but that is false for a column split, where the shard
interleaves its pieces row by row rather than appending them. An implementation built on
the concatenation assumption reproduces row-split layouts correctly and silently
transposes column-split ones.

2.2 Parameter map

A parameter's map pairs each rank's shard shape with the pieces that fill it:

ParamMap = { logical_shape, { rank: (shard_shape, [Piece, ...]) } }

Because both ends of every piece are affine views, the two directions are the same loop
with the copy reversed — extract writes source_view -> dest_view, rebuild writes
dest_view -> source_view. Neither direction needs to know what the parameter means.

Implementation note. Piece offsets are storage offsets, because that is what
as_strided takes. A shard handed in by a caller is frequently a view into a larger
buffer — splitting a tensor produces exactly that — and its elements then begin partway
into that storage. Applying a piece to such a tensor reads from the wrong place, with no
error raised. An implementation must normalise a shard to a buffer starting at its first
element before applying pieces to it.

2.3 The two maps

• M_s — the source map: the layout the checkpoint was written from. Conversion to
  universal form is M_s⁻¹.
• M_t — the target map: the layout being restored into.

Restore is M_t. Phase 2 — moving weights between two live topologies without ever
materialising F — is M_t ∘ M_s⁻¹.

2.4 Why it is a relation, not a function

Under replication, one position in F lives on several ranks, so M_s⁻¹ maps one element
to a set of locations. This is deliberate. A function would have to name a single owner,
and naming an owner is a decision — the cheapest source depends on topology, link
bandwidth, and what else is in flight, none of which the checkpoint format knows. The IR
records that all of these ranks hold the byte; the planner picks one.

This is the same argument that removes the average/combine op (§4.3): a reduce is not
invertible, so keeping it would make M_s⁻¹ undefined for part of the language.

---

3. Properties

For a parameter map to be usable it must satisfy:

P1 — Coverage. The union of all pieces over all ranks covers every element of F
exactly once or more. If some element is covered by no piece, F cannot be rebuilt.

P2 — Consistency under replication. Where pieces from different ranks overlap, the
underlying data is identical. This is what makes picking any one of them safe, and it is
checkable — the current converter already asserts it for replicated parameters
(ds_to_universal.py:341-342).

P3 — Invertibility. Every piece is a pure view of F. No piece is the result of a
computation over several source elements. Given P1–P3, M_s⁻¹ is total and well defined
as a relation.

P4 — Composability. M_t ∘ M_s⁻¹ is again a set of affine pieces, so a transfer plan
can be computed without materialising F. This is not assumed: §8.2 measures it against
the real partition functions for every in-tree layout. It holds, but only for a composer
that folds pieces in N dimensions and groups them by source rank — §8.2 states both as
normative, because a composer missing either produces a correct map whose size grows with
model size.

P3 is the property that forces the average op out of the language, and the property that
vocabulary truncation currently violates (§8).

---

4. Lowering: every current key

This is the compatibility contract. Every key that exists today must have a defined image
in the IR, or be explicitly declared out of scope.

4.1 Convert path — merge_tp_slices, ds_to_universal.py:340-455

The current implementation is a branch chain over regex-matched name categories. Each
branch lowers to a piece list:

Branch (ds_to_universal.py)	Current behaviour	Lowers to
TP_REPLICATED_PARAMETER_PATTERNS (:340)	take slices[0], assert all equal	1 piece covering all of F, locations = every rank
PARAMETER_TO_AVERAGE_PATTERNS (:345)	sum(slices) / len(slices)	not a view — see §4.3
PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0 (:348)	chunk each slice in 2, cat groupwise	2 pieces/rank
PARAMETER_WITH_SUB_PARAMS + SUB_PARAM_SHAPE (:356)	N sub-params, per-rank widths	N pieces/rank
default (:442)	cat(slices, dim=cat_dim) where cat_dim = 1 if PARAMETER_WITH_ROW_PARALLELISM_PATTERNS else 0	1 piece/rank, split along cat_dim
VOCABULARY_PARAMETER_PATTERNS (:448)	param[:original_vocab_size, :]	not a view — see §8

Two keys are written into the produced checkpoint rather than read from metadata, and
both become redundant:

Key	Why it disappears
CAT_DIM	absorbed into strides — a row split and a column split differ only in stride
PARAM_N_SUB_PARAMS	absorbed into the length of the piece list

SUB_PARAM_SHARD_WIDTHS (UCP 0.4, #8185) is the closest thing in tree to the IR already:
it records the physical extent each rank holds of each sub-parameter. It lowers directly
to per-piece shape entries, and it is the key that makes uneven splits expressible.

4.2 Restore path — _resolve_autotp_partition, universal_checkpoint.py:74-148

Meta key	Current use	Lowers to
replicated (:95)	return all of full_hp_param	1 piece, all of F
partition_dim (:79)	axis to narrow on	the axis whose piece offsets vary
logical_shape (:80)	full_hp_param.view(...)	shape of F
sub_param_sizes / sub_param_shape (:82, :81)	_narrow_sub_params	N pieces
sub_param_shard_widths (:83)	per-rank widths inside _narrow_sub_params	per-piece shape
partition_sizes (:84)	narrow(dim, sum(sizes[:rank]), sizes[rank])	1 piece, offset = prefix sum
fallback (:141-148)	full_view.chunk(world)[rank], asserts divisibility	1 piece, even split — the assert disappears, because uneven extents are expressible
unsupported_reason (:91-93)	raises RuntimeError	eliminated — see §5

output_shape and target_partition_shape are carried in the restore meta
(layers.py:68-123) but describe the destination. By the boundary rule they are planner
inputs, not IR. They stay where they are; the IR does not absorb them.

4.3 The average op

PARAMETER_TO_AVERAGE_PATTERNS is the one branch that cannot be a view: an average is a
reduce over several source elements, and no affine view expresses it. It is also the one
branch that would break P3 for the whole language.

Dropping it is safe in tree. The constant has exactly three references repo-wide — its
declaration (constants.py:88), its import, and its .get in the converter — and no
producer anywhere in the repository writes the key. Nothing in tree can reach that
branch. This matches delock's reading in #8230 that it papers over a historical rank-drift
issue rather than describing a real layout.

---

5. Coverage: the layouts #8185 marks unsupported

#8185 introduced AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS — layouts whose partitioning the
current schema cannot describe, which conversion therefore refuses. If the IR is to
replace the semantic categories, it has to cover these. This was tested rather than
assumed.

Method. Per layout, per rank: run the real DeepSpeed partition function on a marker
tensor to recover which element of F each shard element came from; compress that index
map into maximal runs to obtain pieces; then re-run the same partition function on
independent random data and check that those pieces, materialised with
torch.as_strided, reproduce the shard bit-exactly. The last step is the actual test — if
pieces derived from markers reproduce a random tensor's shard, the mapping is a pure view
(P3), not something data-dependent.

Result — all four are affine-expressible.

Layout	Reason string in master	F	tp	Pieces/rank	Covers F
bigcodetype (fused_LinearLayer, layers.py:1078)	"interleaves or replicates blocks … cannot be reassembled from the shards"	48×8	4	2 (1 on last rank)	yes
codegentype (same class)	same	24×8	2	11	yes
Yuan value, dim 0 (Yuan_LinearLayer, layers.py:1171)	"selects noncontiguous head groups … cannot currently describe"	32×8	2	2	yes
Yuan o_proj, dim 1 (Yuan_LinearAllreduce, layers.py:1137)	same	8×32	2	2, strides (32, 1)	yes

Piece count is bounded by block structure, not model size — CodeGen holds at 11
pieces/rank while rows-per-rank grows 32×:

 hidden   F rows   tp  kv_heads  rows/rank  pieces/rank
      8       24    2         8         12           11
     16       48    2         8         24           11
     32       96    2         8         48           11
     64      192    2         8         96           11
    128      384    2         8        192           11
    256      768    2         8        384           11

Three consequences for this spec:

1. No information is lost in any of these cases. All four cover F exactly (P1), and
   Yuan's pieces are disjoint — it was always a clean partition. "Cannot be reassembled" is
   a limit of the description format, not of the data. That is what makes it fixable by
   an IR change alone.

2. bigcodetype is the concrete case for the location set. Its kv block is
   byte-identical on every rank — partial replication within a single parameter. A schema
   with one partition_dim per parameter cannot say "these rows sharded, those rows
   replicated"; a per-piece locations set says it with no extra mechanism. GPTBigCode /
   StarCoder is an in-tree model this unblocks.

3. Strides are load-bearing. Yuan's dim-1 case is strides=(32, 1) — a column
   selection, not a contiguous span. The representation cannot be narrowed to "a list of
   contiguous ranges."

The harness that produced this table is now a test —
tests/unit/checkpoint/test_affine_shard_map.py, 14 cases covering coverage, round-trip
against unseen data, and rebuild-inverts-extract for all four layouts. It needs no process
group or accelerator, because the partition functions take an explicit rank.

---

6. On-disk format

6.1 What is stored, and what is not

The file stores M_s only — how the job that wrote the checkpoint had the parameter
laid out. M_t is not stored and must not be: the job doing the restore derives its own
map from its own layers, and a target map baked into the file would be a map for somebody
else's topology.

Within M_s, §1's boundary rule splits the content in two:

	stable across jobs?	stored	why
geometry (shape, offsets, strides)	yes — a property of the parameter	yes	this is what makes the checkpoint portable
locations (which ranks held a piece)	no — a property of one grid	yes, as provenance	the converter must know which shard file to read

Storing locations is not a violation of the boundary. The file records where the bytes
were, which is a fact; it does not record where they should go next, which is a decision.
A reader that wants a different placement ignores the field entirely.

6.2 The document

Written into UNIVERSAL_CHECKPOINT_INFO under a new affine_map key, keyed by the same
exact-match patterns collect_autotp_universal_checkpoint_info already emits
(layers.py:626 produces rf"^{re.escape(full_name)}$", one per parameter):

{
  "affine_map": {
    "version": 1,
    "^transformer.h.0.attn.c_attn.weight$": {
      "logical_shape": [48, 8],
      "ranks": {
        0: {"shard_shape": [24, 8],
            "pieces": [
              {"shape": [8, 8],  "source": [0,   [8, 1]], "dest": [0,  [8, 1]], "locations": [0]},
              {"shape": [16, 8], "source": [256, [8, 1]], "dest": [64, [8, 1]], "locations": [0,1,2,3]}
            ]},
        ...
      }
    }
  }
}

Each piece is shape plus [offset, strides] on each side — the argument list of
torch.as_strided twice over, so a reader applies it with no interpretation step.

6.3 Rules

• Additive. A checkpoint carrying affine_map also carries today's keys. Old readers
  ignore the new key and take the existing branch chain; new readers prefer affine_map
  when present and fall back otherwise. Nothing written today becomes unreadable, and
  nothing written by this scheme becomes unreadable by an old converter.
• Version. UNIVERSAL_CHECKPOINT_VERSION_VALUE goes 0.4 → 0.5. The inner
  "version": 1 covers the map encoding itself, so the two can move independently.
• Pieces are canonical. A writer must emit maximally merged pieces, folded in N
  dimensions (§8.2). This is normative, not cosmetic: an unmerged map is correct but its
  size grows with model size rather than with block structure.
• Plain scalars only. No tensors, no pickled classes, no SubparamShape objects. The
  map should be readable without importing DeepSpeed, which matters for external tooling
  and for debugging a checkpoint that will not load.
• locations is per piece, not per parameter. This is the field that expresses
  bigcodetype, whose kv block is held identically by every rank (§5).

6.4 Size

Measured on a fused-QKV parameter at TP=8: 96 pieces, 6836 bytes of compact JSON —
about 71 bytes per piece. Extrapolated by tensor count:

model	tensors	affine_map
Llama-3-8B	~291	~1.9 MB
Llama-3-70B	~723	~4.7 MB

Megabytes, not kilobytes, and it scales with tensor count × TP degree rather than with
parameter size — a 70B and a 7B with the same layer count cost the same. Against
checkpoints measured in hundreds of gigabytes this is negligible, but it is large enough
that the per-rank redundancy (most ranks differ only in offset) is worth revisiting if the
map is ever loaded somewhere latency-sensitive. That is an encoding question and is
deliberately outside the semantics.

---

7. Phase 1 — the converter

7.1 The algorithm

With the map present, merge_tp_slices loses its branch chain entirely:

for each parameter:
    F = empty(logical_shape)
    for rank, pieces in M_s:
        shard = flat_buffer(load_shard(rank))
        for piece in pieces:
            F[piece.source_view] = shard[piece.dest_view]
    assert M_s.uncovered_offsets() == []
    if the parameter is a padded vocabulary:
        F = F[:V_orig]                       # outside the IR, see §8.1
    write F

Restore is the same loop with the copy reversed, against the target's own M_t:

for each rank r of this job:
    shard = empty(M_t.shard_shape[r])
    for piece in M_t.pieces[r]:
        shard[piece.dest_view] = F[piece.source_view]

Both directions are driven by the same piece list. That is the structural reason the
current convert/restore asymmetry — two builders in layers.py:42-123 emitting two
different key sets — stops being expressible.

7.2 Two things an implementation must get right

Normalise the shard buffer. Piece offsets are storage offsets, because that is what
as_strided takes. A loaded shard is often a view into a larger buffer whose elements
begin partway into that storage — splitting a tensor produces exactly that — and applying
a piece to it then reads from the wrong place, silently and with no error. flat_buffer
above is not decorative.

Take any one holder of a replicated piece. Where locations has more than one rank,
every listed rank holds identical data by construction (P2), so the converter reads
whichever is cheapest and skips the rest. Phase 1 is the degenerate case of the location
set; phase 2 is where the choice becomes interesting.

7.3 What changes

Removed: the regex categories, the cat_dim decision, the sub-parameter narrowing
arithmetic, the chunk fallback with its divisibility assert
(universal_checkpoint.py:141-148), and the unsupported_reason refusal path
(:91-93) — the layouts it refuses are expressible (§5).

Added: a coverage assertion that is a real correctness property. Today's checks are shape
heuristics — they confirm the numbers line up, not that every element of the parameter has
a home. uncovered_offsets() answers the second question, and it is the check that fails
loudly if a future layout outgrows the language.

Unchanged: parallelism. Pieces write to disjoint regions of F (or identical values where
they overlap), and parameters remain independent, so the existing ProcessPoolExecutor
fan-out over merge_tp_slices carries over untouched.

---

8. Open cases

8.1 Vocabulary truncation. ds_to_universal.py:451-452 does
param = param[:original_vocab_size, :], dropping padding rows. This is not a view, so it
violates P3 and makes M_s⁻¹ non-total over the shards. Two options:

• (a) F carries the padding; truncation becomes a post-step outside the IR. Every
  piece stays a view and M_s⁻¹ stays well defined for the whole language. Cost: the
  padded size must be recorded.
• (b) The IR admits pieces that map to nothing. Cost: invertibility, i.e. P3, for the
  whole language rather than for one case.

Resolved: (a), agreed with delock on #8252. One key is cheaper than one weakened
property, and the case is narrow — only legacy TP needs an even split and therefore
padding at all; AutoTP's uneven sharding does not pad F.

(a) also survives phase 2, which was the real risk. When a source and target topology
pad to different heights (TP=4 → 3 padding rows, TP=2 → 1), the truncate/pad step is the
identity restricted to [0, V_orig). Restricting an affine map to a range leaves it
affine, so M_t ∘ truncate ∘ M_s⁻¹ composes to plain overlap arithmetic and neither
Fsrc nor Ftgt is ever materialised. Verified over all 64 source/target pairs for
TP ∈ {1,2,4,8} and vocab ∈ {13, 32, 100, 50257}, in both the truncating and padding
directions, with source padding poisoned so that reading it would fail the check.

The composed map is partial over the target — target padding rows have no source.
That is much weaker than M_s being non-total: padding is not data, and the target has to
initialise it regardless. Piece count is bounded by the two TP degrees rather than by the vocabulary: across
vocab 13 → 262144 and TP up to 16 it never exceeds src_tp + tgt_tp, though the exact
count varies with how the two padded heights align.

Consequence: legacy TP does not need excluding from phase 2, and padding need not enter
the UC file even in legacy mode — the composition needs only V_orig.

Relevant: PADDED_VOCAB_SIZE is declared at constants.py:68 and referenced nowhere
else in the repository (verified against 92843ad70) — it looks like exactly this slot,
never wired up. Worth confirming before adding a new key.

8.2 Do we ever need explicit indices? — RESOLVED: no. Every layout in tree is
expressible with affine pieces at bounded count (§5), and composition M_t ∘ M_s⁻¹ between
two TP degrees stays affine for all of them, verified against the real partition functions
with F never materialised. No escape hatch to explicit index lists is needed.

Two rules are normative for any composer, because getting either wrong yields a piece
count that grows with model size — which looks like evidence the language is insufficient:

• Compress in N dimensions. A column split emits one run per row with identical strides
  and constant start deltas; that is an outer dimension, and the stack is one piece. In 1-D
  only, Yuan o_proj needs one piece per row (32 → 1024 as rows go 4 → 128); in N-D it is 4.
• Group runs by source rank before folding. A target shard interleaves blocks from
  several source ranks, so runs that stack are usually not adjacent in shard order. For
  Yuan o_proj 4→2 they alternate rank 0/rank 1, and a neighbour-only merge never fires:
  256 pieces instead of 4.

Grouping is only legal because a piece carries its own destination offset (§2.1), so pieces
may be reordered freely.

Measured piece counts are bounded by topology and block structure, not model size: across
an 8× hidden sweep, codegen 4→8 holds at 40, bigcode 2→4 at 7, Yuan o_proj at 4.

8.3 ZeRO and offload placement. delock's extension in #8230 — "a subset of a parameter
combined with a list of ranks holding this subset" — is what §2.1's locations implements.
ZeRO-1/3 partitions and offload replicas should fall out as pieces whose locations
describe the DP group rather than the TP group, but this spec does not yet work through
AutoEP's expert placement, where locations are per-expert rather than per-parameter.

---

9. Staging

1. Done. IR data structure (deepspeed/checkpoint/affine.py: AffinePiece,
   ParamAffineMap), behind no flag — pure addition, nothing reads it yet.
2. Done. The §5 harness as a test, asserting all four layouts lower and round-trip.
3. Converter reads affine_map when present, existing chain otherwise. Parity test:
   both paths produce byte-identical F for every layout in tree.
4. Emit affine_map from collect_autotp_universal_checkpoint_info, and drop
   unsupported_reason for the four layouts §5 covers.
5. Phase 2 — M_t ∘ M_s⁻¹ as a transfer plan — builds on this and is out of scope here.

Steps 1–3 change no behaviour: they add a representation and prove it agrees with the
current one. Only step 4 changes what a checkpoint contains, and only by adding a key.

[delock]

Hi @Achyuthan-S , thanks for your IR spec. Overall the IR spec looks good. I have two questions when studying your IR definition.

1. I see in Yuan conv layer bias data is pre-divided by world_size, the ration should be allow 'sum-allreduce output then add bias' replaced by 'add bias, then sum-allreduce'. Is it a real case we need to consider and do we need to have scale in piece.
   My thinking is framework optimization passes may want to do this optimization (add bias is free op then move before allreduce) and in-memory bias will be divided. scale will help when we write it back to UC or transfer to another node with different TP.

2. I see in two pieces on different rank may fully overlap (so locations can help us to find these ranks), or non-overlap at all, and partial overlap. For partial overlap, read locations is not sufficient to tell whether some data in this shard is also shared by other ranks. So read all mapping on all ranks looks like inevitable if we want to avoid redudant data transferation during planning pass. So my question is: is locations a necessary information in IR or it is just info that reduce planning effort? If the purpose is reduce planning effort, then we need to define partial overlap as illegal otherwise planner still needs to look at all ranks data to avoid redudant data transfer. However, if partial overlap is not legal, then concat q in last rank with kv in MQA as one piece will become illegal IR, because it induce partial overlap.

[Achyuthan-S]

@delock Both good catches. Taking the second one first, because you found a real defect in the spec.

2. Overlap and locations

You are right that the merged q+kv piece is a problem, and it is worse than that: my own §5 table has the bug in it. "2 pieces per rank (1 on the last rank)" — that last rank is exactly the case you describe. Checking it against the real partition function:

rank 0: F[0:64]    rows 0-7    homogeneous
        F[256:384] rows 32-47  homogeneous
rank 3: F[192:384] rows 24-47  *** SPANS 2 HOLDER-SETS ***

On the last rank the q slice ends where the kv block begins, so maximal-run merging fuses them — and the result is one piece whose first half is rank-exclusive and whose second half is on all four ranks. A single locations cannot describe it. So §6.3's "pieces are maximally merged" rule is wrong as written.

But I don't think the fix is to make partial overlap illegal. There is a third option: keep overlap legal, and forbid merging across a holder-set boundary.

Define holders(e) = the set of ranks holding element e of F. Then require every piece to be homogeneous: all its elements share one holder set, and locations is that set. Canonicalisation becomes "merge maximally, subject to homogeneity."

That makes locations exact rather than a hint, which answers your question directly. The argument is one line: if piece P (locations L) and piece Q (locations M) share an element e, then homogeneity gives holders(e) = L and holders(e) = M, so L = M. Two pieces with different location sets cannot share an element. Partial overlap of the kind that would mislead a planner is impossible by construction, so the planner never has to scan other ranks to discover sharing it was not told about.

Overlap between pieces with the same location set stays legal, and is harmless — both already declare the same sharing.

Measured over the four layouts, naive merging vs homogeneity-constrained merging:

layout        naive  homogeneous  cross-location overlaps
bigcode           7            8                        0
codegen          48           48                        0
yuan_value        4            4                        0
yuan_oproj       32           32                        0

The whole cost is one extra piece, on exactly the rank you predicted. Everything else was already homogeneous. So: locations is necessary IR, not a planning hint, and MQA/bigcode stays expressible.

1. Scale

You are right that this is real, and it is in two places, both the same pattern — a row-parallel layer whose output is sum-all-reduced, so the bias is pre-divided and replicated:

• fusedqkv_utils.py:259 — bias = bias / float(world_size) in the Yuan o_proj path
• layers.py:1299 — bias_data = conv.bias.data / float(self.world_size), last conv layer, split input channel

I think scale should be in the piece, for a reason that is different from why we dropped the average op.

A reduce is N -> 1 and destroys information, so M_s⁻¹ is undefined on it. A scale is 1 -> 1 and is a bijection for any non-zero constant, so M_s⁻¹ is just the reciprocal. Admitting scale does not reopen what dropping average closed. The line becomes: the IR may carry an invertible elementwise map; it may not carry a reduction. I find that a more defensible boundary than "pure views", and it is still one that a checkpoint reader can execute without knowing what the parameter means.

The decisive argument for me is the one you gave: the divisor is world_size, so it changes with topology. Going TP=4 to TP=2 is x4 then /2. If the IR does not carry the scale, the only way to know a bias needs that treatment is a rule that says "this is an all-reduce bias" — which is a semantic category, and re-introduces exactly the thing the IR exists to remove. Carrying a scalar on the piece is the non-semantic way to state it.

Two practical notes:

Precision is better than I expected. (b / N) * N is exact for power-of-two N in fp32, bf16 and fp16 alike — dividing by 2^k only shifts the exponent. I checked; it round-trips bit-exactly for N ∈ {2,4,8,16} and is lossy for N ∈ {3,6,12}. So the round-trip is exact for every TP degree anyone actually runs, and the lossy case is worth an explicit note in the spec rather than silence.

Scale should be per piece and a scalar, for the same reason as locations: homogeneity extends to it unchanged — a piece must be uniform in holder set and in scale. A per-element scale vector would not be affine and I would keep it out.

Spec changes

• §2.1 — piece gains scale; the invertible-elementwise vs reduction boundary replaces "pure views"
• §4.3 — restated as "no reductions", with scale as the contrasting invertible case
• §5 — the bigcode row corrected to 2 pieces on every rank including the last, and the "1 on last rank" note removed
• §6.3 — canonicalisation becomes "maximally merged subject to homogeneity in locations and scale"
• §3 — a new property: every piece is homogeneous, with the one-line argument that cross-location overlap is therefore impossible
• §8 — the fp precision note for non-power-of-two TP

I will fold the two bias sites into the coverage test as well, since they are a layout the current spec would have got wrong.

Does it look good now?

[sfc-gh-truwase]

@minjiaz , @xylian86 FYI

[delock]

@delock Both good catches. Taking the second one first, because you found a real defect in the spec.

2. Overlap and locations

You are right that the merged q+kv piece is a problem, and it is worse than that: my own §5 table has the bug in it. "2 pieces per rank (1 on the last rank)" — that last rank is exactly the case you describe. Checking it against the real partition function:

rank 0: F[0:64]    rows 0-7    homogeneous
        F[256:384] rows 32-47  homogeneous
rank 3: F[192:384] rows 24-47  *** SPANS 2 HOLDER-SETS ***

On the last rank the q slice ends where the kv block begins, so maximal-run merging fuses them — and the result is one piece whose first half is rank-exclusive and whose second half is on all four ranks. A single locations cannot describe it. So §6.3's "pieces are maximally merged" rule is wrong as written.

But I don't think the fix is to make partial overlap illegal. There is a third option: keep overlap legal, and forbid merging across a holder-set boundary.

Define holders(e) = the set of ranks holding element e of F. Then require every piece to be homogeneous: all its elements share one holder set, and locations is that set. Canonicalisation becomes "merge maximally, subject to homogeneity."

That makes locations exact rather than a hint, which answers your question directly. The argument is one line: if piece P (locations L) and piece Q (locations M) share an element e, then homogeneity gives holders(e) = L and holders(e) = M, so L = M. Two pieces with different location sets cannot share an element. Partial overlap of the kind that would mislead a planner is impossible by construction, so the planner never has to scan other ranks to discover sharing it was not told about.

Overlap between pieces with the same location set stays legal, and is harmless — both already declare the same sharing.

Measured over the four layouts, naive merging vs homogeneity-constrained merging:

layout        naive  homogeneous  cross-location overlaps
bigcode           7            8                        0
codegen          48           48                        0
yuan_value        4            4                        0
yuan_oproj       32           32                        0

The whole cost is one extra piece, on exactly the rank you predicted. Everything else was already homogeneous. So: locations is necessary IR, not a planning hint, and MQA/bigcode stays expressible.

1. Scale

You are right that this is real, and it is in two places, both the same pattern — a row-parallel layer whose output is sum-all-reduced, so the bias is pre-divided and replicated:

• fusedqkv_utils.py:259 — bias = bias / float(world_size) in the Yuan o_proj path
• layers.py:1299 — bias_data = conv.bias.data / float(self.world_size), last conv layer, split input channel

I think scale should be in the piece, for a reason that is different from why we dropped the average op.

A reduce is N -> 1 and destroys information, so M_s⁻¹ is undefined on it. A scale is 1 -> 1 and is a bijection for any non-zero constant, so M_s⁻¹ is just the reciprocal. Admitting scale does not reopen what dropping average closed. The line becomes: the IR may carry an invertible elementwise map; it may not carry a reduction. I find that a more defensible boundary than "pure views", and it is still one that a checkpoint reader can execute without knowing what the parameter means.

The decisive argument for me is the one you gave: the divisor is world_size, so it changes with topology. Going TP=4 to TP=2 is x4 then /2. If the IR does not carry the scale, the only way to know a bias needs that treatment is a rule that says "this is an all-reduce bias" — which is a semantic category, and re-introduces exactly the thing the IR exists to remove. Carrying a scalar on the piece is the non-semantic way to state it.

Two practical notes:

Precision is better than I expected. (b / N) * N is exact for power-of-two N in fp32, bf16 and fp16 alike — dividing by 2^k only shifts the exponent. I checked; it round-trips bit-exactly for N ∈ {2,4,8,16} and is lossy for N ∈ {3,6,12}. So the round-trip is exact for every TP degree anyone actually runs, and the lossy case is worth an explicit note in the spec rather than silence.

Scale should be per piece and a scalar, for the same reason as locations: homogeneity extends to it unchanged — a piece must be uniform in holder set and in scale. A per-element scale vector would not be affine and I would keep it out.

Spec changes

• §2.1 — piece gains scale; the invertible-elementwise vs reduction boundary replaces "pure views"
• §4.3 — restated as "no reductions", with scale as the contrasting invertible case
• §5 — the bigcode row corrected to 2 pieces on every rank including the last, and the "1 on last rank" note removed
• §6.3 — canonicalisation becomes "maximally merged subject to homogeneity in locations and scale"
• §3 — a new property: every piece is homogeneous, with the one-line argument that cross-location overlap is therefore impossible
• §8 — the fp precision note for non-power-of-two TP

I will fold the two bias sites into the coverage test as well, since they are a layout the current spec would have got wrong.

Does it look good now?

Hi @Achyuthan-S thanks for consider and resolve these two points. The amend looks good to me. Thanks!

[delock]

For staging, I agree that 1-3 could be done in one PR and 4 should be done in a separate PR. I also suggest add spec doc inside deepspeed (deepspeed/checkpoint/affine_ir_spec.md) for better understanding, and reference this doc in affine IR code.

@sfc-gh-truwase for comments.