| """Offline integrity, upstream parity, and real checkpoint inference checks.""" |
| import hashlib |
| import json |
| from pathlib import Path |
| import struct |
| import inspect |
|
|
| import torch |
| from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer |
| from transformers import Gemma4ForCausalLM, Gemma4TextConfig |
| from transformers.models.gemma4 import modeling_gemma4 as upstream_code |
|
|
|
|
| def main(): |
| path = Path(__file__).resolve().parent |
| torch.set_num_threads(8) |
| report = json.loads((path / 'conversion_report.json').read_text()) |
| with (path / 'model.safetensors').open('rb') as f: |
| length = struct.unpack('<Q', f.read(8))[0] |
| header = json.loads(f.read(length)) |
| for key, expected in report['tensor_sha256'].items(): |
| start, end = header[key]['data_offsets'] |
| f.seek(8 + length + start) |
| digest = hashlib.sha256() |
| left = end - start |
| while left: |
| chunk = f.read(min(left, 16 * 1024**2)) |
| assert chunk, key |
| digest.update(chunk) |
| left -= len(chunk) |
| assert digest.hexdigest() == expected, key |
| assert set(header) - {'__metadata__'} == set(report['tensor_sha256']) |
| print(f'All {len(report["tensor_sha256"])} retained tensors match original bytes.', flush=True) |
|
|
| config = AutoConfig.from_pretrained(path, trust_remote_code=True, local_files_only=True) |
| assert not any(k.startswith('mtp_') or k in ['vision_config', 'audio_config'] for k in config.to_dict()) |
| small = json.loads((path / 'config.json').read_text()) |
| small['_name_or_path'] = str(path) |
| small.update(vocab_size=128, vocab_size_per_layer_input=128, hidden_size=64, |
| hidden_size_per_layer_input=8, intermediate_size=96, num_hidden_layers=4, |
| num_attention_heads=2, num_key_value_heads=1, head_dim=16, global_head_dim=32, |
| num_kv_shared_layers=2, sliding_window=4, |
| layer_types=['sliding_attention', 'full_attention'] * 2, |
| eos_token_id=1, pad_token_id=0, dtype='float32') |
| custom = AutoModelForCausalLM.from_config(type(config)(**small), trust_remote_code=True).eval() |
| upstream = Gemma4ForCausalLM(Gemma4TextConfig(**small)).eval() |
| upstream.load_state_dict(custom.state_dict(), strict=True) |
| ids = torch.tensor([[3, 4, 5, 6, 7], [0, 0, 8, 9, 10]]) |
| mask = ids.ne(0).long() |
| with torch.inference_mode(): |
| a = custom(ids, attention_mask=mask, use_cache=True) |
| b = upstream(ids, attention_mask=mask, use_cache=True) |
| torch.testing.assert_close(a.logits, b.logits, rtol=0, atol=0) |
| extended_mask = torch.cat([mask, torch.ones(2, 1, dtype=torch.long)], dim=1) |
| nxt = torch.tensor([[11], [12]]) |
| ac = custom(nxt, attention_mask=extended_mask, past_key_values=a.past_key_values, use_cache=True) |
| bc = upstream(nxt, attention_mask=extended_mask, past_key_values=b.past_key_values, use_cache=True) |
| torch.testing.assert_close(ac.logits, bc.logits, rtol=0, atol=0) |
| full = custom(torch.cat([ids, nxt], dim=1), attention_mask=extended_mask, use_cache=False) |
| torch.testing.assert_close(ac.logits[:, -1], full.logits[:, -1], rtol=1e-4, atol=1e-5) |
| print('Upstream parity: exact prefill and cached decode; cached/full consistency passed.', flush=True) |
|
|
| model, info = AutoModelForCausalLM.from_pretrained( |
| path, trust_remote_code=True, local_files_only=True, dtype=torch.bfloat16, |
| output_loading_info=True, |
| ) |
| model.eval() |
| for key in ['missing_keys', 'unexpected_keys', 'mismatched_keys', 'error_msgs']: |
| assert not info.get(key), info |
| info = {key: sorted(value) if isinstance(value, set) else value for key, value in info.items()} |
| tied = model.lm_head.weight is model.model.embed_tokens.weight |
| assert tied == config.tie_word_embeddings |
| assert tied and "lm_head.weight" not in header |
| assert sum(p.numel() for p in model.parameters()) == report["text_parameters"] |
| assert not any(any(t in n for t in ['vision', 'audio', 'mtp']) for n, _ in model.named_parameters()) |
| tokenizer = AutoTokenizer.from_pretrained(path, local_files_only=True) |
| prompt = tokenizer.apply_chat_template([{'role': 'user', 'content': '请用一句话介绍你自己。'}], |
| tokenize=False, add_generation_prompt=True, enable_thinking=False) |
| inputs = tokenizer(prompt, return_tensors='pt', add_special_tokens=False) |
| with torch.inference_mode(): |
| model_inputs = {k: v for k, v in inputs.items() if k != 'token_type_ids'} |
| logits = model(**model_inputs, logits_to_keep=1).logits |
| assert torch.isfinite(logits).all() |
| generated = model.generate(**model_inputs, max_new_tokens=32, do_sample=False) |
| answer = tokenizer.decode(generated[0, inputs.input_ids.shape[1]:], skip_special_tokens=True) |
| result = {'torch': torch.__version__, 'device': 'cpu', 'dtype': 'bfloat16', |
| 'tensor_integrity': f'{len(report["tensor_sha256"])}/{len(report["tensor_sha256"])} byte-identical to source', |
| 'upstream_small_model_prefill_and_decode': 'exact match', |
| 'cached_vs_full_forward': 'passed (rtol=1e-4, atol=1e-5)', |
| 'loading_info': info, 'tied_embeddings': tied, 'finite_logits': True, |
| 'generation': answer} |
| (path / 'validation_report.json').write_text(json.dumps(result, indent=2, ensure_ascii=False) + '\n') |
| print(json.dumps(result, indent=2, ensure_ascii=False)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
