Describe universal checkpoint shards as affine maps

[Achyuthan-S]

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 (aa3914d).

cc @delock

[chatgpt-codex-connector]

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

[chatgpt-codex-connector]

Transform optimizer states independently of parameter scaling

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 👍 / 👎.

[chatgpt-codex-connector]

Verify replicated pieces before overwriting them

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 👍 / 👎.

[chatgpt-codex-connector]

Reject shards that disagree with their recorded shapes

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 👍 / 👎.

[chatgpt-codex-connector]

Treat zero-extent pieces as covering no offsets

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 👍 / 👎.

[chatgpt-codex-connector]

Add the required Signed-off-by trailer

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 👍 / 👎.

[Copilot]

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.py implementing AffinePiece and ParamAffineMap plus constructors for existing converter layouts.
• Updated merge_tp_slices in deepspeed/checkpoint/ds_to_universal.py to prefer rebuilding parameters via an affine_map entry (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_slices arithmetic.

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.rebuild initializes full_param with torch.empty. If a malformed/partial map leaves any offsets uncovered (and only validate() 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_map does not validate that per_rank_sizes exactly covers the full tensor along partition_dim. If the sizes don't sum to shape[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_map assumes sub_dim_sizes and shard_widths are 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.

[Achyuthan-S]

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 .

[Achyuthan-S]

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

[delock]

Remove Microsoft license header.

[Achyuthan-S]

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.

[delock]

scale should appear in this line.

[Achyuthan-S]

Fixed — it was in the notes below but not in the tuple itself.

[delock]

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!

[Achyuthan-S]

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!

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:

| state      | power | s**power | a rank holds | F recovers to          |
|------------|-------|----------|--------------|------------------------|
| fp32       | 1     | 1/4      | 2.0          | 8.0 — the logical bias |
| exp_avg    | -1    | 4        | 8.0          | 2.0                    |
| exp_avg_sq | -2    | 16       | 32.0         | 2.0                    |

The moments move the opposite way to the parameter, and the second twice as far, because the optimizer trains p = b/4 and ∂L/∂b = (∂L/∂p)·(1/4). Using the parameter's own factor for all three would multiply exp_avg by 4 where it should be divided — off by 16× — and change the trajectory after a resume without any visible error. The numbers in that table are the actual output of ParamAffineMap.rebuild, not hand-derived.

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.

[0z5a]

@Achyuthan-S @delock Following the train→convert→resume gap discussed here and in #8230/#8252, I ran actual checkpoint tests on top of 35eb2e2.

Follow-up commit adds a small column/row AutoTP model with real FP32 ZeRO-1 Adam training: TP2 saves after four steps, converts through either the legacy or affine path, and resumes for four more steps at TP1/TP2. It compares logits, losses, gradients, FP32 parameters, both moments, and step counters with uninterrupted training. Affine metadata is supplied explicitly through checkpoint client state, since the producer is a later step. Native LinearAllreduce adds its bias after reduction, so that bias's scale in this test is 1.

The test found two pre-existing loader problems in both conversion paths:

• TP2→TP1 invoked the Megatron weight merger and failed looking for attention.dense.weight.
• map_to_flat_opt_states rebound the FP32 parameter fragment to a moment buffer, instead of rebinding the corresponding optimizer fragment.

The commit fixes these. All four resume cases passed on 2× RTX 4000 Ada 20 GB, torch 2.13.0+cu130, Python 3.12.14 / NCCL. CPU/Gloo also passed those four cases, 71 affine/AutoTP unit regressions, and four existing fragment/DP-resize regressions. Modified-file pre-commit checks passed.

pytest -sv tests/unit/checkpoint/test_autotp_uc_checkpoint.py::TestAffineUniversalCheckpointResume

There is a separate non-unit-scale result: standalone real-checkpoint reproducer and recorded output. This is an explicitly parameterized p = 0.5*b bias on one CPU/Gloo rank, with clipping disabled and a plain torch Adam oracle; it is not native AutoTP bias. The trained parameter and moments convert/restore correctly with the current powers, but loading also restores source-coordinate lr=0.015, eps=0.1, whereas the target coordinates need lr=0.03, eps=0.05. The next update has max parameter error 0.0087981, growing to 0.0272084 after four updates. Restoring those target-coordinate group hyperparameters gives zero parameter error at every step in this probe.

So this supports the moment powers for that explicit reparameterization, but not general scaled-optimizer resume with unchanged group hyperparameters. Before enabling that path, the contract needs to cover optimizer hyperparameters and replicated-gradient semantics, or reject unsupported scaled optimizer states as you suggested. The follow-up commit is available for cherry-picking; its regression scope is the native scale-one path.

[Achyuthan-S]

@0z5a Thanks for running this — an end-to-end resume was the one thing the spec asserted without measuring, and it's good to have it closed.

I reproduced the hyperparameter result independently with a plain Adam oracle, and it matches yours exactly. Keeping the source lr/eps, max parameter error is 0.0068 after one update growing to 0.0160 after four; transforming to lr / s and eps * s gives exactly zero at every step. Your lr 0.015 → 0.03, eps 0.1 → 0.05 is what the chain rule gives: the optimizer trains p, so matching Δb = Δp / s needs lr_b = lr_p / s and eps_b = eps_p * s.

The part worth stating clearly, since it's easy to read this as the powers being wrong: they aren't. Your probe confirms the parameter and both moments convert and restore correctly under the powers in §2.1. The divergence is entirely from the optimizer's own hyperparameters, which sit a layer below parameter geometry.

That makes this the case I raised with @delock last week — that the powers are derived rather than measured, and that refusing scaled optimizer states is the smaller and more defensible surface until the contract covers them. Your data settles which way to go, so I've taken it: rebuild/extract now raise NotImplementedError on a scaled optimizer state, naming the reason. A scaled parameter still converts normally, since that path is correct. §2.1 records the constraint and the lr / s, eps * s transform a future contract will need. The factors themselves stay pinned in tests — they're what that contract builds on.

On the two loader bugs: those are pre-existing in both conversion paths and independent of this change, and the map_to_flat_opt_states one looks like the more serious of the two — rebinding a parameter fragment to a moment buffer would corrupt a resume with nothing to show for it. I'd suggest sending them as their own PR rather than folding them in here: they're a bug fix in shared code, they'll review faster on their own, and they shouldn't wait on an IR discussion. Happy to review it.

If you'd rather they ride along with this one, I'll cherry-pick with attribution — either way works, and thanks again for the thorough write-up.

[Achyuthan-S]

@delock could you kick off CI when you get a chance? Modal hasn't managed to complete a run on this branch yet, so the suite hasn't actually been exercised.

[0z5a]

@0z5a Thanks for running this — an end-to-end resume was the one thing the spec asserted without measuring, and it's good to have it closed.

I reproduced the hyperparameter result independently with a plain Adam oracle, and it matches yours exactly. Keeping the source lr/eps, max parameter error is 0.0068 after one update growing to 0.0160 after four; transforming to lr / s and eps * s gives exactly zero at every step. Your lr 0.015 → 0.03, eps 0.1 → 0.05 is what the chain rule gives: the optimizer trains p, so matching Δb = Δp / s needs lr_b = lr_p / s and eps_b = eps_p * s.

The part worth stating clearly, since it's easy to read this as the powers being wrong: they aren't. Your probe confirms the parameter and both moments convert and restore correctly under the powers in §2.1. The divergence is entirely from the optimizer's own hyperparameters, which sit a layer below parameter geometry.

That makes this the case I raised with @delock last week — that the powers are derived rather than measured, and that refusing scaled optimizer states is the smaller and more defensible surface until the contract covers them. Your data settles which way to go, so I've taken it: rebuild/extract now raise NotImplementedError on a scaled optimizer state, naming the reason. A scaled parameter still converts normally, since that path is correct. §2.1 records the constraint and the lr / s, eps * s transform a future contract will need. The factors themselves stay pinned in tests — they're what that contract builds on.

On the two loader bugs: those are pre-existing in both conversion paths and independent of this change, and the map_to_flat_opt_states one looks like the more serious of the two — rebinding a parameter fragment to a moment buffer would corrupt a resume with nothing to show for it. I'd suggest sending them as their own PR rather than folding them in here: they're a bug fix in shared code, they'll review faster on their own, and they shouldn't wait on an IR discussion. Happy to review it.

If you'd rather they ride along with this one, I'll cherry-pick with attribution — either way works, and thanks again for the thorough write-up.

Cherry-picking with attribution works for me — thanks!

[delock]

Hi @Achyuthan-S , thanks for provide the example. @0z5a thanks for running the checkpoint test. This discussion reveals scale!=1 should be a case for rollout end checkpoint only. During inference bias could be moved before allreduce and scale could help. In the training end (where optimizer and lr/eps needs to be scaled), keep scale=1 should be better choice. My rational is handling scale with optimizer would increase software complexity, handling scale in inference would catch most potential performance benefit.

In three map, why only contiguous_split_map accept scale as parameter, can you explain?

The intergration test is a good idea, thanks @0z5a for contribution and thanks @Achyuthan-S for cherrypicking.

[delock]

@delock could you kick off CI when you get a chance? Modal hasn't managed to complete a run on this branch yet, so the suite hasn't actually been exercised.

@Achyuthan-S I started one round of test. For modal test, they will be tested when this PR enter merge queue. We will see the result when we start the merge process.

[Achyuthan-S]

@delock Answers below, plus two changes pushed since your approval — flagging so you know a re-approval may be needed.

On the three maps. scale belongs on replicated_map, and I've moved it there. The layouts that pre-divide a value hold it whole on every rank: Yuan's o_proj and the last conv layer both replicate the bias divided by the world size, so the all-reduced sum adds it exactly once. The weight beside them is what gets split, and it's unscaled — so a split constructor has nothing to scale, and contiguous_split_map and sub_param_map now take no scale argument. There's a test covering the replicated case.

On scale != 1 being rollout-only — agreed, and the code already lands there. Refusing scaled optimizer states means a training checkpoint, which carries exp_avg/exp_avg_sq, cannot have scale != 1; an inference or rollout export carries only the parameter, so it can. Happy to state that as an explicit rule in §2.1 if you'd like it written down rather than implied.

Cherry-picked @0z5a's fixes as agreed, with attribution. Worth noting this changes the PR's scope: engine.py and tensor_fragment.py are live paths, so it is no longer purely additive. I've updated the description to say so. Locally on CPU/gloo their four resume cases pass alongside the 40 affine tests — 43 passed.

On the merge queue — both attempts ended with SandboxStartTimeout: Sandbox did not start within 1800s, so no test ran, which the job reports as a capacity problem rather than a test failure. All 12 PR checks are green. Ready to re-queue whenever it's worth another attempt.

[Achyuthan-S]

@delock The current queue entry predates the two commits I pushed just after you queued — its squash is the earlier 5-file version, without the scale move or @0z5a's fixes. Worth dequeuing and re-adding on dd5c57f so those land.

[delock]

Hi @Achyuthan-S ，push new commit should deque a merge request, but for some reason this time it didn't happen. I'll review #8477 and 0z5a's #8474, thank you!