diff --git a/CHANGELOG.md b/CHANGELOG.md index 4f680ad..42ffc5e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,9 @@ All notable changes to this project are documented here. This project adheres to ## [0.4.1] - 2026-07-15 -Distributed checkpoint-load and step-failure hardening. All fixes are backward-compatible; no public API changes. +Distributed checkpoint-load and step-failure hardening. All changes are backward-compatible; the only public API addition is the opt-in `GefenDCPState` resharding wrapper. + +- Add `GefenDCPState`, a standalone DCP Stateful wrapper for **resharding** plain Gefen on one-dimensional default-world FSDP2 `Shard(0)` parameters (`factored_v_2d=False`, `capturable=False`). Save writes shard-addressable dense optimizer tensors so DCP can reshard N ranks to M without a rank-0 full-state gather; load validates before mutation, then re-blocks the resharded dense momentum back into Gefen's compact ~1 byte/param block state (re-running the period search, relearning the codebook, and re-quantizing on each new shard). Resume is a correct continuation within 256-level quantization noise, not bit-exact — use the unchanged native full-state path for same-topology resumes. Correctness — checkpoint loads are now atomic: diff --git a/COMPATIBILITY.md b/COMPATIBILITY.md index 9dc4f94..be322f5 100644 --- a/COMPATIBILITY.md +++ b/COMPATIBILITY.md @@ -103,7 +103,9 @@ Plain Gefen supports FSDP2 training-time CPU offload via `CPUOffloadPolicy` (`fu Native single-process optimizer `state_dict()`/`load_state_dict()` and the explicitly documented GefenMuon `distributed` state path remain separate from the rank-local DCP adapter. For plain Gefen and `GefenMuon(sharded_mode="approx")`, each rank can learn different local codebook and block geometry, so `state_dict()` collectively replaces ordinary rank-local tensors with one tagged, rank-indexed CPU payload that PyTorch full-state DCP preserves. Both optimizers' two-GPU FSDP2 `fully_shard` paths are exercised through `get_optimizer_state_dict(..., full_state_dict=True, cpu_offload=True)` and `set_optimizer_state_dict(..., full_state_dict=True, broadcast_from_rank0=True)`, with exact next-step continuation; the flattened optimizer-state form is also covered. -This format is deliberately same-topology only and currently requires one 1-D DeviceMesh spanning the default process-group world. Every rank must participate in both save and restore; each process temporarily holds all serialized rank payloads on CPU, so the leading checkpoint-time CPU cost approaches `world_size ×` its local optimizer-state size plus local serialization scratch. Loading validates world size, parameter order and names, global and local shapes and dtypes, mesh membership and names, structural placements, rank coordinates, global step, deterministic policy, frozen codebook, and sharded mode before mutation. Multidimensional meshes, subgroups, pipeline-local optimizers, world-size/topology changes, and old unsafe untagged full checkpoints fail closed rather than silently applying rank 0's state to every shard. No optimizer-state reshard portability is claimed, model-only DCP is unaffected, and the full-state DCP support described here does not extend beyond plain Gefen and Muon `approx`. +The native format is deliberately same-topology only and currently requires one 1-D DeviceMesh spanning the default process-group world. It is bit-exact on resume and is the recommended path when the world size does not change. Plain Gefen additionally exposes `GefenDCPState`, a standalone `torch.distributed.checkpoint` Stateful wrapper whose purpose is **resharding**: a checkpoint saved on N ranks can be loaded on M ranks for one-dimensional default-world `Shard(0)` DTensors. Save dequantizes the quantized momentum against each rank's learned codebook into dense `Shard(0)` DTensors, so DCP can reshard them without gathering the full optimizer state onto one rank. Load validates every field before mutating the live optimizer, then re-blocks the resharded dense momentum back into Gefen's compact per-block state on each new local shard — re-running the block-variance period search, relearning the exact codebook, and re-quantizing — so the restored optimizer keeps Gefen's ~1 byte/param footprint rather than expanding to per-element state. + +Because the momentum is routed through a dense reshard and re-blocked against a freshly learned per-shard codebook, resume is a **correct, finite continuation within 256-level quantization noise, not a bit-exact restore** — even at the same world size. Same-topology resumes should therefore use the native `state_dict` path above. `GefenDCPState` requires `factored_v_2d=False` and `capturable=False` (the factored row/column second moment is not shard-addressable, and capturable/compiled counter and seed semantics are not host-serializable); Muon, Hybrid, `sharded_mode`, multidimensional meshes, subgroups, non-`Shard(0)` placements, and the native full-state format remain same-topology and fail closed. ## Transformers Trainer DDP diff --git a/README.md b/README.md index 129be1d..683c4d6 100644 --- a/README.md +++ b/README.md @@ -121,12 +121,12 @@ Gefen drops into standard distributed training like any other PyTorch optimizer, |---|---| | DDP | All optimizers | | FSDP2 | All optimizers; training-time CPU offload (`CPUOffloadPolicy`) validated for plain Gefen, single and multi-GPU | -| FSDP2 checkpoints | Plain Gefen and Muon `approx`; resume needs the same GPU count — scope note below | +| FSDP2 checkpoints | Plain Gefen can use `GefenDCPState` to reshard N ranks to M; same-topology resumes use the bit-exact native full-state path (Muon `approx` too) | | Muon `distributed` checkpoints | Resume on any GPU count, even a single GPU — [details](#experimental-lever-sharded-newton-schulz-under-fsdp2-sharded_mode) | | DeepSpeed ZeRO 1-3 | Plain Gefen (client optimizer); optimizer CPU-offload verified at ZeRO-2/3, parameter offload at ZeRO-3 (full fine-tune and LoRA); use FSDP2 or DDP for the Muon family — config note below | | Megatron-LM | All optimizers, including checkpoint resume — [scope](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#megatron-lm-integration-scope) | -> **FSDP2 checkpoint scope.** Plain Gefen and Muon `approx` save and resume exactly through PyTorch's standard full-state checkpoint calls, as long as the GPU count and sharding layout are unchanged and every GPU joins the save. Anything outside that scope refuses to load instead of corrupting state — [full details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope). +> **FSDP2 checkpoint scope.** `GefenDCPState` is for **resharding** plain Gefen: wrap the optimizer with it for `torch.distributed.checkpoint`, and a checkpoint saved on N ranks loads on M ranks for a one-dimensional default-world `Shard(0)` mesh. Load re-blocks the resharded state back to Gefen's compact ~1 byte/param form, so resume is a correct continuation within quantization noise rather than a bit-exact restore; same-topology resumes should use the bit-exact native full-state path (Muon `approx` too). Requires `factored_v_2d=False` and `capturable=False` — [full details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope). > > **FSDP2 CPU offload.** Training-time CPU offload via `CPUOffloadPolicy` (`fully_shard(module, offload_policy=CPUOffloadPolicy())`) is validated for plain Gefen on single and multiple GPUs: each rank steps its CPU-resident local shard directly (the codebook is learned rank-locally, with no cross-rank codebook collective), and the multi-GPU run completes on an NCCL-only process group. @@ -636,7 +636,7 @@ Measured (Qwen3-0.6B, 2 and 4 GPUs): `"distributed"` matched `"exact"` in the re ## Known limitations - **Hybrid checkpoint schema.** `GefenMuonHybrid`'s `state_dict()` uses its own nested `{"muon": ..., "backup": ..., "backup_optimizer": "gefen" | "adamw"}` layout. Resume from a checkpoint the hybrid itself saved—not one consolidated or converted to the flat torch `{state, param_groups}` layout. Cross-backend loads are rejected before either child is mutated; legacy untagged hybrid checkpoints are interpreted as Gefen-backed. -- **FSDP2 optimizer checkpoints don't reshard.** Plain Gefen and Muon `approx` resume only on the same GPU count and layout; changing either refuses to load. Model weights are unaffected — [details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope). +- **FSDP2 optimizer resharding is explicit.** Use `dcp.save({"optimizer": GefenDCPState(optimizer)}, ...)` and the matching `dcp.load` call to reshard plain Gefen across world sizes on a one-dimensional default-world `Shard(0)` mesh (`factored_v_2d=False`, `capturable=False`). Load re-blocks the state to Gefen's compact form and continues within quantization noise; for a same-topology resume use the bit-exact native path instead. Muon, Hybrid, other placements, and the native full-state format remain same-topology and fail closed — [details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope). - **True-FP16 overflow skips are invisible to Accelerate's `step_was_skipped` flag.** BF16 and standard AMP are unaffected and are the recommended modes in Trainer/Accelerate. ## Troubleshooting diff --git a/src/gefen/__init__.py b/src/gefen/__init__.py index d2eb309..6105828 100644 --- a/src/gefen/__init__.py +++ b/src/gefen/__init__.py @@ -11,6 +11,7 @@ "Gefen", "GefenMuon", "GefenMuonHybrid", + "GefenDCPState", "split_params_for_muon", "validate_split", "kernels", @@ -31,6 +32,10 @@ def __getattr__(name): from .hybrid import GefenMuonHybrid return GefenMuonHybrid + if name == "GefenDCPState": + from .dcp import GefenDCPState + + return GefenDCPState if name in ("split_params_for_muon", "validate_split"): from . import params diff --git a/src/gefen/dcp.py b/src/gefen/dcp.py new file mode 100644 index 0000000..d3c63b0 --- /dev/null +++ b/src/gefen/dcp.py @@ -0,0 +1,820 @@ +"""Standalone DCP adapter for reshardable plain-Gefen FSDP2 state.""" + +from __future__ import annotations + +from dataclasses import dataclass +import math +import re + +import torch + + +# Names Gefen synthesizes positionally when the caller passes no parameter name +# (``group__param_`` for a multi-param group, ``param_`` for an +# implicit/single group). Such names encode *registration order*, not stable +# parameter identity, so two differently ordered reconstructions produce the +# same identity list and would cross-assign momentum on load. DCP resharding +# therefore requires caller-provided names and rejects the synthesized pattern. +_SYNTHESIZED_NAME_RE = re.compile(r"^(group_\d+_param_\d+|param_\d+)$") + + +# Schema history: +# v1 -- counters + dense momentum/second-moment slots only. +# v2 -- adds per-slot parameter identities (name + group + global shape) and +# per-group hyperparameters (lr/betas/eps/weight_decay). Load is +# fail-closed on version, identity, topology, and deterministic mismatch. +_FORMAT_VERSION = 2 + +_GROUP_HYPER_KEYS = ("lr", "beta1", "beta2", "eps", "weight_decay") + + +def _local(value): + if hasattr(value, "to_local"): + value = value.to_local() + if hasattr(value, "wait"): + value = value.wait() + return value + + +def _is_dtensor(value) -> bool: + return ( + hasattr(value, "to_local") + and hasattr(value, "placements") + and hasattr(value, "device_mesh") + ) + + +def _counter(value, name: str) -> int: + if torch.is_tensor(value): + if value.numel() != 1: + raise ValueError("{} must be scalar".format(name)) + value = value.item() + if type(value) not in (int, float) or int(value) != value or value < 0: + raise ValueError("{} must be a nonnegative integer".format(name)) + return int(value) + + +def _dtensor_from_local(local: torch.Tensor, parameter): + from torch.distributed.tensor import DTensor + + return DTensor.from_local( + local, + device_mesh=parameter.device_mesh, + placements=parameter.placements, + shape=parameter.shape, + stride=parameter.stride(), + run_check=False, + ) + + +def _dense_momentum(optimizer, parameter, state) -> torch.Tensor: + indices = state["m_codebook"].reshape(-1).long() + magnitude = state["m_magnitude"].reshape(-1).float() + codebook = optimizer._gefen_codebook + if codebook is None: + raise RuntimeError("Gefen DCP save requires an initialized codebook") + codebook = codebook.detach().to(device=indices.device, dtype=torch.float32) + period = _counter(state["automatic_period"], "automatic_period") + if period < 1 or indices.numel() != magnitude.numel() * period: + raise ValueError("Gefen momentum block geometry is invalid") + return ( + codebook[indices] + .reshape(-1, period) + .mul(magnitude.reshape(-1, 1)) + .reshape(_local(parameter).shape) + ) + + +def _dense_second_moment(parameter, state) -> torch.Tensor: + period = _counter(state["automatic_period"], "automatic_period") + vmean = state["vmean"].reshape(-1).float() + local_numel = _local(parameter).numel() + if period < 1 or vmean.numel() * period != local_numel: + raise ValueError("Gefen second-moment block geometry is invalid") + return ( + vmean.reshape(-1, 1) + .expand(-1, period) + .reshape(_local(parameter).shape) + .clone() + ) + + +def _choose_period(optimizer, name: str, parameter, second: torch.Tensor) -> int: + """Re-derive a compact block period for the resharded local shard. + + Runs Gefen's block-variance period search on the dense per-element second + moment (a grad^2 proxy -- ``vmean`` is itself the EMA of block-mean grad^2), + exactly as the native first step runs it on grad^2. The result is a divisor + of the *new* local numel, so the restored state keeps ~1 byte/param instead + of collapsing to per-element (period one). + + The optimizer's explicit period-one routing gates (``force_1d_period_one``, + ``force_2d_period_one``, ``period_one_substrings``) are honored first, exactly + as :meth:`Gefen._resolve_automatic_period` applies them before the raw search. + Otherwise a checkpoint could restore period>1 into an optimizer configured for + period-one, and the frozen restored codebook would keep the next step from + correcting it. + """ + from gefen.partitioning import find_period_by_block_variance + + if getattr(optimizer, "_force_1d_period_one", False) and parameter.ndim == 1: + return 1 + if getattr(optimizer, "_force_2d_period_one", False) and parameter.ndim == 2: + return 1 + substrings = getattr(optimizer, "_period_one_substrings", ()) + if substrings: + lname = str(name).lower() + if any(sub in lname for sub in substrings): + return 1 + + flat = second.reshape(-1) + if flat.numel() < 8: + return 1 + if flat.device.type == "cuda": + return find_period_by_block_variance( + flat.detach().float(), + print_results=False, + parameter_name=name, + parameter_shape=tuple(parameter.shape), + backend="gpu", + input_is_squared=True, + ) + return find_period_by_block_variance( + flat.detach().float().cpu().numpy(), + print_results=False, + parameter_name=name, + parameter_shape=tuple(parameter.shape), + backend="cpu", + input_is_squared=True, + ) + + +def _learn_codebook(name_flat_period, device): + """Re-learn one exact codebook per rank from the resharded local momentum. + + Mirrors the native per-rank codebook (one histogram across all of the rank's + parameters); the resharded momentum is the distribution being quantized, so + learning from it directly minimizes the re-quantization error on this shard. + Returns ``None`` when there is no initialized momentum to learn from. + """ + from gefen.gefen import learn_gefen_exact_codebook_from_grad_periods + + grad_periods = [ + (name, flat, period, flat) for name, flat, period in name_flat_period + ] + if not grad_periods: + return None + return learn_gefen_exact_codebook_from_grad_periods( + grad_periods=grad_periods, + codebook_device=device, + num_codebooks=256, + force_endpoints=True, + verbose=False, + compute_mse_logging=False, + use_fused_histogram=False, + ) + + +def _reblock(codebook, momentum: torch.Tensor, second: torch.Tensor, period: int): + """Re-quantize a dense fp32 momentum shard into compact Gefen block state. + + Reproduces the native quantize arithmetic (``_automatic_momentum_update``): + per-block max-abs magnitude, normalize into the codebook domain, nearest + codeword indices. ``vmean`` is the per-block mean of the dense second moment + (the block-mean grad^2 for the new blocking). Returns + ``(m_codebook, m_magnitude, vmean)``. + """ + from gefen.gefen import ( + automatic_partition_reduce, + automatic_partition_view, + gefen_nearest_codebook_indices, + ) + + flat = momentum.reshape(-1) + # Co-locate the learned codebook with this shard's operand device. When the + # default-world DTensors span more than one device (e.g. mixed CPU/CUDA + # shards) the codebook was learned on the first slot's device, but + # gefen_nearest_codebook_indices rejects a codebook whose device differs from + # the operand -- native keeps the codebook on the compute device. + codebook = codebook.to(device=flat.device, dtype=torch.float32) + blocks = automatic_partition_view(flat, period) + magnitude = automatic_partition_reduce(flat.abs(), period, reduce_op="max") + nonzero = magnitude > 0 + normalized = torch.where(nonzero, blocks / magnitude, torch.zeros_like(blocks)) + indices = gefen_nearest_codebook_indices(codebook, normalized) + vmean = automatic_partition_reduce(second.reshape(-1), period, reduce_op="mean") + return indices, magnitude.float(), vmean.float() + + +@dataclass(frozen=True) +class _Slot: + index: int + group_index: int + name: str + parameter: torch.Tensor + + +class GefenDCPState: + """DCP ``Stateful`` wrapper for reshardable plain-Gefen FSDP2 state. + + Use this object as the optimizer value passed to + :func:`torch.distributed.checkpoint.save` and ``load``. Save dequantizes the + quantized momentum against the learned per-rank codebook into dense + ``Shard(0)`` DTensors; load reshards those dense tensors and re-blocks them + back into Gefen's compact per-block state (re-running the period search, + relearning the codebook, and re-quantizing on each new local shard), so the + restored optimizer keeps its ~1 byte/param footprint. Because it goes through + a dense reshard this path is for *changing* topology; same-topology resumes + should use the native bit-exact ``state_dict`` path (left unchanged). + + The adapter is intentionally limited to plain :class:`gefen.Gefen` with + ``factored_v_2d=False`` and ``capturable=False``, one-dimensional + default-world DTensors, and one ``Shard(0)`` placement. + """ + + def __init__(self, optimizer): + from gefen.gefen import Gefen + + if type(optimizer) is not Gefen: + raise TypeError("GefenDCPState supports plain Gefen only") + self.optimizer = optimizer + self._slots = self._validate_layout() + + def _validate_layout(self): + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + raise RuntimeError( + "GefenDCPState requires an initialized distributed process group" + ) + import torch.distributed as dist + + if getattr(self.optimizer, "capturable", False): + raise RuntimeError( + "GefenDCPState does not support capturable=True: the step and " + "global-step counters are per-device tensors under capturable " + "(and compiled) execution, so their serialized host-side " + "semantics differ. Checkpoint with capturable=False." + ) + if getattr(self.optimizer, "_factored_v_2d", False): + raise RuntimeError( + "GefenDCPState does not support factored_v_2d=True: the factored " + "row/column second moment is not shard-addressable. Construct the " + "optimizer with factored_v_2d=False for DCP resharding." + ) + + slots = [] + index = 0 + seen_identities = {} + for group_index, group in enumerate(self.optimizer.param_groups): + if group.get("sharded_mode") is not None: + raise RuntimeError("GefenDCPState does not support Muon sharded modes") + names = group.get("param_names") + # Index-addressed slots are only safe to reshard if each carries a + # caller-stable, unique identity. A missing/short param_names list or + # a positionally synthesized name (Gefen's default for unnamed + # parameters) encodes registration order, not identity, so refuse it + # rather than silently cross-assigning momentum on a reordered load. + if names is None or len(names) != len(group["params"]): + raise RuntimeError( + "GefenDCPState requires caller-provided stable param_names on " + "every group (group {} has {} name(s) for {} parameter(s)); " + "pass named parameters so resharded momentum is assigned by " + "identity, not by registration order".format( + group_index, + 0 if names is None else len(names), + len(group["params"]), + ) + ) + for name, parameter in zip(names, group["params"]): + if _SYNTHESIZED_NAME_RE.match(str(name)): + raise RuntimeError( + "GefenDCPState requires caller-provided stable parameter " + "names, but slot {} carries the synthesized positional " + "name {!r}. Construct the optimizer with explicit " + "parameter names so resharded state is addressed by " + "identity, not registration order.".format(index, str(name)) + ) + if not _is_dtensor(parameter): + raise RuntimeError( + "GefenDCPState requires every parameter to be a DTensor" + ) + mesh = parameter.device_mesh + if len(mesh.shape) != 1 or mesh.size() != dist.get_world_size(): + raise RuntimeError( + "GefenDCPState supports only a one-dimensional default-world DeviceMesh" + ) + # A reordered full-world mesh (e.g. ranks [1, 0]) passes the size + # check but permutes shard->rank ownership, so the replicated + # initialized/counter metadata can mis-align with the sharded + # momentum. The documented scope is the canonical default world; + # reject anything else. + if mesh.mesh.flatten().tolist() != list(range(dist.get_world_size())): + raise RuntimeError( + "GefenDCPState supports only the canonical default-world " + "rank order (0..N-1); a reordered mesh is rejected" + ) + if len(parameter.placements) != 1: + raise RuntimeError( + "GefenDCPState supports exactly one DTensor placement" + ) + placement = parameter.placements[0] + if type(placement).__name__ != "Shard" or placement.dim != 0: + raise RuntimeError( + "GefenDCPState supports only one-dimensional Shard(0) parameters" + ) + identity = ( + str(name), + group_index, + tuple(int(dim) for dim in parameter.shape), + ) + if identity in seen_identities: + raise RuntimeError( + "GefenDCPState requires unique parameter identities, but " + "slots {} and {} share identity (name/group/shape) {!r}; " + "resharding cannot disambiguate their momentum.".format( + seen_identities[identity], index, identity + ) + ) + seen_identities[identity] = index + slots.append(_Slot(index, group_index, str(name), parameter)) + index += 1 + if not slots: + raise RuntimeError("GefenDCPState requires at least one parameter") + return tuple(slots) + + @staticmethod + def _key(slot: _Slot, field: str) -> str: + return "slot_{:08d}.{}".format(slot.index, field) + + # Required per-parameter Gefen state fields. A slot that carries all of them + # is materialized; a slot that carries none is fresh (name-only, pre-first + # step); a slot with only some is partial/corrupted and rejected on save. + _REQUIRED_STATE_FIELDS = ( + "automatic_period", + "step", + "m_codebook", + "m_magnitude", + "vmean", + ) + + def _identities(self): + """Stable per-slot identity list (replicated, in slot order). + + Persists the Gefen parameter name, its group membership, and the global + parameter shape so a load can reject a target whose parameters were + registered in a different order (index-only addressing would otherwise + cross-assign each same-shaped parameter the other's momentum). + """ + return [ + { + "name": slot.name, + "group": slot.group_index, + "shape": [int(dim) for dim in slot.parameter.shape], + } + for slot in self._slots + ] + + def _group_hypers(self): + """Per-group hyperparameters (replicated, in group order). + + The compact per-block state is only half of an optimizer resume; the + native full-state path restores the param-group hyperparameters too. A + freshly constructed optimizer resumed after an LR-schedule or hyper change + would otherwise silently keep its constructor values, not the checkpoint's. + """ + hypers = [] + for group in self.optimizer.param_groups: + entry = {} + for key in _GROUP_HYPER_KEYS: + value = group[key] + if torch.is_tensor(value): + value = value.item() + entry[key] = float(value) + hypers.append(entry) + return hypers + + @staticmethod + def _validate_hyper_entry(group_index, entry): + """Parse and range-validate one group's saved hyperparameters. + + Mirrors native ``Gefen._validate_group_options`` (lr>=0, 0<=betas<1, + weight_decay>=0, finite eps>0) and additionally rejects NaN/inf so a + corrupted checkpoint cannot commit values that silently poison the next + update. Returns a fresh ``{key: float}`` dict; raises before any mutation. + """ + if not isinstance(entry, dict): + raise ValueError( + "Gefen DCP checkpoint group {} hyperparameters must be a mapping".format( + group_index + ) + ) + parsed = {} + for key in _GROUP_HYPER_KEYS: + if key not in entry: + raise ValueError( + "Gefen DCP checkpoint group {} is missing hyperparameter {!r}".format( + group_index, key + ) + ) + value = entry[key] + if torch.is_tensor(value): + if value.numel() != 1: + raise ValueError( + "Gefen DCP checkpoint group {} hyperparameter {!r} must be " + "scalar".format(group_index, key) + ) + value = value.item() + try: + value = float(value) + except (TypeError, ValueError): + raise ValueError( + "Gefen DCP checkpoint group {} hyperparameter {!r} is " + "non-numeric".format(group_index, key) + ) + if not math.isfinite(value): + raise ValueError( + "Gefen DCP checkpoint group {} hyperparameter {!r} is not " + "finite ({})".format(group_index, key, value) + ) + parsed[key] = value + if parsed["lr"] < 0.0: + raise ValueError( + "Gefen DCP checkpoint group {} has invalid lr {}".format( + group_index, parsed["lr"] + ) + ) + if not 0.0 <= parsed["beta1"] < 1.0: + raise ValueError( + "Gefen DCP checkpoint group {} has invalid beta1 {}".format( + group_index, parsed["beta1"] + ) + ) + if not 0.0 <= parsed["beta2"] < 1.0: + raise ValueError( + "Gefen DCP checkpoint group {} has invalid beta2 {}".format( + group_index, parsed["beta2"] + ) + ) + if parsed["eps"] <= 0.0: + raise ValueError( + "Gefen DCP checkpoint group {} has invalid eps {} (must be > 0)".format( + group_index, parsed["eps"] + ) + ) + if parsed["weight_decay"] < 0.0: + raise ValueError( + "Gefen DCP checkpoint group {} has invalid weight_decay {}".format( + group_index, parsed["weight_decay"] + ) + ) + return parsed + + def _slot_initialized(self, slot: _Slot) -> bool: + """Classify a slot's live state; reject a partial/malformed materialize. + + Returns True for a fully materialized slot and False for a fresh + name-only slot. A slot that carries some but not all required fields is a + corrupted materialize whose missing history must not be silently + zero-filled, so it raises instead. + """ + state = self.optimizer.state.get(slot.parameter) + if not state: + return False + present = [key for key in self._REQUIRED_STATE_FIELDS if key in state] + if not present: + return False + if len(present) != len(self._REQUIRED_STATE_FIELDS): + missing = [ + key for key in self._REQUIRED_STATE_FIELDS if key not in state + ] + raise ValueError( + "Gefen DCP slot {} ({}) has partial optimizer state; missing " + "{}. Refusing to save it as uninitialized (which would zero the " + "momentum/second-moment history).".format( + slot.index, slot.name, ", ".join(missing) + ) + ) + return True + + def state_dict(self): + # Re-derive the slot layout (and re-run the fail-closed layout gate) so a + # retained wrapper whose optimizer gained parameters via add_param_group + # after construction saves the current set, not the stale snapshot. + self._slots = self._validate_layout() + result = { + "format_version": torch.tensor(_FORMAT_VERSION, dtype=torch.int64), + "global_step": torch.tensor( + _counter(self.optimizer._gefen_global_step, "global_step"), + dtype=torch.int64, + ), + "deterministic": torch.tensor( + int(bool(self.optimizer._deterministic)), dtype=torch.uint8 + ), + "slot_count": torch.tensor(len(self._slots), dtype=torch.int64), + "group_count": torch.tensor( + len(self.optimizer.param_groups), dtype=torch.int64 + ), + "param_identities": self._identities(), + "param_group_hypers": self._group_hypers(), + } + for slot in self._slots: + state = self.optimizer.state.get(slot.parameter) + initialized = self._slot_initialized(slot) + result[self._key(slot, "initialized")] = torch.tensor( + int(initialized), dtype=torch.uint8 + ) + result[self._key(slot, "step")] = torch.tensor( + _counter(state.get("step", 0), "step") if state else 0, + dtype=torch.int64, + ) + result[self._key(slot, "vmean_step")] = torch.tensor( + _counter(state.get("vmean_step", state.get("step", 0)), "vmean_step") + if state + else 0, + dtype=torch.int64, + ) + local = _local(slot.parameter) + momentum = ( + _dense_momentum(self.optimizer, slot.parameter, state) + if initialized + else torch.zeros_like(local, dtype=torch.float32) + ) + second = ( + _dense_second_moment(slot.parameter, state) + if initialized + else torch.zeros_like(local, dtype=torch.float32) + ) + result[self._key(slot, "momentum")] = _dtensor_from_local( + momentum, slot.parameter + ) + result[self._key(slot, "second_moment")] = _dtensor_from_local( + second, slot.parameter + ) + return result + + def load_state_dict(self, state_dict): + # Re-derive the slot layout (and re-run the fail-closed layout gate) so a + # load reflects any add_param_group that ran after construction instead of + # clearing+repopulating from a stale construction-time snapshot. + self._slots = self._validate_layout() + required = { + "format_version", + "global_step", + "deterministic", + "slot_count", + "group_count", + "param_identities", + "param_group_hypers", + } + if not required.issubset(state_dict): + raise ValueError("Gefen DCP checkpoint metadata is incomplete") + version = _counter(state_dict["format_version"], "format_version") + if version != _FORMAT_VERSION: + raise ValueError( + "Unsupported Gefen DCP checkpoint format version {}".format(version) + ) + if _counter(state_dict["slot_count"], "slot_count") != len(self._slots): + raise ValueError("Gefen DCP checkpoint parameter count differs") + + # Fail closed on a parameter-identity or group-topology mismatch BEFORE + # touching live state. Index-only slot addressing would otherwise let a + # target that registered its parameters in a different order load + # "successfully" while cross-assigning each same-shaped parameter the + # other's momentum/second-moment. + if _counter(state_dict["group_count"], "group_count") != len( + self.optimizer.param_groups + ): + raise ValueError("Gefen DCP checkpoint parameter-group count differs") + if list(state_dict["param_identities"]) != self._identities(): + raise ValueError( + "Gefen DCP checkpoint parameter identities (name/group/shape) do " + "not match this optimizer; refusing to load index-addressed state " + "that could cross-assign momentum between parameters" + ) + saved_hypers = list(state_dict["param_group_hypers"]) + if len(saved_hypers) != len(self.optimizer.param_groups): + raise ValueError( + "Gefen DCP checkpoint parameter-group hyperparameters differ in count" + ) + # Parse and range-validate every hyperparameter BEFORE any mutation so a + # missing/non-numeric key or an out-of-range value (lr=NaN, beta1=1, + # negative eps/weight_decay) is rejected fail-atomically instead of being + # committed and silently corrupting the next update. + staged_hypers = [ + self._validate_hyper_entry(group_index, entry) + for group_index, entry in enumerate(saved_hypers) + ] + + # Reject a deterministic-policy mismatch instead of silently overwriting + # it, matching native Gefen.load_state_dict: the flag changes fused-routing + # / replica-determinism semantics, so resuming under a different policy + # requires an intentional migration, not a checkpoint that flips it. + checkpoint_deterministic = bool( + _counter(state_dict["deterministic"], "deterministic") + ) + if checkpoint_deterministic != bool(self.optimizer._deterministic): + raise ValueError( + "Gefen DCP checkpoint deterministic={!r}, but this optimizer was " + "constructed with deterministic={!r}. Resuming under a different " + "replica-determinism policy requires an intentional state " + "migration or a fresh optimizer; refusing to change it " + "silently.".format( + checkpoint_deterministic, bool(self.optimizer._deterministic) + ) + ) + + # Stage, validate, and re-block everything locally BEFORE mutating any + # live optimizer state. The whole preparation runs inside this helper so a + # per-rank failure (a non-finite resharded slice, an allocation failure, + # ...) is caught and synchronized across the process group below: either + # every rank commits or every rank raises. Without the cross-rank + # agreement a slice that only fails on one target rank would leave some + # live optimizers restored and others untouched. + def _stage(): + staged = [] + for slot in self._slots: + keys = { + field: self._key(slot, field) + for field in ( + "initialized", + "step", + "vmean_step", + "momentum", + "second_moment", + ) + } + if any(key not in state_dict for key in keys.values()): + raise ValueError( + "Gefen DCP checkpoint is missing state for slot {}".format( + slot.index + ) + ) + initialized = bool( + _counter(state_dict[keys["initialized"]], "initialized") + ) + step = _counter(state_dict[keys["step"]], "step") + vmean_step = _counter( + state_dict[keys["vmean_step"]], "vmean_step" + ) + momentum = _local(state_dict[keys["momentum"]]).detach().float() + second = _local(state_dict[keys["second_moment"]]).detach().float() + local = _local(slot.parameter) + if momentum.shape != local.shape or second.shape != local.shape: + raise ValueError( + "Gefen DCP checkpoint local shape differs for slot {}".format( + slot.index + ) + ) + if ( + not torch.isfinite(momentum).all() + or not torch.isfinite(second).all() + ): + raise ValueError("Gefen DCP checkpoint contains non-finite state") + if (second < 0).any(): + raise ValueError("Gefen DCP second moment must be nonnegative") + # Reject incoherent initialization metadata instead of loading it + # or silently discarding history, matching native Gefen validation + # (an initialized momentum + its second moment must have positive + # ages; an uninitialized slot must carry no counters or history). + # An initialized slot that reshards to an EMPTY local shard keeps + # its replicated positive counters and is materialized as name-only + # later, exactly as native never materializes an empty local shard. + if initialized: + if step < 1: + raise ValueError( + "Gefen DCP slot {} is initialized but carries step={}; " + "initialized momentum requires step >= 1".format( + slot.index, step + ) + ) + if vmean_step < 1: + raise ValueError( + "Gefen DCP slot {} is initialized but carries vmean_step" + "={}; initialized second moment requires vmean_step >= " + "1".format(slot.index, vmean_step) + ) + else: + if step != 0 or vmean_step != 0: + raise ValueError( + "Gefen DCP slot {} is uninitialized but carries nonzero " + "counters (step={}, vmean_step={})".format( + slot.index, step, vmean_step + ) + ) + if momentum.numel() and ( + bool(momentum.any()) or bool(second.any()) + ): + raise ValueError( + "Gefen DCP slot {} is uninitialized but carries nonzero " + "momentum/second-moment history that would be silently " + "discarded".format(slot.index) + ) + staged.append((slot, initialized, step, vmean_step, momentum, second)) + + global_step = _counter(state_dict["global_step"], "global_step") + + # Re-block the resharded dense momentum back into Gefen's compact + # per-block representation instead of collapsing it to period one. + # Going through a dense DCP reshard is inherently a new blocking (so + # this is NOT bit-exact to a native same-topology resume), but it + # restores the ~1 byte/param memory profile and is a correct, finite + # continuation: per new local shard we re-run the block-variance + # period search, relearn the exact codebook, and re-quantize. All of + # this happens BEFORE the optimizer state is cleared, preserving the + # fail-atomic load contract. + # + # A slot that is initialized in the (replicated) checkpoint metadata + # but reshards to an EMPTY local shard on this rank -- N->M where dim-0 + # < the target world -- carries no local momentum to quantize. Native + # Gefen never materializes state for an empty local shard (the step + # returns early on an empty grad), so it is restored as name-only here + # and left out of the codebook learning; feeding an empty shard through + # the learner would return a None codebook and crash the re-quantize. + def _materialized(initialized, momentum): + return initialized and momentum.numel() > 0 + + device = _local(self._slots[0].parameter).device + periods = {} + for slot, initialized, step, vmean_step, momentum, second in staged: + if _materialized(initialized, momentum): + periods[slot.index] = _choose_period( + self.optimizer, slot.name, slot.parameter, second + ) + codebook = _learn_codebook( + [ + (slot.name, momentum.reshape(-1), periods[slot.index]) + for slot, initialized, _, _, momentum, _ in staged + if _materialized(initialized, momentum) + ], + device, + ) + + new_state = {} + for slot, initialized, step, vmean_step, momentum, second in staged: + if not _materialized(initialized, momentum): + new_state[slot.parameter] = {"name": slot.name} + continue + period = periods[slot.index] + m_codebook, m_magnitude, vmean = _reblock( + codebook, momentum, second, period + ) + new_state[slot.parameter] = { + "name": slot.name, + "automatic_period": period, + "step": step, + "m_codebook": m_codebook, + "m_magnitude": m_magnitude, + "vmean": vmean, + "vmean_step": vmean_step, + } + return new_state, global_step, codebook + + import torch.distributed as dist + from gefen.gefen import _synchronize_step_failure + + local_error = None + prepared = None + try: + prepared = _stage() + except Exception as exc: # resynchronized across the group below + local_error = exc + # Agree on success across the process group before publishing so a + # one-rank failure aborts every rank rather than committing a partial, + # cross-rank-inconsistent restore. The plain-Gefen DTensors live on the + # default world (validated in _validate_layout), so the WORLD group is the + # right scope; _synchronize_step_failure keeps the flag on CPU for + # gloo/CPU-resident state and no-ops for a single-rank group. + if _synchronize_step_failure(local_error is not None, dist.group.WORLD): + if local_error is not None: + raise local_error + raise RuntimeError( + "Gefen DCP load aborted before commit: another rank failed " + "staging/validation/re-block, so no rank commits its restore." + ) + new_state, global_step, codebook = prepared + + self.optimizer.state.clear() + self.optimizer.state.update(new_state) + # Restore the checkpoint's per-group hyperparameters so an LR-schedule or + # hyperparameter change survives the resume, matching the native full-state + # path. A live tensor lr is updated in place to preserve its identity/device + # for any fused kernel holding a reference. + for group, entry in zip(self.optimizer.param_groups, staged_hypers): + for key in _GROUP_HYPER_KEYS: + value = entry[key] + current = group.get(key) + if torch.is_tensor(current): + current.fill_(value) + else: + group[key] = value + self.optimizer._gefen_global_step = global_step + # A relearned codebook reflects this rank's resharded momentum and is kept + # by _maybe_refresh_gefen_codebook on the next step (which then also skips + # re-predicting periods, keeping the restored block geometry consistent). + # With no initialized momentum there is nothing to learn from; leave the + # codebook unset so the next real step learns it cold. + self.optimizer._gefen_codebook = codebook + self.optimizer._gefen_codebook_by_device.clear() + self.optimizer._gefen_codebook_lut_by_device.clear() + self.optimizer._sr_seed_by_device.clear() + self.optimizer._reset_gefen_global_step_devices() + self.optimizer._static_mark_sig = None diff --git a/tests/test_dcp_resharding.py b/tests/test_dcp_resharding.py new file mode 100644 index 0000000..09a1a57 --- /dev/null +++ b/tests/test_dcp_resharding.py @@ -0,0 +1,1418 @@ +"""Real DCP save/load and N-to-M resharding for standalone Gefen state. + +These tests exercise the *real* save path -- a genuine ``Gefen.step`` learns the +exact codebook and picks a block period > 1 -- and check that after a DCP reshard +the restored optimizer (a) keeps Gefen's compact ~1 byte/param block state +instead of collapsing to per-element period one, and (b) continues within a +quantization-noise tolerance of a native run at the *target* topology. The +continuation is tolerance-based, not exact, because routing the momentum through +a dense DCP reshard re-blocks it against a freshly learned per-shard codebook -- +a different (and inherently lossy) blocking than the native target-topology run. +""" + +from __future__ import annotations + +from datetime import timedelta +import multiprocessing as mp +import os +import queue +import socket +import traceback + +import pytest +import torch +import torch.distributed as dist +from torch import nn +from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter +from torch.distributed.tensor import ( + Replicate, + Shard, + distribute_tensor, + init_device_mesh, +) + +from gefen import Gefen, GefenDCPState + + +# A 2D weight large enough that the block-variance period search returns a real +# period > 1 (so the compact vs per-element distinction is actually tested) and +# that 2, 4 both divide the row count for clean N-to-M resharding. +_SHAPE = (256, 128) +_SAVE_STEPS = 3 + + +def _free_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return str(sock.getsockname()[1]) + + +def _local(value): + value = value.to_local() if hasattr(value, "to_local") else value + return value.wait() if hasattr(value, "wait") else value + + +def _global_param(): + numel = _SHAPE[0] * _SHAPE[1] + return torch.linspace(-0.8, 0.7, numel).reshape(_SHAPE) + + +def _global_grad(step, device): + # Deterministic global gradient shared by every topology, so the momentum / + # second-moment history is identical across world sizes up to quantization. + generator = torch.Generator().manual_seed(1000 + step) + grad = torch.randn(_SHAPE, generator=generator) + return grad.to(device) + + +def _build(mesh, device, fused=False): + parameter = nn.Parameter( + distribute_tensor(_global_param().to(device), mesh, [Shard(0)]) + ) + optimizer = Gefen( + [("weight", parameter)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2e-8, + weight_decay=0.03, + fused=fused, + factored_v_2d=False, + deterministic=True, + ) + return parameter, optimizer + + +def _step(parameter, optimizer, mesh, step, device): + parameter.grad = distribute_tensor( + _global_grad(step, device), mesh, [Shard(0)] + ) + optimizer.step() + + +def _period_info(optimizer, parameter): + state = optimizer.state[parameter] + period = int(state["automatic_period"]) + local_numel = int(_local(parameter).numel()) + return { + "period": period, + "vmean_numel": int(state["vmean"].numel()), + "mmag_numel": int(state["m_magnitude"].numel()), + "mcodebook_numel": int(state["m_codebook"].numel()), + "local_numel": local_numel, + # Compact iff the second-moment / magnitude state is one value per block, + # i.e. numel == local_numel / period, not one-per-element (period one). + "compact": ( + period > 1 + and int(state["vmean"].numel()) * period == local_numel + and int(state["m_magnitude"].numel()) * period == local_numel + ), + } + + +def _worker(rank, world, port, checkpoint_dir, mode, device_type, fused, result_queue): + try: + os.environ.update( + MASTER_ADDR="127.0.0.1", + MASTER_PORT=port, + RANK=str(rank), + WORLD_SIZE=str(world), + ) + if device_type == "cuda": + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.device("cuda", rank % torch.cuda.device_count()) + else: + device = torch.device("cpu") + dist.init_process_group( + "cuda:nccl,cpu:gloo" if device_type == "cuda" else "gloo", + rank=rank, + world_size=world, + timeout=timedelta(seconds=120), + ) + mesh = init_device_mesh(device_type, (world,), mesh_dim_names=("dp",)) + parameter, optimizer = _build(mesh, device, fused) + + if mode == "save": + for step in range(_SAVE_STEPS): + _step(parameter, optimizer, mesh, step, device) + saved = _period_info(optimizer, parameter) + torch.distributed.checkpoint.save( + {"optimizer": GefenDCPState(optimizer)}, + storage_writer=FileSystemWriter(checkpoint_dir), + ) + result_queue.put({"rank": rank, "saved": True, **saved}) + return + + # mode == "load": reshard the saved state onto this (possibly different) + # world, then compare a one-step continuation against a native reference + # trained at THIS topology on the identical global gradient history. + torch.distributed.checkpoint.load( + {"optimizer": GefenDCPState(optimizer)}, + storage_reader=FileSystemReader(checkpoint_dir), + ) + restored = _period_info(optimizer, parameter) + + reference_param, reference = _build(mesh, device, fused) + for step in range(_SAVE_STEPS): + _step(reference_param, reference, mesh, step, device) + # Continue both from the identical parameter so the delta isolates the + # restored optimizer STATE (momentum / second moment), not the parameter + # trajectory that already diverged by per-shard quantization noise. + reference_snapshot = _local(reference_param).detach().clone() + with torch.no_grad(): + _local(parameter).copy_(reference_snapshot) + + _step(parameter, optimizer, mesh, _SAVE_STEPS, device) + _step(reference_param, reference, mesh, _SAVE_STEPS, device) + + restored_next = _local(parameter).detach() + reference_next = _local(reference_param).detach() + finite = bool(torch.isfinite(restored_next).all()) + max_abs_diff = float((restored_next - reference_next).abs().max()) + result_queue.put( + { + "rank": rank, + "shape": tuple(_local(parameter).shape), + "restored": restored, + "finite": finite, + "continuation_max_abs_diff": max_abs_diff, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run(world, checkpoint_dir, mode, *, device_type="cpu", fused=False, timeout=240): + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_worker, + args=( + rank, + world, + port, + checkpoint_dir, + mode, + device_type, + fused, + result_queue, + ), + ) + for rank in range(world) + ] + for process in processes: + process.start() + results = [] + try: + for _ in processes: + results.append(result_queue.get(timeout=timeout)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + assert len(results) == world, (results, [p.exitcode for p in processes]) + assert all(p.exitcode == 0 for p in processes), [p.exitcode for p in processes] + return sorted(results, key=lambda item: item["rank"]) + + +# Continuation tolerance. The resharded momentum is re-blocked against a codebook +# relearned on the new local shard -- a different, lossy blocking than the native +# target-topology run -- so the one-step continuation matches only up to +# accumulated 256-level quantization noise (empirically ~1e-3 on these shards), +# not bit-for-bit. +_CONTINUE_TOL = 5e-3 + + +def _assert_loaded(loaded, expected_shapes=None): + assert all("fatal_error" not in item for item in loaded), loaded + if expected_shapes is not None: + assert [item["shape"] for item in loaded] == expected_shapes, loaded + for item in loaded: + assert item["finite"], item + # The whole point of Gefen: the restored state must stay compact + # (~1 byte/param), not collapse to per-element period one. + assert item["restored"]["compact"], item + assert item["restored"]["period"] > 1, item + assert item["continuation_max_abs_diff"] < _CONTINUE_TOL, item + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +def test_dcp_same_topology_stays_compact_and_continues(tmp_path): + checkpoint_dir = str(tmp_path / "gefen-dcp-2-to-2") + saved = _run(2, checkpoint_dir, "save") + assert all(item.get("saved") for item in saved), saved + assert all(item["compact"] and item["period"] > 1 for item in saved), saved + loaded = _run(2, checkpoint_dir, "load") + _assert_loaded(loaded, expected_shapes=[(128, 128), (128, 128)]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +def test_dcp_save_on_two_ranks_loads_and_continues_on_four(tmp_path): + checkpoint_dir = str(tmp_path / "gefen-dcp-2-to-4") + saved = _run(2, checkpoint_dir, "save") + assert all(item.get("saved") for item in saved), saved + loaded = _run(4, checkpoint_dir, "load") + _assert_loaded( + loaded, + expected_shapes=[(64, 128), (64, 128), (64, 128), (64, 128)], + ) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +def test_dcp_save_on_four_ranks_loads_and_continues_on_two(tmp_path): + checkpoint_dir = str(tmp_path / "gefen-dcp-4-to-2") + saved = _run(4, checkpoint_dir, "save") + assert all(item.get("saved") for item in saved), saved + loaded = _run(2, checkpoint_dir, "load") + _assert_loaded(loaded, expected_shapes=[(128, 128), (128, 128)]) + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.device_count() < 4 + or not dist.is_nccl_available(), + reason="GPU DCP resharding coverage requires four CUDA devices and NCCL", +) +def test_dcp_save_on_two_gpus_loads_and_continues_on_four(tmp_path): + # NCCL forbids two ranks per physical device in one communicator, so the + # four-rank load needs four distinct GPUs (rank -> set_device(rank)). + checkpoint_dir = str(tmp_path / "gefen-dcp-gpu-2-to-4") + saved = _run(2, checkpoint_dir, "save", device_type="cuda") + assert all(item.get("saved") for item in saved), saved + loaded = _run(4, checkpoint_dir, "load", device_type="cuda") + _assert_loaded(loaded) + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.device_count() < 2 + or not dist.is_nccl_available(), + reason="fused DCP coverage requires two CUDA devices and NCCL", +) +def test_dcp_fused_state_reshards_and_stays_compact(tmp_path): + # The fused CUDA kernels (fused=True) produce the same optimizer-state layout + # as the decomposed path, so GefenDCPState must save/reshard fused-trained + # state too. Same-topology (2->2) on CUDA is enough to prove the fused state + # round-trips compactly and continues; the reshard math itself is covered by + # the non-fused N-to-M tests and is fused-independent. + checkpoint_dir = str(tmp_path / "gefen-dcp-fused") + saved = _run(2, checkpoint_dir, "save", device_type="cuda", fused=True) + assert all(item.get("saved") for item in saved), saved + assert all(item["compact"] and item["period"] > 1 for item in saved), saved + loaded = _run(2, checkpoint_dir, "load", device_type="cuda", fused=True) + _assert_loaded(loaded) + + +def _fully_shard_worker(rank, world, port, checkpoint_dir, mode, result_queue): + try: + os.environ.update( + MASTER_ADDR="127.0.0.1", + MASTER_PORT=port, + RANK=str(rank), + WORLD_SIZE=str(world), + ) + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + "cuda:nccl,cpu:gloo", + rank=rank, + world_size=world, + timeout=timedelta(seconds=120), + ) + from torch.distributed.fsdp import fully_shard + + torch.manual_seed(0) + model = nn.Linear(128, 64, bias=False).to(device) + fully_shard(model) + parameter = model.weight + optimizer = Gefen( + [("weight", parameter)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2e-8, + fused=False, + factored_v_2d=False, + deterministic=True, + ) + + def real_step(seed): + generator = torch.Generator(device=device).manual_seed(seed) + inputs = torch.randn(32, 128, generator=generator, device=device) + targets = torch.randn(32, 64, generator=generator, device=device) + optimizer.zero_grad() + loss = ((model(inputs) - targets) ** 2).mean() + loss.backward() + optimizer.step() + + if mode == "save": + for step in range(_SAVE_STEPS): + real_step(2000 + step) + info = _period_info(optimizer, parameter) + torch.distributed.checkpoint.save( + {"optimizer": GefenDCPState(optimizer)}, + storage_writer=FileSystemWriter(checkpoint_dir), + ) + result_queue.put({"rank": rank, "saved": True, **info}) + return + + torch.distributed.checkpoint.load( + {"optimizer": GefenDCPState(optimizer)}, + storage_reader=FileSystemReader(checkpoint_dir), + ) + restored = _period_info(optimizer, parameter) + real_step(2000 + _SAVE_STEPS) + finite = bool(torch.isfinite(_local(parameter)).all()) + result_queue.put({"rank": rank, "restored": restored, "finite": finite}) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_fully_shard(world, checkpoint_dir, mode, timeout=240): + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_fully_shard_worker, + args=(rank, world, port, checkpoint_dir, mode, result_queue), + ) + for rank in range(world) + ] + for process in processes: + process.start() + results = [] + try: + for _ in processes: + results.append(result_queue.get(timeout=timeout)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + assert len(results) == world, (results, [p.exitcode for p in processes]) + assert all(p.exitcode == 0 for p in processes), [p.exitcode for p in processes] + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.device_count() < 2 + or not dist.is_nccl_available(), + reason="fully_shard DCP coverage requires two CUDA devices and NCCL", +) +def test_dcp_fully_shard_real_model_stays_compact(tmp_path): + checkpoint_dir = str(tmp_path / "gefen-dcp-fully-shard") + saved = _run_fully_shard(2, checkpoint_dir, "save") + assert all(item.get("saved") for item in saved), saved + assert all(item["compact"] and item["period"] > 1 for item in saved), saved + loaded = _run_fully_shard(2, checkpoint_dir, "load") + assert all("fatal_error" not in item for item in loaded), loaded + for item in loaded: + assert item["finite"] and item["restored"]["compact"], item + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="requires Gloo") +def test_rejects_non_dtensor_optimizer(tmp_path): + init_file = tmp_path / "single-rank-init" + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=0, + world_size=1, + ) + try: + parameter = nn.Parameter(torch.ones(4, 4)) + optimizer = Gefen( + [("weight", parameter)], lr=1e-3, fused=False, factored_v_2d=False + ) + with pytest.raises(RuntimeError, match="every parameter to be a DTensor"): + GefenDCPState(optimizer) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="requires Gloo") +def test_rejects_capturable(tmp_path): + init_file = tmp_path / "capturable-init" + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=0, + world_size=1, + ) + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + parameter = nn.Parameter(distribute_tensor(_global_param(), mesh, [Shard(0)])) + optimizer = Gefen( + [("weight", parameter)], + lr=1e-3, + fused=False, + factored_v_2d=False, + capturable=True, + ) + with pytest.raises(RuntimeError, match="capturable"): + GefenDCPState(optimizer) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="requires Gloo") +def test_rejects_factored_v_2d(tmp_path): + init_file = tmp_path / "factored-init" + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=0, + world_size=1, + ) + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + parameter = nn.Parameter(distribute_tensor(_global_param(), mesh, [Shard(0)])) + optimizer = Gefen( + [("weight", parameter)], + lr=1e-3, + fused=False, + factored_v_2d=True, + ) + with pytest.raises(RuntimeError, match="factored_v_2d"): + GefenDCPState(optimizer) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="requires Gloo") +def test_rejects_non_shard0_placement(tmp_path): + # The PR advertises that non-Shard(0) placements fail closed. + init_file = tmp_path / "replicate-init" + dist.init_process_group( + "gloo", init_method="file://{}".format(init_file), rank=0, world_size=1 + ) + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + parameter = nn.Parameter(distribute_tensor(_global_param(), mesh, [Replicate()])) + optimizer = Gefen( + [("weight", parameter)], lr=1e-3, fused=False, factored_v_2d=False + ) + with pytest.raises(RuntimeError, match="Shard"): + GefenDCPState(optimizer) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="requires Gloo") +def test_rejects_multidimensional_mesh(tmp_path): + # Multidimensional meshes fail closed. + init_file = tmp_path / "mesh2d-init" + dist.init_process_group( + "gloo", init_method="file://{}".format(init_file), rank=0, world_size=1 + ) + try: + mesh = init_device_mesh("cpu", (1, 1), mesh_dim_names=("dp", "tp")) + parameter = nn.Parameter( + distribute_tensor(_global_param(), mesh, [Shard(0), Replicate()]) + ) + optimizer = Gefen( + [("weight", parameter)], lr=1e-3, fused=False, factored_v_2d=False + ) + with pytest.raises(RuntimeError, match="one-dimensional"): + GefenDCPState(optimizer) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="requires Gloo") +def test_rejects_muon_sharded_mode(tmp_path): + # A group carrying a Muon sharded_mode fails closed. + init_file = tmp_path / "shardedmode-init" + dist.init_process_group( + "gloo", init_method="file://{}".format(init_file), rank=0, world_size=1 + ) + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + parameter = nn.Parameter(distribute_tensor(_global_param(), mesh, [Shard(0)])) + optimizer = Gefen( + [("weight", parameter)], lr=1e-3, fused=False, factored_v_2d=False + ) + optimizer.param_groups[0]["sharded_mode"] = "approx" + with pytest.raises(RuntimeError, match="sharded mode"): + GefenDCPState(optimizer) + finally: + dist.destroy_process_group() + + +def _reordered_mesh_worker(rank, world, port, result_queue): + try: + os.environ.update( + MASTER_ADDR="127.0.0.1", + MASTER_PORT=port, + RANK=str(rank), + WORLD_SIZE=str(world), + ) + dist.init_process_group( + "gloo", rank=rank, world_size=world, timeout=timedelta(seconds=60) + ) + from torch.distributed.device_mesh import DeviceMesh + + # A full-world 1-D mesh whose ranks are reordered ([1, 0] on world 2): + # it passes the size check but permutes shard->rank ownership. + mesh = DeviceMesh("cpu", torch.tensor([1, 0])) + parameter = nn.Parameter(distribute_tensor(_global_param(), mesh, [Shard(0)])) + optimizer = Gefen( + [("weight", parameter)], lr=1e-3, fused=False, factored_v_2d=False + ) + raised = False + try: + GefenDCPState(optimizer) + except RuntimeError as exc: + raised = "reordered" in str(exc) or "canonical" in str(exc) + result_queue.put({"rank": rank, "raised": raised}) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="requires Gloo", +) +def test_rejects_reordered_full_world_mesh(): + # A reordered full-world mesh must fail closed on every rank. + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process(target=_reordered_mesh_worker, args=(rank, 2, port, result_queue)) + for rank in range(2) + ] + for process in processes: + process.start() + results = [] + try: + for _ in processes: + results.append(result_queue.get(timeout=120)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + assert all("fatal_error" not in item for item in results), results + assert len(results) == 2 and all(item["raised"] for item in results), results + + +# --- Review-finding coverage (Codex PR #83) ------------------------------- +# +# The tests below cover the fail-closed / semantics fixes. Findings 1, 2, 3 and 5 +# do not need a real topology change, so they run in-process on a one-rank Gloo +# group (save, then load into a differently configured optimizer). Finding 4 +# (an N->M reshard whose dim-0 is smaller than the target world) needs multiple +# ranks and reuses the spawn/worker harness. + + +def _shaped_dparam(shape, mesh, device="cpu"): + numel = 1 + for dim in shape: + numel *= dim + tensor = torch.linspace(-0.8, 0.7, numel).reshape(shape).to(device) + return nn.Parameter(distribute_tensor(tensor, mesh, [Shard(0)])) + + +def _shaped_grad(shape, step): + generator = torch.Generator().manual_seed(1000 + step) + return torch.randn(shape, generator=generator) + + +def _save_optimizer(optimizer, param, shape, mesh, checkpoint_dir): + for step in range(_SAVE_STEPS): + param.grad = distribute_tensor(_shaped_grad(shape, step), mesh, [Shard(0)]) + optimizer.step() + torch.distributed.checkpoint.save( + {"optimizer": GefenDCPState(optimizer)}, + storage_writer=FileSystemWriter(checkpoint_dir), + ) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +def test_dcp_restores_checkpoint_group_hyperparameters(tmp_path): + init_file = tmp_path / "hyper-init" + dist.init_process_group( + "gloo", init_method="file://{}".format(init_file), rank=0, world_size=1 + ) + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + param = _shaped_dparam(_SHAPE, mesh) + optimizer = Gefen( + [("weight", param)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2e-8, + weight_decay=0.03, + fused=False, + factored_v_2d=False, + deterministic=True, + ) + checkpoint_dir = str(tmp_path / "hyper-ckpt") + _save_optimizer(optimizer, param, _SHAPE, mesh, checkpoint_dir) + + # Resume into an optimizer built with DIFFERENT hyperparameters, as if an + # LR scheduler had advanced or the run was reconfigured. The restore must + # adopt the checkpoint's hyperparameters, not silently keep these. + param2 = _shaped_dparam(_SHAPE, mesh) + resumed = Gefen( + [("weight", param2)], + lr=9.9e-9, + betas=(0.1, 0.2), + eps=7e-7, + weight_decay=0.5, + fused=False, + factored_v_2d=False, + deterministic=True, + ) + torch.distributed.checkpoint.load( + {"optimizer": GefenDCPState(resumed)}, + storage_reader=FileSystemReader(checkpoint_dir), + ) + group = resumed.param_groups[0] + assert abs(group["lr"] - 2.5e-3) < 1e-12, group["lr"] + assert abs(group["beta1"] - 0.8) < 1e-12, group["beta1"] + assert abs(group["beta2"] - 0.97) < 1e-12, group["beta2"] + assert abs(group["eps"] - 2e-8) < 1e-15, group["eps"] + assert abs(group["weight_decay"] - 0.03) < 1e-12, group["weight_decay"] + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +def test_dcp_rejects_swapped_parameter_identities(tmp_path): + init_file = tmp_path / "identity-init" + dist.init_process_group( + "gloo", init_method="file://{}".format(init_file), rank=0, world_size=1 + ) + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + shape = (128, 64) + first = _shaped_dparam(shape, mesh) + second = _shaped_dparam(shape, mesh) + optimizer = Gefen( + [("alpha", first), ("beta", second)], + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + for step in range(_SAVE_STEPS): + for param in (first, second): + param.grad = distribute_tensor( + _shaped_grad(shape, step), mesh, [Shard(0)] + ) + optimizer.step() + checkpoint_dir = str(tmp_path / "identity-ckpt") + torch.distributed.checkpoint.save( + {"optimizer": GefenDCPState(optimizer)}, + storage_writer=FileSystemWriter(checkpoint_dir), + ) + + # The same two same-shaped parameters, registered in SWAPPED name order. + # Slot-index addressing alone would load "successfully" and hand each + # parameter the other's momentum; the identity guard must reject it. + first_swapped = _shaped_dparam(shape, mesh) + second_swapped = _shaped_dparam(shape, mesh) + swapped = Gefen( + [("beta", second_swapped), ("alpha", first_swapped)], + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + with pytest.raises(ValueError, match="parameter identities"): + torch.distributed.checkpoint.load( + {"optimizer": GefenDCPState(swapped)}, + storage_reader=FileSystemReader(checkpoint_dir), + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +def test_dcp_rejects_deterministic_mismatch(tmp_path): + init_file = tmp_path / "deterministic-init" + dist.init_process_group( + "gloo", init_method="file://{}".format(init_file), rank=0, world_size=1 + ) + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + param = _shaped_dparam(_SHAPE, mesh) + optimizer = Gefen( + [("weight", param)], + lr=1e-3, + fused=False, + factored_v_2d=False, + deterministic=True, + ) + checkpoint_dir = str(tmp_path / "deterministic-ckpt") + _save_optimizer(optimizer, param, _SHAPE, mesh, checkpoint_dir) + + param2 = _shaped_dparam(_SHAPE, mesh) + other = Gefen( + [("weight", param2)], + lr=1e-3, + fused=False, + factored_v_2d=False, + deterministic=False, + ) + with pytest.raises(ValueError, match="deterministic"): + torch.distributed.checkpoint.load( + {"optimizer": GefenDCPState(other)}, + storage_reader=FileSystemReader(checkpoint_dir), + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +def test_dcp_restore_honors_force_2d_period_one(tmp_path): + init_file = tmp_path / "period-init" + dist.init_process_group( + "gloo", init_method="file://{}".format(init_file), rank=0, world_size=1 + ) + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + param = _shaped_dparam(_SHAPE, mesh) + optimizer = Gefen( + [("weight", param)], lr=1e-3, fused=False, factored_v_2d=False + ) + checkpoint_dir = str(tmp_path / "period-ckpt") + _save_optimizer(optimizer, param, _SHAPE, mesh, checkpoint_dir) + + # Control: an unconstrained resume re-derives the compact block period. + control_param = _shaped_dparam(_SHAPE, mesh) + control = Gefen( + [("weight", control_param)], lr=1e-3, fused=False, factored_v_2d=False + ) + torch.distributed.checkpoint.load( + {"optimizer": GefenDCPState(control)}, + storage_reader=FileSystemReader(checkpoint_dir), + ) + assert int(control.state[control_param]["automatic_period"]) > 1 + + # A resume that explicitly forces 2D params to period one must honor that + # gate rather than running the raw search and restoring period > 1. + forced_param = _shaped_dparam(_SHAPE, mesh) + forced = Gefen( + [("weight", forced_param)], + lr=1e-3, + fused=False, + factored_v_2d=False, + force_2d_period_one=True, + ) + torch.distributed.checkpoint.load( + {"optimizer": GefenDCPState(forced)}, + storage_reader=FileSystemReader(checkpoint_dir), + ) + assert int(forced.state[forced_param]["automatic_period"]) == 1, forced.state[ + forced_param + ]["automatic_period"] + finally: + dist.destroy_process_group() + + +def _empty_shard_worker(rank, world, port, checkpoint_dir, mode, result_queue): + try: + os.environ.update( + MASTER_ADDR="127.0.0.1", + MASTER_PORT=port, + RANK=str(rank), + WORLD_SIZE=str(world), + ) + dist.init_process_group( + "gloo", rank=rank, world_size=world, timeout=timedelta(seconds=120) + ) + mesh = init_device_mesh("cpu", (world,), mesh_dim_names=("dp",)) + # dim-0 == 2: fits two ranks fully, but reshards to empty local shards on + # ranks >= 2 of a four-rank target. + shape = (2, 128) + parameter = nn.Parameter( + distribute_tensor( + torch.linspace(-0.8, 0.7, shape[0] * shape[1]).reshape(shape), + mesh, + [Shard(0)], + ) + ) + optimizer = Gefen( + [("weight", parameter)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2e-8, + fused=False, + factored_v_2d=False, + deterministic=True, + ) + + if mode == "save": + for step in range(_SAVE_STEPS): + parameter.grad = distribute_tensor( + _shaped_grad(shape, step), mesh, [Shard(0)] + ) + optimizer.step() + torch.distributed.checkpoint.save( + {"optimizer": GefenDCPState(optimizer)}, + storage_writer=FileSystemWriter(checkpoint_dir), + ) + result_queue.put({"rank": rank, "saved": True}) + return + + # The load must not crash on ranks whose local shard resharded to empty + # (a None per-rank codebook previously dereferenced codebook.device). + torch.distributed.checkpoint.load( + {"optimizer": GefenDCPState(optimizer)}, + storage_reader=FileSystemReader(checkpoint_dir), + ) + local_numel = int(_local(parameter).numel()) + state = optimizer.state.get(parameter, {}) + materialized = all( + key in state + for key in ("automatic_period", "m_codebook", "m_magnitude", "vmean") + ) + finite = True + if materialized: + finite = bool( + torch.isfinite(state["m_magnitude"]).all() + and torch.isfinite(state["vmean"]).all() + ) + result_queue.put( + { + "rank": rank, + "local_numel": local_numel, + "materialized": materialized, + "finite": finite, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_empty(world, checkpoint_dir, mode, timeout=240): + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_empty_shard_worker, + args=(rank, world, port, checkpoint_dir, mode, result_queue), + ) + for rank in range(world) + ] + for process in processes: + process.start() + results = [] + try: + for _ in processes: + results.append(result_queue.get(timeout=timeout)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + assert len(results) == world, (results, [p.exitcode for p in processes]) + assert all(p.exitcode == 0 for p in processes), [p.exitcode for p in processes] + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +def test_dcp_reshard_to_empty_target_shard_does_not_crash(tmp_path): + checkpoint_dir = str(tmp_path / "gefen-dcp-empty-shard") + saved = _run_empty(2, checkpoint_dir, "save") + assert all(item.get("saved") for item in saved), saved + loaded = _run_empty(4, checkpoint_dir, "load") + assert all("fatal_error" not in item for item in loaded), loaded + by_rank = {item["rank"]: item for item in loaded} + # dim-0 == 2 over world 4: ranks 0 and 1 own a row, ranks 2 and 3 reshard to + # empty local shards. The empty ranks must load as unmaterialized state (like + # native, which never materializes an empty local shard) instead of crashing. + assert by_rank[0]["materialized"] and by_rank[1]["materialized"], loaded + assert by_rank[0]["local_numel"] == 128 and by_rank[1]["local_numel"] == 128, loaded + assert by_rank[2]["local_numel"] == 0 and by_rank[3]["local_numel"] == 0, loaded + assert not by_rank[2]["materialized"] and not by_rank[3]["materialized"], loaded + for item in loaded: + assert item["finite"], item + + +# --- Second review round (CodeRabbit + Codex on 7ea3f52) ------------------- +# +# These harden the v2 additions themselves: unique/caller-stable identities +# (finding 1), pre-commit hyperparameter validation (finding 2), stale-slot +# refresh after add_param_group (finding 3), mixed-device codebook co-location +# (finding 4), cross-rank commit agreement (finding 5), and coherent +# initialization metadata (finding 6). Most run on a single-rank Gloo group and, +# where a full distributed repro is impractical, exercise the validation via a +# direct GefenDCPState.state_dict() -> mutate -> load_state_dict() round trip +# (which is the exact code path DCP drives per rank). + + +def _single_rank_group(init_file): + dist.init_process_group( + "gloo", init_method="file://{}".format(init_file), rank=0, world_size=1 + ) + + +def _live_state_signature(optimizer, param): + state = optimizer.state.get(param, {}) + return { + "period": int(state["automatic_period"]) if "automatic_period" in state else None, + "step": int(state["step"]) if "step" in state else None, + "vmean_step": int(state["vmean_step"]) if "vmean_step" in state else None, + "m_magnitude_sum": ( + float(state["m_magnitude"].double().sum()) + if "m_magnitude" in state + else None + ), + } + + +def _group_hyper_signature(optimizer): + return [ + { + key: float( + group[key].item() + if torch.is_tensor(group[key]) + else group[key] + ) + for key in ("lr", "beta1", "beta2", "eps", "weight_decay") + } + for group in optimizer.param_groups + ] + + +# --- Finding 1: caller-stable, UNIQUE parameter identities ----------------- + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +def test_dcp_rejects_unnamed_synthesized_identities(tmp_path): + _single_rank_group(tmp_path / "unnamed-init") + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + # Bare (unnamed) parameters: Gefen synthesizes positional names + # (group_0_param_0/1) that encode registration order, not identity, so + # two reorderings would cross-assign momentum. Construction must refuse. + first = _shaped_dparam((128, 64), mesh) + second = _shaped_dparam((128, 64), mesh) + optimizer = Gefen( + [first, second], lr=1e-3, fused=False, factored_v_2d=False + ) + with pytest.raises(RuntimeError, match="synthesized"): + GefenDCPState(optimizer) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +def test_dcp_rejects_duplicate_identities(tmp_path): + _single_rank_group(tmp_path / "dup-init") + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + first = _shaped_dparam((128, 64), mesh) + second = _shaped_dparam((128, 64), mesh) + optimizer = Gefen( + [("alpha", first), ("beta", second)], + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + # Force two identical (name, group, shape) identities, as a corrupted or + # adversarial caller could. Same-shaped same-group same-named slots cannot + # be disambiguated on a resharded load, so construction must reject them. + optimizer.param_groups[0]["param_names"] = ["dup", "dup"] + with pytest.raises(RuntimeError, match="unique parameter identities"): + GefenDCPState(optimizer) + finally: + dist.destroy_process_group() + + +# --- Finding 2: validate group hyperparameters BEFORE committing ----------- + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +@pytest.mark.parametrize( + "key,bad_value,match", + [ + ("lr", float("nan"), "not\\s+finite"), + ("beta1", 1.0, "invalid beta1"), + ("eps", -1e-8, "invalid eps"), + ("weight_decay", -0.1, "invalid weight_decay"), + ], +) +def test_dcp_rejects_corrupt_hyper_fail_atomic(tmp_path, key, bad_value, match): + _single_rank_group(tmp_path / "corrupt-hyper-init") + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + param = _shaped_dparam(_SHAPE, mesh) + optimizer = Gefen( + [("weight", param)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2e-8, + weight_decay=0.03, + fused=False, + factored_v_2d=False, + deterministic=True, + ) + for step in range(_SAVE_STEPS): + param.grad = distribute_tensor(_shaped_grad(_SHAPE, step), mesh, [Shard(0)]) + optimizer.step() + + saved = GefenDCPState(optimizer).state_dict() + # Corrupt one committed-in hyper. It is finite-but-invalid or NaN, and it + # would silently poison the next update if committed. + saved["param_group_hypers"][0][key] = bad_value + + state_before = _live_state_signature(optimizer, param) + hypers_before = _group_hyper_signature(optimizer) + with pytest.raises(ValueError, match=match): + GefenDCPState(optimizer).load_state_dict(saved) + # Fail-atomic: neither optimizer state nor param-group hypers moved. + assert _live_state_signature(optimizer, param) == state_before + assert _group_hyper_signature(optimizer) == hypers_before + finally: + dist.destroy_process_group() + + +# --- Finding 3: refresh cached _slots before save/load --------------------- + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +def test_dcp_retained_wrapper_saves_added_param_group(tmp_path): + _single_rank_group(tmp_path / "addgroup-init") + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + shape = (128, 64) + first = _shaped_dparam(shape, mesh) + optimizer = Gefen( + [("alpha", first)], lr=1e-3, fused=False, factored_v_2d=False + ) + # Retain the wrapper (as a training loop would), THEN grow the optimizer. + wrapper = GefenDCPState(optimizer) + second = _shaped_dparam(shape, mesh) + optimizer.add_param_group({"params": [("beta", second)]}) + for step in range(_SAVE_STEPS): + for param in (first, second): + param.grad = distribute_tensor( + _shaped_grad(shape, step), mesh, [Shard(0)] + ) + optimizer.step() + + checkpoint_dir = str(tmp_path / "addgroup-ckpt") + # Saving through the RETAINED wrapper must include the added parameter, + # not the construction-time single-slot snapshot. + saved = wrapper.state_dict() + assert int(saved["slot_count"]) == 2, saved["slot_count"] + torch.distributed.checkpoint.save( + {"optimizer": wrapper}, + storage_writer=FileSystemWriter(checkpoint_dir), + ) + + # The resumed optimizer must mirror the saved group topology (alpha in + # group 0, beta added as group 1) for the identity/group guard to accept. + first2 = _shaped_dparam(shape, mesh) + second2 = _shaped_dparam(shape, mesh) + resumed = Gefen( + [("alpha", first2)], lr=1e-3, fused=False, factored_v_2d=False + ) + resumed.add_param_group({"params": [("beta", second2)]}) + torch.distributed.checkpoint.load( + {"optimizer": GefenDCPState(resumed)}, + storage_reader=FileSystemReader(checkpoint_dir), + ) + for param in (first2, second2): + assert "automatic_period" in resumed.state[param], resumed.state[param] + assert int(resumed.state[param]["automatic_period"]) >= 1 + finally: + dist.destroy_process_group() + + +# --- Finding 4: co-locate the codebook with each slot's device ------------- + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="mixed-device codebook coverage requires CUDA", +) +def test_reblock_relocates_codebook_to_operand_device(): + # Direct unit test (no distributed needed): a codebook learned on CUDA must + # re-quantize a CPU-resident momentum shard. gefen_nearest_codebook_indices + # rejects a codebook whose device differs from the operand, so without the + # per-slot co-location _reblock would abort on mixed CPU/CUDA shards. + import gefen.dcp as dcp_mod + + torch.manual_seed(0) + period = 4 + cuda_momentum = torch.randn(256, device="cuda") + codebook = dcp_mod._learn_codebook( + [("w", cuda_momentum.reshape(-1), period)], torch.device("cuda") + ) + assert codebook is not None and codebook.is_cuda + + cpu_momentum = torch.randn(256) + cpu_second = torch.rand(256) + m_codebook, m_magnitude, vmean = dcp_mod._reblock( + codebook, cpu_momentum, cpu_second, period + ) + assert m_codebook.device.type == "cpu" + assert m_magnitude.device.type == "cpu" and vmean.device.type == "cpu" + assert bool(torch.isfinite(m_magnitude).all()) + assert bool(torch.isfinite(vmean).all()) + # The original CUDA codebook must be untouched (co-location returns a copy). + assert codebook.is_cuda + + +# --- Finding 5: cross-rank commit agreement -------------------------------- + + +def _fail_sync_worker(rank, world, port, checkpoint_dir, mode, result_queue): + try: + os.environ.update( + MASTER_ADDR="127.0.0.1", + MASTER_PORT=port, + RANK=str(rank), + WORLD_SIZE=str(world), + ) + dist.init_process_group( + "gloo", rank=rank, world_size=world, timeout=timedelta(seconds=120) + ) + mesh = init_device_mesh("cpu", (world,), mesh_dim_names=("dp",)) + device = torch.device("cpu") + parameter, optimizer = _build(mesh, device) + + if mode == "save": + for step in range(_SAVE_STEPS): + _step(parameter, optimizer, mesh, step, device) + torch.distributed.checkpoint.save( + {"optimizer": GefenDCPState(optimizer)}, + storage_writer=FileSystemWriter(checkpoint_dir), + ) + result_queue.put({"rank": rank, "saved": True}) + return + + # Give the live optimizer real state, then inject a one-rank re-block + # failure. The cross-rank success sync must make EVERY rank raise and + # leave EVERY rank's live state untouched (fail-atomic, no partial + # restore) rather than committing on the ranks that did not fail. + _step(parameter, optimizer, mesh, 0, device) + before = _live_state_signature(optimizer, parameter) + + import gefen.dcp as dcp_mod + + if rank == 0: + def _boom(*args, **kwargs): + raise RuntimeError("injected one-rank re-block failure") + + dcp_mod._reblock = _boom + + raised = False + try: + torch.distributed.checkpoint.load( + {"optimizer": GefenDCPState(optimizer)}, + storage_reader=FileSystemReader(checkpoint_dir), + ) + except Exception: + raised = True + after = _live_state_signature(optimizer, parameter) + result_queue.put( + { + "rank": rank, + "raised": raised, + "state_unchanged": after == before, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_fail_sync(world, checkpoint_dir, mode, timeout=240): + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_fail_sync_worker, + args=(rank, world, port, checkpoint_dir, mode, result_queue), + ) + for rank in range(world) + ] + for process in processes: + process.start() + results = [] + try: + for _ in processes: + results.append(result_queue.get(timeout=timeout)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + assert len(results) == world, (results, [p.exitcode for p in processes]) + assert all(p.exitcode == 0 for p in processes), [p.exitcode for p in processes] + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +def test_dcp_one_rank_failure_aborts_every_rank(tmp_path): + checkpoint_dir = str(tmp_path / "gefen-dcp-fail-sync") + saved = _run_fail_sync(2, checkpoint_dir, "save") + assert all(item.get("saved") for item in saved), saved + loaded = _run_fail_sync(2, checkpoint_dir, "load") + assert all("fatal_error" not in item for item in loaded), loaded + # A failure injected only on rank 0 must abort BOTH ranks and commit on + # neither: every rank raised and every rank's live state is unchanged. + for item in loaded: + assert item["raised"], item + assert item["state_unchanged"], item + + +# --- Finding 6: reject inconsistent initialization metadata/counters ------- + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +def test_dcp_rejects_initialized_slot_with_zero_step(tmp_path): + _single_rank_group(tmp_path / "coherent-step-init") + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + param = _shaped_dparam(_SHAPE, mesh) + optimizer = Gefen( + [("weight", param)], + lr=1e-3, + fused=False, + factored_v_2d=False, + deterministic=True, + ) + for step in range(_SAVE_STEPS): + param.grad = distribute_tensor(_shaped_grad(_SHAPE, step), mesh, [Shard(0)]) + optimizer.step() + + saved = GefenDCPState(optimizer).state_dict() + # An initialized slot whose age counter collapsed to zero is incoherent + # (native requires initialized momentum to carry a positive step). + saved["slot_00000000.step"] = torch.tensor(0, dtype=torch.int64) + + before = _live_state_signature(optimizer, param) + with pytest.raises(ValueError, match="initialized momentum requires step"): + GefenDCPState(optimizer).load_state_dict(saved) + assert _live_state_signature(optimizer, param) == before + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DCP resharding coverage requires Gloo", +) +def test_dcp_rejects_uninitialized_slot_with_nonzero_counters(tmp_path): + _single_rank_group(tmp_path / "coherent-uninit-init") + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + param = _shaped_dparam(_SHAPE, mesh) + optimizer = Gefen( + [("weight", param)], + lr=1e-3, + fused=False, + factored_v_2d=False, + deterministic=True, + ) + for step in range(_SAVE_STEPS): + param.grad = distribute_tensor(_shaped_grad(_SHAPE, step), mesh, [Shard(0)]) + optimizer.step() + + saved = GefenDCPState(optimizer).state_dict() + # Flip the slot to uninitialized while it still carries a nonzero step and + # dense momentum: native never materializes an uninitialized slot with + # history, so the load must reject rather than silently discard it. + saved["slot_00000000.initialized"] = torch.tensor(0, dtype=torch.uint8) + + before = _live_state_signature(optimizer, param) + with pytest.raises(ValueError, match="uninitialized"): + GefenDCPState(optimizer).load_state_dict(saved) + assert _live_state_signature(optimizer, param) == before + finally: + dist.destroy_process_group()