From 6eeb740a491dc3b0ae833bda67307c970de37e21 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Thu, 16 Jul 2026 12:25:30 -0700 Subject: [PATCH 1/7] Add standalone DCP resharding for FSDP2 Gefen state --- CHANGELOG.md | 2 + COMPATIBILITY.md | 2 +- README.md | 6 +- src/gefen/__init__.py | 5 + src/gefen/dcp.py | 305 +++++++++++++++++++++++++++++++++++ tests/test_dcp_resharding.py | 282 ++++++++++++++++++++++++++++++++ 6 files changed, 598 insertions(+), 4 deletions(-) create mode 100644 src/gefen/dcp.py create mode 100644 tests/test_dcp_resharding.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4f680ad..76fa1d5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ All notable changes to this project are documented here. This project adheres to Distributed checkpoint-load and step-failure hardening. All fixes are backward-compatible; no public API changes. +- Add `GefenDCPState`, a standalone DCP Stateful wrapper for plain Gefen on one-dimensional default-world FSDP2 `Shard(0)` parameters. It saves shard-addressable dense optimizer tensors, supports N-to-M DCP resharding without a rank-0 full-state gather, validates before mutation, and leaves the existing native checkpoint format unchanged. + Correctness — checkpoint loads are now atomic: - `Gefen.load_state_dict` and `GefenMuonHybrid.load_state_dict` stage the full restore on an isolated shadow, validate it, then publish through non-throwing dict swaps, so a malformed, truncated, or mismatched checkpoint — or a CUDA OOM while re-staging tensors — leaves the live optimizer untouched instead of half-loaded and corrupted (#69). The hybrid previously committed the Muon child before validating the backup half and recovered a backup rejection with a second full Muon reload that could itself raise and strand a new-Muon/old-backup optimizer; the shadow-then-publish path makes a rejection a true no-op. diff --git a/COMPATIBILITY.md b/COMPATIBILITY.md index 9dc4f94..c1aa7ca 100644 --- a/COMPATIBILITY.md +++ b/COMPATIBILITY.md @@ -103,7 +103,7 @@ 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. Plain Gefen additionally exposes `GefenDCPState`, a standalone `torch.distributed.checkpoint` Stateful wrapper for one-dimensional default-world `Shard(0)` DTensors. It writes dense momentum and second-moment DTensors without gathering the complete optimizer state onto one rank; DCP can therefore reshard a checkpoint from N ranks to M ranks. Load validates every field before mutation and projects the restored state into an exact period-one Gefen representation. Muon, Hybrid, 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..c1a3e4d 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` for DCP resharding; the native full-state path and Muon `approx` remain same-topology | | 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.** Plain Gefen can save shard-addressable optimizer state through `torch.distributed.checkpoint` by wrapping it with `GefenDCPState`; a checkpoint saved on N ranks can load on M ranks for a one-dimensional default-world `Shard(0)` mesh. The existing native full-state path and Muon `approx` remain same-topology — [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 for plain Gefen on a one-dimensional default-world `Shard(0)` mesh. 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..89c0596 --- /dev/null +++ b/src/gefen/dcp.py @@ -0,0 +1,305 @@ +"""Standalone DCP adapter for reshardable plain-Gefen FSDP2 state.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +_FORMAT_VERSION = 1 +_CODEBOOK = torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + + +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() + ) + + +@dataclass(frozen=True) +class _Slot: + 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``. The adapter is + intentionally limited to plain :class:`gefen.Gefen`, one-dimensional + default-world DTensors, and one ``Shard(0)`` placement. Native + ``state_dict`` checkpoints remain unchanged. + """ + + 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 + + slots = [] + index = 0 + for group in 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", ()) + if len(names) != len(group["params"]): + names = tuple("param_{}".format(index + offset) for offset in range(len(group["params"]))) + for name, parameter in zip(names, group["params"]): + 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" + ) + 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" + ) + slots.append(_Slot(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) + + def state_dict(self): + 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), + } + for slot in self._slots: + state = self.optimizer.state.get(slot.parameter) + initialized = bool( + state + and all( + key in state + for key in ( + "automatic_period", + "step", + "m_codebook", + "m_magnitude", + "vmean", + ) + ) + ) + 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): + required = {"format_version", "global_step", "deterministic", "slot_count"} + 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") + + 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") + staged.append((slot, initialized, step, vmean_step, momentum, second)) + + global_step = _counter(state_dict["global_step"], "global_step") + deterministic = bool( + _counter(state_dict["deterministic"], "deterministic") + ) + new_state = {} + codebook = _CODEBOOK.to( + device=_local(self._slots[0].parameter).device, + dtype=torch.float32, + ) + for slot, initialized, step, vmean_step, momentum, second in staged: + if not initialized: + new_state[slot.parameter] = {"name": slot.name} + continue + flat = momentum.reshape(-1) + indices = torch.where( + torch.signbit(flat), + torch.zeros(flat.numel(), dtype=torch.uint8, device=flat.device), + torch.full( + (flat.numel(),), 255, dtype=torch.uint8, device=flat.device + ), + ) + new_state[slot.parameter] = { + "name": slot.name, + "automatic_period": 1, + "step": step, + "m_codebook": indices.reshape(-1, 1), + "m_magnitude": flat.abs().reshape(-1, 1), + "vmean": second.reshape(-1, 1).clone(), + "vmean_step": vmean_step, + } + + self.optimizer.state.clear() + self.optimizer.state.update(new_state) + self.optimizer._gefen_global_step = global_step + self.optimizer._deterministic = deterministic + 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..973ea44 --- /dev/null +++ b/tests/test_dcp_resharding.py @@ -0,0 +1,282 @@ +"""Real DCP save/load and 2-to-4 resharding for standalone Gefen state.""" + +from __future__ import annotations + +from datetime import timedelta +import multiprocessing as mp +import os +import queue +import socket +import tempfile +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 Shard, distribute_tensor, init_device_mesh + +from gefen import Gefen, GefenDCPState + + +_SHAPE = (8, 6) + + +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 _full_values(): + numel = _SHAPE[0] * _SHAPE[1] + parameter = torch.linspace(-0.8, 0.7, numel).reshape(_SHAPE) + momentum = torch.linspace(-1.25, 1.5, numel).reshape(_SHAPE) + second = torch.arange(1, numel + 1, dtype=torch.float32).reshape(_SHAPE) / 19 + gradient = torch.arange(numel, dtype=torch.float32).reshape(_SHAPE).sin() + return parameter, momentum, second, gradient + + +def _build(mesh, full_parameter): + parameter = nn.Parameter( + distribute_tensor(full_parameter.clone(), mesh, [Shard(0)]) + ) + optimizer = Gefen( + [("weight", parameter)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2e-8, + weight_decay=0.03, + fused=False, + factored_v_2d=False, + deterministic=True, + ) + return parameter, optimizer + + +def _local_slice(full, rank, world): + return torch.chunk(full, world, dim=0)[rank] + + +def _seed_projected(optimizer, parameter, momentum, second, rank, world): + local_momentum = _local_slice(momentum, rank, world).clone() + local_second = _local_slice(second, rank, world).clone() + flat = local_momentum.reshape(-1) + optimizer._gefen_codebook = torch.linspace( + -1, 1, 256, device=flat.device + ) + optimizer._gefen_global_step = 9 + optimizer.state[parameter].update( + { + "name": "weight", + "automatic_period": 1, + "step": 7, + "m_codebook": torch.where( + torch.signbit(flat), + torch.zeros( + flat.numel(), dtype=torch.uint8, device=flat.device + ), + torch.full( + (flat.numel(),), + 255, + dtype=torch.uint8, + device=flat.device, + ), + ).reshape(-1, 1), + "m_magnitude": flat.abs().reshape(-1, 1), + "vmean": local_second.reshape(-1, 1), + "vmean_step": 6, + } + ) + + +def _snapshot(optimizer, parameter): + state = optimizer.state[parameter] + codebook = optimizer._gefen_codebook + momentum = codebook[state["m_codebook"].long()].mul(state["m_magnitude"]) + return { + "parameter": _local(parameter).detach().clone(), + "momentum": momentum.reshape(_local(parameter).shape).clone(), + "second": state["vmean"].reshape(_local(parameter).shape).clone(), + "step": state["step"], + "vmean_step": state["vmean_step"], + "global_step": optimizer._gefen_global_step, + } + + +def _worker(rank, world, port, checkpoint_dir, mode, device_type, 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) + dist.init_process_group( + "cuda:nccl,cpu:gloo" if device_type == "cuda" else "gloo", + rank=rank, + world_size=world, + timeout=timedelta(seconds=90), + ) + mesh = init_device_mesh(device_type, (world,), mesh_dim_names=("dp",)) + parameter_value, momentum, second, gradient = _full_values() + parameter_value = parameter_value.to(device_type) + momentum = momentum.to(device_type) + second = second.to(device_type) + gradient = gradient.to(device_type) + parameter, optimizer = _build(mesh, parameter_value) + wrapper = GefenDCPState(optimizer) + if mode == "save": + _seed_projected( + optimizer, parameter, momentum, second, rank, world + ) + torch.distributed.checkpoint.save( + {"optimizer": wrapper}, + storage_writer=FileSystemWriter(checkpoint_dir), + ) + result_queue.put({"rank": rank, "saved": True}) + return + + torch.distributed.checkpoint.load( + {"optimizer": wrapper}, + storage_reader=FileSystemReader(checkpoint_dir), + ) + reference_parameter, reference = _build(mesh, parameter_value) + _seed_projected( + reference, reference_parameter, momentum, second, rank, world + ) + before = _snapshot(optimizer, parameter) + expected_before = _snapshot(reference, reference_parameter) + parameter.grad = distribute_tensor(gradient.clone(), mesh, [Shard(0)]) + reference_parameter.grad = distribute_tensor( + gradient.clone(), mesh, [Shard(0)] + ) + optimizer.step() + reference.step() + after = _snapshot(optimizer, parameter) + expected_after = _snapshot(reference, reference_parameter) + result_queue.put( + { + "rank": rank, + "restored_exact": all( + torch.equal(before[key], expected_before[key]) + if torch.is_tensor(before[key]) + else before[key] == expected_before[key] + for key in before + ), + "continued_exact": all( + torch.equal(after[key], expected_after[key]) + if torch.is_tensor(after[key]) + else after[key] == expected_after[key] + for key in after + ), + "shape": tuple(_local(parameter).shape), + } + ) + 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", timeout=120): + 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, + 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_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 all("fatal_error" not in item for item in loaded), loaded + assert [item["shape"] for item in loaded] == [ + (2, 6), + (2, 6), + (2, 6), + (2, 6), + ] + assert all(item["restored_exact"] for item in loaded), loaded + assert all(item["continued_exact"] for item in loaded), loaded + + +@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): + 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 all("fatal_error" not in item for item in loaded), loaded + assert all(item["restored_exact"] for item in loaded), loaded + assert all(item["continued_exact"] for item in loaded), loaded + + +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) + with pytest.raises(RuntimeError, match="every parameter to be a DTensor"): + GefenDCPState(optimizer) + finally: + dist.destroy_process_group() From 02adaecb42bf937b560f60f7018026428e83e101 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Thu, 16 Jul 2026 13:07:04 -0700 Subject: [PATCH 2/7] DCP resharding: re-block on load, real reshard tests, fail-closed guards Address the PR review on GefenDCPState. - Load path re-blocks instead of collapsing to period one. After a DCP reshard, re-derive Gefen's compact per-rank state from the resharded dense fp32 momentum: re-run the block-variance period search on the dense second moment (a grad^2 proxy), relearn the exact codebook on the new local shard, and re-quantize into blocks (per-block max-abs magnitude, nearest codeword indices). vmean is reconstructed as the per-block mean of the dense second moment for the new blocking. This restores the ~1 byte/param footprint that period-one re-encoding destroyed (e.g. a 2048x vmean/m_magnitude blow-up). Re-blocking happens before optimizer.state.clear(), preserving the fail-atomic load. - Fail closed on capturable=True and factored_v_2d=True: capturable (and compiled) counter/seed semantics are not host-serializable, and the factored row/column second moment is not shard-addressable. - Replace the tautological period-one tests with real ones: a genuine Gefen.step learns the codebook and picks period > 1; save, reshard (2->2, 2->4, 4->2, GPU/NCCL 2->4, fully_shard), and assert the restored state stays compact and continues within quantization-noise tolerance of a native run at the target topology. Add capturable/factored guard tests. - Docs (COMPATIBILITY, README, CHANGELOG): drop the "exact period-one" overclaim; state that GefenDCPState is for resharding, that resume is a correct continuation within quantization noise (not bit-exact), steer same-topology users to the native path, and note the compact re-block and the factored_v_2d/capturable requirements. --- CHANGELOG.md | 2 +- COMPATIBILITY.md | 4 +- README.md | 6 +- src/gefen/dcp.py | 167 ++++++++++++-- tests/test_dcp_resharding.py | 412 ++++++++++++++++++++++++++--------- 5 files changed, 459 insertions(+), 132 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 76fa1d5..8c175fa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ All notable changes to this project are documented here. This project adheres to Distributed checkpoint-load and step-failure hardening. All fixes are backward-compatible; no public API changes. -- Add `GefenDCPState`, a standalone DCP Stateful wrapper for plain Gefen on one-dimensional default-world FSDP2 `Shard(0)` parameters. It saves shard-addressable dense optimizer tensors, supports N-to-M DCP resharding without a rank-0 full-state gather, validates before mutation, and leaves the existing native checkpoint format unchanged. +- 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 c1aa7ca..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. -The native format is deliberately same-topology only and currently requires one 1-D DeviceMesh spanning the default process-group world. Plain Gefen additionally exposes `GefenDCPState`, a standalone `torch.distributed.checkpoint` Stateful wrapper for one-dimensional default-world `Shard(0)` DTensors. It writes dense momentum and second-moment DTensors without gathering the complete optimizer state onto one rank; DCP can therefore reshard a checkpoint from N ranks to M ranks. Load validates every field before mutation and projects the restored state into an exact period-one Gefen representation. Muon, Hybrid, multidimensional meshes, subgroups, non-`Shard(0)` placements, and the native full-state format remain same-topology and fail closed. +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 c1a3e4d..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 can use `GefenDCPState` for DCP resharding; the native full-state path and Muon `approx` remain same-topology | +| 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 can save shard-addressable optimizer state through `torch.distributed.checkpoint` by wrapping it with `GefenDCPState`; a checkpoint saved on N ranks can load on M ranks for a one-dimensional default-world `Shard(0)` mesh. The existing native full-state path and Muon `approx` remain same-topology — [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 resharding is explicit.** Use `dcp.save({"optimizer": GefenDCPState(optimizer)}, ...)` and the matching `dcp.load` call for plain Gefen on a one-dimensional default-world `Shard(0)` mesh. 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). +- **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/dcp.py b/src/gefen/dcp.py index 89c0596..98da247 100644 --- a/src/gefen/dcp.py +++ b/src/gefen/dcp.py @@ -8,7 +8,6 @@ _FORMAT_VERSION = 1 -_CODEBOOK = torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) def _local(value): @@ -82,6 +81,90 @@ def _dense_second_moment(parameter, state) -> torch.Tensor: ) +def _choose_period(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). + """ + from gefen.partitioning import find_period_by_block_variance + + 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) + 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 @@ -93,10 +176,18 @@ 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``. The adapter is - intentionally limited to plain :class:`gefen.Gefen`, one-dimensional - default-world DTensors, and one ``Shard(0)`` placement. Native - ``state_dict`` checkpoints remain unchanged. + :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): @@ -114,6 +205,20 @@ def _validate_layout(self): ) 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 for group in self.optimizer.param_groups: @@ -266,30 +371,47 @@ def load_state_dict(self, state_dict): deterministic = bool( _counter(state_dict["deterministic"], "deterministic") ) - new_state = {} - codebook = _CODEBOOK.to( - device=_local(self._slots[0].parameter).device, - dtype=torch.float32, + + # 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. + device = _local(self._slots[0].parameter).device + periods = {} + for slot, initialized, step, vmean_step, momentum, second in staged: + if initialized: + periods[slot.index] = _choose_period( + slot.name, slot.parameter, second + ) + codebook = _learn_codebook( + [ + (slot.name, momentum.reshape(-1), periods[slot.index]) + for slot, initialized, _, _, momentum, _ in staged + if initialized + ], + device, ) + + new_state = {} for slot, initialized, step, vmean_step, momentum, second in staged: if not initialized: new_state[slot.parameter] = {"name": slot.name} continue - flat = momentum.reshape(-1) - indices = torch.where( - torch.signbit(flat), - torch.zeros(flat.numel(), dtype=torch.uint8, device=flat.device), - torch.full( - (flat.numel(),), 255, dtype=torch.uint8, device=flat.device - ), + period = periods[slot.index] + m_codebook, m_magnitude, vmean = _reblock( + codebook, momentum, second, period ) new_state[slot.parameter] = { "name": slot.name, - "automatic_period": 1, + "automatic_period": period, "step": step, - "m_codebook": indices.reshape(-1, 1), - "m_magnitude": flat.abs().reshape(-1, 1), - "vmean": second.reshape(-1, 1).clone(), + "m_codebook": m_codebook, + "m_magnitude": m_magnitude, + "vmean": vmean, "vmean_step": vmean_step, } @@ -297,6 +419,11 @@ def load_state_dict(self, state_dict): self.optimizer.state.update(new_state) self.optimizer._gefen_global_step = global_step self.optimizer._deterministic = deterministic + # 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() diff --git a/tests/test_dcp_resharding.py b/tests/test_dcp_resharding.py index 973ea44..4ccabba 100644 --- a/tests/test_dcp_resharding.py +++ b/tests/test_dcp_resharding.py @@ -1,4 +1,14 @@ -"""Real DCP save/load and 2-to-4 resharding for standalone Gefen state.""" +"""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 @@ -7,7 +17,6 @@ import os import queue import socket -import tempfile import traceback import pytest @@ -20,7 +29,11 @@ from gefen import Gefen, GefenDCPState -_SHAPE = (8, 6) +# 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(): @@ -34,18 +47,22 @@ def _local(value): return value.wait() if hasattr(value, "wait") else value -def _full_values(): +def _global_param(): numel = _SHAPE[0] * _SHAPE[1] - parameter = torch.linspace(-0.8, 0.7, numel).reshape(_SHAPE) - momentum = torch.linspace(-1.25, 1.5, numel).reshape(_SHAPE) - second = torch.arange(1, numel + 1, dtype=torch.float32).reshape(_SHAPE) / 19 - gradient = torch.arange(numel, dtype=torch.float32).reshape(_SHAPE).sin() - return parameter, momentum, second, gradient + 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, full_parameter): +def _build(mesh, device): parameter = nn.Parameter( - distribute_tensor(full_parameter.clone(), mesh, [Shard(0)]) + distribute_tensor(_global_param().to(device), mesh, [Shard(0)]) ) optimizer = Gefen( [("weight", parameter)], @@ -60,53 +77,30 @@ def _build(mesh, full_parameter): return parameter, optimizer -def _local_slice(full, rank, world): - return torch.chunk(full, world, dim=0)[rank] - - -def _seed_projected(optimizer, parameter, momentum, second, rank, world): - local_momentum = _local_slice(momentum, rank, world).clone() - local_second = _local_slice(second, rank, world).clone() - flat = local_momentum.reshape(-1) - optimizer._gefen_codebook = torch.linspace( - -1, 1, 256, device=flat.device - ) - optimizer._gefen_global_step = 9 - optimizer.state[parameter].update( - { - "name": "weight", - "automatic_period": 1, - "step": 7, - "m_codebook": torch.where( - torch.signbit(flat), - torch.zeros( - flat.numel(), dtype=torch.uint8, device=flat.device - ), - torch.full( - (flat.numel(),), - 255, - dtype=torch.uint8, - device=flat.device, - ), - ).reshape(-1, 1), - "m_magnitude": flat.abs().reshape(-1, 1), - "vmean": local_second.reshape(-1, 1), - "vmean_step": 6, - } +def _step(parameter, optimizer, mesh, step, device): + parameter.grad = distribute_tensor( + _global_grad(step, device), mesh, [Shard(0)] ) + optimizer.step() -def _snapshot(optimizer, parameter): +def _period_info(optimizer, parameter): state = optimizer.state[parameter] - codebook = optimizer._gefen_codebook - momentum = codebook[state["m_codebook"].long()].mul(state["m_magnitude"]) + period = int(state["automatic_period"]) + local_numel = int(_local(parameter).numel()) return { - "parameter": _local(parameter).detach().clone(), - "momentum": momentum.reshape(_local(parameter).shape).clone(), - "second": state["vmean"].reshape(_local(parameter).shape).clone(), - "step": state["step"], - "vmean_step": state["vmean_step"], - "global_step": optimizer._gefen_global_step, + "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 + ), } @@ -119,66 +113,63 @@ def _worker(rank, world, port, checkpoint_dir, mode, device_type, result_queue): WORLD_SIZE=str(world), ) if device_type == "cuda": - torch.cuda.set_device(rank) + 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=90), + timeout=timedelta(seconds=120), ) mesh = init_device_mesh(device_type, (world,), mesh_dim_names=("dp",)) - parameter_value, momentum, second, gradient = _full_values() - parameter_value = parameter_value.to(device_type) - momentum = momentum.to(device_type) - second = second.to(device_type) - gradient = gradient.to(device_type) - parameter, optimizer = _build(mesh, parameter_value) - wrapper = GefenDCPState(optimizer) + parameter, optimizer = _build(mesh, device) + if mode == "save": - _seed_projected( - optimizer, parameter, momentum, second, rank, world - ) + for step in range(_SAVE_STEPS): + _step(parameter, optimizer, mesh, step, device) + saved = _period_info(optimizer, parameter) torch.distributed.checkpoint.save( - {"optimizer": wrapper}, + {"optimizer": GefenDCPState(optimizer)}, storage_writer=FileSystemWriter(checkpoint_dir), ) - result_queue.put({"rank": rank, "saved": True}) + 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": wrapper}, + {"optimizer": GefenDCPState(optimizer)}, storage_reader=FileSystemReader(checkpoint_dir), ) - reference_parameter, reference = _build(mesh, parameter_value) - _seed_projected( - reference, reference_parameter, momentum, second, rank, world - ) - before = _snapshot(optimizer, parameter) - expected_before = _snapshot(reference, reference_parameter) - parameter.grad = distribute_tensor(gradient.clone(), mesh, [Shard(0)]) - reference_parameter.grad = distribute_tensor( - gradient.clone(), mesh, [Shard(0)] - ) - optimizer.step() - reference.step() - after = _snapshot(optimizer, parameter) - expected_after = _snapshot(reference, reference_parameter) + restored = _period_info(optimizer, parameter) + + reference_param, reference = _build(mesh, device) + 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, - "restored_exact": all( - torch.equal(before[key], expected_before[key]) - if torch.is_tensor(before[key]) - else before[key] == expected_before[key] - for key in before - ), - "continued_exact": all( - torch.equal(after[key], expected_after[key]) - if torch.is_tensor(after[key]) - else after[key] == expected_after[key] - for key in after - ), "shape": tuple(_local(parameter).shape), + "restored": restored, + "finite": finite, + "continuation_max_abs_diff": max_abs_diff, } ) except BaseException: @@ -188,7 +179,7 @@ def _worker(rank, world, port, checkpoint_dir, mode, device_type, result_queue): dist.destroy_process_group() -def _run(world, checkpoint_dir, mode, *, device_type="cpu", timeout=120): +def _run(world, checkpoint_dir, mode, *, device_type="cpu", timeout=240): context = mp.get_context("spawn") result_queue = context.Queue() port = _free_port() @@ -229,6 +220,40 @@ def _run(world, checkpoint_dir, mode, *, device_type="cpu", timeout=120): 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", @@ -238,15 +263,22 @@ def test_dcp_save_on_two_ranks_loads_and_continues_on_four(tmp_path): saved = _run(2, checkpoint_dir, "save") assert all(item.get("saved") for item in saved), saved loaded = _run(4, checkpoint_dir, "load") - assert all("fatal_error" not in item for item in loaded), loaded - assert [item["shape"] for item in loaded] == [ - (2, 6), - (2, 6), - (2, 6), - (2, 6), - ] - assert all(item["restored_exact"] for item in loaded), loaded - assert all(item["continued_exact"] for item in loaded), loaded + _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( @@ -256,13 +288,130 @@ def test_dcp_save_on_two_ranks_loads_and_continues_on_four(tmp_path): 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) + + +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 - assert all(item["restored_exact"] for item in loaded), loaded - assert all(item["continued_exact"] for item in loaded), loaded + for item in loaded: + assert item["finite"] and item["restored"]["compact"], item def test_rejects_non_dtensor_optimizer(tmp_path): @@ -275,8 +424,57 @@ def test_rejects_non_dtensor_optimizer(tmp_path): ) try: parameter = nn.Parameter(torch.ones(4, 4)) - optimizer = Gefen([("weight", parameter)], lr=1e-3, fused=False) + 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() + + +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() + + +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() From bd7cfc60abbb51907deff10959cf58b3acd4f195 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Thu, 16 Jul 2026 13:27:07 -0700 Subject: [PATCH 3/7] Cover the fused path in DCP resharding tests GefenDCPState only reads optimizer state (m_codebook / m_magnitude / vmean / automatic_period), which is layout-identical whether the step ran through the fused CUDA kernels or the decomposed path, so the adapter is fused-agnostic -- but every test built with fused=False. Thread a fused flag through the test harness (default False, so existing tests are unchanged) and add a GPU/NCCL test that trains with fused=True and asserts the fused-produced state saves, reshards, stays compact, and continues within tolerance. --- tests/test_dcp_resharding.py | 33 +++++++++++++++++++++++++++------ 1 file changed, 27 insertions(+), 6 deletions(-) diff --git a/tests/test_dcp_resharding.py b/tests/test_dcp_resharding.py index 4ccabba..92acb44 100644 --- a/tests/test_dcp_resharding.py +++ b/tests/test_dcp_resharding.py @@ -60,7 +60,7 @@ def _global_grad(step, device): return grad.to(device) -def _build(mesh, device): +def _build(mesh, device, fused=False): parameter = nn.Parameter( distribute_tensor(_global_param().to(device), mesh, [Shard(0)]) ) @@ -70,7 +70,7 @@ def _build(mesh, device): betas=(0.8, 0.97), eps=2e-8, weight_decay=0.03, - fused=False, + fused=fused, factored_v_2d=False, deterministic=True, ) @@ -104,7 +104,7 @@ def _period_info(optimizer, parameter): } -def _worker(rank, world, port, checkpoint_dir, mode, device_type, result_queue): +def _worker(rank, world, port, checkpoint_dir, mode, device_type, fused, result_queue): try: os.environ.update( MASTER_ADDR="127.0.0.1", @@ -124,7 +124,7 @@ def _worker(rank, world, port, checkpoint_dir, mode, device_type, result_queue): timeout=timedelta(seconds=120), ) mesh = init_device_mesh(device_type, (world,), mesh_dim_names=("dp",)) - parameter, optimizer = _build(mesh, device) + parameter, optimizer = _build(mesh, device, fused) if mode == "save": for step in range(_SAVE_STEPS): @@ -146,7 +146,7 @@ def _worker(rank, world, port, checkpoint_dir, mode, device_type, result_queue): ) restored = _period_info(optimizer, parameter) - reference_param, reference = _build(mesh, device) + 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 @@ -179,7 +179,7 @@ def _worker(rank, world, port, checkpoint_dir, mode, device_type, result_queue): dist.destroy_process_group() -def _run(world, checkpoint_dir, mode, *, device_type="cpu", timeout=240): +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() @@ -193,6 +193,7 @@ def _run(world, checkpoint_dir, mode, *, device_type="cpu", timeout=240): checkpoint_dir, mode, device_type, + fused, result_queue, ), ) @@ -297,6 +298,26 @@ def test_dcp_save_on_two_gpus_loads_and_continues_on_four(tmp_path): _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( From 7ea3f52e0222d17e659cd37d747057cda87d332d Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Thu, 16 Jul 2026 13:59:09 -0700 Subject: [PATCH 4/7] Harden GefenDCPState resharding load against Codex review findings Address six correctness findings on the DCP resharding adapter and bump the on-disk schema to v2 (fail-closed on version/identity/topology mismatch): - Serialize and restore per-group hyperparameters (lr/betas/eps/weight_decay) so a resume after an LR-schedule or hyperparameter change adopts the checkpoint's values instead of silently keeping the freshly-constructed optimizer's, matching the native full-state path. - Persist stable per-slot parameter identities (name + group + global shape) and validate them before accepting a load, so a target that registered same-shaped parameters in a different order is rejected instead of cross-assigning each parameter the other's momentum/second-moment. - Reject a deterministic-policy mismatch before mutating live state, matching native Gefen.load_state_dict, instead of silently overwriting the flag. - Treat an initialized slot that reshards to an empty local shard (N->M where dim-0 < the target world) as unmaterialized, name-only state and leave it out of codebook learning, so it no longer feeds a None per-rank codebook into the re-quantize and dereferences codebook.device. - Route the restore period search through the same force_1d_period_one / force_2d_period_one / period_one_substrings gates as _resolve_automatic_period, so an explicit period-one config is honored on resume instead of restoring a frozen period>1 codebook. - Reject a slot carrying only some required optimizer fields instead of treating it as uninitialized and zeroing the momentum/second-moment history. Extend the resharding tests with per-fix coverage (changed-lr restore, swapped-identity rejection, deterministic-mismatch rejection, force-period-one restore, and an N->M reshard to empty target shards) and correct the 0.4.1 changelog summary, which claimed no public API changes despite adding GefenDCPState. --- CHANGELOG.md | 2 +- src/gefen/dcp.py | 210 +++++++++++++++++--- tests/test_dcp_resharding.py | 359 +++++++++++++++++++++++++++++++++++ 3 files changed, 544 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8c175fa..42ffc5e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ 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. diff --git a/src/gefen/dcp.py b/src/gefen/dcp.py index 98da247..6f31c43 100644 --- a/src/gefen/dcp.py +++ b/src/gefen/dcp.py @@ -7,7 +7,14 @@ import torch -_FORMAT_VERSION = 1 +# 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): @@ -81,7 +88,7 @@ def _dense_second_moment(parameter, state) -> torch.Tensor: ) -def _choose_period(name: str, parameter, second: torch.Tensor) -> int: +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 @@ -89,9 +96,26 @@ def _choose_period(name: str, parameter, second: torch.Tensor) -> int: 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 @@ -168,6 +192,7 @@ def _reblock(codebook, momentum: torch.Tensor, second: torch.Tensor, period: int @dataclass(frozen=True) class _Slot: index: int + group_index: int name: str parameter: torch.Tensor @@ -221,7 +246,7 @@ def _validate_layout(self): slots = [] index = 0 - for group in self.optimizer.param_groups: + 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", ()) @@ -246,7 +271,7 @@ def _validate_layout(self): raise RuntimeError( "GefenDCPState supports only one-dimensional Shard(0) parameters" ) - slots.append(_Slot(index, str(name), parameter)) + slots.append(_Slot(index, group_index, str(name), parameter)) index += 1 if not slots: raise RuntimeError("GefenDCPState requires at least one parameter") @@ -256,6 +281,80 @@ def _validate_layout(self): 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 + + 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): result = { "format_version": torch.tensor(_FORMAT_VERSION, dtype=torch.int64), @@ -267,22 +366,15 @@ def state_dict(self): 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 = bool( - state - and all( - key in state - for key in ( - "automatic_period", - "step", - "m_codebook", - "m_magnitude", - "vmean", - ) - ) - ) + initialized = self._slot_initialized(slot) result[self._key(slot, "initialized")] = torch.tensor( int(initialized), dtype=torch.uint8 ) @@ -316,7 +408,15 @@ def state_dict(self): return result def load_state_dict(self, state_dict): - required = {"format_version", "global_step", "deterministic", "slot_count"} + 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") @@ -327,6 +427,45 @@ def load_state_dict(self, state_dict): 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" + ) + + # 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) + ) + ) + staged = [] for slot in self._slots: keys = { @@ -368,9 +507,6 @@ def load_state_dict(self, state_dict): staged.append((slot, initialized, step, vmean_step, momentum, second)) global_step = _counter(state_dict["global_step"], "global_step") - deterministic = bool( - _counter(state_dict["deterministic"], "deterministic") - ) # Re-block the resharded dense momentum back into Gefen's compact # per-block representation instead of collapsing it to period one. Going @@ -380,25 +516,36 @@ def load_state_dict(self, state_dict): # 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 initialized: + if _materialized(initialized, momentum): periods[slot.index] = _choose_period( - slot.name, slot.parameter, second + 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 initialized + if _materialized(initialized, momentum) ], device, ) new_state = {} for slot, initialized, step, vmean_step, momentum, second in staged: - if not initialized: + if not _materialized(initialized, momentum): new_state[slot.parameter] = {"name": slot.name} continue period = periods[slot.index] @@ -417,8 +564,19 @@ def load_state_dict(self, state_dict): 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, saved_hypers): + for key in _GROUP_HYPER_KEYS: + value = float(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 - self.optimizer._deterministic = deterministic # 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). diff --git a/tests/test_dcp_resharding.py b/tests/test_dcp_resharding.py index 92acb44..83c2f4b 100644 --- a/tests/test_dcp_resharding.py +++ b/tests/test_dcp_resharding.py @@ -499,3 +499,362 @@ def test_rejects_factored_v_2d(tmp_path): GefenDCPState(optimizer) finally: dist.destroy_process_group() + + +# --- 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 From ffbe7ad7c83178d3201943c069531f79186070c1 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Thu, 16 Jul 2026 14:31:57 -0700 Subject: [PATCH 5/7] Harden GefenDCPState resharding against second-round review findings Address the CodeRabbit + Codex critiques of the v2 identity/hyper/init additions: 1. Require caller-stable, unique parameter identities. Reject at construction when param_names is missing/short, matches Gefen's synthesized positional pattern (group_N_param_M / param_N), or yields a duplicate (name, group, shape) triple, instead of silently cross-assigning momentum on a reordered load. 2. Parse and range-validate every group hyperparameter (finite, lr>=0, 0<=betas<1, eps>0, wd>=0) BEFORE any mutation, so a missing/non-numeric key or a poisoning value (lr=NaN, beta1=1, negative eps/wd) is rejected fail-atomically rather than committed. 3. Recompute the slot layout (re-run the fail-closed layout gate) at save and load time so a retained wrapper reflects add_param_group instead of a stale construction-time snapshot. 4. Co-locate the learned codebook with each slot's operand device inside the re-block so mixed CPU/CUDA shards do not abort on a device mismatch. 5. Synchronize staging/validation/re-block success across the process group before committing (all_reduce of a local-ok flag via _synchronize_step_failure) so either every rank commits or every rank raises; gloo-safe for CPU-resident state. 6. Reject incoherent initialization metadata: an initialized slot must carry step>=1 and vmean_step>=1; an uninitialized slot must carry no counters and no dense history that would be silently discarded. Serialized schema is unchanged, so _FORMAT_VERSION stays at 2. Extends tests/test_dcp_resharding.py with per-finding coverage (unnamed/duplicate rejection, fail-atomic corrupt-hyper load, add_param_group save/load, mixed-device codebook, one-rank-failure abort, incoherent-counter rejection). --- src/gefen/dcp.py | 408 +++++++++++++++++++++++++-------- tests/test_dcp_resharding.py | 425 +++++++++++++++++++++++++++++++++++ 2 files changed, 739 insertions(+), 94 deletions(-) diff --git a/src/gefen/dcp.py b/src/gefen/dcp.py index 6f31c43..89e0859 100644 --- a/src/gefen/dcp.py +++ b/src/gefen/dcp.py @@ -3,10 +3,21 @@ 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 @@ -180,6 +191,12 @@ def _reblock(codebook, momentum: torch.Tensor, second: torch.Tensor, period: int ) 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 @@ -246,13 +263,36 @@ def _validate_layout(self): 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", ()) - if len(names) != len(group["params"]): - names = tuple("param_{}".format(index + offset) for offset in range(len(group["params"]))) + 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" @@ -271,6 +311,20 @@ def _validate_layout(self): 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: @@ -328,6 +382,82 @@ def _group_hypers(self): 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. @@ -356,6 +486,10 @@ def _slot_initialized(self, slot: _Slot) -> bool: 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( @@ -408,6 +542,10 @@ def state_dict(self): 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", @@ -447,6 +585,14 @@ def load_state_dict(self, state_dict): 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 @@ -466,101 +612,175 @@ def load_state_dict(self, state_dict): ) ) - 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 + # 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", ) - ) - 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 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") ) - 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") - 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 + step = _counter(state_dict[keys["step"]], "step") + vmean_step = _counter( + state_dict[keys["vmean_step"]], "vmean_step" ) - codebook = _learn_codebook( - [ - (slot.name, momentum.reshape(-1), periods[slot.index]) - for slot, initialized, _, _, momentum, _ in staged - if _materialized(initialized, momentum) - ], - device, - ) + 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 - 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 + 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[slot.parameter] = { - "name": slot.name, - "automatic_period": period, - "step": step, - "m_codebook": m_codebook, - "m_magnitude": m_magnitude, - "vmean": vmean, - "vmean_step": vmean_step, - } + new_state, global_step, codebook = prepared self.optimizer.state.clear() self.optimizer.state.update(new_state) @@ -568,9 +788,9 @@ def _materialized(initialized, momentum): # 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, saved_hypers): + for group, entry in zip(self.optimizer.param_groups, staged_hypers): for key in _GROUP_HYPER_KEYS: - value = float(entry[key]) + value = entry[key] current = group.get(key) if torch.is_tensor(current): current.fill_(value) diff --git a/tests/test_dcp_resharding.py b/tests/test_dcp_resharding.py index 83c2f4b..1b5313a 100644 --- a/tests/test_dcp_resharding.py +++ b/tests/test_dcp_resharding.py @@ -858,3 +858,428 @@ def test_dcp_reshard_to_empty_target_shard_does_not_crash(tmp_path): 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() From ea89e85a2f0e0211ab91bcb63f3b328d0c6f87d6 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Thu, 16 Jul 2026 14:45:17 -0700 Subject: [PATCH 6/7] Test the advertised fail-closed layout scope The PR states that multidimensional meshes, subgroups, and non-Shard(0) placements fail closed, but only the non-DTensor / capturable / factored_v_2d rejections had tests. Add reject tests for a non-Shard(0) (Replicate) placement, a multidimensional DeviceMesh, and a group carrying a Muon sharded_mode. --- tests/test_dcp_resharding.py | 64 +++++++++++++++++++++++++++++++++++- 1 file changed, 63 insertions(+), 1 deletion(-) diff --git a/tests/test_dcp_resharding.py b/tests/test_dcp_resharding.py index 1b5313a..34a0361 100644 --- a/tests/test_dcp_resharding.py +++ b/tests/test_dcp_resharding.py @@ -24,7 +24,12 @@ import torch.distributed as dist from torch import nn from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter -from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh +from torch.distributed.tensor import ( + Replicate, + Shard, + distribute_tensor, + init_device_mesh, +) from gefen import Gefen, GefenDCPState @@ -501,6 +506,63 @@ def test_rejects_factored_v_2d(tmp_path): dist.destroy_process_group() +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() + + +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() + + +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() + + # --- Review-finding coverage (Codex PR #83) ------------------------------- # # The tests below cover the fail-closed / semantics fixes. Findings 1, 2, 3 and 5 From bd47e80b6a12009a662880f85ea84318e13ce017 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Thu, 16 Jul 2026 15:05:53 -0700 Subject: [PATCH 7/7] Reject reordered full-world meshes; guard Gloo-only reject tests Codex: `_validate_layout` checked mesh size but not rank order, so a reordered full-world 1-D mesh (e.g. ranks [1, 0]) passed while permuting shard->rank ownership -- misaligning the replicated initialized/counter metadata with the sharded momentum. Now reject any mesh whose ranks are not the canonical 0..N-1 order. Test: test_rejects_reordered_full_world_mesh (2-rank). CodeRabbit: the layout-reject tests called init_process_group("gloo") unconditionally; add the same `skipif(not gloo available)` guard the reshard tests use so unsupported builds skip instead of failing. --- src/gefen/dcp.py | 10 +++++ tests/test_dcp_resharding.py | 71 ++++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+) diff --git a/src/gefen/dcp.py b/src/gefen/dcp.py index 89e0859..d3c63b0 100644 --- a/src/gefen/dcp.py +++ b/src/gefen/dcp.py @@ -302,6 +302,16 @@ def _validate_layout(self): 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" diff --git a/tests/test_dcp_resharding.py b/tests/test_dcp_resharding.py index 34a0361..09a1a57 100644 --- a/tests/test_dcp_resharding.py +++ b/tests/test_dcp_resharding.py @@ -440,6 +440,7 @@ def test_dcp_fully_shard_real_model_stays_compact(tmp_path): 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( @@ -459,6 +460,7 @@ def test_rejects_non_dtensor_optimizer(tmp_path): 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( @@ -483,6 +485,7 @@ def test_rejects_capturable(tmp_path): 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( @@ -506,6 +509,7 @@ def test_rejects_factored_v_2d(tmp_path): 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" @@ -524,6 +528,7 @@ def test_rejects_non_shard0_placement(tmp_path): 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" @@ -544,6 +549,7 @@ def test_rejects_multidimensional_mesh(tmp_path): 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" @@ -563,6 +569,71 @@ def test_rejects_muon_sharded_mode(tmp_path): 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