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