Describe universal checkpoint shards as affine maps - #8385
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b4dc5e37ce
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # category branches below are consulted. A checkpoint converted this way carries | ||
| # the map rather than the per-category keys those branches write, because the | ||
| # geometry is what the restoring side needs and it is not tied to a category. | ||
| param = matched_affine_map.rebuild(dict(enumerate(slices))) |
There was a problem hiding this comment.
When an affine piece has a non-unit scale, this branch invokes the same rebuild operation for fp32, exp_avg, and exp_avg_sq. The scale describes the relationship between parameter values, whereas Adam's first and second moments transform differently under a parameter-coordinate change; in particular, applying one division by the parameter scale to both moments cannot be correct. Converting a scaled row-parallel bias therefore changes its saved optimizer state and the trajectory after resuming, so scaling needs to be state-aware or restricted to the parameter value.
Useful? React with 👍 / 👎.
| for rank, pieces in self.pieces_by_rank.items(): | ||
| flat_shard = _flat_buffer(shards[rank]) | ||
| for piece in pieces: | ||
| target = piece.source_view(full_param) | ||
| target.copy_(piece.dest_view(flat_shard)) |
There was a problem hiding this comment.
When multiple ranks hold a replicated piece and their values differ because of rank drift or checkpoint corruption, this loop copies every replica into the same destination and silently lets the last rank win. The existing replicated conversion branch explicitly checks that all slices are equal, so switching that layout to an affine map removes a corruption guard. Check overlapping replicas for equality, or select a single holder only after validating consistency.
Useful? React with 👍 / 👎.
| for rank, pieces in self.pieces_by_rank.items(): | ||
| flat_shard = _flat_buffer(shards[rank]) | ||
| for piece in pieces: |
There was a problem hiding this comment.
If an affine entry is stale or selected by an overlapping regex and records a shard smaller than the actual slice, validate() only compares its pieces with the metadata's own shard_shapes; this line then flattens the real tensor and reads the described prefix while silently ignoring its remaining values. Assert that every supplied shard's shape or element count matches self.shard_shapes[rank] before applying pieces so mismatched metadata cannot produce a plausible but incomplete parameter.
Useful? React with 👍 / 👎.
| def _offsets(self, base, strides): | ||
| if not self.shape: | ||
| yield base | ||
| return |
There was a problem hiding this comment.
For layouts where a TP rank receives a zero-width sub-parameter, a piece can have a shape such as (0, 8). _offsets() currently enters the iteration and yields eight offsets even though numel is zero, allowing an entirely empty piece to make uncovered_offsets() return an empty result and validate_coverage() approve a map that contains no data for those elements. Return immediately when any dimension is zero.
Useful? React with 👍 / 👎.
| @@ -0,0 +1,420 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
There was a problem hiding this comment.
This non-merge commit has no Signed-off-by trailer, so it violates the repository's mandatory commit-signing requirement and is liable to fail the corresponding CI/DCO check. Recreate the commit with --signoff using the configured Git identity.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Pull request overview
This PR introduces an affine (geometry-based) intermediate representation for describing how tensor-parallel shards map to a full logical parameter tensor, and wires the universal-checkpoint converter to read this representation when present. The goal is to eventually remove reliance on semantic regex categories for conversion by making shard layouts executable from recorded offsets/strides.
Changes:
- Added
deepspeed/checkpoint/affine.pyimplementingAffinePieceandParamAffineMapplus constructors for existing converter layouts. - Updated
merge_tp_slicesindeepspeed/checkpoint/ds_to_universal.pyto prefer rebuilding parameters via anaffine_mapentry (when present) and added new constants for the UC-info key shape. - Added a spec document and unit tests validating coverage, invertibility, and parity with today’s
merge_tp_slicesarithmetic.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
tests/unit/checkpoint/test_affine_shard_map.py |
Adds pytest coverage to prove specific previously-unsupported AutoTP layouts are representable as affine views and round-trip correctly. |
deepspeed/checkpoint/ds_to_universal.py |
Adds optional affine-map-based merge path in merge_tp_slices and validates affine map version presence. |
deepspeed/checkpoint/constants.py |
Introduces UC-info keys for storing affine-map metadata (affine_map, version, params). |
deepspeed/checkpoint/affine.py |
Implements the affine IR (pieces + per-rank maps), serialization, and constructors for common sharding layouts. |
deepspeed/checkpoint/affine_ir_spec.md |
Documents the IR, intended properties, and staging plan for adoption. |
Suppressed comments (3)
deepspeed/checkpoint/affine.py:217
ParamAffineMap.rebuildinitializesfull_paramwithtorch.empty. If a malformed/partial map leaves any offsets uncovered (and onlyvalidate()is called), the result can contain uninitialized garbage values. Initializing with zeros makes failure modes deterministic and avoids propagating uninitialized data.
self.validate()
any_shard = next(iter(shards.values()))
full_param = torch.empty(self.numel, dtype=any_shard.dtype, device=any_shard.device)
deepspeed/checkpoint/affine.py:336
contiguous_split_mapdoes not validate thatper_rank_sizesexactly covers the full tensor alongpartition_dim. If the sizes don't sum toshape[partition_dim],rebuild()can silently leave parts of the full tensor unwritten.
shape = tuple(shape)
source_strides = _row_major_strides(shape)
pieces_by_rank = {}
shard_shapes = {}
start = 0
deepspeed/checkpoint/affine.py:366
sub_param_mapassumessub_dim_sizesandshard_widthsare consistent (counts, tp degree, and per-sub-param sums) but doesn't validate them. If these inputs are inconsistent, the resulting map can be structurally valid per-rank while leaving holes/overlaps in the logical tensor.
shape = tuple(shape)
source_strides = _row_major_strides(shape)
tp_degree = len(shard_widths[0])
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| affine_map_version = affine_map_info.get(AFFINE_MAP_VERSION, AFFINE_MAP_FORMAT_VERSION) | ||
| assert affine_map_version <= AFFINE_MAP_FORMAT_VERSION, ( | ||
| f"Checkpoint records affine map format version {affine_map_version}, but this DeepSpeed understands " | ||
| f"up to {AFFINE_MAP_FORMAT_VERSION}. Reading it could misinterpret fields added since.") | ||
| affine_params = affine_map_info.get(AFFINE_MAP_PARAMS, {}) |
| self.scale = float(scale) | ||
| assert self.scale != 0.0, 'A zero scale is not invertible, so the full tensor could not be rebuilt.' |
| for pattern_, entry_ in affine_params.items(): | ||
| if re.match(pattern_, name_): | ||
| return ParamAffineMap.from_dict(entry_) | ||
| return None |
| # The pieces say where every element of the parameter lives, so none of the | ||
| # category branches below are consulted. A checkpoint converted this way carries | ||
| # the map rather than the per-category keys those branches write, because the | ||
| # geometry is what the restoring side needs and it is not tied to a category. |
afab92b to
5769544
Compare
|
The modal-torch-latest failure is a dependency-install error, not a test failure — THESE PACKAGES DO NOT MATCH THE HASHES FROM THE REQUIREMENTS FILE while downloading torch/CUDA wheels, before any test runs. This branch passed the same workflow two hours ago (run 33527949501) and the diff touches no requirements or CI files. Looks transient — could someone re-run it? @delock . |
|
@delock ., The modal-torch-latest failures aren't from this change. The first was a pip hash mismatch while installing torch/CUDA wheels (THESE PACKAGES DO NOT MATCH THE HASHES FROM THE REQUIREMENTS FILE), before any test ran; the second was the job being cancelled at the 75-minute mark. Neither reported a test failure, and the diff touches no requirements or CI files. This branch passed the same workflow green in 62m on the first push (run 33527949501). Happy to rebase or re-run if that helps. Also correcting the description: it says 31 test cases, now 39 after the review fixes. |
| @@ -0,0 +1,455 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
There was a problem hiding this comment.
Remove Microsoft license header.
There was a problem hiding this comment.
Done. Switched both new files to # Copyright (c) DeepSpeed Team. to match what files added recently carry. Confirmed with scripts/check-license.py — the Microsoft line is commented out in its COPYRIGHT list, so only the SPDX and DeepSpeed Team lines are actually required.
| where it sits in the shard: | ||
|
|
||
| ```python | ||
| Piece = (shape, source_offset, source_strides, dest_offset, dest_strides, locations) |
There was a problem hiding this comment.
scale should appear in this line.
There was a problem hiding this comment.
Fixed — it was in the notes below but not in the tuple itself.
|
Hi @Achyuthan-S thanks for your PR. I have read the spec part and left my comments. Nice catch to optimizer states scaling where we didn't covered early. I like the json structure and how that simplify the code in 7.1 Modal failure is due to a different reason and we are working on it and hopefully it won't block PRs again. I'll come back to review the rest part next week. Also need to ponder on the optimizer state scaling part a little more. The rest part of the spec document itself looks fine to me. Its better if you have an example on optimizer state scaling to help understanding. Thanks! |
Describe a parameter's tensor-parallel layout as affine views of the full tensor instead of matching its name against regex categories. Each piece records where a block sits in the full parameter and in the shard, so both conversion directions are the same copy with the ends swapped. Covers the fused QKV and Yuan shared-QK layouts that AutoTP currently marks unsupported, with tests showing their shards do cover the full parameter. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
A row-parallel layer pre-divides its replicated bias by the world size, so a piece carries the factor its shard holds the block by. Scaling is invertible where a reduction is not, so this keeps conversion reversible in both directions. A piece must also cover elements held by the same set of ranks. Merging by adjacency alone fuses a rank-private block onto a replicated one where they happen to be neighbours, leaving a piece whose own locations is wrong for half of it. Adds the specification the module implements. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Add constructors for the layouts the universal checkpoint converter already handles: replicated parameters, contiguous splits along either axis, and parameters holding several sub-parameters split unevenly across ranks. Row and column parallelism differ only in stride, so one constructor covers both and the recorded concat dimension becomes redundant. Tests require each constructor to reproduce merge_tp_slices' own arithmetic exactly, so a map can replace a branch without changing what a checkpoint converts to. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Read the geometric description from universal checkpoint info and rebuild the parameter from it, falling back to the existing category branches when a parameter has no map. Nothing writes a map yet, so this changes no conversion. Add the on-disk form, which holds plain scalars so the map can be read without importing DeepSpeed, and omits the scale factor where it is 1. Tests require each constructor to reproduce merge_tp_slices' own arithmetic, including the uneven sub-parameter widths that earlier metadata could not describe. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
AutoTP now derives kv-head and grain values into an AutoTPMeta per model instead of process-wide globals, so the tests construct one and pass it to the partition functions they exercise. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
The stored map declares its own encoding version, separate from the universal checkpoint version so the two can move independently. A reader that predates a version would otherwise misinterpret fields added since, so refuse instead. Keep the stride helper private, since nothing outside the module builds a piece by hand yet. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Nothing outside the module builds a piece by hand yet, so exporting the helper widens the public surface for no caller. Step 4 can promote it when it needs it. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Adam's moments live in the parameter's scaled coordinate, so scaling a parameter by s scales its gradient by 1/s. Applying the parameter's factor to the moments would corrupt the optimizer state and change the trajectory after a resume, so the caller now says which power of the scale applies. Refuse rather than guess in four more places: replicas that disagree, a shard whose size contradicts the map, a map format newer than this reader, and a zero scale. Empty pieces no longer report covering elements they do not hold. Raise instead of asserting where the check guards checkpoint compatibility, since asserts are stripped under python -O. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Use the DeepSpeed Team copyright line that new files carry; the license check does not ask for the Microsoft one. Record scale in the piece definition itself, not only in the notes below it, and add a worked example of how each optimizer state recovers from a scaled shard. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
5769544 to
35eb2e2
Compare
Thanks @delock. All three addressed in the latest push, and good to know the modal failure is being handled separately. Optimizer state example — added to §2.1 as a worked table, for a row-parallel bias at world size 4 (scale = 1/4). A shard holds full * spower and conversion recovers full = shard / spower: One thing worth flagging while you ponder it: those powers are derived, not measured. They follow from the chain rule, and the round-trip is tested, but nothing yet trains a model with a scaled bias, converts, resumes, and checks the trajectory matches. §5 and §8.2 are verified against the real partition functions; this part is reasoned. If you'd rather the IR simply refused a scaled optimizer state until that end-to-end check exists, that's a smaller and more defensible surface — happy to go that way. Also correcting the description: it said 31 test cases, now 39 after the review fixes. |
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_infois 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_slicesis 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_PATTERNSrecords. 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.mdis the specification, developed in #8252 and #8230.affine.pyreferences it.What is here
AffinePiece— a block of elements, recording where it sits in the full tensor and where it sits in the shard, withshapeshared. Each side istorch.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 ownerscale, 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 -> 1and reverses, a reduce isN -> 1and does not.ParamAffineMap—extractandrebuild, 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 (wheretorch.chunkand AutoTP disagree) and sub-parameters of different sizes, none dividing evenly by the tp degree.Coverage — the four layouts
AUTOTP_UNSUPPORTED_PARAMETER_PATTERNScurrently 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_stridedtakes. 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_buffernormalises 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 (aa3914d).
cc @delock