#!/usr/bin/env python3 """Prefix-MRL student trainer with ARD. One maximum projection is initialized from Preview d128. A frozen EVIE-8B teacher transfers row/column MaxSim relations and hard-negative margins; the adapter-disabled Preview path anchors d128. """ from __future__ import annotations import argparse import hashlib import json import os import time import sys from pathlib import Path _SHARED = Path(__file__).resolve().parents[2] / "shared" if str(_SHARED) not in sys.path: sys.path.insert(0, str(_SHARED)) import torch from peft import LoraConfig from torch.distributed.elastic.multiprocessing.errors import record from transformers import TrainingArguments, set_seed from data_loader import ( ALLOWED_SOURCES, _default_dataset_cache_dir, build_hardneg_dataset, build_train_dataset, ) from paths import forbid_venv_path from colpali_engine.loss.late_interaction_losses import ColbertLoss, ColbertNegativeCELoss from colpali_engine.loss.ard import ARDLoss from colpali_engine.models import ColQwen3_5, ColQwen3_5Processor from transformers.models.qwen3_5 import Qwen3_5Config from colpali_engine.trainer.colmodel_training import ColModelTraining, ColModelTrainingConfig from colpali_engine.trainer.ard_trainer import ARDTrainer TARGET_MODULES = ( r"(.*(model)(?!.*visual).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj|" r"in_proj_qkv|in_proj_z|in_proj_b|in_proj_a|out_proj).*$)" ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--base-model", default="") parser.add_argument("--data-root", default="./data") parser.add_argument("--output-dir", required=True) parser.add_argument("--sources", nargs="*", default=[], choices=ALLOWED_SOURCES, help="Optional ablation filter on the source column; empty = full corpus.") parser.add_argument("--seed", type=int, default=42) parser.add_argument("--epochs", type=float, default=1.0) parser.add_argument("--max-steps", type=int, default=-1) parser.add_argument("--max-samples-per-source", type=int, default=0) parser.add_argument("--per-device-batch-size", type=int, default=2) parser.add_argument("--grad-accum", type=int, default=1) parser.add_argument("--learning-rate", type=float, default=1.5e-5, help="Warm-start default. Teacher from-scratch uses 4.57e-5.") parser.add_argument("--weight-decay", type=float, default=0.02) parser.add_argument("--warmup-ratio", type=float, default=0.08) parser.add_argument("--max-visual-tokens", type=int, default=1024) parser.add_argument("--col-dim", type=int, default=512, help="custom_text_proj output dim (unused when --head-dims is set).") parser.add_argument("--dataloader-workers", type=int, default=2) parser.add_argument( "--dataloader-prefetch-factor", type=int, default=4, help="Batches prefetched by each DataLoader worker; ignored when workers=0.", ) parser.add_argument("--save-steps", type=int, default=500) parser.add_argument("--logging-steps", type=int, default=10) parser.add_argument("--lora-r", type=int, default=32) parser.add_argument("--lora-alpha", type=int, default=128) parser.add_argument("--lora-dropout", type=float, default=0.197) parser.add_argument("--loss-temperature", type=float, default=0.02) parser.add_argument("--hardneg-in-batch-weight", type=float, default=0.5) parser.add_argument("--resume-from-checkpoint", default="", help="Checkpoint path, or 'latest' to resume the newest output checkpoint.") parser.add_argument("--attn", choices=("flash_attention_2", "sdpa", "eager"), default="flash_attention_2") parser.add_argument("--grad-checkpointing", choices=("on", "off"), default="off") parser.add_argument( "--bidirectional-attention", choices=("on", "off"), default="on", help="on = encoder-ize full-attention layers (ColEmbed V2).", ) parser.add_argument( "--hardneg-root", default="", help="Hardneg output with queries/. corpus/ is optional; without it images load from --data-root.", ) parser.add_argument("--num-hard-negs", type=int, default=2) parser.add_argument( "--use-hardnegatives", choices=("on", "off"), default="on", help="When --hardneg-root is set: on=ColbertNegativeCELoss; off=same rows, in-batch only.", ) parser.add_argument( "--report-to", default="none", help="Comma-separated metric trackers, e.g. wandb,tensorboard.", ) parser.add_argument("--logging-dir", default=None, help="TensorBoard event output directory.") parser.add_argument("--run-name", default=None, help="Run name for the tracker (e.g. wandb).") parser.add_argument( "--head-dims", default="", help="Comma-separated Prefix-MRL widths, e.g. '64,128,256,512,1024,2048'. " "Empty turns off prefixes and the teacher.", ) parser.add_argument( "--anchor-dim", type=int, default=128, help="Preview projection width copied into the Prefix-MRL prefix rows.", ) parser.add_argument( "--teacher-dir", default="", help="Directory of a frozen full-weight teacher (EVIE-8B). " "Empty together with --head-dims is retrieval-only.", ) parser.add_argument( "--teacher-md5", default="", help="Expected md5 of the teacher model.safetensors; verified before training starts.", ) parser.add_argument( "--kd-dims", default="64,128,256,512,1024,2048", help="Prefixes that receive relation distillation.", ) parser.add_argument("--calibration-dim", type=int, default=128) parser.add_argument("--teacher-temperature", type=float, default=0.13) parser.add_argument( "--student-temperatures", default="64:0.13,128:0.13,256:0.13,512:0.13,1024:0.13,2048:0.13", ) parser.add_argument("--relation-weight", type=float, default=1.0) parser.add_argument("--margin-weight", type=float, default=0.25) parser.add_argument("--anchor-weight", type=float, default=0.25) parser.add_argument("--column-weight", type=float, default=1.0) parser.add_argument("--confidence-floor", type=float, default=0.1) parser.add_argument("--teacher-wrong-factor", type=float, default=0.25) parser.add_argument("--head-weights", default="") parser.add_argument("--kd-head-weights", default="") parser.add_argument( "--kd-directions", choices=("both", "row", "column", "none"), default="both", ) parser.add_argument("--kd-include-hardnegs", choices=("on", "off"), default="on") parser.add_argument("--anchor-teacher", choices=("on", "off"), default="on") parser.add_argument("--task-consistent-batches", choices=("on", "off"), default="on") parser.add_argument("--gradient-target-ratio", type=float, default=0.5) parser.add_argument("--gradient-calibration-steps", type=int, default=100) parser.add_argument("--gradient-calibration-interval", type=int, default=10) parser.add_argument("--gradient-scale-min", type=float, default=0.05) parser.add_argument("--gradient-scale-max", type=float, default=20.0) parser.add_argument("--gradient-scale-ema", type=float, default=0.9) parser.add_argument("--gradient-diagnostics", choices=("on", "off"), default="on") parser.add_argument("--gradient-diagnostic-steps", type=int, default=100) parser.add_argument("--gradient-diagnostic-interval", type=int, default=10) parser.add_argument("--head-warmup-steps", type=int, default=100) return parser.parse_args() def resolve_pretrained(spec: str) -> str: path = Path(spec) return str(path.resolve()) if path.exists() else spec def parse_head_dims(spec: str) -> tuple[int, ...]: if not spec.strip(): return () dims = tuple(sorted({int(x) for x in spec.replace(" ", "").split(",") if x})) if not dims or dims[0] <= 0: raise ValueError(f"--head-dims must be positive integers, got {spec!r}") return dims def parse_dim_map(spec: str, allowed: tuple[int, ...], name: str) -> dict[int, float]: if not spec.strip(): return {} values: dict[int, float] = {} for item in spec.replace(" ", "").split(","): if not item: continue try: raw_dim, raw_value = item.split(":", 1) dim, value = int(raw_dim), float(raw_value) except ValueError as exc: raise ValueError(f"{name} expects d:value entries, got {item!r}") from exc if dim not in allowed: raise ValueError(f"{name} dim {dim} is not in {allowed}") values[dim] = value return values def file_md5(path: Path) -> str: digest = hashlib.md5() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(16 * 1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def load_teacher(teacher_dir: Path, expected_md5: str, attn: str, bidirectional: str): """Load the frozen teacher at full bf16 weight, eval + no_grad. Full weights only (merged soup, not base + adapter). Optional md5 check rejects a wrong teacher before training starts. Returns ``(teacher, verified_md5)``. The digest is reused for ``run_config.json`` so every rank does not hash the same file twice. """ weights = teacher_dir / "model.safetensors" shards = sorted(teacher_dir.glob("model-*-of-*.safetensors")) if not weights.is_file() and not shards: raise FileNotFoundError(f"teacher has no model*.safetensors: {teacher_dir}") verified_md5 = None if expected_md5: if not weights.is_file(): raise ValueError( f"--teacher-md5 given but teacher is sharded ({len(shards)} shards); " "cannot verify a single-file md5" ) verified_md5 = file_md5(weights) if verified_md5 != expected_md5: raise ValueError( f"teacher md5 mismatch: expected {expected_md5}, got {verified_md5} ({weights})" ) print(f"[teacher] md5 verified: {verified_md5}") config = Qwen3_5Config.from_pretrained(teacher_dir) if getattr(config, "head_dims", None): raise ValueError("teacher must be single-head; found head_dims in its config") teacher = ColQwen3_5.from_pretrained( teacher_dir, config=config, torch_dtype=torch.bfloat16, attn_implementation=attn, ) try: teacher.rope_deltas = None except AttributeError: pass if bidirectional == "on": teacher.enable_bidirectional_attention() teacher.eval() for p in teacher.parameters(): p.requires_grad_(False) print( f"[teacher] {teacher_dir.name}: dim={teacher.dim} " f"bidir={bidirectional} params={sum(p.numel() for p in teacher.parameters()):,}" ) return teacher, verified_md5 def assert_shared_batch(student, teacher, processor) -> None: checks = { "patch_size": (student.patch_size, teacher.patch_size, processor.image_processor.patch_size), "spatial_merge_size": ( student.spatial_merge_size, teacher.spatial_merge_size, processor.image_processor.merge_size, ), } for name, (s, t, p) in checks.items(): if not (int(s) == int(t) == int(p)): raise ValueError( f"{name} disagrees: student={s} teacher={t} processor={p}" ) for name in ("image_token_id", "video_token_id", "vision_start_token_id", "vocab_size"): s = getattr(student.config, name, None) t = getattr(teacher.config, name, None) if s != t: raise ValueError(f"{name} disagrees: student={s} teacher={t}") print( f"[teacher] shared-batch OK (patch={student.patch_size} " f"merge={student.spatial_merge_size} image_token_id={student.config.image_token_id})" ) def prepare_output(path: Path, resume_requested: bool) -> None: path = forbid_venv_path(path, "output-dir") prepared_by_launcher = os.environ.get("EVIE_OUTPUT_PREPARED") == "1" if path.exists() and any(path.iterdir()) and not resume_requested and not prepared_by_launcher: raise FileExistsError(f"Refusing to overwrite a non-empty output directory: {path}") path.mkdir(parents=True, exist_ok=True) def main() -> None: args = parse_args() head_dims = parse_head_dims(args.head_dims) kd_dims = parse_head_dims(args.kd_dims) if head_dims: if args.anchor_dim not in head_dims: raise ValueError(f"--anchor-dim {args.anchor_dim} must be in --head-dims {head_dims}") if not kd_dims or not set(kd_dims).issubset(head_dims): raise ValueError(f"--kd-dims {kd_dims} must be a non-empty subset of {head_dims}") if args.calibration_dim not in kd_dims: raise ValueError( f"--calibration-dim {args.calibration_dim} must be in --kd-dims {kd_dims}" ) student_temperatures = parse_dim_map( args.student_temperatures, kd_dims, "--student-temperatures" ) head_weights = parse_dim_map(args.head_weights, head_dims, "--head-weights") kd_head_weights = parse_dim_map(args.kd_head_weights, kd_dims, "--kd-head-weights") else: student_temperatures = {} head_weights = {} kd_head_weights = {} teacher_dir = Path(args.teacher_dir).resolve() if args.teacher_dir else None if teacher_dir is not None: if not head_dims: raise ValueError("--teacher-dir requires --head-dims") if not teacher_dir.is_dir(): raise FileNotFoundError(f"teacher dir is missing: {teacher_dir}") if args.margin_weight > 0 and args.kd_include_hardnegs != "on": raise ValueError("--margin-weight > 0 requires --kd-include-hardnegs on") output_dir = Path(args.output_dir).resolve() base_model = resolve_pretrained(args.base_model) data_root = Path(args.data_root).resolve() if not data_root.is_dir(): raise FileNotFoundError(f"Data root is missing: {data_root}") if args.hardneg_root: hn = Path(args.hardneg_root).resolve() subdir = os.environ.get("HARDNEG_SUBDIR", "allpos") if not (hn / subdir).is_dir(): raise FileNotFoundError(f"hardneg-root needs {subdir}/: {hn}") prepare_output(output_dir, resume_requested=bool(args.resume_from_checkpoint)) os.environ.setdefault("HF_DATASETS_CACHE", str(_default_dataset_cache_dir())) set_seed(args.seed) print("== building train-only query/image pairs ==") use_hardneg = bool(args.hardneg_root) and args.use_hardnegatives == "on" if teacher_dir is not None and args.margin_weight > 0 and not use_hardneg: raise ValueError("--margin-weight > 0 requires --hardneg-root and --use-hardnegatives on") if args.hardneg_root: train_dataset = build_hardneg_dataset( hardneg_root=args.hardneg_root, data_root=data_root, num_negatives=args.num_hard_negs, max_samples=args.max_samples_per_source, use_negatives=use_hardneg, ) else: train_dataset = build_train_dataset( data_root=data_root, sources=args.sources, max_samples_per_source=args.max_samples_per_source, ) print("== loading processor and bf16 base model ==") processor = ColQwen3_5Processor.from_pretrained( base_model, max_num_visual_tokens=args.max_visual_tokens, ) # Load Preview in its original d128 shape first, then expand to Prefix-MRL. # Setting head_dims before from_pretrained would shape-mismatch and drop # the deployed 128-d projection. model_config = Qwen3_5Config.from_pretrained(base_model) if not head_dims: model_config.dim = args.col_dim model = ColQwen3_5.from_pretrained( base_model, config=model_config, torch_dtype=torch.bfloat16, attn_implementation=args.attn, ) if head_dims: model.enable_prefix_mrl(head_dims, anchor_dim=args.anchor_dim) print( f"[model] prefix MRL = Linear({model.hidden_size_for_heads}->{max(head_dims)}), " f"dims={list(head_dims)}, copied Preview rows [0:{args.anchor_dim})" ) else: print(f"[model] custom_text_proj = Linear(-> {args.col_dim}), fresh full-rank head") try: model.rope_deltas = None except AttributeError: pass if args.bidirectional_attention == "on": model.enable_bidirectional_attention() print("[model] bidirectional attention enabled on full-attention layers") use_gc = args.grad_checkpointing == "on" if use_gc: model.enable_input_require_grads() report_to = [item.strip() for item in args.report_to.split(",") if item.strip()] if not report_to or report_to == ["none"]: report_to = [] training_args = TrainingArguments( output_dir=str(output_dir), num_train_epochs=args.epochs, max_steps=args.max_steps, per_device_train_batch_size=args.per_device_batch_size, gradient_accumulation_steps=args.grad_accum, gradient_checkpointing=use_gc, gradient_checkpointing_kwargs={"use_reentrant": False} if use_gc else None, dataloader_num_workers=args.dataloader_workers, dataloader_pin_memory=True, dataloader_persistent_workers=args.dataloader_workers > 0, dataloader_prefetch_factor=args.dataloader_prefetch_factor if args.dataloader_workers > 0 else None, dataloader_drop_last=True, save_steps=args.save_steps, save_total_limit=2, logging_steps=args.logging_steps, learning_rate=args.learning_rate, lr_scheduler_type="cosine", warmup_ratio=args.warmup_ratio, weight_decay=args.weight_decay, bf16=True, seed=args.seed, data_seed=args.seed, ddp_find_unused_parameters=False, report_to=report_to, logging_dir=args.logging_dir, run_name=args.run_name, ) head_modules = model.head_module_names() lora = LoraConfig( r=args.lora_r, lora_alpha=args.lora_alpha, lora_dropout=args.lora_dropout, init_lora_weights="gaussian", bias="none", task_type="FEATURE_EXTRACTION", target_modules=TARGET_MODULES, modules_to_save=head_modules, ) print(f"[lora] modules_to_save={head_modules}") teacher = None teacher_md5 = None trainer_cls = None trainer_kwargs: dict = {} if head_dims: if teacher_dir is not None: print(f"== loading frozen teacher: {teacher_dir} ==") teacher, teacher_md5 = load_teacher( teacher_dir, args.teacher_md5, args.attn, args.bidirectional_attention ) assert_shared_batch(model, teacher, processor) else: print("== no teacher: Prefix-MRL retrieval-only ==") print( f"== Loss: ARD (prefixes={list(head_dims)}, kd_dims={list(kd_dims)}, " f"tau_T={args.teacher_temperature}, tau_S={student_temperatures}, " f"relation={args.relation_weight}, margin={args.margin_weight}, " f"anchor={args.anchor_weight}, dirs={args.kd_directions}, " f"global_hardnegs={args.kd_include_hardnegs}, " f"task_consistent={args.task_consistent_batches}) ==" ) loss_func = ARDLoss( head_dims=head_dims, kd_dims=kd_dims, temperature=args.loss_temperature, teacher_temperature=args.teacher_temperature, student_temperatures=student_temperatures, relation_weight=args.relation_weight if teacher is not None else 0.0, margin_weight=args.margin_weight if teacher is not None else 0.0, anchor_weight=args.anchor_weight if args.anchor_teacher == "on" else 0.0, column_weight=args.column_weight, in_batch_term_weight=args.hardneg_in_batch_weight, kd_directions=args.kd_directions if teacher is not None else "none", confidence_floor=args.confidence_floor, teacher_wrong_factor=args.teacher_wrong_factor, head_weights=head_weights, kd_head_weights=kd_head_weights, pos_aware_negative_filtering=True, ) trainer_cls = ARDTrainer trainer_kwargs = { "teacher_model": teacher, "head_dims": head_dims, "anchor_dim": args.anchor_dim, "use_anchor_teacher": args.anchor_teacher == "on", "kd_include_hardnegs": args.kd_include_hardnegs == "on", "task_consistent_batches": args.task_consistent_batches == "on", "gradient_target_ratio": args.gradient_target_ratio, "gradient_calibration_steps": args.gradient_calibration_steps, "gradient_calibration_interval": args.gradient_calibration_interval, "gradient_scale_min": args.gradient_scale_min, "gradient_scale_max": args.gradient_scale_max, "gradient_scale_ema": args.gradient_scale_ema, "calibration_dim": args.calibration_dim, "gradient_diagnostics": args.gradient_diagnostics == "on", "gradient_diagnostic_steps": args.gradient_diagnostic_steps, "gradient_diagnostic_interval": args.gradient_diagnostic_interval, "head_warmup_steps": args.head_warmup_steps, } elif use_hardneg: judged = bool(getattr(train_dataset, "judged_pos", False)) all_pos = bool(getattr(train_dataset, "all_pos", False)) mode = ", judged-pos" if judged else (", all-pos" if all_pos else "") print( f"== Loss: ColbertNegativeCELoss (hard_negs={args.num_hard_negs}{mode}) ==" ) loss_func = ColbertNegativeCELoss( temperature=args.loss_temperature, normalize_scores=True, use_smooth_max=False, pos_aware_negative_filtering=True, in_batch_term_weight=args.hardneg_in_batch_weight, ) else: print("== Loss: ColbertLoss (in-batch only) ==") loss_func = ColbertLoss( temperature=args.loss_temperature, normalize_scores=True, use_smooth_max=False, ) trainer = ColModelTraining( ColModelTrainingConfig( output_dir=str(output_dir), processor=processor, model=model, train_dataset=train_dataset, eval_dataset=None, run_eval=False, loss_func=loss_func, tr_args=training_args, peft_config=lora, trainer_cls=trainer_cls, trainer_kwargs=trainer_kwargs, ) ) print("== training ==") started = time.time() resume = args.resume_from_checkpoint or None if resume == "latest": checkpoints = sorted( output_dir.glob("checkpoint-*"), key=lambda p: int(p.name.rsplit("-", 1)[-1]), ) if not checkpoints: raise FileNotFoundError(f"no checkpoint-* found under {output_dir}") resume = str(checkpoints[-1]) training_args.resume_from_checkpoint = resume trainer.train() trainer.save() world_size = int(os.environ.get("WORLD_SIZE", 1)) try: n_samples = len(train_dataset) except TypeError: n_samples = None config = vars(args) | { "base_model": str(base_model), "data_root": str(data_root), "framework": "evie-ard-prefix-mrl", "hardneg_subdir": ( os.environ.get("HARDNEG_SUBDIR", "allpos") if args.hardneg_root else None ), "judged_pos": bool(getattr(train_dataset, "judged_pos", False)), "all_pos": bool(getattr(train_dataset, "all_pos", False)), "head_dims": list(head_dims) if head_dims else None, "mrl_prefix": bool(head_dims), "anchor_dim": args.anchor_dim if head_dims else None, "kd_dims": list(kd_dims) if head_dims else None, "calibration_dim": args.calibration_dim if head_dims else None, "student_temperatures_parsed": ( loss_func.student_temperatures if head_dims else None ), "col_dim": max(head_dims) if head_dims else args.col_dim, "teacher_dir": str(teacher_dir) if teacher_dir else None, "teacher_md5": teacher_md5, "teacher_col_dim": int(teacher.dim) if teacher is not None else None, "kd_active": teacher is not None, "world_size": world_size, "effective_batch_size": args.per_device_batch_size * args.grad_accum * world_size, "train_samples": n_samples, "train_runtime_seconds": round(time.time() - started, 1), "total_params": sum(p.numel() for p in trainer.model.parameters()), "trainable_params": sum( p.numel() for p in trainer.model.parameters() if p.requires_grad ), } if int(os.environ.get("RANK", "0")) == 0: (output_dir / "run_config.json").write_text( json.dumps(config, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) if torch.distributed.is_available() and torch.distributed.is_initialized(): torch.distributed.barrier() print(f"== complete: {output_dir} ==") @record def run_with_error_recording() -> None: """Persist the original failing DDP rank's traceback for torchrun.""" try: main() finally: if torch.distributed.is_available() and torch.distributed.is_initialized(): torch.distributed.destroy_process_group() if __name__ == "__main__": run_with_error_recording()