| """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 |
|
|
