diff --git a/recipe/cua_s1/export_multimodal_reference.py b/recipe/cua_s1/export_multimodal_reference.py new file mode 100644 index 0000000..89961a6 --- /dev/null +++ b/recipe/cua_s1/export_multimodal_reference.py @@ -0,0 +1,334 @@ +"""Export the #12 reference's actual multimodal forward inputs and readout.""" + +from __future__ import annotations + +import argparse +import base64 +import hashlib +import importlib.metadata +import importlib.util +import json +import os +import platform +import subprocess +from pathlib import Path + +from PIL import Image, ImageDraw + +from models.cua_s1.multimodal.model import ( + ADAPTER_REVISION, + BASE_REVISION, + REFERENCE_REVISION, + WEIGHTS_MANIFEST_SHA256, + MultimodalEngine, + letter_ids, +) +from models.cua_s1.multimodal.protocol import ( + answer, + build_messages, + decode_request, + parse_request, +) + +SCHEMA = "cua-s1-multimodal-reference-v1" +PACKAGES = { + "torch": "2.14.0", + "torchvision": "0.29.0", + "transformers": "5.17.0", + "peft": "0.21.0", + "accelerate": "1.15.0", + "Pillow": "11.3.0", + "safetensors": "0.8.0", + "huggingface-hub": "1.32.0", + "tokenizers": "0.23.2", + "numpy": "2.5.3", +} + + +def sha256(raw): + return hashlib.sha256(raw).hexdigest() + + +def write_json(path, value): + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + "\n" + ) + + +def make_cases(folder): + folder.mkdir(parents=True) + cases = [] + for name, size, fmt in [ + ("small", (320, 240), "PNG"), + ("wide", (640, 320), "PNG"), + ("portrait", (320, 640), "PNG"), + ("jpeg", (640, 480), "JPEG"), + ("single-option", (320, 240), "PNG"), + ("26-options", (256, 256), "PNG"), + ("two-questions", (320, 240), "PNG"), + ]: + image = Image.new("RGB", size, "#f4f6f8") + draw = ImageDraw.Draw(image) + width, height = size + draw.rectangle( + (16, 16, width - 16, height - 16), fill="white", outline="#8899aa" + ) + draw.text((24, 24), "Account settings", fill="black") + draw.text((24, 48), "Display name: Alice", fill="black") + draw.rectangle((24, height // 2, width // 2, height // 2 + 32), fill="#1460b4") + draw.text((28, height // 2 + 8), "Save", fill="white") + draw.text((width // 2 + 16, height // 2 + 8), "Cancel", fill="black") + image_path = folder / (name + (".jpg" if fmt == "JPEG" else ".png")) + image.save(image_path, format=fmt) + criteria = {"save": "Click Save", "cancel": "Click Cancel", "wait": "Wait"} + if name == "single-option": + criteria = {"save": "Click Save"} + elif name == "26-options": + criteria = {f"option-{i}": f"Choose action {i}" for i in range(26)} + questions = { + "next": { + "type": "choice", + "instructions": "Save the changed display name.", + "criteria": criteria, + } + } + if name == "two-questions": + questions["second"] = { + "type": "choice", + "instructions": {"goal": "保存名称"}, + "criteria": {"continue": {"label": "Save"}, "cancel": None}, + } + mime = "jpeg" if fmt == "JPEG" else "png" + request = { + "model": "cua-s1-4b-0.2", + "state": { + "image": f"data:image/{mime};base64," + + base64.b64encode(image_path.read_bytes()).decode() + }, + "questions": questions, + } + write_json(folder / f"{name}.json", request) + cases.append({"name": name, "image": image_path.name, "request": request}) + return cases + + +def tensor_info(tensor): + import torch + + value = tensor.detach().cpu().contiguous() + return { + "shape": list(value.shape), + "dtype": str(value.dtype).removeprefix("torch."), + "sha256": sha256(value.view(torch.uint8).numpy().tobytes()), + } + + +def save_tensors(path, tensors): + from safetensors.torch import save_file + + # Clone individually: safetensors refuses shared storage, even for equal inputs. + tensors = { + name: value.detach().cpu().contiguous().clone() + for name, value in tensors.items() + } + save_file(tensors, str(path), metadata={"schema": SCHEMA}) + return {name: tensor_info(value) for name, value in tensors.items()} + + +def capture(engine, inputs, question): + import torch + + tensors = {} + + def keep(name, value): + if name in tensors: + raise RuntimeError(f"expected one forward per question: duplicate {name}") + tensors[name] = value.detach().cpu().contiguous().clone() + + def vision_hook(module, args, output): + keep("image_features", output.pooler_output) + + def language_pre_hook(module, args, kwargs): + keep("inputs_embeds", kwargs["inputs_embeds"]) + keep("position_ids", kwargs["position_ids"]) + + def language_hook(module, args, output): + keep("last_hidden_state", output.last_hidden_state[:, -1, :]) + + core = engine.model.get_base_model().model + hooks = [ + core.visual.register_forward_hook(vision_hook), + core.language_model.register_forward_pre_hook( + language_pre_hook, with_kwargs=True + ), + core.language_model.register_forward_hook(language_hook), + ] + try: + with torch.no_grad(): + output = engine.model( + **{ + name: value.to(engine.model.device) + for name, value in inputs.items() + } + ) + ids = torch.tensor( + letter_ids(engine.tokenizer, len(question.keys)), + device=output.logits.device, + ) + logits = output.logits[0, -1, ids] + keep("candidate_token_ids", ids) + keep("candidate_logits", logits) + keep("candidate_probabilities", torch.softmax(logits.float(), dim=-1)) + finally: + for handle in hooks: + handle.remove() + for name, value in inputs.items(): + keep(name, value) + keep("rope_deltas", core.rope_deltas) + indices = ( + (inputs["input_ids"][0] == engine.model.config.image_token_id) + .nonzero() + .flatten() + ) + keep("image_token_indices", indices) + if not torch.equal(tensors["inputs_embeds"][0, indices], tensors["image_features"]): + raise RuntimeError("image feature insertion differs from language input") + return tensors + + +def environment(): + import torch + from transformers.models.qwen3_5 import modeling_qwen3_5 + + packages = {name: importlib.metadata.version(name) for name in PACKAGES} + for name, expected in PACKAGES.items(): + if packages[name].split("+")[0] != expected: + raise ValueError(f"{name} must be {expected}, got {packages[name]}") + if any( + importlib.util.find_spec(name) is not None for name in ("fla", "causal_conv1d") + ): + raise ValueError( + "reference export requires the PyTorch DeltaNet path, without FLA" + ) + return { + "python": platform.python_version(), + "torch_build": str(torch.__version__), + "torch_num_threads": torch.get_num_threads(), + "torch_num_interop_threads": torch.get_num_interop_threads(), + "packages": packages, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "compute_capability": list(torch.cuda.get_device_capability()), + "driver": subprocess.check_output( + ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"], + text=True, + ).strip(), + "transformers_source_sha256": sha256( + Path(modeling_qwen3_5.__file__).read_bytes() + ), + "dtype": "bfloat16", + "adapter_merged": False, + "tf32": False, + "deterministic_algorithms": True, + "cublas_workspace_config": os.environ["CUBLAS_WORKSPACE_CONFIG"], + } + + +def export(weights, output): + if output.exists(): + raise FileExistsError(f"output must be a new directory: {output}") + # Must be set before Torch initializes CUDA, including model construction. + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + import torch + + torch.manual_seed(0) + torch.use_deterministic_algorithms(True) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + report = { + "schema": SCHEMA, + "reference_revision": REFERENCE_REVISION, + "base_revision": BASE_REVISION, + "adapter_revision": ADAPTER_REVISION, + "weights_manifest_sha256": WEIGHTS_MANIFEST_SHA256, + "environment": environment(), + "files": {}, + "questions": [], + } + engine = MultimodalEngine( + str(weights / "Qwen3.5-4B"), str(weights / "cua-s1-4b-0.2/multimodal") + ) + core = engine.model.get_base_model().model + report["execution"] = { + "visual_attention": core.visual.config._attn_implementation, + "text_attention": core.language_model.config._attn_implementation, + "processor_class": type(engine.processor).__name__, + "image_processor_class": type(engine.processor.image_processor).__name__, + } + output.mkdir(parents=True) + cases = make_cases(output / "inputs") + (output / "tensors").mkdir() + (output / "configs").mkdir() + for label, path in [ + ("base", weights / "Qwen3.5-4B/config.json"), + ("processor", weights / "Qwen3.5-4B/preprocessor_config.json"), + ("adapter", weights / "cua-s1-4b-0.2/multimodal/adapter_config.json"), + ]: + (output / "configs" / f"{label}.json").write_bytes(path.read_bytes()) + root = Path(__file__).resolve().parents[2] + report["source_sha256"] = { + name: sha256((root / name).read_bytes()) + for name in [ + "recipe/cua_s1/export_multimodal_reference.py", + "src/models/cua_s1/multimodal/model.py", + "src/models/cua_s1/multimodal/protocol.py", + ] + } + for case in cases: + request_path = output / "inputs" / f"{case['name']}.json" + request = parse_request(decode_request(request_path.read_bytes())) + for index, question in enumerate(request.questions): + inputs = engine.prepare(request.image, question) + tensors = capture(engine, inputs, question) + probabilities = tensors["candidate_probabilities"].tolist() + if probabilities != engine.score(inputs, question): + raise RuntimeError("hooked readout differs from ordinary #12 score") + relative = f"tensors/{case['name']}-{index}.safetensors" + entry = { + "case": case["name"], + "question": question.name, + "request": request_path.relative_to(output).as_posix(), + "image": f"inputs/{case['image']}", + "image_size_wh": list(request.image.size), + "option_keys": list(question.keys), + "prompt": engine.processor.apply_chat_template( + build_messages(question), tokenize=False, add_generation_prompt=True + ), + "tensors_file": relative, + "tensors": save_tensors(output / relative, tensors), + "answer": answer(question, probabilities), + "ordinary_score_equal": True, + } + report["questions"].append(entry) + print(f"exported {case['name']}/{question.name}", flush=True) + for path in sorted(output.rglob("*")): + if path.is_file(): + report["files"][path.relative_to(output).as_posix()] = { + "sha256": sha256(path.read_bytes()), + "size": path.stat().st_size, + } + # A manifest is written only after every forward and ordinary-score check passed. + write_json(output / "manifest.json", report) + return report + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--weights", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + export(args.weights, args.output) + + +if __name__ == "__main__": + main() diff --git a/recipe/cua_s1/verify_multimodal_reference.py b/recipe/cua_s1/verify_multimodal_reference.py new file mode 100644 index 0000000..0f9a7a2 --- /dev/null +++ b/recipe/cua_s1/verify_multimodal_reference.py @@ -0,0 +1,239 @@ +"""Verify export integrity and multimodal tensor relations; optionally compare runs.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch +from export_multimodal_reference import SCHEMA, sha256, tensor_info +from safetensors.torch import load_file + +from models.cua_s1.multimodal.protocol import answer, decode_request, parse_request + +REQUIRED = { + "input_ids", + "attention_mask", + "mm_token_type_ids", + "pixel_values", + "image_grid_thw", + "image_features", + "image_token_indices", + "inputs_embeds", + "position_ids", + "rope_deltas", + "last_hidden_state", + "candidate_token_ids", + "candidate_logits", + "candidate_probabilities", +} + + +def require(condition, message): + if not condition: + raise ValueError(message) + + +def check_files(folder, manifest): + for name, info in manifest["files"].items(): + path = folder / name + require( + path.resolve().is_relative_to(folder.resolve()), f"unsafe file path: {name}" + ) + require(path.is_file(), f"missing file: {name}") + require( + path.stat().st_size == info["size"] + and sha256(path.read_bytes()) == info["sha256"], + f"file checksum: {name}", + ) + + actual = { + path.relative_to(folder).as_posix() + for path in folder.rglob("*") + if path.is_file() and path != folder / "manifest.json" + } + require(actual == set(manifest["files"]), "file inventory mismatch") + + +def check_tensors(tensors, entry, configs): + require( + set(tensors) == REQUIRED, + f"missing tensors or unknown keys: {set(tensors) ^ REQUIRED}", + ) + require( + {name: tensor_info(value) for name, value in tensors.items()} + == entry["tensors"], + "tensor fingerprint mismatch", + ) + ids = tensors["input_ids"] + sequence = ids.shape[-1] + hidden = configs["base"]["text_config"]["hidden_size"] + vision = configs["base"]["vision_config"] + merge = vision["spatial_merge_size"] + patch = vision["patch_size"] + temporal = vision["temporal_patch_size"] + grid = tensors["image_grid_thw"] + require( + grid.shape == (1, 3) and grid[0, 0].item() == 1, "expected one still-image grid" + ) + require( + grid[0, 1].item() % merge == grid[0, 2].item() % merge == 0, + "grid merge alignment", + ) + image_tokens = grid.prod().item() // merge**2 + expected_indices = (ids[0] == configs["base"]["image_token_id"]).nonzero().flatten() + shapes = { + "input_ids": (1, sequence), + "attention_mask": (1, sequence), + "mm_token_type_ids": (1, sequence), + "pixel_values": (grid.prod().item(), 3 * temporal * patch**2), + "image_features": (image_tokens, hidden), + "image_token_indices": (image_tokens,), + "inputs_embeds": (1, sequence, hidden), + "position_ids": (3, 1, sequence), + "rope_deltas": (1, 1), + "last_hidden_state": (1, hidden), + "candidate_token_ids": (len(entry["option_keys"]),), + "candidate_logits": (len(entry["option_keys"]),), + "candidate_probabilities": (len(entry["option_keys"]),), + } + for name, shape in shapes.items(): + value = tensors[name] + require(tuple(value.shape) == shape, f"shape mismatch: {name}") + require(bool(torch.isfinite(value).all()), f"nonfinite tensor: {name}") + for name in [ + "input_ids", + "attention_mask", + "mm_token_type_ids", + "image_grid_thw", + "image_token_indices", + "position_ids", + "rope_deltas", + "candidate_token_ids", + ]: + require(tensors[name].dtype == torch.int64, f"expected int64: {name}") + for name in [ + "image_features", + "inputs_embeds", + "last_hidden_state", + "candidate_logits", + ]: + require(tensors[name].dtype == torch.bfloat16, f"expected bfloat16: {name}") + require( + tensors["pixel_values"].dtype + == tensors["candidate_probabilities"].dtype + == torch.float32, + "expected float32 pixels/probabilities", + ) + require(bool((tensors["attention_mask"] == 1).all()), "expected unpadded prompt") + require( + torch.equal(expected_indices, tensors["image_token_indices"]), + "image token indices mismatch", + ) + require( + torch.equal( + expected_indices, (tensors["mm_token_type_ids"][0] == 1).nonzero().flatten() + ), + "image token types mismatch", + ) + require( + torch.equal( + tensors["inputs_embeds"][0, expected_indices], tensors["image_features"] + ), + "image insertion mismatch", + ) + require( + tensors["rope_deltas"].item() + == tensors["position_ids"].max().item() + 1 - sequence, + "rope delta mismatch", + ) + require( + torch.equal( + tensors["candidate_token_ids"], + torch.arange(32, 32 + len(entry["option_keys"])), + ), + "candidate token ordering", + ) + require( + torch.allclose( + torch.softmax(tensors["candidate_logits"].float(), dim=-1), + tensors["candidate_probabilities"], + atol=1e-7, + rtol=0, + ), + "readout mismatch", + ) + + +def verify(folder): + manifest = json.loads((folder / "manifest.json").read_text()) + require(manifest["schema"] == SCHEMA, "unsupported schema") + check_files(folder, manifest) + configs = { + name: json.loads((folder / f"configs/{name}.json").read_text()) + for name in ["base", "processor", "adapter"] + } + seen_images = {} + count = 0 + for entry in manifest["questions"]: + require(entry["tensors_file"] in manifest["files"], "unlisted tensor file") + require( + entry["request"] in manifest["files"] + and entry["image"] in manifest["files"], + "unlisted input file", + ) + tensors = load_file(str(folder / entry["tensors_file"])) + check_tensors(tensors, entry, configs) + request = parse_request( + decode_request((folder / entry["request"]).read_bytes()) + ) + question = next(q for q in request.questions if q.name == entry["question"]) + require( + list(request.image.size) == entry["image_size_wh"], "image size mismatch" + ) + require(list(question.keys) == entry["option_keys"], "option ordering mismatch") + require( + answer(question, tensors["candidate_probabilities"].tolist()) + == entry["answer"], + "answer mismatch", + ) + require(entry["ordinary_score_equal"] is True, "ordinary score check missing") + image = manifest["files"][entry["image"]]["sha256"] + features = { + name: entry["tensors"][name] + for name in ["pixel_values", "image_grid_thw", "image_features"] + } + require( + image not in seen_images or seen_images[image] == features, + "same-image vision mismatch", + ) + seen_images[image] = features + count += len(tensors) + require(len(manifest["questions"]) == 8, "expected eight questions") + return manifest, { + "questions": 8, + "tensors": count, + "files": len(manifest["files"]), + "integrity_and_relations": "pass", + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("bundle", type=Path) + parser.add_argument("--compare", type=Path) + args = parser.parse_args() + manifest, summary = verify(args.bundle) + if args.compare: + other, _ = verify(args.compare) + require( + manifest == other, + "exports differ (environment, metadata, files or tensors)", + ) + summary["independent_export_equality"] = "pass" + print(json.dumps(summary, indent=2)) + + +if __name__ == "__main__": + main()