training/prepare_observation_export.py

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