Add decision core: autojev/optim.py

autojev/optim.py

@@ -0,0 +1,102 @@
+"""AdamW with exact FP32 master weights and moments kept in host memory."""
+
+from collections.abc import Iterable
+from typing import TypedDict, cast
+
+import torch
+
+
+class ParameterGroup(TypedDict):
+    params: list[torch.Tensor]
+    lr: float
+
+
+class OffloadState(TypedDict):
+    format_version: int
+    parameter_names: list[str]
+    masters: list[torch.Tensor]
+    optimizer: dict[str, object]
+
+
+class CPUOffloadAdamW:
+    def __init__(
+        self, named_parameters: Iterable[tuple[str, torch.nn.Parameter]], lr: float,
+        weight_decay: float = 0.01,
+    ) -> None:
+        parameters = [(name, parameter) for name, parameter in named_parameters if parameter.requires_grad]
+        if not parameters:
+            raise ValueError("The optimizer needs trainable parameters.")
+        self.names = [name for name, _ in parameters]
+        if len(set(self.names)) != len(self.names):
+            raise ValueError("Parameter names must be unique.")
+        self.parameters = [parameter for _, parameter in parameters]
+        self.masters = [torch.nn.Parameter(parameter.detach().to(device="cpu", dtype=torch.float32, copy=True))
+                        for parameter in self.parameters]
+        self.gradients = [torch.empty_like(master) for master in self.masters]
+        self.buffers = {
+            dtype: torch.empty(max(parameter.numel() for parameter in self.parameters if parameter.dtype == dtype),
+                               dtype=dtype, device="cpu", pin_memory=any(parameter.is_cuda for parameter in self.parameters))
+            for dtype in {parameter.dtype for parameter in self.parameters}
+        }
+        self.optimizer = torch.optim.AdamW(self.masters, lr=lr, weight_decay=weight_decay, fused=True)
+
+    @property
+    def param_groups(self) -> list[ParameterGroup]:
+        return cast(list[ParameterGroup], self.optimizer.param_groups)
+
+    def zero_grad(self, set_to_none: bool = True) -> None:
+        self.optimizer.zero_grad(set_to_none=set_to_none)
+        for parameter in self.parameters:
+            if set_to_none:
+                parameter.grad = None
+            elif parameter.grad is not None:
+                parameter.grad.zero_()
+
+    @torch.no_grad()
+    def step(self) -> None:
+        for parameter, master, gradient in zip(self.parameters, self.masters, self.gradients):
+            if parameter.grad is None:
+                master.grad = None
+                continue
+            buffer = self.buffers[parameter.dtype][:parameter.numel()].view_as(parameter)
+            buffer.copy_(parameter.grad, non_blocking=True)
+            if parameter.is_cuda:
+                torch.cuda.synchronize(parameter.device)
+            gradient.copy_(buffer)
+            master.grad = gradient
+            parameter.grad = None
+        self.optimizer.step()
+        for parameter, master in zip(self.parameters, self.masters):
+            if master.grad is not None:
+                buffer = self.buffers[parameter.dtype][:parameter.numel()].view_as(parameter)
+                buffer.copy_(master)
+                parameter.copy_(buffer, non_blocking=True)
+                if parameter.is_cuda:
+                    torch.cuda.synchronize(parameter.device)
+                master.grad = None
+
+    def state_dict(self) -> OffloadState:
+        """Return CPU tensor references for torch.save; never clone the large states."""
+        return {"format_version": 1, "parameter_names": list(self.names),
+                "masters": [master.detach() for master in self.masters],
+                "optimizer": cast(dict[str, object], self.optimizer.state_dict())}
+
+    @torch.no_grad()
+    def load_state_dict(self, state: OffloadState) -> None:
+        if state["format_version"] != 1 or state["parameter_names"] != self.names:
+            raise ValueError("Optimizer checkpoint does not match model parameter names and order.")
+        saved_masters = state["masters"]
+        if len(saved_masters) != len(self.masters):
+            raise ValueError("Optimizer checkpoint has the wrong number of master tensors.")
+        for master, saved in zip(self.masters, saved_masters):
+            if saved.device.type != "cpu" or saved.dtype != torch.float32 or saved.shape != master.shape:
+                raise ValueError("Master tensors must be CPU FP32 with matching parameter shapes.")
+        self.optimizer.load_state_dict(state["optimizer"])
+        for parameter, master, saved in zip(self.parameters, self.masters, saved_masters):
+            master.copy_(saved)
+            parameter.copy_(master)
+        for parameter in self.parameters:
+            if parameter.is_cuda:
+                torch.cuda.synchronize(parameter.device)
+                break
+        self.zero_grad()