#!/usr/bin/env python3 """Compare a privileged teacher with the ordinary student on gold answer tokens. Gold labels are appended only as teacher-forced diagnostic continuations after both prompts have been constructed. They are never visible while building the student prompt, privileged teacher prompt, or independent teacher-gate answer. Launch with ``torchrun`` to shard rows over all selected GPUs. """ from __future__ import annotations import argparse import json import math import os import sys import time from pathlib import Path from typing import Any, Mapping, Sequence ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from data_utils.chart.teacher_gate import ( chartqa_reference_answer, chartqa_sample_fingerprint, default_manifest_path, read_jsonl, read_manifest, validate_teacher_gate_cache, ) from data_utils.paths import local_pretrained_kwargs, resolve_model_path, validate_local_model_dir from opsd_utils.opsd_loss import ( _teacher_logits_with_oom_retry, build_opsd_answer_diagnostic_masks, ) from opsd_utils.teacher_batching import process_teacher_sample, stack_teacher_processor_batches from opsd_utils.vocab_align import align_cross_model_logits from reward_utils.teacher_generate import _align_stacked_batch from scripts.preprocess_chartqa_teacher_gate import ( _build_job, _load_rows, _prepare_sample, _selected_indices, ) def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--input", required=True) parser.add_argument("--output", required=True) parser.add_argument("--student-model", required=True) parser.add_argument("--teacher-model", required=True) parser.add_argument("--student-processor", default="") parser.add_argument("--teacher-processor", default="") parser.add_argument("--teacher-gate-cache", required=True) parser.add_argument("--providers", default="visual_facts_deplot") parser.add_argument("--max-samples", type=int, default=0) parser.add_argument("--sample-strategy", choices=("head", "spread"), default="spread") parser.add_argument("--top-k", type=int, default=16) parser.add_argument("--max-relative-change", type=float, default=0.05) parser.add_argument("--human-only", action=argparse.BooleanOptionalAction, default=False) parser.add_argument( "--attn-implementation", choices=("sdpa", "flash_attention_2"), default="flash_attention_2", ) parser.add_argument("--log-every", type=int, default=20) return parser.parse_args(argv) def distributed_context() -> tuple[Any, int, int, int]: import torch import torch.distributed as dist rank = int(os.environ.get("RANK", "0")) world_size = int(os.environ.get("WORLD_SIZE", "1")) local_rank = int(os.environ.get("LOCAL_RANK", "0")) if not torch.cuda.is_available(): raise RuntimeError("privileged-teacher alignment analysis requires CUDA") torch.cuda.set_device(local_rank) if world_size > 1 and not dist.is_initialized(): dist.init_process_group(backend="nccl") return torch, rank, world_size, local_rank def barrier(world_size: int) -> None: if world_size > 1: import torch.distributed as dist dist.barrier() def load_model_processor( model_source: str, processor_source: str, *, torch: Any, local_rank: int, attn_implementation: str, ) -> tuple[Any, Any, str, str]: from transformers import AutoProcessor, LlavaOnevisionForConditionalGeneration model_path = validate_local_model_dir(resolve_model_path(model_source), role="model") processor_path = resolve_model_path(processor_source or model_source) if not os.path.isdir(processor_path): raise FileNotFoundError(f"processor directory is incomplete or missing: {processor_path}") processor = AutoProcessor.from_pretrained( processor_path, **local_pretrained_kwargs(processor_path), ) processor.tokenizer.padding_side = "left" model = LlavaOnevisionForConditionalGeneration.from_pretrained( model_path, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, attn_implementation=attn_implementation, **local_pretrained_kwargs(model_path), ).to(torch.device("cuda", local_rank)) model.eval() model.requires_grad_(False) return model, processor, model_path, processor_path def prepared_prompt(processor: Any, model: Any, text: str, images: list[Any]) -> dict[str, Any]: from opsd_utils.teacher_batching import ( model_inference_device, move_batch_num_images_to_model_device, move_pixel_values_to_model_device, ) processed = process_teacher_sample(processor, text, images) stacked = stack_teacher_processor_batches(processor, [processed]) ids, mask, pixel_values, image_sizes, batch_num_images = _align_stacked_batch( model, processor, stacked ) if isinstance(pixel_values, list): raise RuntimeError("single diagnostic row retained incompatible visual tensors") device = model_inference_device(model) return { "prompt_ids": ids.to(device), "prompt_mask": mask.to(device), "pixel_values": move_pixel_values_to_model_device(model, pixel_values), "image_sizes": image_sizes.to(device) if hasattr(image_sizes, "to") else image_sizes, "batch_num_images": move_batch_num_images_to_model_device(model, batch_num_images), } def completion_ids_for_gold(tokenizer: Any, gold: str, torch: Any, device: Any) -> Any: token_ids = [ int(value) for value in tokenizer.encode( f"\nAnswer: {str(gold).strip()}", add_special_tokens=False ) ] eos_token_id = getattr(tokenizer, "eos_token_id", None) if eos_token_id is None: raise ValueError("diagnostic tokenizer has no eos_token_id") token_ids.append(int(eos_token_id)) return torch.tensor([token_ids], dtype=torch.long, device=device) def score_completion( *, torch: Any, model: Any, processor: Any, prompt: Mapping[str, Any], completion_ids: Any, ) -> Any: completion_mask = torch.ones_like(completion_ids) return _teacher_logits_with_oom_retry( model, processor, prompt["prompt_ids"], prompt["prompt_mask"], completion_ids, completion_mask, prompt["pixel_values"], prompt["image_sizes"], int(completion_ids.size(1)), teacher_batch_num_images=prompt["batch_num_images"], ) def mean_selected(values: Any, mask: Any) -> float: selected = values[mask] return float(selected.mean().item()) if selected.numel() else 0.0 def pair_metrics( *, torch: Any, student_logits: Any, teacher_logits: Any, completion_ids: Any, tokenizer: Any, top_k: int, ) -> dict[str, float]: import torch.nn.functional as F student_logits, teacher_logits = align_cross_model_logits( student_logits.float(), teacher_logits.float() ) teacher_logits = teacher_logits.to(student_logits.device) completion_ids = completion_ids.to(student_logits.device) completion_mask = torch.ones_like(completion_ids) masks = build_opsd_answer_diagnostic_masks( completion_ids, completion_mask, tokenizer, answer_flag="Answer:", ) answer_mask = masks["answer_span_mask"].bool().to(student_logits.device) eos_mask = masks["eos_mask"].bool().to(student_logits.device) first_position = int(masks["first_answer_positions"][0].item()) if first_position < 0 or not bool(answer_mask.any().item()) or not bool(eos_mask.any().item()): raise RuntimeError("gold diagnostic continuation lacks answer span or EOS") student_log = F.log_softmax(student_logits, dim=-1) teacher_log = F.log_softmax(teacher_logits, dim=-1) target = completion_ids.unsqueeze(-1) student_gold_logp = student_log.gather(-1, target).squeeze(-1) teacher_gold_logp = teacher_log.gather(-1, target).squeeze(-1) beta = torch.tensor(0.5, device=student_logits.device) mixture = torch.logaddexp( student_log + torch.log1p(-beta), teacher_log + torch.log(beta), ) jsd = 0.5 * ( F.kl_div(mixture, teacher_log, reduction="none", log_target=True).sum(-1) + F.kl_div(mixture, student_log, reduction="none", log_target=True).sum(-1) ) jsd = jsd.clamp_min(0.0) selected_s_log = student_log[answer_mask] selected_t_log = teacher_log[answer_mask] selected_s_prob = selected_s_log.exp() selected_t_prob = selected_t_log.exp() k = min(max(1, int(top_k)), int(selected_s_log.size(-1))) student_topk = torch.topk(selected_s_log, k=k, dim=-1).indices teacher_topk = torch.topk(selected_t_log, k=k, dim=-1).indices topk_overlap = ( student_topk.unsqueeze(-1) == teacher_topk.unsqueeze(-2) ).any(dim=-1).float().sum(dim=-1) / float(k) shared_mass = torch.minimum(selected_s_prob, selected_t_prob).sum(dim=-1) student_entropy = -(selected_s_prob * selected_s_log).sum(dim=-1) teacher_entropy = -(selected_t_prob * selected_t_log).sum(dim=-1) entropy_gap = teacher_entropy - student_entropy first_student_prob = selected_s_prob[0] first_teacher_prob = selected_t_prob[0] first_student_topk = torch.topk(selected_s_log[0], k=k).indices first_teacher_topk = torch.topk(selected_t_log[0], k=k).indices first_overlap = ( first_student_topk.unsqueeze(-1) == first_teacher_topk.unsqueeze(-2) ).any(dim=-1).float().sum() / float(k) first_shared = torch.minimum(first_student_prob, first_teacher_prob).sum() first_student_entropy = -(first_student_prob * selected_s_log[0]).sum() first_teacher_entropy = -(first_teacher_prob * selected_t_log[0]).sum() return { "first_answer_student_gold_logp": float(student_gold_logp[0, first_position].item()), "first_answer_teacher_gold_logp": float(teacher_gold_logp[0, first_position].item()), "first_answer_teacher_gold_logp_advantage": float( (teacher_gold_logp[0, first_position] - student_gold_logp[0, first_position]).item() ), "answer_span_student_gold_logp": mean_selected(student_gold_logp, answer_mask), "answer_span_teacher_gold_logp": mean_selected(teacher_gold_logp, answer_mask), "answer_span_teacher_gold_logp_advantage": mean_selected( teacher_gold_logp - student_gold_logp, answer_mask ), "first_answer_jsd": float(jsd[0, first_position].item()), "answer_span_jsd": mean_selected(jsd, answer_mask), "eos_jsd": mean_selected(jsd, eos_mask), "first_answer_topk_overlap": float(first_overlap.item()), "answer_span_topk_overlap": float(topk_overlap.mean().item()), "first_answer_shared_probability_mass": float(first_shared.item()), "answer_span_shared_probability_mass": float(shared_mass.mean().item()), "first_answer_teacher_minus_student_entropy": float( (first_teacher_entropy - first_student_entropy).item() ), "answer_span_teacher_minus_student_entropy": float(entropy_gap.mean().item()), "answer_span_entropy_abs_gap": float(entropy_gap.abs().mean().item()), "answer_token_count": float(answer_mask.sum().item()), } def atomic_jsonl(path: Path, rows: Sequence[Mapping[str, Any]]) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}") with temporary.open("w", encoding="utf-8") as handle: for row in rows: handle.write(json.dumps(dict(row), ensure_ascii=False, sort_keys=True) + "\n") os.replace(temporary, path) def atomic_json(path: Path, payload: Mapping[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}") with temporary.open("w", encoding="utf-8") as handle: json.dump(dict(payload), handle, ensure_ascii=False, indent=2, sort_keys=True) handle.write("\n") os.replace(temporary, path) def summarize(rows: Sequence[Mapping[str, Any]]) -> dict[str, Any]: numeric_keys = sorted( { key for row in rows for key, value in row.items() if isinstance(value, (int, float)) and key not in {"source_index"} } ) def group_summary(group: Sequence[Mapping[str, Any]]) -> dict[str, Any]: means: dict[str, float] = {} for key in numeric_keys: values = [float(row[key]) for row in group if key in row and math.isfinite(float(row[key]))] if values: means[key] = sum(values) / len(values) return {"count": len(group), "means": means} correct = [row for row in rows if row.get("teacher_correct") is True] incorrect = [row for row in rows if row.get("teacher_correct") is False] return { "all": group_summary(rows), "teacher_correct": group_summary(correct), "teacher_incorrect": group_summary(incorrect), "teacher_correct_rate": len(correct) / len(rows) if rows else 0.0, } def main(argv: Sequence[str] | None = None) -> int: args = parse_args(argv) if args.max_samples < 0 or args.top_k <= 0 or args.log_every <= 0: raise ValueError("max-samples must be non-negative and top-k/log-every positive") providers = [value.strip() for value in str(args.providers).split(",") if value.strip()] if providers != ["visual_facts_deplot"]: raise ValueError("this gold-blind diagnostic currently requires visual_facts_deplot only") raw_rows = _load_rows(args.input, human_only=bool(args.human_only)) gate_manifest = read_manifest(default_manifest_path(args.teacher_gate_cache)) if int(gate_manifest.get("record_count", -1)) == len(raw_rows): gate_records, gate_manifest = validate_teacher_gate_cache( raw_rows, cache_path=args.teacher_gate_cache, require_validation_calibration=False, ) else: gate_records = read_jsonl(args.teacher_gate_cache) if int(gate_manifest.get("source_row_count", -1)) != len(raw_rows): raise ValueError("subset teacher-gate manifest has the wrong source row count") manifest_indices = gate_manifest.get("selected_source_indices") record_indices = [int(record.get("source_index", -1)) for record in gate_records] if not isinstance(manifest_indices, list) or record_indices != [ int(index) for index in manifest_indices ]: raise ValueError("subset teacher-gate source indices do not match its manifest") for record in gate_records: source_index = int(record["source_index"]) if source_index < 0 or source_index >= len(raw_rows): raise ValueError("subset teacher-gate source index is outside the dataset") expected_fingerprint = chartqa_sample_fingerprint(raw_rows[source_index]) if str(record.get("sample_fingerprint") or "") != expected_fingerprint: raise ValueError( f"subset teacher-gate fingerprint mismatch at source row {source_index}" ) gate_by_source_index = { int(record["source_index"]): record for record in gate_records } gate_teacher = str(Path(str(gate_manifest.get("teacher_model") or "")).resolve()) requested_teacher = str(Path(resolve_model_path(args.teacher_model)).resolve()) if gate_teacher != requested_teacher: raise ValueError( "teacher-gate cache model does not match diagnostic teacher: " f"cache={gate_teacher}, requested={requested_teacher}" ) torch, rank, world_size, local_rank = distributed_context() selected = _selected_indices(len(raw_rows), int(args.max_samples), args.sample_strategy) missing_gate_rows = [index for index in selected if index not in gate_by_source_index] if missing_gate_rows: raise ValueError( "teacher-gate cache does not cover selected diagnostic rows: " f"count={len(missing_gate_rows)}, examples={missing_gate_rows[:8]}" ) local_indices = selected[rank::world_size] student_model, student_processor, student_model_path, student_processor_path = ( load_model_processor( args.student_model, args.student_processor, torch=torch, local_rank=local_rank, attn_implementation=args.attn_implementation, ) ) same_model = str(Path(student_model_path).resolve()) == requested_teacher if same_model and (not args.teacher_processor or str(Path(student_processor_path).resolve()) == str(Path(resolve_model_path(args.teacher_processor or args.teacher_model)).resolve())): teacher_model = student_model teacher_processor = student_processor teacher_model_path = student_model_path teacher_processor_path = student_processor_path else: teacher_model, teacher_processor, teacher_model_path, teacher_processor_path = ( load_model_processor( args.teacher_model, args.teacher_processor, torch=torch, local_rank=local_rank, attn_implementation=args.attn_implementation, ) ) if len(student_processor.tokenizer) != len(teacher_processor.tokenizer): raise ValueError("student and teacher tokenizers have different vocabulary sizes") records: list[dict[str, Any]] = [] started = time.monotonic() for local_position, source_index in enumerate(local_indices, start=1): sample = _prepare_sample(raw_rows[source_index]) teacher_job = _build_job( sample, providers=providers, source_index=source_index, ) from opsd_utils.privileged.image_utils import load_rgb student_image = load_rgb(sample.get("image")) if student_image is None: raise ValueError(f"could not load student image for source row {source_index}") student_prompt = prepared_prompt( student_processor, student_model, str(sample["prompt"]), [student_image], ) teacher_prompt = prepared_prompt( teacher_processor, teacher_model, str(teacher_job["prompt"]), list(teacher_job["images"]), ) gold = chartqa_reference_answer(sample) student_completion = completion_ids_for_gold( student_processor.tokenizer, gold, torch, student_prompt["prompt_ids"].device, ) teacher_completion = completion_ids_for_gold( teacher_processor.tokenizer, gold, torch, teacher_prompt["prompt_ids"].device, ) if student_completion.detach().cpu().tolist() != teacher_completion.detach().cpu().tolist(): raise ValueError("student and teacher tokenize the gold continuation differently") with torch.inference_mode(): student_logits = score_completion( torch=torch, model=student_model, processor=student_processor, prompt=student_prompt, completion_ids=student_completion, ) teacher_logits = score_completion( torch=torch, model=teacher_model, processor=teacher_processor, prompt=teacher_prompt, completion_ids=teacher_completion, ) metric_values = pair_metrics( torch=torch, student_logits=student_logits, teacher_logits=teacher_logits, completion_ids=student_completion, tokenizer=student_processor.tokenizer, top_k=int(args.top_k), ) gate = gate_by_source_index[source_index] records.append( { "source_index": int(source_index), "teacher_correct": bool(gate.get("teacher_correct") is True), "teacher_output": str(gate.get("teacher_output") or ""), "gold_answer": gold, **metric_values, } ) if local_position == len(local_indices) or local_position % int(args.log_every) == 0: elapsed = time.monotonic() - started print( f"[teacher-alignment][rank={rank}/{world_size}] " f"{local_position}/{len(local_indices)} elapsed={elapsed:.1f}s " f"rate={local_position / max(elapsed, 1e-9):.3f}/s", flush=True, ) output = Path(args.output).resolve() parts = Path(f"{output}.parts") part = parts / f"world{world_size:03d}-rank{rank:03d}.jsonl" atomic_jsonl(part, records) barrier(world_size) if rank == 0: merged: list[dict[str, Any]] = [] for part_rank in range(world_size): merged.extend( json.loads(line) for line in ( parts / f"world{world_size:03d}-rank{part_rank:03d}.jsonl" ).read_text(encoding="utf-8").splitlines() if line.strip() ) merged.sort(key=lambda row: int(row["source_index"])) if [int(row["source_index"]) for row in merged] != selected: raise ValueError("merged alignment rows are incomplete or duplicated") atomic_jsonl(output, merged) summary = { "input": str(Path(args.input).resolve()), "output": str(output), "student_model": str(Path(student_model_path).resolve()), "teacher_model": str(Path(teacher_model_path).resolve()), "student_processor": str(Path(student_processor_path).resolve()), "teacher_processor": str(Path(teacher_processor_path).resolve()), "teacher_gate_cache": str(Path(args.teacher_gate_cache).resolve()), "teacher_gate_manifest": str( Path(default_manifest_path(args.teacher_gate_cache)).resolve() ), "providers": providers, "top_k": int(args.top_k), "world_size": int(world_size), "selected_samples": len(selected), "gold_usage": "teacher_forced_diagnostics_only_never_prompt_or_routing", "summary": summarize(merged), } atomic_json(Path(f"{output}.summary.json"), summary) print( "[teacher-alignment][complete] " + json.dumps(summary, ensure_ascii=False, sort_keys=True), flush=True, ) barrier(world_size) return 0 if __name__ == "__main__": raise SystemExit(main())