royguw/ground-qwen3-5

ground_qwen3_5

Research code tree for grounded segmentation / referring-expression work built on Qwen3-VL, including the video extension, ablation variants, baselines and analysis tooling.

This is a working research tree, published as-is rather than as a packaged library. It contains many parallel variant directories (each a self-contained copy of the model, training loop and eval scripts) that differ in a single design choice.

What's here

path	contents
ground_qwen35*/	model + training + eval variants (one directory per design choice)
ground_qwen35_video*/	video extension: tube assignment, temporal action tokens
baselines/	MM-GroundingDINO / mmdetection and other baseline checkouts
analysis/	bridge-attention and multi-scale analysis outputs
ConCor-1-release/	the ConCor-1 public release staging area
data_scripts/, scripts/	dataset construction and job launch scripts
docs/	design notes
wandb/	training run logs

Trained checkpoints for these runs are published separately at royguw/ground-ckpts.

Packaging note

HuggingFace caps a repository at 20,000 files; the raw tree has ~175,800. To fit, these directories are published as tar archives rather than loose files:

archive	contents
wandb.tar	wandb/ training run logs
analysis.tar	analysis/ result dumps
baselines/mm_grounding_dino/venv.tar	that baseline's .venv
ground_qwen35_video/baselines/*/venv.tar	each video baseline's .venv

Extract in place, e.g. tar -xf wandb.tar.

Not included: .git/ history, and __pycache__/ / *.pyc bytecode.

WANDB_API_KEY and --hf-token are read from the environment; set them yourself.

Notes

• WANDB_API_KEY is read from the environment; set it yourself before training.
• Git history is not included.
• baselines/ and the nested ground_qwen35_video/baselines/ vendor third-party repositories (mmdetection and others), which carry their own licences.

---

Project README

GroundQwen35 — Instance Segmentation on Qwen3.5-0.8B

Instance-level image segmentation using configurable bridge (query) tokens on top of Qwen3.5-0.8B's hybrid-attention multimodal backbone.

Ported from ground_olmo.

Overview

GroundQwen35 is a grounding-only model (not autoregressive). It appends learnable bridge tokens after text in the sequence, whose hidden states are fed to grounding heads that predict:

• Segmentation masks: Per-bridge, per-visual-token predictions (seg_logits)
• Presence: Which bridges are "occupied" by an instance (presence_logits)
• Text grounding (optional via config): Which category text tokens align with each instance (text_logits), using a second bilinear head when query_text_grounding is enabled

Typical configuration

Many training runs use together:

Setting	Role
query_text_grounding=True (CLI default; see below)	Instantiates text_head and builds text labels; forward requires non-empty text_grounding_mask.
--spatial_grid_assignment	Hungarian assignment of instances to bridge slots on an N×N grid (num_bridges must be N²).
--presence_mlp	SwiGLU-style MLP presence head instead of a single linear layer.

Optional / less central for that path: --reset_position, --mask_incomplete_presence, --mask_out_incomplete_presence.

Architecture

                    Qwen3.5-0.8B Backbone
                    ┌─────────────────┐
  image pixels ──►  │  Vision Encoder  │ ──► visual tokens (N_vis)
                    │  (12-layer ViT)  │
                    └────────┬────────┘
                             ▼
  input_ids ──────► ┌─────────────────┐
  (vis + text +     │  Text Model     │ ──► hidden_states (B, S, 1024)
   bridge tokens)   │  (24 hybrid     │
                    │   layers)       │
                    └────────┬────────┘
                             ▼
                    ┌─────────────────┐
                    │ extract_tokens  │
                    │ _by_mask        │
                    └──┬──────────┬───┬┘
                       ▼          ▼   ▼
                 query_hidden  visual_hidden  text_hidden *
                  (B,Q,D)      (B,N_vis,D)     (B,N_txt,D)
                       │          │            │
                       └────┬─────┘            │
                            ▼                  │
                     ┌─────────────┐           │
                     │  seg_head   │           │
                     │ (Visual     │           │
                     │ SingleSlot  │           │
                     │ Gate)       │           │
                     └──────┬──────┘           │
                            ▼                  │
                     seg_logits                │
                   (B,Q,N_vis)                 │
                                               ▼
                     ┌─────────────────────────┴────────────┐
                     │  text_head (VisualSingleSlotGate) *   │
                     │  inputs: text_hidden + query_hidden    │
                     └──────────────────┬─────────────────────┘
                                        ▼
                                 text_logits *
                              (B,Q,N_txt)

                            query_hidden
                                  │
                                  ▼
                           ┌─────────────┐
                           │ presence_   │
                           │ head        │
                           │ (MLP or     │
                           │  Linear)    │
                           └──────┬──────┘
                                  ▼
                           presence_logits
                               (B,Q)

* text_hidden, text_head, and text_logits exist only when query_text_grounding is True (ground_qwen35/model.py). N_txt is the number of tokenizer outputs for the category string segment. seg_head uses visual + query hidden states; text_head uses category text + query; presence_head uses query only.

Sequence format

<|vision_start|> <|image_pad|>×N_vis <|vision_end|> category_text bridge_token×Q

Here Q = num_bridges (default 225 in train.py / GroundQwen35Config). Bridge token IDs are consecutive unused embedding rows starting at 248077; if Q > 243, the embedding table is resized (init_bridge_embeddings).

No system/assistant wrapping. Causal attention is used as in the backbone.

Data pipeline

flowchart LR
  A[GroundingSourceDataset] --> B[DeterministicGroundingDataset]
  B --> C[QwenExamplePreprocessor]
  C --> D[QwenGroundingCollator]
  D --> E[IterableGroundingMixture]

• Sources (data/dataset.py): ground_olmo dataset classes (COCONut / LVIS / ADE20K / RefCOCO / baselines, etc.); empty-segment rows removed.
• Deterministic wrapper: deterministic RNG + QwenExamplePreprocessor; _validate_sample (image, instances, text labels, text_grounding_mask.any()); up to 50 retries with rerouted indices.
• Preprocessor (data/preprocessor.py): HF image processor → merged patch labels (data/segmentation.py); optional compute_spatial_grid_assignment (perfect-square num_bridges); format_grounding_sequence (data/formatter.py); optional build_instance_text_labels (first token subsequence match for category + char spans).
• Collator (data/collator.py): pad sequences and labels; presence_loss_mask; optional slot_assignment.
• Iterable mixture (data/iterable_mixture.py): mixture probabilities, per-dataset epoch shuffle, global batch sharding across workers and ranks.

Training loss (loss.py)

All three terms are computed every step; train.py adds a zero multiple of head outputs so DDP always touches all parameters (DDP-safe empty batches).

• Presence: BCE with optional weighting (presence_loss_type: raw / weight / neg_down_sample) and presence_loss_mask.
• Segmentation: Dice + BCE (or BCE only) on active bridge predictions vs instance masks; with slot_assignment, maps slots back to the correct instance row (--grounding_loss_func: dice_bce or bce).
• Text grounding: BCE on active slots vs padded text labels (-100 ignored).

Training script notes (train.py)

• Monkey patch Qwen3_5TextModel.forward (mRoPE / position_ids handling) — applied before model construction.
• Optimizer: AdamW; seg_head and text_head use --head_lr; presence_head and backbone use --backbone_lr.
• Schedule: Linear warmup then cosine decay to alpha_f × peak LR. global_step counts optimizer steps (after grad_accum_steps micro-batches).
• Checkpoints: model.pt, optimizer.pt, per-rank train/rank*.pt (RNG, global_train_examples_seen, etc.).
• Weights & Biases: _init_wandb currently assigns WANDB_API_KEY inside the script — remove hardcoded credentials for any shared or public repo; use wandb login or environment variables instead.

num_bridges constraints

Situation	Requirement
--spatial_grid_assignment	num_bridges must be N² (asserted in data/segmentation.py::compute_spatial_grid_assignment). Examples: 64, 100, 225. Not 32. Training startup does not re-assert this (unlike --reset_position).
--reset_position	Also requires a perfect square; asserted in train.py.
Neither	Any positive Q works for sequential slot filling (instance i → bridge i, capped by loss logic when n_inst > Q).
Reserved tokenizer IDs	243 slots without tokenizer entries; larger Q triggers resize_token_embeddings.

Spatial grid assignment

Hungarian matching assigns each instance (after reading-order sort) to the nearest cell on an N×N grid in normalized coordinates. If more instances than slots, some rows stay -1 in slot_assignment and are skipped where the loss checks slot indices.

Dataset mixtures

Configured in ground_qwen35/train.py:get_training_data(), not CLI.

Mode	Content (summary)
debug	Single group single_ref (weight 0.2): refcoco_clean / refcoco_plus_clean / refclef with sampling_rate 0.4 / 0.4 / 0.2. (Commented blocks show older vocab / pos_neg debug.)
full	allvocab 0.1 (coconut_vocab_all_cat, lvis_vocab_all_cat; ade20k_vocab_all_cat commented), pos_neg 0.2, panoptic_caption 0.4 (coconut_pancap), reference 0.3 (refcoco_clean, refcoco_plus_clean, refclef with 0.45 / 0.45 / 0.1).
debug_single	coconut_vocab_all_cat only.
debug_mixed	0.5 coconut_pos_neg + 0.5 Ref clean splits (for DDP / complete vs incomplete testing).

Per-dataset sampling uses sqrt(n) or sampling_rate * sqrt(n) within a group, then group weights; probabilities are normalized globally.

Registry also includes baselines and cleaned Ref variants — see DATASET_REGISTRY in data/dataset.py.

Key training arguments

Arg	Default	Description
--mode	debug	Mixture selector (debug, full, debug_single, debug_mixed).
--num_bridges	225	Bridge / query count.
--grounding_loss_func	dice_bce	dice_bce or bce.
--presence_loss_type	weight	raw, weight, neg_down_sample.
--freeze_vision	True	Freeze vision tower.
--freeze_backbone	False	Freeze full backbone.
--query_text_grounding	effectively on	Implemented as action="store_true", default=True — the flag does not toggle off from CLI; to disable, change train.py / argparse.
--spatial_grid_assignment	False	N×N Hungarian slot assignment (needs num_bridges square).
--presence_mlp	False	SwiGLU MLP presence head.
--reset_position	False	Override bridge MRoPE to interpolated grid (perfect-square num_bridges).
--head_lr / --backbone_lr	1e-4 / 1e-5	LR groups.
--warmup_steps	200	Warmup length (optimizer steps).
--alpha_f	0.1	Cosine floor factor.
--max_grad_norm	1.0	Clipping (0 = off).
--grad_accum_steps	1	Micro-steps per optimizer step.
--eval_interval	0	Eval every N optimizer steps (0 = off).
--eval_max_examples	512	Cap per eval dataset.
--detail_debug	False	Verbose debug dumps.

Sample validation

Every sample passes _validate_sample() (data/dataset.py):

1. pixel_values present
2. At least one instance and N_vis > 0
3. instance_text_labels not None
4. text_grounding_mask has at least one True

Implication: text-only preprocessor paths are not accepted by the current training loop validation.

Validation eval loop

With --eval_interval N, eval runs on: EVAL_DATASETS (coconut_pancap, coconut_pos_neg, ade20k_pos_neg, lvis_pos_neg), plus REFCOCO_EVAL_DATASETS (refcoco, refcoco_plus), plus CAPTION_BASELINE_EVAL_DATASETS (refcocog_baseline, grefcoco_baseline, coconut_pancap_baseline), each capped by --eval_max_examples.

Known limitations (current code)

• _build_presence_labels (loss.py): when slot_assignment is set, the loop uses n = min(num_instances, Q, max_instances), so if num_instances > Q some instance indices are never scanned for presence targets (Hungarian may still assign higher indices to slots). Be aware when using dense scenes with small Q.
• Collator assumes the same Q for every sample in a batch (bridge_token_mask.sum() from the first example).
• --query_text_grounding: cannot be turned off via CLI with the current argparse pattern (see table above).

Quick start

cd /weka-mm/royg/ground_qwen3_5

# 1. Run debug test (single sample, full intermediate logging)
python -m ground_qwen35.debug_test

# 2. Full training (8 GPUs, example: global batch 64 = 8 × 4 × 2 accum)
/opt/conda/bin/torchrun --nproc-per-node 8 -m ground_qwen35.train \
    --mode full --batch_size 4 --grad_accum_steps 2 --max_steps 10000 \
    --spatial_grid_assignment --presence_mlp \
    --num_bridges 225 \
    --save_dir /path/to/checkpoints --eval_interval 500 --wandb

# 3. Debug training (single GPU)
python -m ground_qwen35.train --mode debug

# 4. Inference demo (Gradio web UI)
python demo/segmentation_demo.py /weka-mm/royg/ground_ckpts/qwen3_5/mar_10_g8/step3000

# 5. Integration test (2 GPUs, damaged samples, checkpoint resume)
bash ground_qwen35/run_test.sh

Inference demo

python demo/segmentation_demo.py /path/to/checkpoint [--port 7860] [--share]

Three modes: Image + Text (masks + text highlights), Panoptic (empty text), Text-only (no image).

Features: Threshold sliders, per-bridge overlays, bridge checkboxes, presence / coverage / text grounding readouts.

Integration test

bash ground_qwen35/run_test.sh

2-GPU end-to-end: short training with damaged-sample injection, checkpoint save/resume. Verifies DDP, validation, eval, auto-resume.

Verification / smoke testing

# With spatial grid (use perfect-square num_bridges, e.g. 64 or 100)
/opt/conda/bin/torchrun --nproc-per-node 2 -m ground_qwen35.train \
    --mode debug --batch_size 2 --max_steps 3 --num_bridges 64 \
    --spatial_grid_assignment \
    --eval_interval 1 --eval_max_examples 8 --detail_debug --log_interval 1

# Without spatial grid
/opt/conda/bin/torchrun --nproc-per-node 2 -m ground_qwen35.train \
    --mode debug --batch_size 2 --max_steps 3 \
    --eval_interval 1 --eval_max_examples 8 --detail_debug --log_interval 1

LR schedule

Cosine with warmup (same spirit as ground_olmo CosWithWarmup): linear 0 → peak over warmup_steps, then cosine to alpha_f × peak by max_steps.

Param groups: seg_head, text_head (head_lr); presence_head, backbone (backbone_lr).

File structure

ground_qwen3_5/
    Qwen3.5-0.8B/                  # Pretrained weights
    docs/                          # Architecture tutorial
    demo/
        segmentation_demo.py       # Gradio inference demo
    ground_qwen35/
        __init__.py
        config.py                  # GroundQwen35Config
        model.py                   # GroundQwen35 model
        grounding_heads.py         # VisualSingleSlotGate
        token_utils.py             # extract_tokens_by_mask
        loss.py                    # Dice+BCE, presence BCE, text BCE (slot_assignment)
        train.py                   # Training, mixtures, monkey patch, CLI
        test_training.py           # Integration test helpers
        run_test.sh                # 2-GPU integration test runner
        debug_test.py              # Single-sample debug test
        debug_logger.py            # Detailed debug logging
        data/
            __init__.py
            segmentation.py        # Mask → merged-token labels, spatial Hungarian
            formatter.py           # Sequence construction
            preprocessor.py        # Example preprocessing
            collator.py            # Batch collation
            dataset.py             # Source + validation + deterministic wrapper
            iterable_mixture.py    # Iterable mixture + sharding
    README.md                      # This file

Dependencies

• Python 3.11+
• PyTorch 2.x
• transformers (Qwen3.5 support); model uses Flash Attention 2 where configured
• scipy (Hungarian matching)
• pycocotools (via ground_olmo / COCO)
• gradio (demo)
• ground_olmo at /weka-mm/royg/ground_olmo (datasets, coco_instances_to_bool_masks, etc.)