Harden convention codebook-scope preambles and contract validation - #76
Conversation
Address a round of Codex + CodeRabbit findings on the slimmed optimizer convention (PR #67). All findings are convention-specific. Codebook-scope collective symmetry: - initialize_codebook() now captures the failure-vote binding and synchronizes a one-sided preamble failure (finalized layout / runtime binding / capture-readiness) through the scope before the scoped "initialize" operation header, so a member that fails its preamble no longer strands peers inside the header all_gather. - plain Gefen.step() rejects a closure that replaces or clears the runtime codebook binding between capture and the header: the captured-binding recheck raises and is synchronized through the captured scope, so every rank follows the same header collective instead of one rank skipping it. Checkpoint and contract validation: - CheckpointProcessGroupBinding no longer rejects a CUDA collective device on Gloo/MPI (both support CUDA tensors); such a device now falls through to the existing device-availability check. - Contract dataclasses enforce exact runtime types: StateField.checkpointed, StateVariant.initialized/migration_only, and CheckpointSupport.requires_collective/atomic_load must be real bools, and OptimizerContract.schema_version must be a real int, so truthy strings or bool/float look-alikes can no longer advertise false guarantees. - The hybrid finalized-layout fast token folds in the _state_param_owner registry contents (keys and values), so an in-place owner replacement that preserves the dict identity and length invalidates the cached verdict. Tests: - Snapshot per-device cache membership so a cache clear/removal is detected. - Verify the backup parameter is untouched in hybrid scoped-failure tests. - Guard the remaining Gloo-only distributed tests with the dist/Gloo skipif. - Documented that only ordinary nested GefenMuonHybrid.load_state_dict() reports atomic_load=False; portable composite imports keep the guarantee. - New gloo regression tests for the initialize-preamble and step group-swap synchronization, plus unit tests for the checkpoint device, contract type, and hybrid owner-token fixes.
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![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: 57d72f58c4
ℹ️ 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".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # Gloo (and MPI) support CUDA tensors in addition to CPU, so no | ||
| # backend-level device restriction applies here; a CUDA collective | ||
| # device on those backends still falls through to the availability | ||
| # check below. |
There was a problem hiding this comment.
On CUDA hosts using the MPI backend without CUDA-aware MPI, this now lets validate_runtime() succeed for a CUDA checkpoint device just because the device index exists. The portable checkpoint path then moves its control/metadata tensors to binding.collective_device before dist.all_gather/broadcast (for example in src/gefen/portable_collective.py), so these bindings fail later with backend errors instead of being rejected during validation; keep MPI CPU-only unless CUDA-aware support can be detected explicitly.
Useful? React with 👍 / 👎.
|
@coderabbitai review |
✅ Action performedReview finished.
|
Narrow the F4 relaxation: Gloo genuinely supports CUDA tensors, but MPI moves GPU tensors only when built CUDA-aware, which PyTorch cannot reliably detect. Reject a CUDA collective device on the MPI backend during validate_runtime() instead of deferring a backend error to the portable collective's all_gather/broadcast (Codex P2 on #76).
a1bd0d1
into
feat/platform-optimizer-contracts
1 participant
Addresses a review round (Codex P1/P2 + CodeRabbit review
4707322434) on the slimmed optimizer convention. Base isfeat/platform-optimizer-contracts— this PR stacks into #67, not main. Every finding is convention-specific (present only on #67).Disposition
Gefen.step()preamble not synchronized before scoped headergefen.py_synchronize_prevalidated_codebook_scope_failure("step preamble", scope_binding)before the header. Reconciled with F3 (below).initialize_codebook()preamble not synchronizedgefen.pyinitializeheader.gefen.pyscope_binding(swap to None or another binding) and routes it through the synchronized scope-failure path, so no rank skips/enters a different header collective. This is the F1/F3 reconciliation: F1 = preamble/closure failure sync (already present); F3 = closure replacing the group (added here)._validate_collective_devicerejects valid Gloo+CUDAcheckpoint.pyschema_versionaccept non-exact typescontracts.pyStateField.checkpointed,StateVariant.initialized/migration_only,CheckpointSupport.requires_collective/atomic_loadrequiretype(...) is bool;OptimizerContract.schema_versionrequirestype(...) is int. (TrainingSupportandOptimizerCapabilitiesalready enforced this.)hybrid.py_state_param_ownerkeys and values (mirroring the base_param_namessnapshot), so an in-place entry replacement preserving dict identity+length invalidates the cached verdict._gefen_codebook_by_devicemembershiptests/_state_snapshot.pyclear()/removal is detected even when retained tensors still match their clones.tests/test_codebook_scope_distributed.pydist.is_available()/dist.is_gloo_available()skipif to the three unguarded Gloo tests._untouchednever checksbackup_parametertests/test_hybrid_scoped_failure_protocol.py_untouchedtakesbackup_shardand asserts the backup parameter equals its_backup_initial()slice; all callers updated.atomic_load=Falsecontradicts portable guaranteedocs/optimizer_contracts.mdGefenMuonHybrid.load_state_dict()reportsatomic_load=False; portable composite imports retain the fail-before-mutation guarantee.Notes on collective symmetry
F2/F3 follow the existing convention: capture a validated failure-vote binding first, wrap rank-local checks (and user closure) in try/except, and vote through
_synchronize_prevalidated_codebook_scope_failureon the captured binding before any scoped collective. No collective is entered/skipped on only some ranks;@torch._dynamo.disableand binding guards on the sibling primitives are untouched. F3 is scoped to plainGefenper the finding;GefenMuon/Hybridwere left as-is.Tests
New gloo regression tests (with bounded no-hang/liveness guards, mirroring
_closure_preamble_worker):test_initialize_codebook_preamble_failure_raises_symmetrically_across_the_scope(F2) — one-rank preamble failure makes both ranks raise fast.test_plain_gefen_scoped_step_closure_group_swap_raises_symmetrically_across_the_scope(F3, newswap_groupfailure mode) — one-rank binding clear makes both ranks raise fast.New/updated unit tests:
test_validate_collective_device_allows_gloo_cuda_but_requires_available_device+ reworked distributed binding case (F4).test_contract_flags_and_schema_version_reject_non_exact_types(F5).test_composite_registry_in_place_entry_swap_is_detected_with_warm_verdict(F6).test_deep_snapshot_detects_per_device_cache_membership_clear(F7).Each new/changed test was verified to FAIL without its fix (F2/F3 hang → bounded liveness guard fails the run; F4–F7 assert) and PASS with it.
Validation (CPU gloo venv,
CUDA_VISIBLE_DEVICES="")test_codebook_scope_distributed.py,test_codebook_scope_cpu.py,test_hybrid_scoped_failure_protocol.py,test_optimizer_contracts.py,test_checkpoint_binding.py,test_native_load_atomicity.py,test_hybrid_load_atomicity.py,test_shard_identity_contracts.py,test_hybrid_layout_cache.py:140 passed, 2 skipped (skips are NCCL/2-CUDA-device gated). 0 failures.