Replace with ContrastGround-trained SelectGround-8B

Plain-base single-stage SFT plus auxiliary selection loss checkpoint, exact recipe, checksums, and evaluation code.

README.md

@@ -1,27 +1,93 @@
 ---
-license: apache-2.0
-library_name: peft
 base_model: Qwen/Qwen3-VL-8B-Instruct
+library_name: peft
 pipeline_tag: image-text-to-text
+license: apache-2.0
+datasets:
+- ruotian/ContrastGround
 tags:
 - gui-grounding
 - computer-use
 - qwen3-vl
-datasets:
-- ruotian/ClickContrast
+- lora
+- selectground
 ---
 
 # SelectGround-8B
 
-SelectGround-8B is a GUI grounding model built from Qwen3-VL-8B-Instruct with standard SFT and a target–distractor selection loss. It returns one click point as `[x, y]` in normalized `[0, 1000]` coordinates.
+SelectGround-8B maps a screenshot and a natural-language instruction to one
+click point. This release replaces the earlier ClickContrast-trained checkpoint
+with a checkpoint trained from the pinned plain Qwen3-VL-8B-Instruct base on
+the `selectground-8b` configuration of
+[`ruotian/ContrastGround`](https://huggingface.co/datasets/ruotian/ContrastGround).
+It is a single directly trained LoRA checkpoint, not a model soup or weight
+aggregate.
+
+## Direct grounding results
+
+| Benchmark | Accuracy | Semantic error |
+|---|---:|---:|
+| ScreenSpot-Pro | 65.09 | 29.35 |
+| UI-Vision | 37.12 | 49.10 |
+| OSWorld-G | 69.41 | 20.00 |
+
+UI-Vision is the equal-weight macro over its basic, functional, and spatial
+element-grounding subsets. OSWorld-G uses its 510 target-bearing examples;
+refusal-only rows are excluded. These public benchmarks were used during model
+selection, so results are test-tuned rather than held-out validation estimates.
+
+## Direct inference
+
+The repository includes the exact loader and evaluator. `visual_merger.pt` must
+be loaded in addition to the PEFT adapter; `selectground.py` does this.
+
+```bash
+python evaluate.py \
+  --model ruotian/SelectGround-8B \
+  --benchmark screenspot_pro \
+  --data data/screenspot-pro \
+  --output outputs/screenspot-pro.jsonl
+```
+
+Inference uses the full native screenshot, the prompt in `selectground.py`,
+Qwen smart resize with `min_pixels=3136` and `max_pixels=8847360`, greedy
+decoding for at most 32 tokens, and normalized 0–1000 point coordinates.
+
+## Reproduce training from the plain base
+
+```bash
+hf download ruotian/ContrastGround --repo-type dataset \
+  --local-dir data/ContrastGround
+
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+accelerate launch --mixed_precision bf16 --num_processes 2 train.py \
+  --model 8b \
+  --data data/ContrastGround \
+  --pairs-file data/ContrastGround/data/selectground-8b/train_pairs.jsonl \
+  --replay-file data/ContrastGround/data/selectground-8b/train_replays.jsonl \
+  --output outputs/SelectGround-8B \
+  --steps 240 --gpus 2 --accumulation 64 \
+  --learning-rate 3e-5 --selector-learning-rate 1e-4 \
+  --aux-weight 0.1 --ground-coordinate-weight 1.0 \
+  --margin 0.3 --pair-weight 0.5 \
+  --warmup-steps 10 --scheduler-steps 384 \
+  --holdout-fraction 0.02 --seed 20260819
+```
 
-This repository contains the LoRA checkpoint and the visual-merger state used for exact reproduction. `selection_head.pt` stores the training-only layer/head aggregation state; inference uses only the trained grounding model.
+This is SFT coordinate cross-entropy on pair and replay rows plus the paper's
+auxiliary competitor-selection loss on pair rows. Pair and replay microbatches
+alternate. LoRA uses rank 64, alpha 128, dropout 0.05 on
+`q/k/v/o/gate/up/down` projections. The selector reads semantic attention from
+layers 18–23. See `training_manifest.json` for the complete recipe and artifact
+SHA-256 checksums.
 
-| Method | ScreenSpot-Pro | UI-Vision | OSWorld-G |
-|---|---:|---:|---:|
-| SelectGround-8B | 64.96 | 38.68 | 70.00 |
-| + CVL | 72.36 | 45.10 | 72.55 |
+The reference environment used PyTorch 2.11.0+cu128, Transformers 4.57.1,
+PEFT 0.19.1, Accelerate 1.13.0, and qwen-vl-utils 0.0.14. CUDA kernels are not
+bitwise deterministic; clean runs should be expected to be close rather than
+byte-identical.
 
-OSWorld-G excludes refusal cases. UI-Vision is the equal-weight mean over basic, functional, and spatial element grounding.
+## License and data
 
-Code and usage are in the [ClickContrast repository](https://github.com/zhangruotian/ClickContrast). Training data are in [ClickContrast](https://huggingface.co/datasets/ruotian/ClickContrast).
+The adapter follows the Apache-2.0 license of the base model. Dataset assets
+retain their upstream terms; consult the ContrastGround data card and its
+row-level provenance.

adapter_config.json

@@ -1,15 +1,48 @@
 {
+  "alora_invocation_tokens": null,
+  "alpha_pattern": {},
+  "arrow_config": null,
+  "auto_mapping": null,
   "base_model_name_or_path": "Qwen/Qwen3-VL-8B-Instruct",
   "bias": "none",
+  "corda_config": null,
+  "ensure_weight_tying": false,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
   "inference_mode": true,
   "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
   "lora_alpha": 128,
+  "lora_bias": false,
   "lora_dropout": 0.05,
+  "lora_ga_config": null,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
   "peft_type": "LORA",
+  "peft_version": "0.19.1",
+  "qalora_group_size": 16,
   "r": 64,
+  "rank_pattern": {},
   "revision": "0c351dd01ed87e9c1b53cbc748cba10e6187ff3b",
-  "target_modules": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+  "target_modules": [
+    "q_proj",
+    "v_proj",
+    "gate_proj",
+    "down_proj",
+    "k_proj",
+    "up_proj",
+    "o_proj"
+  ],
+  "target_parameters": null,
   "task_type": "CAUSAL_LM",
+  "trainable_token_indices": null,
+  "use_bdlora": null,
   "use_dora": false,
+  "use_qalora": false,
   "use_rslora": false
 }

adapter_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:aa33306498c323fea5405798a8feded09f4e5dfda386dd412c48cf174d1581cd
+oid sha256:6ff35d55e9000af46eb6e134b78c7c0029369bf3ec5fe74d2dfb29f3a09e923c
 size 349251816

evaluate.py

@@ -0,0 +1,166 @@
+import argparse
+import io
+import itertools
+import json
+from pathlib import Path
+
+from PIL import Image
+
+from selectground import SelectGround
+
+
+def load_cases(name: str, root: Path):
+    if name == "screenspot_pro":
+        parquet_files = sorted((root / "data").glob("*.parquet"))
+        if parquet_files:
+            import pyarrow.parquet as parquet
+
+            for path in parquet_files:
+                for batch in parquet.ParquetFile(path).iter_batches(batch_size=1):
+                    row = batch.to_pylist()[0]
+                    encoded = row["image"]
+                    image = Image.open(io.BytesIO(encoded["bytes"])).convert("RGB")
+                    yield {
+                        "id": row["id"],
+                        "image": image,
+                        "image_name": encoded.get("path") or row["id"],
+                        "instruction": row["instruction"],
+                        "target": row["bbox"],
+                        "type": "xyxy",
+                        "group": row.get("group"),
+                    }
+        else:
+            for annotation in sorted((root / "annotations").glob("*.json")):
+                for row in json.loads(annotation.read_text()):
+                    yield {
+                        "id": row["id"],
+                        "image": root / "images" / row["img_filename"],
+                        "image_name": row["img_filename"],
+                        "instruction": row["instruction"],
+                        "target": row["bbox"],
+                        "type": "xyxy",
+                        "group": row.get("group"),
+                    }
+    elif name == "ui_vision":
+        for split in ("basic", "functional", "spatial"):
+            path = root / "annotations" / "element_grounding" / f"element_grounding_{split}.json"
+            for index, row in enumerate(json.loads(path.read_text())):
+                yield {
+                    "id": f"{split}-{index}",
+                    "image": root / "images" / row["image_path"],
+                    "image_name": row["image_path"],
+                    "instruction": row["prompt_to_evaluate"],
+                    "target": row["bbox"],
+                    "type": "xyxy",
+                    "group": split,
+                }
+    else:
+        benchmark = root / "benchmark" if (root / "benchmark").is_dir() else root
+        for row in json.loads((benchmark / "OSWorld-G.json").read_text()):
+            if row["box_type"] == "refusal":
+                continue
+            yield {
+                "id": row["id"],
+                "image": benchmark / "images" / row["image_path"],
+                "image_name": row["image_path"],
+                "instruction": row["instruction"],
+                "target": row["box_coordinates"],
+                "type": row["box_type"],
+                "group": None,
+            }
+
+
+def contains(point, target, target_type):
+    if point is None:
+        return False
+    x, y = point
+    if target_type in {"bbox", "xyxy"}:
+        if target_type == "xyxy":
+            left, top, right, bottom = target
+        else:
+            left, top, width, height = target[:4]
+            right, bottom = left + width, top + height
+        center_x, center_y = (left + right) / 2, (top + bottom) / 2
+        half_width, half_height = abs(right - left) / 2, abs(bottom - top) / 2
+        return (
+            center_x - half_width <= x <= center_x + half_width
+            and center_y - half_height <= y <= center_y + half_height
+        )
+    vertices = list(zip(target[0::2], target[1::2]))
+    previous, inside = vertices[-1], False
+    for current in vertices:
+        x1, y1 = current
+        x2, y2 = previous
+        cross = (x - x1) * (y2 - y1) - (y - y1) * (x2 - x1)
+        if abs(cross) <= 1e-7 and min(x1, x2) - 1e-7 <= x <= max(x1, x2) + 1e-7 and min(y1, y2) - 1e-7 <= y <= max(y1, y2) + 1e-7:
+            return True
+        if (y1 > y) != (y2 > y) and x < (x2 - x1) * (y - y1) / (y2 - y1) + x1:
+            inside = not inside
+        previous = current
+    return inside
+
+
+def metrics(rows, benchmark):
+    if benchmark != "ui_vision":
+        return {"total": len(rows), "correct": sum(row["correct"] for row in rows), "accuracy": 100 * sum(row["correct"] for row in rows) / len(rows)}
+    splits = {}
+    for split in ("basic", "functional", "spatial"):
+        selected = [row for row in rows if row["group"] == split]
+        if selected:
+            splits[split] = 100 * sum(row["correct"] for row in selected) / len(selected)
+    return {"total": len(rows), "accuracy": sum(splits.values()) / len(splits), "splits": splits}
+
+
+parser = argparse.ArgumentParser(description="Evaluate SelectGround on a GUI grounding benchmark.")
+parser.add_argument("--model", default="ruotian/SelectGround-8B")
+parser.add_argument("--benchmark", choices=("screenspot_pro", "ui_vision", "osworld_g"), required=True)
+parser.add_argument("--data", type=Path, required=True)
+parser.add_argument("--output", type=Path, required=True)
+parser.add_argument("--lcr", action="store_true")
+parser.add_argument(
+    "--lcr-variant",
+    choices=("full", "no_competitor", "one_competitor", "no_incumbent"),
+    default="full",
+)
+parser.add_argument("--limit", type=int)
+parser.add_argument("--num-shards", type=int, default=1)
+parser.add_argument("--shard", type=int, default=0)
+args = parser.parse_args()
+
+cases = (
+    case for index, case in enumerate(load_cases(args.benchmark, args.data))
+    if index % args.num_shards == args.shard
+)
+if args.limit is not None:
+    cases = itertools.islice(cases, args.limit)
+existing = []
+if args.output.exists():
+    existing = [json.loads(line) for line in args.output.read_text().splitlines() if line.strip()]
+done = {row["id"] for row in existing}
+grounder = SelectGround(args.model)
+args.output.parent.mkdir(parents=True, exist_ok=True)
+with args.output.open("a") as output:
+    for number, case in enumerate(cases, 1):
+        if case["id"] in done:
+            continue
+        prediction = grounder.predict(
+            case["image"],
+            case["instruction"],
+            lcr=args.lcr,
+            benchmark=args.benchmark,
+            lcr_variant=args.lcr_variant,
+        )
+        row = {
+            "id": case["id"],
+            "instruction": case["instruction"],
+            "image": case["image_name"],
+            "point": prediction["point"],
+            "correct": contains(prediction["point"], case["target"], case["type"]),
+            "group": case["group"],
+            "prediction": prediction,
+        }
+        output.write(json.dumps(row) + "\n")
+        output.flush()
+        existing.append(row)
+        print(f"[{number}] {case['id']} correct={int(row['correct'])}", flush=True)
+print(json.dumps(metrics(existing, args.benchmark), indent=2))

selectground.py

@@ -0,0 +1,465 @@
+import json
+import math
+import re
+from itertools import combinations
+from pathlib import Path
+from typing import Any
+
+import torch
+from PIL import Image
+from huggingface_hub import snapshot_download
+from peft import PeftModel
+from transformers import AutoModelForImageTextToText, AutoProcessor
+from transformers.cache_utils import DynamicCache
+
+
+PROMPT = """You are an expert GUI grounding model.
+Given a screenshot and an instruction, point to the UI element that should be clicked.
+Return only one point as [x, y], where x and y are normalized integers from 0 to 1000 relative to the full image.
+For an element with area, return the center point.
+Instruction: {instruction}"""
+
+LCR_PROMPT = """You are an expert UI element locator. Given a GUI image and a user's element description, provide the coordinates of the specified element as a single (x,y) point. The image resolution is height {height} and width {width}. For elements with area, return the center point.
+
+Output the coordinate pair exactly:
+(x,y)"""
+
+
+class SelectGround:
+    """SelectGround direct grounding and LCR test-time inference."""
+
+    def __init__(self, checkpoint: str = "ruotian/SelectGround-8B") -> None:
+        checkpoint_path = Path(checkpoint)
+        if not checkpoint_path.exists():
+            checkpoint_path = Path(snapshot_download(checkpoint))
+        adapter_config = json.loads((checkpoint_path / "adapter_config.json").read_text())
+        base_model = adapter_config["base_model_name_or_path"]
+        revision = adapter_config["revision"]
+        model = AutoModelForImageTextToText.from_pretrained(
+            base_model,
+            revision=revision,
+            dtype=torch.bfloat16,
+            device_map="auto",
+            attn_implementation="sdpa",
+        )
+        self.model = PeftModel.from_pretrained(model, checkpoint_path)
+        merger = torch.load(checkpoint_path / "visual_merger.pt", map_location="cpu")
+        parameters = dict(self.model.named_parameters())
+        with torch.no_grad():
+            for name, value in merger.get("state_dict", merger).items():
+                parameters[name].copy_(value.to(parameters[name].device, parameters[name].dtype))
+        self.model.eval()
+        self.model.config.use_cache = True
+        self.processor = AutoProcessor.from_pretrained(
+            base_model,
+            revision=revision,
+            min_pixels=3136,
+            max_pixels=8847360,
+        )
+        self.device = next(self.model.parameters()).device
+        self.core = self.model.get_base_model().model
+        self.vision_start_token_id = int(self.core.config.vision_start_token_id)
+        self.vision_end_token_id = int(self.core.config.vision_end_token_id)
+        self.merge_size = int(self.core.config.vision_config.spatial_merge_size)
+        self.large_lcr = len(self.core.language_model.layers) > 36
+
+    def predict(
+        self,
+        image: str | Path | Image.Image,
+        instruction: str,
+        *,
+        lcr: bool = False,
+        benchmark: str | None = None,
+        lcr_variant: str = "full",
+    ) -> dict[str, Any]:
+        """Return a source-image click; benchmark only selects UI-Vision's crop size."""
+        source = Image.open(image).convert("RGB") if not isinstance(image, Image.Image) else image.convert("RGB")
+        size = source.size
+        p0, attention, grid = self._observe(
+            source,
+            instruction,
+            capture_attention=lcr,
+            lcr_prompt=lcr and not self.large_lcr,
+        )
+        if not lcr:
+            return {"method": "SelectGround", **p0}
+        if p0["point"] is None:
+            return {"method": "SelectGround+LCR", **p0}
+
+        attention_boxes = (
+            _attention_crops(attention, grid, size, benchmark)
+            if attention is not None
+            else []
+        )
+        if lcr_variant not in {"full", "no_competitor", "one_competitor", "no_incumbent"}:
+            raise ValueError(f"unknown LCR variant: {lcr_variant}")
+        if lcr_variant == "no_competitor":
+            attention_boxes = []
+        elif lcr_variant == "one_competitor":
+            attention_boxes = attention_boxes[:1]
+        if self.large_lcr:
+            views = [
+                *((box, 2.0) for box in attention_boxes[:1]),
+                (_fraction_crop(p0["point"], size, 0.25, 256), 2.5),
+                (_fraction_crop(p0["point"], size, 0.40, 320), 2.0),
+            ]
+        else:
+            views = [
+                *((box, 2.0) for box in attention_boxes),
+                (_pixel_budget_crop(p0["point"], size, 501_760), 1.5),
+            ]
+        if lcr_variant == "no_incumbent":
+            views = views[:-2] if self.large_lcr else views[:-1]
+        observations = [(p0, (0, 0, size[0], size[1]))]
+        for box, scale in views:
+            crop = source.crop(box)
+            view = crop.resize(
+                (round(crop.width * scale), round(crop.height * scale)),
+                Image.Resampling.LANCZOS if self.large_lcr else Image.Resampling.BICUBIC,
+            )
+            prediction, _, _ = self._observe(
+                view, instruction, lcr_prompt=not self.large_lcr
+            )
+            if prediction["point"] is not None:
+                observations.append((_map_crop(prediction, box, size, scale), box))
+        if len(observations) == 1:
+            return {"method": "SelectGround+LCR", **p0}
+        first, second = min(
+            combinations(range(len(observations)), 2),
+            key=lambda pair: (
+                _distance(
+                    observations[pair[0]][0]["point"],
+                    observations[pair[1]][0]["point"],
+                    size,
+                ),
+                pair,
+            ),
+        )
+        selected = min(
+            (first, second),
+            key=lambda index: (_area(observations[index][1]), index),
+        )
+        result = observations[selected][0]
+        return {
+            "method": "SelectGround+LCR",
+            "point": result["point"],
+            "normalized_point": result["normalized_point"],
+            "raw_response": result["raw_response"],
+        }
+
+    def _observe(
+        self,
+        image: Image.Image,
+        instruction: str,
+        *,
+        capture_attention: bool = False,
+        lcr_prompt: bool = False,
+    ) -> tuple[dict[str, Any], torch.Tensor | None, tuple[int, int]]:
+        inputs = self._inputs(image, instruction, lcr_prompt)
+        input_ids = inputs["input_ids"]
+        length = int(input_ids.shape[1])
+        image_grid = inputs["image_grid_thw"][0].detach().cpu().long()
+        grid = (int(image_grid[1]) // self.merge_size, int(image_grid[2]) // self.merge_size)
+        position_ids, _ = self.core.get_rope_index(
+            input_ids,
+            inputs.get("image_grid_thw"),
+            inputs.get("video_grid_thw"),
+            attention_mask=inputs.get("attention_mask"),
+        )
+        cache = DynamicCache(config=self.core.language_model.config)
+        attention = (
+            _Attention(
+                self.core,
+                input_ids,
+                self.vision_start_token_id,
+                self.vision_end_token_id,
+            )
+            if capture_attention
+            else None
+        )
+        with torch.inference_mode():
+            output = self.model(
+                **inputs,
+                past_key_values=cache,
+                position_ids=position_ids,
+                cache_position=torch.arange(length, device=self.device),
+                use_cache=True,
+                logits_to_keep=1,
+            )
+            raw = self._decode(
+                output.logits[:, -1, :],
+                cache,
+                position_ids[:, :, -1:] + 1,
+                attention,
+            )
+        return (
+            _prediction(raw, image.size, integer=lcr_prompt),
+            attention.scores if attention else None,
+            grid,
+        )
+
+    def _decode(
+        self,
+        logits: torch.Tensor,
+        cache: DynamicCache,
+        position_ids: torch.Tensor,
+        attention: "_Attention | None",
+    ) -> str:
+        stop = {
+            token
+            for token in (
+                self.processor.tokenizer.eos_token_id,
+                self.processor.tokenizer.pad_token_id,
+            )
+            if token is not None
+        }
+        generated = []
+        for _ in range(32):
+            token = int(logits.argmax())
+            if token in stop:
+                break
+            generated.append(token)
+            token_text = self.processor.decode([token])
+            if attention is not None and ("," in token_text or "，" in token_text):
+                attention.install(cache)
+            try:
+                output = self.model(
+                    input_ids=torch.tensor([[token]], device=self.device),
+                    past_key_values=cache,
+                    position_ids=position_ids,
+                    cache_position=torch.tensor(
+                        [cache.get_seq_length()], device=self.device
+                    ),
+                    use_cache=True,
+                    logits_to_keep=1,
+                )
+            finally:
+                if attention is not None and attention.handle is not None:
+                    attention.remove()
+            cache = output.past_key_values
+            logits = output.logits[:, -1, :]
+            position_ids = position_ids + 1
+        return self.processor.decode(
+            generated,
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False,
+        ).strip()
+
+    def _inputs(
+        self, image: Image.Image, instruction: str, lcr_prompt: bool
+    ) -> Any:
+        from qwen_vl_utils import process_vision_info, smart_resize
+
+        if lcr_prompt:
+            resized_height, resized_width = smart_resize(
+                image.height,
+                image.width,
+                factor=(
+                    self.processor.image_processor.patch_size
+                    * self.processor.image_processor.merge_size
+                ),
+                min_pixels=self.processor.image_processor.min_pixels,
+                max_pixels=self.processor.image_processor.max_pixels,
+            )
+            image = image.resize((resized_width, resized_height))
+            messages = [
+                {
+                    "role": "system",
+                    "content": LCR_PROMPT.format(
+                        height=resized_height, width=resized_width
+                    ),
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image", "image": image},
+                        {"type": "text", "text": instruction},
+                    ],
+                },
+            ]
+        else:
+            messages = [{
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": image},
+                    {"type": "text", "text": PROMPT.format(instruction=instruction)},
+                ],
+            }]
+        text = self.processor.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        image_inputs, video_inputs = process_vision_info(messages)
+        return self.processor(
+            text=[text],
+            images=image_inputs,
+            videos=video_inputs,
+            padding=True,
+            return_tensors="pt",
+        ).to(self.device)
+
+
+class _Attention:
+    def __init__(
+        self,
+        model: Any,
+        input_ids: torch.Tensor,
+        vision_start: int,
+        vision_end: int,
+    ) -> None:
+        self.model = model
+        self.handle: Any = None
+        start = int(torch.nonzero(input_ids[0] == vision_start)[0]) + 1
+        end = int(torch.nonzero(input_ids[0] == vision_end)[0])
+        self.visual = torch.arange(start, end, device=input_ids.device)
+        self.cache: DynamicCache | None = None
+        self.scores: torch.Tensor | None = None
+
+    def install(self, cache: DynamicCache) -> None:
+        self.cache = cache
+        layers = self.model.language_model.layers
+        layer = layers[2 * len(layers) // 3].self_attn
+        self.handle = layer.register_forward_hook(self._hook, with_kwargs=True)
+
+    def remove(self) -> None:
+        self.handle.remove()
+        self.handle = None
+
+    def _hook(
+        self,
+        module: Any,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+        output: Any,
+    ) -> None:
+        from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+            apply_rotary_pos_emb,
+            repeat_kv,
+        )
+
+        hidden = kwargs.get("hidden_states", args[0] if args else None)
+        shape = (*hidden.shape[:-1], -1, int(module.head_dim))
+        query = module.q_norm(module.q_proj(hidden).view(shape)).transpose(1, 2)
+        query, _ = apply_rotary_pos_emb(
+            query, query, *kwargs["position_embeddings"]
+        )
+        keys = repeat_kv(
+            self.cache.layers[module.layer_idx].keys,
+            int(module.num_key_value_groups),
+        )
+        weights = torch.matmul(query, keys.transpose(-2, -1)) * module.scaling
+        weights = weights.squeeze(2).softmax(-1).max(1).values[0]
+        self.scores = (
+            weights.index_select(0, self.visual.to(weights.device))
+            .detach()
+            .float()
+            .cpu()
+        )
+
+
+def _prediction(
+    raw: str, size: tuple[int, int], *, integer: bool = False
+) -> dict[str, Any]:
+    match = re.search(
+        r"[\[(（]\s*(?:x\s*=\s*)?(-?\d+(?:\.\d+)?)\s*[,，]\s*"
+        r"(?:y\s*=\s*)?(-?\d+(?:\.\d+)?)\s*[\])）]",
+        raw,
+        flags=re.IGNORECASE,
+    )
+    normalized = [float(match.group(1)), float(match.group(2))] if match else None
+    point = [normalized[0] / 1000 * size[0], normalized[1] / 1000 * size[1]] if normalized else None
+    if point is not None and integer:
+        point = [int(value) for value in point]
+    return {"point": point, "normalized_point": normalized, "raw_response": raw}
+
+
+def _map_crop(
+    prediction: dict[str, Any],
+    box: tuple[int, int, int, int],
+    size: tuple[int, int],
+    scale: float,
+) -> dict[str, Any]:
+    point = prediction["point"]
+    mapped = [box[0] + point[0] / scale, box[1] + point[1] / scale]
+    normalized = [mapped[0] / size[0] * 1000, mapped[1] / size[1] * 1000]
+    raw = f"[{round(normalized[0])},{round(normalized[1])}]"
+    return {"point": mapped, "normalized_point": normalized, "raw_response": raw}
+
+
+def _pixel_budget_crop(
+    point: list[float], size: tuple[int, int], pixels: int
+) -> tuple[int, int, int, int]:
+    width, height = size
+    fraction = min(1.0, math.sqrt(pixels / (width * height)))
+    crop_width = max(1, round(fraction * width))
+    crop_height = max(1, round(fraction * height))
+    left = round(min(max(0.0, point[0] - crop_width / 2), width - crop_width))
+    top = round(min(max(0.0, point[1] - crop_height / 2), height - crop_height))
+    return left, top, left + crop_width, top + crop_height
+
+
+def _fraction_crop(
+    point: list[float],
+    size: tuple[int, int],
+    fraction: float,
+    minimum: int,
+) -> tuple[int, int, int, int]:
+    width, height = size
+    crop_width = min(width, max(minimum, round(fraction * width)))
+    crop_height = min(height, max(minimum, round(fraction * height)))
+    left = round(min(max(0.0, point[0] - crop_width / 2), width - crop_width))
+    top = round(min(max(0.0, point[1] - crop_height / 2), height - crop_height))
+    return left, top, left + crop_width, top + crop_height
+
+
+def _attention_crops(
+    attention: torch.Tensor,
+    grid: tuple[int, int],
+    size: tuple[int, int],
+    benchmark: str | None,
+) -> list[tuple[int, int, int, int]]:
+    width, height = size
+    window = (1280, 720) if benchmark == "ui_vision" else (1288, 728)
+    crop_width, crop_height = min(window[0], width), min(window[1], height)
+    top = attention.topk(min(100, attention.numel())).indices.tolist()
+    positions = [
+        ((index % grid[1] + 0.5) / grid[1] * width,
+         (index // grid[1] + 0.5) / grid[0] * height)
+        for index in top
+    ]
+    ranked = []
+    for x, y in positions:
+        left = min(max(0.0, x - crop_width / 2), width - crop_width)
+        upper = min(max(0.0, y - crop_height / 2), height - crop_height)
+        box = (
+            int(left),
+            int(upper),
+            int(left + crop_width),
+            int(upper + crop_height),
+        )
+        coverage = sum(
+            left <= px <= left + crop_width
+            and upper <= py <= upper + crop_height
+            for px, py in positions
+        )
+        ranked.append((coverage, box))
+    ranked.sort(key=lambda row: row[0], reverse=True)
+    selected = []
+    for _, box in ranked:
+        if box not in selected:
+            selected.append(box)
+        if len(selected) == 2:
+            break
+    return selected
+
+
+def _distance(
+    first: list[float], second: list[float], size: tuple[int, int]
+) -> float:
+    return math.hypot(
+        (first[0] - second[0]) / size[0],
+        (first[1] - second[1]) / size[1],
+    )
+
+
+def _area(box: tuple[int, int, int, int]) -> int:
+    return (box[2] - box[0]) * (box[3] - box[1])

selection_head.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:456e706f68b76c64f3a7db24ee2e220711d7b16adffee1f4e7b02b938c3f9d77
-size 2586
+oid sha256:e74db8e7fc9344ccc477d66d2439a53e67a646eca57bf1fabc7b790bdb0ab420
+size 2458

train.py

@@ -0,0 +1,663 @@
+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()

training_manifest.json

@@ -0,0 +1,53 @@
+{
+  "schema_version": "selectground.release.v2",
+  "model": "SelectGround-8B",
+  "base_model": "Qwen/Qwen3-VL-8B-Instruct",
+  "base_revision": "0c351dd01ed87e9c1b53cbc748cba10e6187ff3b",
+  "plain_base_start": true,
+  "aggregate": false,
+  "method": "SFT plus auxiliary selection loss",
+  "training": {
+    "stages": 1,
+    "steps": 240,
+    "seed": 20260819,
+    "gpus": 2,
+    "gradient_accumulation": 64,
+    "effective_global_batch": 128,
+    "learning_rate": 0.00003,
+    "selector_learning_rate": 0.0001,
+    "aux_weight": 0.1,
+    "coordinate_weight": 1.0,
+    "margin": 0.3,
+    "pair_weight": 0.5,
+    "warmup_steps": 10,
+    "scheduler_steps": 384,
+    "holdout_fraction": 0.02,
+    "precision": "bf16",
+    "attention": "sdpa",
+    "lora": {
+      "rank": 64,
+      "alpha": 128,
+      "dropout": 0.05,
+      "targets": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+    }
+  },
+  "dataset": {
+    "repo": "ruotian/ContrastGround",
+    "config": "selectground-8b",
+    "pairs": 4942,
+    "replay": 4103,
+    "source_pairs_sha256": "bd9fcb0197d3799fd9d1a6f1178f3e0adaddd8a88d032939101789b259ac6d68",
+    "source_replay_sha256": "e94f589efe8ff197778fe81c33beefd75b59af962d77e17f0fb29596baf03d47"
+  },
+  "checkpoint_sha256": {
+    "adapter_config.json": "6d19205e4597f233d22dfdb37c089d39ef284c6582269f41bd42a58715335c37",
+    "adapter_model.safetensors": "6ff35d55e9000af46eb6e134b78c7c0029369bf3ec5fe74d2dfb29f3a09e923c",
+    "visual_merger.pt": "125b9b37b243724cee09eb064650a6e14d23f186097738b5144170e870da2fab",
+    "selection_head.pt": "e74db8e7fc9344ccc477d66d2439a53e67a646eca57bf1fabc7b790bdb0ab420"
+  },
+  "direct_accuracy_pct": {
+    "screenspot_pro": 65.0853889943074,
+    "ui_vision": 37.123840466010655,
+    "osworld_g": 69.41176470588235
+  }
+}