"""Verify trained HF reload and build a lossless Reasoner-only Omni overlay. The original snapshot/export are immutable. Generator tensors are preserved; this overlay is evaluated only in reasoner mode, not for generation/actions. """ import argparse import gc import json from pathlib import Path import shutil import sys import time sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from training.observation_baseline import digest, write_json, verify_bundle BUNDLE_SHA = '8aab3db21d4775afcb1067328622dd1685c40cdb96d8598a0f748da70f6f7a70' INDEX_SHA = 'da25fad9eb3354f3f3b78777662e9d678ad6718bcd7e55435fdfa1a96c83cdd8' AUDIT_SHA = '96f09363a609872cbd4b6e2e43293ab6231b35dde79748ad5bc5620f16cb8e7e' def indexed_paths(root, index): paths = {} for relative in set(index.values()): path = (root / relative).resolve() if not path.is_relative_to(root.resolve()) or not path.is_file(): raise ValueError('Invalid indexed shard: ' + relative) paths[relative] = path return paths def tensor_hash(tensor): import hashlib import torch tensor = tensor.detach().contiguous().cpu() h = hashlib.sha256(str((tuple(tensor.shape), str(tensor.dtype))).encode()) flat = tensor.reshape(-1).view(torch.uint8).numpy() for i in range(0, flat.size, 1024 * 1024): h.update(memoryview(flat[i:i + 1024 * 1024])) return h.hexdigest() def main(): p = argparse.ArgumentParser(description=__doc__) for name in ('base', 'export', 'training-root', 'bundle', 'overlay', 'output'): p.add_argument('--' + name, type=Path, required=True) p.add_argument('--expected-index-sha256', default=INDEX_SHA) p.add_argument('--expected-audit-sha256', default=AUDIT_SHA) p.add_argument('--training-bundle-sha256', default=BUNDLE_SHA) p.add_argument('--evaluation-bundle-sha256', default=BUNDLE_SHA) p.add_argument('--training-audit', type=Path) p.add_argument('--completed-updates', type=int, default=40) a = p.parse_args(); a.output.mkdir(parents=True, exist_ok=False) import torch from safetensors import safe_open from safetensors.torch import load_file, save_file from transformers import AutoModelForImageTextToText import cosmos_framework.model.generator.reasoner.cosmos3_edge # registers Auto class from cosmos_framework.data.generator.processors import build_processor from cosmos_framework.inference.inference import _decode_reasoner_video from cosmos_framework.model.generator.utils.safetensors_loader import ( _detect_indexed_snapshot, _verify_edge_vision_shard_hash) from tancho.edge_chat import install, CHAT_MODE install(); torch.manual_seed(42) manifest, _ = verify_bundle(a.bundle) assert digest(a.bundle / 'manifest.json') == a.evaluation_bundle_sha256 provenance = json.loads((a.base / 'tancho-model-provenance.json').read_text()) assert provenance['revision'] == manifest['model_revision'] for item in provenance['files']: path = a.base / item['path'] assert path.stat().st_size == item['bytes'] and digest(path) == item['sha256'] assert digest(a.export / 'model.safetensors.index.json') == a.expected_index_sha256 training_audit = a.training_audit or a.training_root / 'results/finish/audit/audit.json' assert digest(training_audit) == a.expected_audit_sha256 export_record = json.loads((a.export / 'tancho-sft-provenance.json').read_text()) assert export_record['completed_updates'] == a.completed_updates assert export_record['bundle_manifest_sha256'] == a.training_bundle_sha256 assert export_record['training_audit_sha256'] == a.expected_audit_sha256 assert export_record['index_sha256'] == a.expected_index_sha256 inventory = json.loads((a.training_root / 'results/checkpoint-inventory.json').read_text()) checked = [] for item in inventory: path = Path(item['path']) if path.is_relative_to(a.export): assert path.stat().st_size == item['bytes'] and digest(path) == item['sha256'], path checked.append(item) indexed = _detect_indexed_snapshot(str(a.base)); assert indexed is not None _verify_edge_vision_shard_hash(str(a.base), indexed.files) base_index = json.loads((a.base / 'model.safetensors.index.json').read_text())['weight_map'] export_index = json.loads((a.export / 'model.safetensors.index.json').read_text())['weight_map'] indexed_paths(a.base, base_index); export_paths = indexed_paths(a.export, export_index) assert set(export_paths.values()) <= {Path(item['path']).resolve() for item in checked}, 'Unverified export shard' assert set(indexed.key_map.values()) == set(export_index) and len(export_index) == 670 weights = {} for relative, path in export_paths.items(): shard = load_file(str(path)) assert all(export_index[k] == relative and k not in weights for k in shard) weights.update(shard) assert set(weights) == set(export_index) assert all(v.dtype == torch.bfloat16 and torch.isfinite(v).all() for v in weights.values()) hashes = {k: tensor_hash(v) for k, v in weights.items()} # Native HF reload: no missing/unexpected tensors or silent initialization. started = time.monotonic() model, info = AutoModelForImageTextToText.from_pretrained( str(a.export), torch_dtype=torch.bfloat16, device_map='cuda:0', attn_implementation='sdpa', local_files_only=True, output_loading_info=True) model.eval() assert not any(info.get(k) for k in ('missing_keys', 'unexpected_keys', 'mismatched_keys', 'error_msgs')), info state = model.state_dict() assert set(state) == set(weights) assert all(tensor_hash(state[k]) == hashes[k] for k in weights), 'HF reload changed weights' # The same wrapper and decoded four frames used in SFT and Omni inference. processor = build_processor(tokenizer_type=str(a.base), config_variant='hf') conversation = json.loads((a.bundle / 'train/text/A01.json').read_text())['conversations'] video = _decode_reasoner_video(str(a.bundle / 'media/A01.mp4'), 0.5) messages = [{'role': 'user', 'content': [ {'type': 'video', 'video': video['frames'], 'fps': video['fps']}, {'type': 'text', 'text': conversation[0]['content'][1]['text']}]}] inputs = processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors='pt') input_hashes = {k: tensor_hash(v) for k, v in inputs.items() if torch.is_tensor(v)} inputs = {k: (v.unsqueeze(0) if k in ('input_ids', 'attention_mask') else v).to('cuda:0') for k, v in inputs.items() if torch.is_tensor(v)} with torch.inference_mode(): generated = model.generate(**inputs, max_new_tokens=16, do_sample=False, use_cache=False, eos_token_id=11, pad_token_id=0) tokens = generated[0, inputs['input_ids'].shape[-1]:].tolist() assert tokens hf_report = {'passed': True, 'weights_verified': len(weights), 'loading_info': info, 'train_probe_id': 'A01', 'generated_token_ids': tokens, 'text': processor.processor.tokenizer.decode(tokens, skip_special_tokens=True), 'input_tensor_hashes': input_hashes, 'chat_mode': CHAT_MODE, 'elapsed_seconds_including_load_and_hash': time.monotonic() - started} write_json(a.output / 'hf-reload.json', hf_report) del state, model, inputs; gc.collect(); torch.cuda.empty_cache() print('HF_RELOAD_VERIFIED ' + json.dumps(hf_report, ensure_ascii=False), flush=True) # Reuse unchanged files with symlinks, replace all three weight shards. a.overlay.mkdir(parents=True, exist_ok=False) for path in a.base.rglob('*'): if not path.is_file(): continue relative = path.relative_to(a.base) if str(relative) in indexed.files or relative.name == 'tancho-model-provenance.json': continue target = a.overlay / relative; target.parent.mkdir(parents=True, exist_ok=True) target.symlink_to(path.resolve()) (a.overlay / 'tancho-base-model-provenance.json').symlink_to((a.base / 'tancho-model-provenance.json').resolve()) base_hashes = {}; changed = []; frozen = [] output_shards = [] for relative in indexed.files: source = a.base / relative state = load_file(str(source)) original_hashes = {k: tensor_hash(v) for k, v in state.items()} replacement = {raw: canonical for raw, canonical in indexed.key_map.items() if base_index[raw] == relative} for raw, canonical in replacement.items(): assert state[raw].shape == weights[canonical].shape and state[raw].dtype == weights[canonical].dtype base_hashes[canonical] = original_hashes[raw] if canonical.startswith('model.visual.'): assert original_hashes[raw] == hashes[canonical], 'Frozen vision changed' frozen.append(canonical) elif original_hashes[raw] != hashes[canonical]: changed.append(canonical) state[raw] = weights[canonical] target = a.overlay / relative; target.parent.mkdir(parents=True, exist_ok=True) save_file(state, str(target), metadata={'format': 'pt'}) # Read saved bytes again: all replacements and all generator weights. with safe_open(str(target), framework='pt', device='cpu') as saved: assert set(saved.keys()) == set(state) for key in saved.keys(): expected = hashes[replacement[key]] if key in replacement else original_hashes[key] assert tensor_hash(saved.get_tensor(key)) == expected, key output_shards.append({'path': str(target), 'bytes': target.stat().st_size, 'sha256': digest(target)}) del state; gc.collect() assert len(base_hashes) == 670 and len(frozen) == 437 and changed for variant, values in (('base', base_hashes), ('sft', hashes)): write_json(a.output / (variant + '-expected-weights.json'), { 'variant': variant, 'canonical_tensor_hashes': values, 'raw_to_canonical': indexed.key_map, 'checkpoint': str(a.base if variant == 'base' else a.overlay)}) report = {'passed': True, 'bundle_sha256': a.evaluation_bundle_sha256, 'training_bundle_sha256': a.training_bundle_sha256, 'source_export_index_sha256': a.expected_index_sha256, 'source_training_audit_sha256': a.expected_audit_sha256, 'source_export_files_verified': checked, 'replaced_tensors': 670, 'frozen_vision_tensors_identical': len(frozen), 'trainable_tensors_changed_after_bf16_export': len(changed), 'all_other_shard_tensors_preserved': True, 'shards': output_shards, 'reasoner_only': True, 'generator_or_action_model_validated': False} write_json(a.overlay / 'tancho-sft-overlay-provenance.json', report) write_json(a.output / 'overlay-verification.json', report) print('OVERLAY_VERIFIED ' + json.dumps({k: v for k, v in report.items() if k not in ('source_export_files_verified', 'shards')}), flush=True) if __name__ == '__main__': main()