Module neuralops.core.snapshot_restore

Model State Snapshot & Exact Bit-Identical Restorer.

Preserves the exact structural integrity, module graph, parameter buffers, and random state so attention modules can be dynamically swapped in-memory with 100% exact bit-identical restoration and zero weight/state drift.

Classes

class ModelStateSnapshot (model: nn.Module)

Expand source code

class ModelStateSnapshot:
    """Exact bit-identical snapshot of an nn.Module."""

    def __init__(self, model: nn.Module):
        self.model = model
        # 1. Structural snapshot: save references to all attention modules and their parent dotted paths
        self.swapped_layer_paths: list[str] = []
        self.original_submodules: dict[str, nn.Module] = {}
        for name, sub in model.named_modules():
            # Match transformer attention layers
            lname = name.lower()
            if (
                lname.endswith(".attention")
                or lname.endswith(".attn")
                or lname.endswith(".self")
                or lname.endswith(".self_attn")
            ):
                self.swapped_layer_paths.append(name)
                self.original_submodules[name] = sub

        # 2. Tensor snapshot: exact deep-copied state_dict in CPU host memory to preserve VRAM
        # and prevent any in-place mutation or optimizer gradient corruption.
        self.clean_state_dict: dict[str, torch.Tensor] = {
            k: v.detach().clone().cpu() for k, v in model.state_dict().items()
        }

        # 3. Extra state: config attributes, training mode
        self.was_training = model.training

    def restore(self) -> None:
        """Restores the model structure and tensor values to exact bit-identical snapshot state."""
        # 1. Structural restore: re-attach the original submodules back into the graph
        for path, original_sub in self.original_submodules.items():
            try:
                replace_module(self.model, path, original_sub)
            except Exception:
                pass

        # 2. Tensor restore: load exact state_dict back onto parameters
        device = next(self.model.parameters()).device
        restored_state = {k: v.to(device) for k, v in self.clean_state_dict.items()}
        self.model.load_state_dict(restored_state, strict=False)

        # 3. Restore training state
        self.model.train(self.was_training)

        # Clean any compression/evaluation attributes attached dynamically
        for attr in ("compression_stats", "is_custom_attention", "_injected_attention"):
            if hasattr(self.model, attr):
                try:
                    delattr(self.model, attr)
                except Exception:
                    pass

        # Reset LoRA merge state
        for m in self.model.modules():
            if hasattr(m, "merged"):
                m.merged = False

        # Clean VRAM
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def verify_bit_identical(self) -> tuple[bool, float]:
        """Verify that current model state is bit-identical to snapshot. Returns (is_identical, max_diff)."""
        max_diff = 0.0
        current_st = self.model.state_dict()
        for k, v_orig in self.clean_state_dict.items():
            if k not in current_st:
                return False, float("inf")
            v_curr = current_st[k].detach().cpu()
            if v_orig.shape != v_curr.shape:
                return False, float("inf")
            diff = (v_orig - v_curr).abs().max().item()
            if diff > max_diff:
                max_diff = diff
        return (max_diff == 0.0), max_diff

Exact bit-identical snapshot of an nn.Module.

Methods

def restore(self) ‑> None

Expand source code

def restore(self) -> None:
    """Restores the model structure and tensor values to exact bit-identical snapshot state."""
    # 1. Structural restore: re-attach the original submodules back into the graph
    for path, original_sub in self.original_submodules.items():
        try:
            replace_module(self.model, path, original_sub)
        except Exception:
            pass

    # 2. Tensor restore: load exact state_dict back onto parameters
    device = next(self.model.parameters()).device
    restored_state = {k: v.to(device) for k, v in self.clean_state_dict.items()}
    self.model.load_state_dict(restored_state, strict=False)

    # 3. Restore training state
    self.model.train(self.was_training)

    # Clean any compression/evaluation attributes attached dynamically
    for attr in ("compression_stats", "is_custom_attention", "_injected_attention"):
        if hasattr(self.model, attr):
            try:
                delattr(self.model, attr)
            except Exception:
                pass

    # Reset LoRA merge state
    for m in self.model.modules():
        if hasattr(m, "merged"):
            m.merged = False

    # Clean VRAM
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

Restores the model structure and tensor values to exact bit-identical snapshot state.

def verify_bit_identical(self) ‑> tuple[bool, float]

Expand source code

def verify_bit_identical(self) -> tuple[bool, float]:
    """Verify that current model state is bit-identical to snapshot. Returns (is_identical, max_diff)."""
    max_diff = 0.0
    current_st = self.model.state_dict()
    for k, v_orig in self.clean_state_dict.items():
        if k not in current_st:
            return False, float("inf")
        v_curr = current_st[k].detach().cpu()
        if v_orig.shape != v_curr.shape:
            return False, float("inf")
        diff = (v_orig - v_curr).abs().max().item()
        if diff > max_diff:
            max_diff = diff
    return (max_diff == 0.0), max_diff

Verify that current model state is bit-identical to snapshot. Returns (is_identical, max_diff).