Emit affine maps from AutoTP layers - #8519
Achyuthan-S wants to merge 2 commits into
Conversation
Derive a parameter's affine map where the layer already knows its per-rank extents. The conversion metadata does not carry them -- they are resolved by _freeze_partition_sizes while the layer is built and are not recoverable from a shape alone -- so the map is built at mark time rather than at collection. Publish the maps alongside the existing pattern lists, so a converter that predates them still reads the checkpoint through the categories. The converter that prefers a map therefore has to mark those patterns as superseded, or a strict conversion rejects them as unused. The resume fixture now takes its layout from the producer instead of supplying one, so the four save-convert-resume cases exercise the metadata a real job writes. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
![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: 129087a461
ℹ️ 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".
| 'original_shape': _normalize_uc_shape(original_shape), | ||
| 'is_bias': is_bias, | ||
| 'replicated': replicated, | ||
| 'affine_map': affine_map, |
There was a problem hiding this comment.
Adding affine_map unconditionally changes every _build_param_uc_restore_meta() conversion dictionary, but tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py:235-244 still compares that dictionary exactly without this key, so the existing unit test now fails even when the argument uses its default. Update that test—preferably to assert only the stable fields rather than the entire private dictionary—or omit the key when no map is supplied.
AGENTS.md reference: AGENTS.md:L30-L32
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
🟡 Changes recommended
Schema compatibility and embedding-path coverage issues remain, and the producer test does not enforce exact output.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Adds affine-map emission for AutoTP checkpoint metadata and uses producer-generated maps during conversion while preserving legacy categories.
Changes:
- Derives maps for supported contiguous and replicated layouts.
- Updates converter compatibility handling.
- Expands producer and resume test coverage.
File summaries
| File | Summary | Final review comments |
|---|---|---|
tests/unit/checkpoint/test_autotp_uc_checkpoint.py |
Tests producer-generated layouts and resume behavior. | Nit (2 votes): Assert the emitted pattern set exactly matches expectations. |
deepspeed/module_inject/layers.py |
Produces and collects affine maps. | Moderate (2 votes): affine_map=None changes the metadata schema and breaks an existing test. Moderate (1 vote): The embedding path does not provide affine maps for contiguous row-sharded layouts. |
deepspeed/checkpoint/ds_to_universal.py |
Consumes affine maps alongside legacy patterns. | No final comments. |
Review details
Suppressed comments (1)
deepspeed/module_inject/layers.py:713
- This collector only serializes
conversion_meta['affine_map'], but the AutoTP embedding path still constructs metadata directly inauto_tp.py:_slice_embeddingwithout supplying an affine map. Consequently a contiguous row-sharded embedding continues through the category fallback and this producer does not cover all contiguous AutoTP layouts described by the PR; derive and pass the map for that path as well.
affine_map = conversion_meta.get('affine_map')
if affine_map is not None:
affine_maps[pattern] = affine_map.to_dict()
- Files reviewed: 3/3 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| 'original_shape': _normalize_uc_shape(original_shape), | ||
| 'is_bias': is_bias, | ||
| 'replicated': replicated, | ||
| 'affine_map': affine_map, |
| for pattern, want in expected.items(): | ||
| assert pattern in maps, f"producer emitted no affine map for {pattern}" | ||
| assert maps[pattern] == want, (f"emitted map for {pattern} differs from the layout the resume " | ||
| f"fixture verifies:\n emitted {maps[pattern]}\n expected {want}") |
Adding the key unconditionally changed the conversion schema for every parameter, including layouts with no map to describe. Emit it only where there is one, so an existing layer publishes exactly what it published before, and store it as plain scalars like the rest of the schema. Require the producer test to match the expected pattern set exactly, so an unintended extra map is a failure rather than something the test ignores. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Step 4 of the staging plan in #8252: AutoTP layers now emit an affine map, and the converter reads it on a path a real job takes rather than one a test supplies.
Follows #8385 (the IR, the lowering constructors, the converter branch) and #8477 (the scale fix). Scope here is contiguous splits and replicated parameters; the fused QKV and Yuan shared-QK layouts are next and still carry
unsupported_reason, so they emit no map and are unaffected.Where a map is built, and why it is not where you would expect
The producer cannot build a map from what
collect_autotp_universal_checkpoint_infosees. The conversion metadata carries eight fields and per-rank partition sizes are not among them — those were deliberately restore-only. But the layer does have them:_freeze_partition_sizesresolves them viaget_shard_size_listwhile the layer is built, and they are not recoverable later from a shape alone.So the map is derived at mark time, in
_set_param_uc_meta, which already receivespartition_sizes,sub_param_shard_widths,logical_shapeandreplicated, and hastp_world_sizein scope. That puts the derivation in one place rather than in each of the seven_mark_uc_metadataimplementations, andcollect_...then gathers what the layers produced.A layout that is not describable yet returns
None, so conversion falls back to the categories.The map and the categories have to coexist
A checkpoint carrying an affine map also carries the existing pattern lists, so a converter predating the map still reads it — that is the additivity §6.3 promises.
The consequence is that the converter which prefers the map never consults those branches, which leaves their patterns looking unused and fails a strict conversion. They are superseded, not unused, so the affine branch now marks them consumed. This is the compatibility question raised as "one boundary for review" in #8385; it turns out both have to be present and the reader has to account for the other.
This only surfaced end to end. The producer's output was correct in isolation, and the emitted maps matched a layout already verified by a passing resume — the failure was a strict-mode assert on rank 0 during conversion, which presented as rank 1 blocking on the following barrier.
Tests
TestAffineMapProducerbuilds a real AutoTP engine at TP2 and requires the producer to emit exactly the four maps the resume fixture previously supplied by hand. Those maps are not a guess: a full train → save → convert → resume cycle reproduces uninterrupted training through them, so matching them is evidence rather than self-consistency.affine_resume_checkpointnow takes its layout from the producer instead of injecting one, so the four save-convert-resume cases exercise the metadata a real job writes. All four pass — TP2 → TP1 and TP2 → TP2, through both the legacy and affine paths, compared against uninterrupted training on logits, losses, gradients, FP32 weights, both Adam moments and step counters.Validated on CPU/gloo (
DS_ACCELERATOR=cpu LOCAL_SIZE=4): 45 passed across the resume matrix, the producer test and the affine unit suite. The resume harness is @0z5a's, from #8474.Related: #8252, #8230.
cc @delock