File size: 2,163 Bytes
93ebaf2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | """Constant-tensor checkpoint reader."""
import json
from pathlib import Path
import torch
from safetensors.torch import load_file
from .configuration_omni import OmniConfig
from .modeling_omni import OmniModel
class ConstantCheckpoint:
def __init__(self, directory):
self.directory = Path(directory)
self.shapes = json.loads((self.directory / "tensor_shapes.json").read_text())
self.values = load_file(self.directory / "initialization.st")
if self.shapes.keys() != self.values.keys():
raise ValueError("Shape and value keys differ")
self.config = OmniConfig.from_pretrained(self.directory, local_files_only=True)
def slice(self, name, start, length):
size = 1
for dimension in self.shapes[name]:
size *= dimension
if (
type(start) is not int
or type(length) is not int
or start < 0
or length < 0
or start + length > size
):
raise ValueError("Invalid tensor slice")
if length > 1_000_000:
raise ValueError("Read at most one million elements per slice")
return self.values[name].expand(length).clone()
def meta_model(self):
with torch.device("meta"):
model = OmniModel(self.config)
parameters = dict(model.named_parameters())
if parameters.keys() != self.shapes.keys():
raise ValueError("Checkpoint and model parameter keys differ")
for name, parameter in parameters.items():
if list(parameter.shape) != self.shapes[name]:
raise ValueError(f"Shape mismatch: {name}")
return model
def materialize(self, max_parameters=10_000_000):
model = self.meta_model()
total = sum(p.numel() for p in model.parameters())
if total > max_parameters:
raise ValueError(f"{total:,} parameters exceed the materialization limit")
model.to_empty(device="cpu")
with torch.no_grad():
for name, parameter in model.named_parameters():
parameter.fill_(self.values[name].item())
return model
|