| import os |
| import shutil |
|
|
| import torch |
| from transformers import AutoModelForCausalLM |
|
|
| from configuration_minigpt import MiniGPTConfig |
| from modeling_minigpt import MiniGPTForCausalLM |
|
|
| OUT_DIR = "runs/check_model" |
| EXPECTED_PARAMS = 31_433_920 |
|
|
|
|
| def param_breakdown(config): |
| embed = config.vocab_size * config.d_model |
| qkvo = 4 * config.d_model * config.d_model |
| mlp = 2 * config.d_model * config.d_ff |
| norms = (2 * config.n_layer + 1) * config.d_model |
| return embed, qkvo + mlp, norms |
|
|
|
|
| def main(): |
| torch.manual_seed(0) |
| config = MiniGPTConfig() |
| model = MiniGPTForCausalLM(config) |
| model.eval() |
|
|
| embed, block, norms = param_breakdown(config) |
| n_params = model.num_parameters() |
| print(f"embedding {embed:>10,}") |
| print(f"blocks {block:>10,}") |
| print(f"norms {norms:>10,}") |
| print(f"total {n_params:>10,}") |
| assert n_params == EXPECTED_PARAMS, f"{n_params} != {EXPECTED_PARAMS}" |
|
|
| assert model.lm_head.weight.data_ptr() == model.model.embed_tokens.weight.data_ptr() |
|
|
| input_ids = torch.randint(0, config.vocab_size, (2, 32)) |
| with torch.no_grad(): |
| logits = model(input_ids).logits |
| assert logits.shape == (2, 32, config.vocab_size), logits.shape |
|
|
| edited = input_ids.clone() |
| edited[:, 17] = (edited[:, 17] + 1) % config.vocab_size |
| with torch.no_grad(): |
| logits_edited = model(edited).logits |
| assert torch.equal(logits[:, :17], logits_edited[:, :17]) |
| assert not torch.equal(logits[:, 17], logits_edited[:, 17]) |
|
|
| if os.path.isdir(OUT_DIR): |
| shutil.rmtree(OUT_DIR) |
| model.save_pretrained(OUT_DIR) |
|
|
| for name in [ |
| "config.json", |
| "model.safetensors", |
| "configuration_minigpt.py", |
| "modeling_minigpt.py", |
| ]: |
| path = os.path.join(OUT_DIR, name) |
| assert os.path.isfile(path), f"missing {name}" |
|
|
| reloaded = AutoModelForCausalLM.from_pretrained(OUT_DIR, trust_remote_code=True) |
| reloaded.eval() |
| with torch.no_grad(): |
| reloaded_logits = reloaded(input_ids).logits |
| assert torch.allclose(logits, reloaded_logits, atol=1e-5), "logits differ after reload" |
|
|
| print("PASS") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
