File size: 3,271 Bytes
8c00b68 | 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | #!/usr/bin/env python3
"""Validate remote-code loading and checkpoint shapes without allocating weights."""
from __future__ import annotations
from pathlib import Path
from accelerate import init_empty_weights
from transformers import AutoConfig, AutoModel, AutoTokenizer
from transformers.dynamic_module_utils import get_class_from_dynamic_module
from validate_repo import EXPECTED_SPECIAL_TOKENS, load_json, read_safetensors_header
ROOT = Path(__file__).resolve().parent
def main() -> None:
config = AutoConfig.from_pretrained(ROOT, trust_remote_code=True)
assert type(config).__name__ == "RaonConfig"
assert type(config.text_model_config).__name__ == "AXK1Config"
assert config.max_position_embeddings == config.text_model_config.max_position_embeddings
tokenizer = AutoTokenizer.from_pretrained(ROOT, trust_remote_code=True)
assert len(tokenizer) == config.text_model_config.vocab_size
for token_id, token_text in EXPECTED_SPECIAL_TOKENS.items():
assert tokenizer.encode(token_text, add_special_tokens=False) == [token_id]
processor_class = get_class_from_dynamic_module("modeling_raon.RaonProcessor", ROOT)
processor = processor_class.from_pretrained(ROOT)
assert processor.sampling_rate == 24_000
with init_empty_weights(include_buffers=True):
model = AutoModel.from_config(config, trust_remote_code=True)
assert type(model).__name__ == "RaonModel"
assert type(model.text_model).__name__ == "AXK1Model"
assert next(model.parameters()).device.type == "meta"
model_shapes = {name: tuple(tensor.shape) for name, tensor in model.state_dict().items()}
index = load_json("model.safetensors.index.json")
checkpoint_shapes: dict[str, tuple[int, ...]] = {}
for shard_name in sorted(set(index["weight_map"].values())):
header = read_safetensors_header(ROOT / shard_name)
checkpoint_shapes.update(
{
name: tuple(info["shape"])
for name, info in header.items()
if name != "__metadata__"
}
)
model_keys = set(model_shapes)
checkpoint_keys = set(checkpoint_shapes)
missing = sorted(model_keys - checkpoint_keys)
unexpected = sorted(checkpoint_keys - model_keys)
mismatched = sorted(
name
for name in model_keys & checkpoint_keys
if model_shapes[name] != checkpoint_shapes[name]
)
assert not missing, f"Checkpoint is missing {len(missing)} model tensors: {missing[:20]}"
assert not unexpected, f"Checkpoint has {len(unexpected)} unexpected tensors: {unexpected[:20]}"
assert not mismatched, (
f"Checkpoint has {len(mismatched)} shape mismatches: "
f"{[(name, model_shapes[name], checkpoint_shapes[name]) for name in mismatched[:20]]}"
)
parameter_count = sum(parameter.numel() for parameter in model.parameters())
assert parameter_count == index["metadata"]["total_parameters"]
print(
"Model validation passed: "
f"{type(model).__name__} with {type(model.text_model).__name__}, "
f"{len(model_shapes):,} state entries, "
f"{parameter_count:,} parameters, "
"zero missing/unexpected/mismatched tensors."
)
if __name__ == "__main__":
main()
|