[Achyuthan-S]

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):

1. Move rank 0's value piece onto its key columns.
2. validate() and validate_coverage() both pass.
3. rebuild(extract(F)) differs from F in 8 of 48 elements.
4. 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

[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]

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]

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

[Achyuthan-S]

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.

[jinyouzhi]

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.

[jinyouzhi]

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

[Achyuthan-S]

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.