Buckets:
| """Immutable run-start parameter view for reference-policy scoring.""" | |
| from __future__ import annotations | |
| import contextlib | |
| import hashlib | |
| import json | |
| import torch | |
| from .training import finite_tensors | |
| class FrozenReference: | |
| """Evaluate an actor through an immutable snapshot of trainable parameters.""" | |
| STATE_SCHEMA = "shared_frozen_reference_v1" | |
| UPDATE_INVARIANT = "never" | |
| def __init__(self, actor) -> None: | |
| named = tuple( | |
| (name, parameter) | |
| for name, parameter in actor.named_parameters() | |
| if parameter.requires_grad | |
| ) | |
| if not named: | |
| raise ValueError("frozen reference requires trainable actor parameters") | |
| self.parameter_names = tuple(name for name, _parameter in named) | |
| self._actor = actor | |
| self._parameters = tuple(parameter for _name, parameter in named) | |
| self._modules = tuple(actor.modules()) | |
| self._parameter_specs = { | |
| name: { | |
| "shape": list(parameter.shape), | |
| "native_dtype": str(parameter.dtype).removeprefix("torch."), | |
| } | |
| for name, parameter in named | |
| } | |
| self._shadow = { | |
| name: parameter.detach().float().clone() for name, parameter in named | |
| } | |
| if not finite_tensors(self._shadow.values()): | |
| raise ValueError("frozen reference run-start snapshot is non-finite") | |
| self._native_by_name = { | |
| name: self._shadow[name].to(dtype=parameter.dtype).clone() | |
| for name, parameter in named | |
| } | |
| self._native = tuple(self._native_by_name.values()) | |
| self._snapshot_sha256 = self._snapshot_hash( | |
| self.parameter_names, self._shadow | |
| ) | |
| self._active = False | |
| def _snapshot_hash(names, shadow) -> str: | |
| digest = hashlib.sha256() | |
| for name in names: | |
| value = shadow[name].detach().float().cpu().contiguous() | |
| digest.update(name.encode()) | |
| digest.update(str(tuple(value.shape)).encode()) | |
| digest.update(value.numpy().tobytes()) | |
| return digest.hexdigest() | |
| def _descriptor(self, snapshot_sha256: str) -> dict: | |
| return { | |
| "schema": self.STATE_SCHEMA, | |
| "update_invariant": self.UPDATE_INVARIANT, | |
| "parameter_names": list(self.parameter_names), | |
| "parameter_specs": { | |
| name: { | |
| "shape": list(spec["shape"]), | |
| "native_dtype": spec["native_dtype"], | |
| } | |
| for name, spec in self._parameter_specs.items() | |
| }, | |
| "snapshot_sha256": snapshot_sha256, | |
| } | |
| def _descriptor_hash(descriptor: dict) -> str: | |
| encoded = json.dumps( | |
| descriptor, sort_keys=True, separators=(",", ":") | |
| ).encode() | |
| return hashlib.sha256(encoded).hexdigest() | |
| def descriptor(self) -> dict: | |
| """Return a tensor-free identity for the immutable reference state.""" | |
| descriptor = self._descriptor(self._snapshot_sha256) | |
| return descriptor | {"descriptor_hash": self._descriptor_hash(descriptor)} | |
| def state_dict(self) -> dict: | |
| descriptor = self.descriptor() | |
| return descriptor | { | |
| "shadow": { | |
| name: value.detach().cpu().clone() | |
| for name, value in self._shadow.items() | |
| } | |
| } | |
| def load_state_dict(self, state: dict) -> None: | |
| """Restore an immutable run-start snapshot after strict validation.""" | |
| if self._active: | |
| raise RuntimeError("cannot load reference state while its view is active") | |
| self._validate_actor(self._actor) | |
| expected_keys = { | |
| "schema", | |
| "update_invariant", | |
| "parameter_names", | |
| "parameter_specs", | |
| "snapshot_sha256", | |
| "descriptor_hash", | |
| "shadow", | |
| } | |
| if not isinstance(state, dict) or set(state) != expected_keys: | |
| raise ValueError("invalid frozen-reference state keys") | |
| if state["schema"] != self.STATE_SCHEMA: | |
| raise ValueError("unsupported frozen-reference state schema") | |
| if state["update_invariant"] != self.UPDATE_INVARIANT: | |
| raise ValueError("frozen reference must retain the never-update invariant") | |
| if tuple(state["parameter_names"]) != self.parameter_names: | |
| raise ValueError("frozen-reference parameter names differ") | |
| if state["parameter_specs"] != self._parameter_specs: | |
| raise ValueError("frozen-reference parameter specs differ") | |
| incoming = state["shadow"] | |
| if not isinstance(incoming, dict) or set(incoming) != set(self.parameter_names): | |
| raise ValueError("frozen-reference shadow parameter keys differ") | |
| validated = {} | |
| for name in self.parameter_names: | |
| value = incoming[name] | |
| if not isinstance(value, torch.Tensor): | |
| raise TypeError(f"frozen-reference shadow {name!r} is not a tensor") | |
| if value.dtype != torch.float32: | |
| raise ValueError("frozen-reference shadows must use FP32") | |
| if tuple(value.shape) != tuple(self._shadow[name].shape): | |
| raise ValueError(f"frozen-reference shadow shape differs for {name}") | |
| if not bool(torch.isfinite(value).all()): | |
| raise ValueError(f"frozen-reference shadow is non-finite for {name}") | |
| validated[name] = value.detach().float().cpu().clone() | |
| snapshot_sha256 = self._snapshot_hash(self.parameter_names, validated) | |
| descriptor = self._descriptor(snapshot_sha256) | |
| if snapshot_sha256 != state["snapshot_sha256"]: | |
| raise ValueError("frozen-reference snapshot hash differs") | |
| if self._descriptor_hash(descriptor) != state["descriptor_hash"]: | |
| raise ValueError("frozen-reference descriptor hash differs") | |
| for name in self.parameter_names: | |
| self._shadow[name].copy_(validated[name].to(self._shadow[name])) | |
| self._native_by_name[name].copy_( | |
| self._shadow[name].to(self._native_by_name[name]) | |
| ) | |
| self._snapshot_sha256 = snapshot_sha256 | |
| def _validate_actor(self, actor) -> None: | |
| if actor is not self._actor: | |
| raise RuntimeError("frozen reference is bound to a different actor") | |
| current = tuple( | |
| (name, parameter) | |
| for name, parameter in actor.named_parameters() | |
| if parameter.requires_grad | |
| ) | |
| if ( | |
| tuple(name for name, _parameter in current) != self.parameter_names | |
| or len(current) != len(self._parameters) | |
| or any( | |
| current_parameter is not original_parameter | |
| for (_name, current_parameter), original_parameter in zip( | |
| current, self._parameters, strict=True | |
| ) | |
| ) | |
| or any( | |
| list(parameter.shape) != self._parameter_specs[name]["shape"] | |
| or str(parameter.dtype).removeprefix("torch.") | |
| != self._parameter_specs[name]["native_dtype"] | |
| for name, parameter in current | |
| ) | |
| ): | |
| raise RuntimeError("reference actor trainable parameter schema changed") | |
| def apply(self, actor): | |
| """Temporarily evaluate ``actor`` at its immutable run-start state.""" | |
| if self._active: | |
| raise RuntimeError("frozen reference parameter context is not reentrant") | |
| self._validate_actor(actor) | |
| original_data = tuple(parameter.data for parameter in self._parameters) | |
| training_modes = tuple(module.training for module in self._modules) | |
| self._active = True | |
| try: | |
| actor.eval() | |
| for parameter, native in zip( | |
| self._parameters, self._native, strict=True | |
| ): | |
| parameter.data = native | |
| with torch.no_grad(): | |
| yield actor | |
| finally: | |
| for parameter, original in zip( | |
| self._parameters, original_data, strict=True | |
| ): | |
| parameter.data = original | |
| for module, training in zip(self._modules, training_modes, strict=True): | |
| module.training = training | |
| self._active = False | |
| __all__ = ["FrozenReference"] | |
Xet Storage Details
- Size:
- 8.51 kB
- Xet hash:
- 89303e75cae79ab5d19b0e11e80387448fa9960f268afe5603af1ac50a764c3c
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.