Add DTensor optimizer contracts and sharded portable resharding

[thad0ctor]

Summary

• Add stable post-sharding rebinding, local training, and exact same-topology native checkpoint continuation for narrow one-dimensional default-world DTensors in Gefen and GefenMuon, including uneven and empty shards.
• Extend GefenMuonHybrid with atomic AdamW-backed and mixed-DTensor rebinding and native load staging while preserving nested child state and fail-before-mutation behavior.
• Add representation-qualified checkpoint transitions, including the directional factored-to-block live-fp32 period-one projection.
• Replace dense-everywhere portable DCP saves with a load-compatible sharded envelope using DCP-native semantic fields, bounded indexed routing, integrity verification, and derived Hybrid child namespaces.
• Update capability declarations and durable integration documentation to state the exact supported and unsupported topology transitions.

Why

This implements the four remaining roadmap slices from the platform optimizer contracts handoff: the DTensor data plane, broader Hybrid rebinding, qualified topology-changing canonical projection, and scalable portable resharding. Training, native same-topology checkpointing, and portable topology-changing support remain separate explicit claims.

Contract boundaries

Muon DTensor support is limited to training and same-topology native checkpoints; portable Muon resharding remains replicated/whole-owner. AdamW-backed Hybrid supports rebinding and native same-topology checkpoints but not portable global state. Block-to-factored conversion, multidimensional meshes, subgroups, asynchronous DCP, and mixed model/optimizer Stateful composition remain unsupported and fail closed.

Validation

• Current Python 3.12 / PyTorch 2.12 CPU suite: 1,203 passed, 333 skipped.
• Python 3.10 / PyTorch 2.5.0 floor suite: 1,190 passed, 346 skipped.
• Installed-wheel two-GPU release gate: 262 passed with zero skips.
• Installed-wheel portable DCP NCCL coverage: passed.
• Ruff, current/floor byte-compilation, git diff checks, wheel build, twine check, isolated wheel smoke, and independent requirement/collective audits: clean.
• Wheel SHA-256: 26ad47132426a316fb7d677fbc7aadfebc3af0ffb71d9c887ef5d93c98c5f7d8.

[coderabbitai]

Warning

Review limit reached

You’ve reached a temporary PR review limit under our Fair Usage Limits Policy.

Your recent review volume is higher than typical usage, so adaptive limits are currently applied.

Next review available in: 40 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: d6b9b3d6-9104-48fd-a280-471d7c3730fc

📥 Commits

Reviewing files that changed from the base of the PR and between 6feefa6 and 106e178.

📒 Files selected for processing (42)

• CHANGELOG.md
• COMPATIBILITY.md
• README.md
• docs/optimizer_contracts.md
• src/gefen/__init__.py
• src/gefen/_adamw_adapter.py
• src/gefen/contracts.py
• src/gefen/dtensor.py
• src/gefen/gefen.py
• src/gefen/gefen_muon.py
• src/gefen/hybrid.py
• src/gefen/portable.py
• src/gefen/portable_dcp.py
• src/gefen/portable_dcp_hybrid_sharded.py
• src/gefen/portable_dcp_sharded.py
• src/gefen/portable_runtime.py
• src/gefen/portable_state.py
• tests/test_capturable_fsdp2.py
• tests/test_checkpoint_projection_contracts.py
• tests/test_codebook_scope_cpu.py
• tests/test_dtensor_rebinding.py
• tests/test_dtensor_rebinding_acceptance.py
• tests/test_gefen_fsdp2_checkpoint.py
• tests/test_hybrid_adamw_dtensor.py
• tests/test_hybrid_adamw_layouts_distributed.py
• tests/test_hybrid_adamw_staging.py
• tests/test_hybrid_mixed_dtensor.py
• tests/test_hybrid_rebinding.py
• tests/test_hybrid_scoped_failure_protocol.py
• tests/test_muon_distributed_checkpoint_safety.py
• tests/test_muon_grad_presence.py
• tests/test_optimizer_contracts.py
• tests/test_portable_dcp.py
• tests/test_portable_dcp_sharded_segments.py
• tests/test_portable_dcp_sharded_validation.py
• tests/test_portable_dcp_topologies.py
• tests/test_portable_dtensor.py
• tests/test_portable_hybrid_distributed.py
• tests/test_portable_hybrid_runtime.py
• tests/test_portable_runtime.py
• tests/test_portable_state.py
• tests/test_portable_state_math.py

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/dtensor-portable-resharding

---

_{Comment @coderabbitai help to get the list of available commands.}

[thad0ctor]

@coderabbitai review
@codex

[coderabbitai]

@thad0ctor, I'll review the changes now.

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[thad0ctor]

Superseded by #82. The replacement PR is based directly on current main, carries forward this PR’s DTensor and sharded portable DCP implementation, adds the real disk-backed 2→4 FSDP2 DCP test, and reconciles #74/#76/#80.