from __future__ import annotations import argparse import json import math import random import signal import shutil import time from contextlib import nullcontext from pathlib import Path from typing import Any import numpy as np import torch import torch.nn.functional as F from accelerate import Accelerator from accelerate.utils import DistributedDataParallelKwargs, set_seed from peft import LoraConfig, PeftModel, get_peft_model from torch import nn from torch.utils.data import DataLoader, Dataset from transformers import AutoConfig, AutoModelForImageTextToText, AutoProcessor, get_scheduler from selectground import PROMPT RECIPES = { "8b": { "base": "Qwen/Qwen3-VL-8B-Instruct", "revision": "0c351dd01ed87e9c1b53cbc748cba10e6187ff3b", "data": "ruotian/ContrastGround", "steps": 135, "gpus": 2, "accumulation": 64, "learning_rate": 5e-5, }, "30b": { "base": "Qwen/Qwen3-VL-30B-A3B-Instruct", "revision": "9c4b90e1e4ba969fd3b5378b57d966d725f1b86c", "data": "ruotian/ContrastGround", "steps": 200, "gpus": 4, "accumulation": 4, "learning_rate": 4e-5, }, } LAYERS = list(range(18, 24)) SEED = 20260625 PREEMPT_REQUESTED = False def request_preemption(_signum: int, _frame: Any) -> None: global PREEMPT_REQUESTED PREEMPT_REQUESTED = True def replace_with_retry(source: Path, destination: Path, attempts: int = 5) -> None: for attempt in range(attempts): try: source.replace(destination) return except OSError: if attempt + 1 == attempts: raise time.sleep(2 ** attempt) def _value(obj: Any, name: str, default: Any = None) -> Any: return obj.get(name, default) if isinstance(obj, dict) else getattr(obj, name, default) def _find_config(model: Any) -> Any: config = getattr(model, "config", None) if config is None or _value(config, "vision_config") is None: raise ValueError("Could not find the Qwen3-VL model config") return config def _repeat_key_value_heads(key_states: torch.Tensor, groups: int) -> torch.Tensor: if groups == 1: return key_states batch, heads, sequence, head_dim = key_states.shape return ( key_states[:, :, None, :, :] .expand(batch, heads, groups, sequence, head_dim) .reshape(batch, heads * groups, sequence, head_dim) ) def _attention_logits( attention: Any, hidden_states: torch.Tensor, query_position: int, visual_positions: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None, ) -> torch.Tensor: if position_embeddings is None: raise RuntimeError("Qwen3-VL semantic logits require position_embeddings") head_dim = int(attention.head_dim) hidden_shape = (*hidden_states.shape[:-1], -1, head_dim) query_states = attention.q_norm(attention.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) key_states = attention.k_norm(attention.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) from transformers.models.qwen3_vl.modeling_qwen3_vl import apply_rotary_pos_emb query_states, key_states = apply_rotary_pos_emb(query_states, key_states, *position_embeddings) key_states = _repeat_key_value_heads(key_states, int(getattr(attention, "num_key_value_groups", 1))) positions = visual_positions.to(device=hidden_states.device, dtype=torch.long) query = query_states[:, :, int(query_position), :] visual_keys = key_states.index_select(2, positions) logits = (query.unsqueeze(2) * visual_keys).sum(dim=-1) * float(getattr(attention, "scaling", 1.0)) return logits.squeeze(0) def load_visual_merger(model: Any, checkpoint: Path) -> None: merger_path = checkpoint / "visual_merger.pt" if not merger_path.is_file(): raise FileNotFoundError(f"Missing visual merger checkpoint: {merger_path}") merger = torch.load(merger_path, map_location="cpu", weights_only=False) parameters = dict(model.named_parameters()) state = merger.get("state_dict", merger) missing = sorted(set(state) - set(parameters)) if missing: raise KeyError(f"Visual merger parameters missing from model: {missing[:3]}") with torch.no_grad(): for name, value in state.items(): parameters[name].copy_(value.to(parameters[name].device, parameters[name].dtype)) class Rows(Dataset): def __init__(self, rows: list[dict[str, Any]], root: Path) -> None: self.rows, self.root = rows, root def __len__(self) -> int: return len(self.rows) def __getitem__(self, index: int) -> dict[str, Any]: return {**self.rows[index], "image": str(self.root / self.rows[index]["image"])} class HeadSelector(nn.Module): def __init__(self, heads: int) -> None: super().__init__() self.layer_head_weights = nn.Parameter(torch.zeros(len(LAYERS), heads)) def forward(self, values: list[torch.Tensor]) -> torch.Tensor: stacked = torch.stack([value.float() for value in values]) weights = self.layer_head_weights.flatten().softmax(0).view_as(self.layer_head_weights).to(stacked.device) return (stacked * weights[:, :, None]).sum(dim=(0, 1)) class Attention: def __init__(self, model: Any, query: int, visual: torch.Tensor) -> None: self.model, self.query, self.visual = model, query, visual self.values: dict[int, torch.Tensor] = {} self.handles: list[Any] = [] def __enter__(self) -> "Attention": for module in self.model.modules(): layer = getattr(module, "layer_idx", None) if layer in LAYERS and hasattr(module, "q_proj"): self.handles.append(module.register_forward_hook(self._hook(int(layer)), with_kwargs=True)) return self def __exit__(self, *_: Any) -> None: for handle in self.handles: handle.remove() def _hook(self, layer: int): def hook(module: Any, args: tuple[Any, ...], kwargs: dict[str, Any], output: Any) -> None: hidden = kwargs.get("hidden_states", args[0] if args else None) if hidden is not None and hidden.shape[1] > self.query: self.values[layer] = _attention_logits( module, hidden, self.query, self.visual, kwargs["position_embeddings"] ) return hook def ordered(self) -> list[torch.Tensor]: return [self.values[layer] for layer in LAYERS] def read_rows(path: Path) -> list[dict[str, Any]]: return [json.loads(line) for line in path.read_text().splitlines() if line.strip()] def stage_files(stage: str) -> tuple[str, str]: if stage == "main": return "train_pairs.jsonl", "train_replay.jsonl" if stage == "refinement": return "refinement_pairs.jsonl", "refinement_replay.jsonl" raise ValueError(f"Unknown training stage: {stage}") def loader(rows: list[dict[str, Any]], root: Path, seed: int) -> DataLoader: return DataLoader( Rows(rows, root), batch_size=1, shuffle=True, collate_fn=lambda batch: batch[0], generator=torch.Generator().manual_seed(seed), ) def next_row(data_loader: DataLoader, iterator: Any): try: return next(iterator), iterator except StopIteration: iterator = iter(data_loader) return next(iterator), iterator def encode(processor: Any, row: dict[str, Any], device: torch.device): from qwen_vl_utils import process_vision_info user = { "role": "user", "content": [ {"type": "image", "image": row["image"]}, {"type": "text", "text": PROMPT.format(instruction=row["instruction"])}, ], } prompt = [user] full = [user, {"role": "assistant", "content": [{"type": "text", "text": row["response"]}]}] def process(messages: list[dict[str, Any]], generation_prompt: bool): text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=generation_prompt) images, videos = process_vision_info(messages) kwargs = {"text": [text], "images": images, "padding": True, "return_tensors": "pt"} if videos is not None: kwargs["videos"] = videos return processor(**kwargs).to(device) inputs = process(full, False) prompt_length = int(process(prompt, True)["attention_mask"].sum()) labels = inputs["input_ids"].clone() labels[:, :prompt_length] = -100 return inputs, labels, prompt_length - 1 def coordinate_loss(logits: torch.Tensor, input_ids: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: start = input_ids.shape[1] - logits.shape[1] targets = input_ids[:, start + 1 :] mask = labels[:, start + 1 :].ne(-100) token_logps = logits[:, :-1].float().log_softmax(-1).gather(-1, targets.unsqueeze(-1)).squeeze(-1) return -(token_logps * mask).sum().to(logits.dtype) / mask.sum() def coordinate_weight(row: dict[str, Any], ground_weight: float) -> float: component = str(row.get("source_ref", {}).get("component") or "") return ground_weight if component.startswith("ground_") else 1.0 def box_mask(box: list[float], row: dict[str, Any], grid: torch.Tensor, config: Any) -> torch.Tensor: vision = _value(config, "vision_config") patch, merge = int(_value(vision, "patch_size", 16)), int(_value(vision, "spatial_merge_size", 2)) grid = grid.detach().cpu().long() height, width = int(grid[1]) // merge, int(grid[2]) // merge resized_width, resized_height = int(grid[2]) * patch, int(grid[1]) * patch x1, y1, x2, y2 = box left, right = sorted((x1 / row["image_width"] * resized_width, x2 / row["image_width"] * resized_width)) top, bottom = sorted((y1 / row["image_height"] * resized_height, y2 / row["image_height"] * resized_height)) rows = torch.arange(height)[:, None] columns = torch.arange(width)[None, :] return ( (left < (columns + 1) * resized_width / width) & (right > columns * resized_width / width) & (top < (rows + 1) * resized_height / height) & (bottom > rows * resized_height / height) ).flatten() def region_score(scores: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: selected = scores[mask.to(scores.device)] return torch.logsumexp(selected.float(), 0) - math.log(selected.numel()) def selection_loss( scores: torch.Tensor, row: dict[str, Any], grid: torch.Tensor, config: Any, margin: float, pair_weight: float, ) -> torch.Tensor: target = box_mask(row["target_bbox"], row, grid, config) distractor = box_mask(row["distractor_bbox"], row, grid, config) overlap = target & distractor target, distractor = target & ~overlap, distractor & ~overlap if not target.any() or not distractor.any(): return scores.sum() * 0 target_score, distractor_score = region_score(scores, target), region_score(scores, distractor) candidates = [target, distractor] extras = [] for index, box in enumerate(row.get("candidate_bboxes", [])): mask = box_mask(box, row, grid, config) if mask.any() and not (mask & target).any() and not (mask & distractor).any(): extras.append((float(region_score(scores, mask).detach()), -index, mask)) extras.sort(reverse=True, key=lambda item: item[:2]) candidates.extend(item[2] for item in extras[:3]) listwise = torch.logsumexp(torch.stack([region_score(scores, mask) for mask in candidates]), 0) - target_score pair = F.softplus(scores.new_tensor(margin) - target_score + distractor_score) return listwise + pair_weight * pair def warmup_cosine(step: int, warmup: int, total: int) -> float: if step < warmup: return step / warmup if step >= total: return 0.0 return .5 * (1 + math.cos(math.pi * (step - warmup) / (total - warmup))) def scheduler_for(optimizer: torch.optim.Optimizer, warmup_steps: int, training_steps: int): return get_scheduler( "cosine", optimizer=optimizer, num_warmup_steps=warmup_steps, num_training_steps=training_steps, ) def paper_scheduler_for( optimizer: torch.optim.Optimizer, phase_a_steps: int, phase_a_warmup_steps: int, phase_a_scheduler_steps: int, phase_b_warmup_steps: int, phase_b_scheduler_steps: int, phase_b_learning_rate: float, phase_b_selector_learning_rate: float, ): base_lrs = [group["lr"] for group in optimizer.param_groups] target_lrs = [phase_b_learning_rate, phase_b_selector_learning_rate] functions = [] for base, target in zip(base_lrs, target_lrs): def schedule(step: int, base=base, target=target): if step < phase_a_steps + 1: return warmup_cosine(step, phase_a_warmup_steps, phase_a_scheduler_steps) return target / base * warmup_cosine( step - phase_a_steps, phase_b_warmup_steps, phase_b_scheduler_steps, ) functions.append(schedule) return torch.optim.lr_scheduler.LambdaLR(optimizer, functions) def save( accelerator: Accelerator, model: Any, selector: Any, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LRScheduler, processor: Any, output: Path, revision: str, completed: int, stage: str, micro_step: int, ) -> None: atomic = output.name.startswith("step-") target = output.with_name(f"{output.name}.incomplete") if atomic else output accelerator.wait_for_everyone() if accelerator.is_main_process: if atomic and target.exists(): shutil.rmtree(target) target.mkdir(parents=True, exist_ok=True) (target / "checkpoint_complete").unlink(missing_ok=True) unwrapped = accelerator.unwrap_model(model) unwrapped.save_pretrained(target, safe_serialization=True) config_path = target / "adapter_config.json" config = json.loads(config_path.read_text()) config["revision"] = revision config_path.write_text(json.dumps(config, indent=2) + "\n") merger = {name: value.detach().cpu() for name, value in unwrapped.named_parameters() if ".visual.merger." in f".{name}"} torch.save({"state_dict": merger}, target / "visual_merger.pt") head = accelerator.unwrap_model(selector) torch.save({"layers": LAYERS, "layer_head_weights": head.layer_head_weights.detach().cpu()}, target / "selection_head.pt") torch.save( { "completed": completed, "micro_step": micro_step, "stage": stage, "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(), }, target / "training_state.pt", ) processor.save_pretrained(target) accelerator.wait_for_everyone() rng = { "python": random.getstate(), "numpy": np.random.get_state(), "torch": torch.get_rng_state(), "cuda": torch.cuda.get_rng_state_all(), } torch.save(rng, target / f"rng_state_rank_{accelerator.process_index}.pt") accelerator.wait_for_everyone() if accelerator.is_main_process: (target / "checkpoint_complete").write_text("complete\n", encoding="utf-8") accelerator.wait_for_everyone() if accelerator.is_main_process and atomic: if output.exists(): shutil.rmtree(output) replace_with_retry(target, output) accelerator.wait_for_everyone() if accelerator.is_main_process and atomic: for previous in output.parent.glob("step-*"): if previous != output and (previous / "checkpoint_complete").is_file(): shutil.rmtree(previous) accelerator.wait_for_everyone() def restore( checkpoint: Path, optimizer: torch.optim.Optimizer, scheduler: Any, accelerator: Accelerator, stage: str, accumulation: int, ) -> tuple[int, int]: state = torch.load(checkpoint / "training_state.pt", map_location="cpu", weights_only=False) optimizer.load_state_dict(state["optimizer"]) scheduler.load_state_dict(state["scheduler"]) rng = torch.load( checkpoint / f"rng_state_rank_{accelerator.process_index}.pt", map_location="cpu", weights_only=False, ) random.setstate(rng["python"]) np.random.set_state(rng["numpy"]) torch.set_rng_state(rng["torch"]) torch.cuda.set_rng_state_all(rng["cuda"]) completed = int(state["completed"]) saved_stage = state.get("stage") micro_step = int(state.get("micro_step", completed * accumulation)) if saved_stage == stage else 0 return completed, micro_step def main() -> None: signal.signal(signal.SIGUSR1, request_preemption) parser = argparse.ArgumentParser(description="Train SelectGround on local paired and replay JSONL files.") parser.add_argument("--model", choices=("8b", "30b"), default="8b") parser.add_argument("--data", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--checkpoint", type=Path) parser.add_argument("--initialize-from", type=Path) parser.add_argument("--stage", choices=("main", "refinement"), default="main") parser.add_argument("--pairs-file", type=Path) parser.add_argument("--replay-file", type=Path) parser.add_argument("--steps", type=int, required=True) parser.add_argument("--gpus", type=int, default=4) parser.add_argument("--accumulation", type=int, default=32) parser.add_argument("--learning-rate", type=float, default=5e-5) parser.add_argument("--selector-learning-rate", type=float, default=1e-4) parser.add_argument("--aux-weight", type=float, default=0.1) parser.add_argument("--ground-coordinate-weight", type=float, default=1.0) parser.add_argument("--margin", type=float, default=0.3) parser.add_argument("--pair-weight", type=float, default=0.5) parser.add_argument("--warmup-steps", type=int, default=10) parser.add_argument("--scheduler-steps", type=int) parser.add_argument("--paper-two-stage", action="store_true") parser.add_argument("--phase-a-steps", type=int) parser.add_argument("--phase-b-warmup-steps", type=int, default=10) parser.add_argument("--phase-b-scheduler-steps", type=int, default=25) parser.add_argument("--phase-b-learning-rate", type=float, default=1e-6) parser.add_argument("--phase-b-selector-learning-rate", type=float, default=1e-4) parser.add_argument("--holdout-fraction", type=float, default=0.0) parser.add_argument("--max-pixels", type=int, default=8847360) parser.add_argument("--seed", type=int, default=SEED) parser.add_argument("--save-every", type=int, default=25) args = parser.parse_args() if args.checkpoint is not None and args.initialize_from is not None: raise ValueError("Use only one of --checkpoint and --initialize-from") if (args.pairs_file is None) != (args.replay_file is None): raise ValueError("--pairs-file and --replay-file must be used together") if args.paper_two_stage and args.phase_a_steps is None: raise ValueError("--paper-two-stage requires --phase-a-steps") recipe = RECIPES[args.model] accelerator = Accelerator( gradient_accumulation_steps=args.accumulation, kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=False)], ) if accelerator.num_processes != args.gpus: raise ValueError(f"Expected {args.gpus} processes, got {accelerator.num_processes}") set_seed(args.seed + accelerator.process_index) data = args.data pair_file, replay_file = stage_files(args.stage) pairs_path = args.pairs_file or data / "data" / pair_file replay_path = args.replay_file or data / "data" / replay_file pairs = read_rows(pairs_path) replay = read_rows(replay_path) if args.stage == "main" and args.holdout_fraction > 0: random.Random(args.seed).shuffle(pairs) pairs = pairs[max(1, round(args.holdout_fraction * len(pairs))) :] seed_offset = 1000 if args.stage == "main" else 3000 pair_seed, replay_seed = args.seed + seed_offset + 1, args.seed + seed_offset + 1001 pair_loader, replay_loader = loader(pairs, data, pair_seed), loader(replay, data, replay_seed) processor = AutoProcessor.from_pretrained( recipe["base"], revision=recipe["revision"], min_pixels=3136, max_pixels=args.max_pixels ) base_config = AutoConfig.from_pretrained(recipe["base"], revision=recipe["revision"]) model = AutoModelForImageTextToText.from_pretrained( recipe["base"], revision=recipe["revision"], config=base_config, dtype=torch.bfloat16, attn_implementation="sdpa" ) source_checkpoint = args.checkpoint or args.initialize_from if source_checkpoint is None: model = get_peft_model(model, LoraConfig( r=64, lora_alpha=128, lora_dropout=.05, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], task_type="CAUSAL_LM", )) else: import transformers.integrations.tensor_parallel as tensor_parallel if not hasattr(tensor_parallel, "EmbeddingParallel"): tensor_parallel.EmbeddingParallel = type("EmbeddingParallel", (), {}) model = PeftModel.from_pretrained(model, source_checkpoint, is_trainable=True) load_visual_merger(model, source_checkpoint) for parameter in model.parameters(): if parameter.requires_grad: parameter.data = parameter.data.to(torch.bfloat16) model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) model.enable_input_require_grads() model.config.use_cache = False config = _find_config(model) heads = int(_value(_value(config, "text_config", config), "num_attention_heads")) selector = HeadSelector(heads) if source_checkpoint is not None: head = torch.load(source_checkpoint / "selection_head.pt", map_location="cpu", weights_only=False) selector.layer_head_weights.data.copy_(head["layer_head_weights"]) optimizer = torch.optim.AdamW([ {"params": [parameter for parameter in model.parameters() if parameter.requires_grad], "lr": args.learning_rate}, {"params": selector.parameters(), "lr": args.selector_learning_rate}, ], weight_decay=0.0) if args.paper_two_stage: scheduler = paper_scheduler_for( optimizer, phase_a_steps=args.phase_a_steps, phase_a_warmup_steps=args.warmup_steps, phase_a_scheduler_steps=args.scheduler_steps or args.steps, phase_b_warmup_steps=args.phase_b_warmup_steps, phase_b_scheduler_steps=args.phase_b_scheduler_steps, phase_b_learning_rate=args.phase_b_learning_rate, phase_b_selector_learning_rate=args.phase_b_selector_learning_rate, ) else: scheduler = scheduler_for(optimizer, args.warmup_steps, args.scheduler_steps or args.steps) model, selector, optimizer, pair_loader, replay_loader = accelerator.prepare( model, selector, optimizer, pair_loader, replay_loader ) model.train() selector.train() iterators = [iter(pair_loader), iter(replay_loader)] completed, micro_step = ( restore(args.checkpoint, optimizer, scheduler, accelerator, args.stage, args.accumulation) if args.checkpoint else (0, 0) ) if micro_step: for skipped in range(micro_step): index = 0 if skipped % 2 == 1 else 1 _, iterators[index] = next_row((pair_loader, replay_loader)[index], iterators[index]) target = args.steps optimizer.zero_grad(set_to_none=True) while completed < target: active_loaders, active_iterators = (pair_loader, replay_loader), iterators competitor_paired = micro_step % 2 == 1 index = 0 if competitor_paired else 1 row, active_iterators[index] = next_row(active_loaders[index], active_iterators[index]) with accelerator.accumulate(model, selector): inputs, labels, query = encode(processor, row, accelerator.device) visual = torch.nonzero(inputs["input_ids"][0] == int(_value(config, "image_token_id")), as_tuple=False).flatten() keep = int(labels.ne(-100).sum()) + 1 context = Attention(model, query, visual) if competitor_paired else nullcontext() with context as attention: output = model(**inputs, use_cache=False, logits_to_keep=keep) if competitor_paired: scores = selector(attention.ordered()) selection_term = selection_loss( scores, row, inputs["image_grid_thw"][0], config, margin=args.margin, pair_weight=args.pair_weight, ) else: selection_term = output.logits.sum() * 0 coord_loss = coordinate_loss(output.logits, inputs["input_ids"], labels) coord_scale = coordinate_weight(row, args.ground_coordinate_weight) loss = coord_scale * coord_loss + args.aux_weight * selection_term accelerator.backward(loss) if accelerator.sync_gradients: accelerator.clip_grad_norm_(list(model.parameters()) + list(selector.parameters()), 1.0) optimizer.step() scheduler.step() optimizer.zero_grad(set_to_none=True) micro_step += 1 if accelerator.sync_gradients: completed += 1 if accelerator.is_main_process: print(f"step={completed} loss={float(loss):.4f} coord={float(coord_loss):.4f} coord_scale={coord_scale:.2f} selection={float(selection_term):.4f}", flush=True) if args.save_every > 0 and completed < target and completed % args.save_every == 0: save( accelerator, model, selector, optimizer, scheduler, processor, args.output / "checkpoints" / f"step-{completed}", recipe["revision"], completed, args.stage, micro_step, ) if PREEMPT_REQUESTED: save( accelerator, model, selector, optimizer, scheduler, processor, args.output / "checkpoints" / f"step-{completed}", recipe["revision"], completed, args.stage, micro_step, ) raise SystemExit(85) save( accelerator, model, selector, optimizer, scheduler, processor, args.output, recipe["revision"], completed, args.stage, micro_step, ) if accelerator.is_main_process: run_config = vars(args) | { "base_model": recipe["base"], "base_revision": recipe["revision"], } (args.output / "run_config.json").write_text( json.dumps(run_config, default=str, indent=2, sort_keys=True) + "\n", encoding="utf-8", ) if __name__ == "__main__": main()