Check where column-split pieces write in a shard by Achyuthan-S · Pull Request #8922 · deepspeedai/DeepSpeed

Achyuthan-S · 2026-10-27T02:26:06Z

ParamAffineMap.validate() checks where a piece writes in its shard only when the piece's destination is one dense interval. That check came from #8624. A column split has no such interval: each piece writes a few columns of every shard row. So for fused QKV split along columns, Yuan's o_proj, or any sub_param_map / segmented_map / block_gather_map with partition_dim=1, two pieces can write the same shard columns and pass every check.

Reproduction on master (fused QKV (4, 12) split along columns at TP 2):

Move rank 0's value piece onto its key columns.
validate() and validate_coverage() both pass.
rebuild(extract(F)) differs from F in 8 of 48 elements.
Two columns of rank 0's shard are never written by extract, so they hold uninitialised memory.

Changes

validate() now also checks strided pieces, still from metadata only, without walking elements.
- Bounds: each strided piece's lowest and highest shard address must lie inside the shard.
- Overlap: a column-split piece is a box in its shard's row-major grid, so overlap between such pieces follows from their corners. A strided view that wraps past the end of a shard row is valid but is not a box; the full check below covers it.
- The box pass runs only on ranks that hold a strided piece. Ranks with only dense pieces cost the same as before.
validate_coverage(), the full O(n) check, now walks every shard address and requires each one to be written exactly once. Together with the element-count check, that makes "every shard address is written by exactly one element" a complete check, for any stride, including strided pieces that are not boxes.

Tests

tests/unit/checkpoint/test_affine_shard_map.py:

Test	Contract	Fails on `master`
`test_strided_destination_overlap_is_refused`	column-split pieces that collide are refused	yes (did not raise)
`test_strided_destination_outside_the_shard_is_refused`	a column-split piece past the shard's width is refused	yes (did not raise)
`test_full_check_refuses_a_destination_collision_of_any_stride`	`validate_coverage()` refuses a collision that no box can describe	yes (did not raise)
`test_strided_destination_validation_does_not_enumerate_elements`	`validate()` stays metadata-only for a `10^9 × 8` column split	no; it guards cost
`test_strided_destination_that_wraps_a_shard_row_round_trips`	a strided view that wraps a shard row but stays inside the shard is accepted and round-trips	no; it guards against refusing a valid map

Validation

macOS 26.5 (arm64, CPU), Python 3.9.6, torch 2.8.0, DS_ACCELERATOR=cpu LOCAL_SIZE=4.

test_affine_shard_map.py: 76 passed.
test_affine_shard_map.py, test_affine_transfer_plan.py, test_autoep_affine.py, test_autoep_affine_universal.py and test_autotp_uc_checkpoint.py: 276 passed, 5 failed. The failures are identical on master. The 5 tests build the CPU Adam op, which fails on macOS arm64 (cpuinfo reports no flags). They were not run here.
In those suites the new box pass checks 28 strided pieces from the existing tests' layouts and rejects none of them.
End to end, to confirm the converter path is unaffected: AutoTP training, save, ds_to_universal, then a universal load at another TP degree. Runs: TP 2→4 under ZeRO 1 and 2, including an uneven split at TP 4; TP 4→2 under ZeRO 1. Parameters, exp_avg and exp_avg_sq are restored bit for bit, and the resumed losses match the uninterrupted run to within 1.2e-7.
yapf 0.43.0 and flake8 are clean on both files.

Refs #8252, #8624.

@delock

validate() checked destination bounds and overlap only for pieces whose destination is one dense interval, so pieces of a column split could write the same shard columns and pass every check. Check those pieces as boxes in the shard's grid from their metadata, and walk every shard address in validate_coverage() so the full check covers any stride. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.

chatgpt-codex-connector · 2026-10-27T02:26:22Z

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-27T02:32:43.789041Z	`5f5af9f`	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5f5af9f49c

ℹ️ 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 · 2026-10-27T02:32:41Z

+                if _dest_interval(piece) is None:
+                    raise ValueError(f'Rank {rank} piece {piece} writes outside the shard of shape '
+                                     f'{tuple(shard_shape)}.')


Keep wrapped strided views valid

When a legitimate strided destination starts near the end of a shard row but all of its storage addresses remain in bounds, this branch rejects it. For example, in a (3, 3) shard, an AffinePiece with shape (2, 2), dest_offset=2, and dest_strides=(3, 1) visits flat addresses 2, 3, 5, 6, but _dest_box computes a column stop of 4, _dest_interval returns None, and validate() now raises. Because destinations are arbitrary torch.as_strided views over _flat_buffer, serialized or custom maps using such wrapped views can no longer rebuild; bounds should be checked against the actual flat address range while box coordinates are used only when the view fits the shard grid.

Useful? React with 👍 / 👎.

Fixed in the latest commit. Bounds for a strided piece now come from its actual lowest and highest shard address, and box coordinates are used only when the view fits the shard grid. A view that wraps a row is left to the address walk in validate_coverage(). Your (3, 3) example is now test_strided_destination_that_wraps_a_shard_row_round_trips.

Bound a strided piece by its lowest and highest shard address, and use box coordinates only when the view fits the shard grid. A view that wraps a row stays valid and is checked by the address walk in validate_coverage(). Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

jinyouzhi · 2026-10-27T09:28:05Z

    return piece.dest_offset, piece.dest_offset + piece.numel


+def _dest_box(piece, shard_shape):


Is these similar to?

DeepSpeed/deepspeed/checkpoint/affine_transfer.py

Lines 93 to 120 in 7501610

@dataclass(frozen=True)

class LogicalBox:

"""An axis-aligned block of the logical parameter: its first coordinate and its extents."""

origin: Tuple[int, ...]

extent: Tuple[int, ...]

@property

def volume(self) -> int:

return _product(self.extent)

@property

def is_empty(self) -> bool:

return any(size == 0 for size in self.extent)

def upper(self) -> Tuple[int, ...]:

return tuple(low + size for low, size in zip(self.origin, self.extent))

def intersection(self, other: 'LogicalBox') -> 'LogicalBox':

low = tuple(max(a, b) for a, b in zip(self.origin, other.origin))

high = tuple(min(a, b) for a, b in zip(self.upper(), other.upper()))

return LogicalBox(low, tuple(max(0, h - l) for h, l in zip(high, low)))

def contains(self, other: 'LogicalBox') -> bool:

return all(other_low >= low and other_high <= high

for low, high, other_low, other_high in zip(self.origin, self.upper(), other.origin, other.upper()))

def __repr__(self) -> str:

return f'Box({self.origin}+{self.extent})'

Maybe re-use/abstract as common utils.

Also, I think there's some follow-up to optimize (I list in #8885 (comment)). cc @0z5a btw, if you both no bandwidth on it, I'd like to takeover. Thanks.

I prefer to keep it not block merging. Refactor it later as seperate follow-up.

Thanks @jinyouzhi! Yes, it's the same idea: _dest_box is a box in a shard's coordinates, and LogicalBox is one in the logical parameter's. +1 to the shared affine_geometry.py you proposed on Slack: affine_transfer.py already imports from affine.py, so a module both import is the right home. Agreed it shouldn't block this PR.

The planner optimizations you listed in #8885 sit in the planner from #8623, so I'm happy for you and @0z5a to take them. One data point for that work, since it overlaps your grid fast path: for CodeGen's fused QKV the planner emits 44 / 88 / 96 copies for TP 2→4 / 4→8 / 2→8 on current master. The counts are flat across hidden sizes, so they're bounded and the plans are exact, but folding in N dimensions as in spec §8.2 gives 40 for 4→8. Same-shape runs at constant offsets aren't stacked into an outer dimension yet.

Happy to share the script.

@0z5a

…deepspeedai#8975) Conversion trusts a map further than it checks it. - `rebuild` (what `ds_to_universal` calls on convert) runs only `validate()`. That checks each rank's element count and where pieces write in the shard, from metadata, but it never looks at the source side. So a map can pass it while two ranks claim the same block of the parameter and no rank holds another. - `rebuild` then returns whatever `torch.empty` left in the elements nobody holds. - It also lets one rank's block overwrite the other's with no comparison. - `extract` (what the AutoTP restore path calls) runs no check at all. - A destination that wraps past the end of a shard row is not a box, so `validate()` cannot place it. deepspeedai#8922 made `validate_coverage()` complete for such destinations, but conversion does not call `validate_coverage()`, because that check walks elements in Python. Reproduction on `master`: 1. Take two ranks with shards of 4 elements each, for an `(8,)` parameter. 2. Point rank 1's piece at source offset 0 instead of 4. 3. `validate()` passes. `rebuild` returns rank 1's values in elements 0–3 and uninitialised memory in 4–7, with no error. How often this happens: - I fuzzed the in-tree constructors (contiguous split, replicated, segmented, block gather). Each case is a random 2-D shape, TP 2–4 with uneven widths, and one offset or stride of one piece corrupted. - On `master`, **671 of 2,334 corrupted maps converted without an error and produced a wrong tensor**. With this change, none do: each one is refused or still round-trips exactly. - Every map that passes `validate_coverage()` round-trips exactly on both (801 of 801). - With the stricter oracle from @0z5a's review (shards built from the uncorrupted map, then rebuilt through the corrupted one; 2,000 cases): master converts 421 silently wrong, this PR none. - One limit: a corruption that leaves a *valid but different* layout, for example two equal-shaped pieces swapping destinations, passes every map check, because the map is self-consistent and disagrees only with the memory it describes. Catching that takes a runtime digest of the data, not a check on the map. ## Changes - **`rebuild`** - Records who wrote each element of the full tensor (nobody, a single holder, or a replica), one byte per element, and refuses: - an element no piece writes; - an element written by two pieces unless both are replicas, whichever rank is visited first; - replicas that disagree on any element. That includes replicas cut at different boundaries on different ranks. Before, only a piece starting exactly where an earlier one started was compared, so the rest of the replica overwrote it unchecked. - Also tracks the shard addresses it reads on each rank, so a destination collision is refused for any stride. - **`extract`** - Runs `rebuild`'s metadata check for its own rank: `validate()` restricted to one rank, through a new `_validate_rank`. - Uses the same shard-address mask, so a restored shard never holds uninitialised memory. - **Both, through `_validate_rank` (and so `validate()`)** - Refuse a piece whose source or destination view reaches one address twice within itself. - Such a view loses a shard value, and which value survives depends on the device. - The check is metadata only: taking the axes that move in order of stride, each stride must clear what the axes before it reach. Every in-tree layout satisfies it. - A zero stride now raises a `ValueError` instead of torch's own `RuntimeError`. - **No change for valid maps.** Error types and messages follow the existing ones (`ValueError`, "twice", "uncovered", "different data"), plus "itself" for the new check. Cost: - One byte per element of the full tensor and one per shard element, with vectorised tensor ops and no per-element Python. - `rebuild` already materialises the full tensor, so the fp32 states that `ds_to_universal` converts need at most a quarter more memory, for the duration of one parameter's conversion. ## Tests `tests/unit/checkpoint/test_affine_shard_map.py`: | Test | Contract | Fails on `master` | |---|---|---| | `test_conversion_refuses_a_destination_collision_of_any_stride` | `extract` and `rebuild` refuse the collision from the full-check test of deepspeedai#8922, which `validate()` accepts | yes (did not raise) | | `test_replicas_cut_at_different_boundaries_are_compared_in_full` | a replica that one rank holds whole and another holds in halves round-trips, and is refused when the second half drifts | yes (did not raise) | | `test_conversion_refuses_a_map_that_does_not_cover_the_parameter` | `rebuild` refuses a map that leaves an element unheld, and one where two ranks own one block | yes (did not raise) | | `test_a_replica_overlapping_a_single_holder_is_refused_in_either_order` | a block declared both replicated and single-held is refused whichever rank comes first | yes (accepted in both orders) | | `test_a_view_that_overlaps_itself_is_refused` | `extract` and `rebuild` refuse a source view that reads one address twice, a zero stride, and a destination view that writes one address twice | yes (silently wrong) | | `test_a_corrupted_map_is_refused_or_still_exact` | fixed seed: 300 random layouts, with shards built before one offset or stride is corrupted, are either refused or rebuilt bit-exact | yes (silently wrong tensor) | The colliding map of deepspeedai#8922's full-check test is now a helper shared by both tests. ## Validation macOS 26.5 (arm64, CPU), Python 3.9.6, torch 2.8.0, `DS_ACCELERATOR=cpu LOCAL_SIZE=4`. - `test_affine_shard_map.py`, `test_affine_transfer_plan.py`, `test_zero_affine.py`, `test_autoep_affine.py` and `test_autoep_affine_universal.py`: **280 passed**. - `test_universal_checkpoint.py`, `test_autotp_uc_checkpoint.py`, `test_autoep_checkpoint.py`, `test_autoep_unit.py` and `test_autoep_affine_universal.py`: **344 passed**. - I ran `master` (`ba0d63ce2`) side by side, and it fails exactly the same 57 tests. - 53 of them build the CPU Adam op, which fails on macOS arm64 (`cpuinfo` reports no `flags`). - The other 4 are `test_storage_maps_match_runtime_ownership_and_plan_into_a_new_placement`. - End to end, to confirm no valid map is refused: AutoTP training, save, `ds_to_universal`, then a universal load at another TP degree, ZeRO 1, gloo. - Models: an MLP block, CodeGen (fused QKV), GPTBigCode (multi-query, replicated K/V), a Yuan-style shared-QK attention (replicated `o_proj` bias), and Llama with a tied embedding and `vocab_parallel_lm_head` (vocabulary of 67, uneven at every degree). - TP changes: 2→4, 4→2 and 1→2. - Parameters, `exp_avg` and `exp_avg_sq` are restored bit for bit, and the resumed losses match the uninterrupted run to within 4.8e-7. - On GPU and XPU, @jinyouzhi ran the first commit on 8× NVIDIA L20 (torch 2.14.1+cu130) and on 8× Intel Arc Pro B70. - The same affine files: 283 passed on each. - On CUDA, rebuilding a 4096×4096 fp32 column split at TP 8 takes 1.2 ms, and 8 extracts take 0.9 ms. - The second commit adds only metadata checks and a one-byte mask in place of a bool one. It has not been rerun on GPU. - `yapf` 0.43.0 and `flake8` are clean on both files. Refs deepspeedai#8252, deepspeedai#8922. @delock --------- Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

Achyuthan-S requested review from loadams, tjruwase and tohtana as code owners October 27, 2026 02:26

Copilot AI balanced review requested due to automatic review settings October 27, 2026 02:26

Copilot AI reviewed Oct 27, 2026

chatgpt-codex-connector Bot reviewed Oct 27, 2026

View reviewed changes

jinyouzhi reviewed Oct 27, 2026

View reviewed changes

delock self-requested a review October 28, 2026 01:16

delock approved these changes Oct 29, 2026

View reviewed changes

delock added this pull request to the merge queue Oct 29, 2026

Merged via the queue into deepspeedai:master with commit 8d4f29c Oct 29, 2026
19 checks passed

Achyuthan-S deleted the check-strided-destinations branch October 30, 2026 02:38

Achyuthan-S mentioned this pull request Nov 1, 2026

Check that converting through an affine map writes every element once #8975

Merged

This was referenced Nov 9, 2026

[AffineIR] Check strided shard destinations without pairwise comparisons #9039

Open

[AffineIR] Check affine holders and source regions without pairwise comparisons #9011

Open

gpu-test-labeler Bot added the skip-gpu-test Requires positive evidence that every changed behavior can be validated without a GPU. label Nov 14, 2026

Navigation Menu

Repository navigation

Check where column-split pieces write in a shard - #8922

Check where column-split pieces write in a shard#8922
delock merged 2 commits into
deepspeedai:masterfrom
Achyuthan-S:check-strided-destinations

Achyuthan-S commented Oct 27, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

chatgpt-codex-connector Bot commented Oct 27, 2026 •

edited

Loading

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Oct 27, 2026

Uh oh!

Achyuthan-S Oct 27, 2026

Uh oh!

jinyouzhi Oct 27, 2026

Uh oh!

jinyouzhi Oct 27, 2026

Uh oh!

Achyuthan-S Oct 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		return piece.dest_offset, piece.dest_offset + piece.numel


		def _dest_box(piece, shard_shape):

	@dataclass(frozen=True)
	class LogicalBox:
	"""An axis-aligned block of the logical parameter: its first coordinate and its extents."""
	origin: Tuple[int, ...]
	extent: Tuple[int, ...]

	@property
	def volume(self) -> int:
	return _product(self.extent)

	@property
	def is_empty(self) -> bool:
	return any(size == 0 for size in self.extent)

	def upper(self) -> Tuple[int, ...]:
	return tuple(low + size for low, size in zip(self.origin, self.extent))

	def intersection(self, other: 'LogicalBox') -> 'LogicalBox':
	low = tuple(max(a, b) for a, b in zip(self.origin, other.origin))
	high = tuple(min(a, b) for a, b in zip(self.upper(), other.upper()))
	return LogicalBox(low, tuple(max(0, h - l) for h, l in zip(high, low)))

	def contains(self, other: 'LogicalBox') -> bool:
	return all(other_low >= low and other_high <= high
	for low, high, other_low, other_high in zip(self.origin, self.upper(), other.origin, other.upper()))

	def __repr__(self) -> str:
	return f'Box({self.origin}+{self.extent})'

Conversation

Achyuthan-S commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

Tests

Validation

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codex Review Summary

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

Achyuthan-S Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

jinyouzhi Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

jinyouzhi Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

Achyuthan-S Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Achyuthan-S commented Oct 27, 2026 •

edited

Loading

chatgpt-codex-connector Bot commented Oct 27, 2026 •

edited

Loading