Describe universal checkpoint shards as affine maps (#8385)
Implements steps 1–3 of the staging plan in #8252: the IR structure, the
lowering from today's metadata, and the converter reading it. Emitting
the map from `collect_autotp_universal_checkpoint_info` is step 4 and
will be a separate PR, per @delock's suggestion.
**No conversion changes.** Nothing writes an affine map yet, so the new
branch in `merge_tp_slices` is never taken and every checkpoint converts
exactly as it does today. The tests are what give the work its value at
this stage: they require the map to reproduce the existing arithmetic
before anything depends on it.
### Why
Universal checkpoint decides how to merge a parameter by matching its
*name* against regex categories — vocabulary, row-parallel, fused
sub-parameters. A layout no category describes cannot be converted at
all, which is what `AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS` records. This
describes the layout geometrically instead, so the question becomes
where the bytes are rather than what the parameter means.
`deepspeed/checkpoint/affine_ir_spec.md` is the specification, developed
in #8252 and #8230. `affine.py` references it.
### What is here
**`AffinePiece`** — a block of elements, recording where it sits in the
full tensor **and** where it sits in the shard, with `shape` shared.
Each side is `torch.as_strided`'s argument list, so a piece is
executable with no interpretation step.
A piece also carries:
- `locations`, the ranks holding it, so a converter can read a
replicated block from whichever rank is cheapest rather than a
designated owner
- `scale`, the factor the shard holds the block by. A row-parallel layer
pre-divides its replicated bias by the world size so the all-reduced sum
adds the bias once; the divisor changes with the world size, so a
checkpoint that cannot record it cannot be restored at a different TP
degree without a rule naming which parameters are biases.
Scaling is admitted where averaging is not, and the line is
invertibility rather than arithmetic: a scale is `1 -> 1` and reverses,
a reduce is `N -> 1` and does not.
**`ParamAffineMap`** — `extract` and `rebuild`, which are the same loop
with the copy reversed, plus coverage and homogeneity validation and the
on-disk form.
**Lowering constructors** for the layouts the converter already handles:
`replicated_map`, `contiguous_split_map`, `sub_param_map`. Row and
column parallelism differ only in stride, so one constructor covers both
— which is why the recorded concat dimension becomes redundant.
### Tests
39 cases, all plain pytest: the partition functions take an explicit
rank, so none of this needs a process group or an accelerator, and the
module runs in about two seconds on any runner.
**Parity** — each constructor must reproduce `merge_tp_slices`' own
arithmetic exactly, including an uneven `[3, 3, 2, 2]` split (where
`torch.chunk` and AutoTP disagree) and sub-parameters of different
sizes, none dividing evenly by the tp degree.
**Coverage** — the four layouts `AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS`
currently refuses (`bigcodetype`, `codegentype`, Yuan shared-QK value
and o_proj) are all describable. Pieces are derived by running the
**real** partition functions on a marker tensor, then validated against
independent random data. That second step is the actual test: if pieces
derived from markers reproduce a random tensor's shard bit-exactly, the
layout is a pure view rather than something data-dependent.
### Two things worth knowing about
**A shard is not its pieces concatenated in order.** For a column split
the shard interleaves them row by row, so an implementation built on the
concatenation assumption reproduces row-split layouts correctly and
silently transposes column-split ones. A round-trip test on the Yuan
o_proj case caught this, and it is why both ends of a piece are
recorded.
**Piece offsets are storage offsets**, because that is what `as_strided`
takes. A loaded shard is frequently a *view* into a larger buffer —
CodeGen's rank 1 begins at offset 96 of a 192-element buffer — so
applying a piece to it directly reads from the wrong place, silently.
`_flat_buffer` normalises this.
### One boundary for review
The category branches write per-category keys into the converted
checkpoint (`CAT_DIM`, `PARAM_N_SUB_PARAMS`, `SUB_PARAM_SHAPE`) that the
restore path reads. The affine branch cannot reconstruct those, and
arguably should not: the geometry is what a restoring job needs and it
is not tied to a category. So a checkpoint converted through the map
carries the map instead. This is inert until step 4, but it is the
compatibility question I would most like a second opinion on.
Related: #8252, #8230. Builds on #8185 (aa3914df8).
cc @delock
---------
Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
