What this repository is. Training and evaluation code for an autoregressive polygon-grounding baseline, plus the eval results and run logs from the experiments — not packaged model weights. There is no config.json/safetensors pair here, so from_pretrained("royguw/ft_qwen3_5") will not work. Use it to reproduce or inspect the baseline; train with ft_qwen35/train.py against the released data (royguw/grounding-data and royguw/mm-olmo-images).

Contents: ft_qwen35/ (the package), scripts/, ft_eval_results*/ (metrics JSON), ft_eval_cache*/ (cached eval tensors, .pt), and wandb/ (offline run logs). Paths inside the code and logs are absolute cluster paths and will need rewriting.

WANDB_API_KEY was removed from ft_qwen35/train.py for this release; export your own.

ft_qwen3_5 — Autoregressive Polygon Grounding Baseline

Autoregressive fine-tuning baseline for Qwen3.5-0.8B that predicts instance segmentation as polygon coordinate text, paired with char span references into the grounding text.

Parent project: /weka-mm/royg/ground_qwen3_5 (bridge-token grounding heads approach)

Architecture

Uses Qwen3_5ForConditionalGeneration directly (no custom heads). The model generates polygon coordinates as text tokens via the standard lm_head (tied to embed_tokens).

• Base model: Qwen3.5-0.8B (hidden_size=1024, 24 layers, hybrid linear+full attention)
• Vision: Qwen3.5 ViT (smart resize, 16×16 patches, 2×2 spatial merge → 32×32 px per token)
• Frozen vision: ViT + PatchMerger frozen by default (--freeze_vision)
• Trainable params: ~752M (LLM layers + embed_tokens/lm_head)

Differences from ground_qwen35

Aspect	ground_qwen35	ft_qwen35
Model class	Custom GroundQwen35 wrapper	Direct Qwen3_5ForConditionalGeneration
Output	seg_head + presence_head + text_head	lm_head (text generation)
Bridge tokens	256 bridge tokens	None
Loss	Dice+BCE + presence BCE + text BCE	Cross-entropy on generated text
Sequence	Flat: image + text + bridges	Chat: user (image+text) → assistant (polygon text)
Target	Merged-token soft masks (32×32 blocks)	Normalized polygon coordinates as text

---

Sequence Format

Qwen3.5 native chat template:

<|im_start|>user\n
<|vision_start|><|image_pad|>×N_vis<|vision_end|>\n
{grounding_text}<|im_end|>\n
<|im_start|>assistant\n
{target_text}<|im_end|>

Label masking: labels[:prompt_len] = -100. Only target_text + trailing <|im_end|> contribute to loss.

---

Target Text Format

One instance per line. Polygon coordinates normalized to [0, 999]. Multiple polygon segments per instance separated by |. Char spans are character offsets into grounding_text.

x1 y1 x2 y2 ... xn yn | x1 y1 x2 y2 ; s1 e1 s2 e2
x1 y1 x2 y2 ... xn yn ; s1 e1

Parsing at inference:

1. Split by \n → one string per instance
2. Split each line by ; → (polygon_part, spans_part)
3. Split polygon_part by | → list of segment strings
4. Split each segment by spaces, pair up → [(x1,y1), (x2,y2), ...]
5. Split spans_part by spaces, pair up → [(s1,e1), (s2,e2), ...]

---

Dataset Mixture

Same ground_olmo datasets as ground_qwen35, loaded directly via GroundingSourceDataset. Training does fully online preprocessing — no pre-converted JSONL needed.

Modes

debug — 1 group:

• single_ref: refcoco (0.4), refcoco_plus (0.4), refclef (0.2) — weight 0.2

full — 4 groups:

• allvocab (0.2): coconut_vocab_all_cat, lvis_vocab_all_cat
• pos_neg (0.2): coconut_pos_neg, ade20k_pos_neg, lvis_pos_neg
• panoptic_caption (0.3): coconut_pancap
• reference (0.3): refcoco (0.4), refcoco_plus (0.4), refclef (0.2)

debug_single — 1 group:

• single: coconut_vocab_all_cat — weight 1.0

Available Datasets (19 total)

Name	Type	Mask Format
coconut	PerCategory	RLE
coconut_all_cat	AllCat	RLE
coconut_vocab_all_cat	VocabAllCat	RLE
coconut_pos_neg	PosNeg (stochastic)	RLE
coconut_pancap	PanCap (dense caption)	RLE
lvis / lvis_all_cat / lvis_vocab_all_cat / lvis_pos_neg	Various	Mixed (polygon/RLE)
ade20k / ade20k_all_cat / ade20k_vocab_all_cat / ade20k_pos_neg	Various	RLE
refcoco / refcoco_plus	Reference	Native polygon
refclef	Reference	RLE
refcocog_baseline / grefcoco_baseline / coconut_pancap_baseline	Eval baselines	Various

Eval Datasets

• RefCOCO: refcoco, refcoco_plus — metrics: mIoU, oIoU, P@{0.5,0.6,0.7,0.8,0.9}
• Caption baselines: refcocog_baseline, grefcoco_baseline, coconut_pancap_baseline — metrics: mSpanIoU, mMaskIoU, Text/Mask/Joint P/R/F1@{0.5,0.75}

---

RNG Determinism

Training uses the same deterministic RNG system as ground_qwen35:

sample_rng = make_random_state(seed, idx, epoch, attempt + 1)
example = raw_dataset.get(idx, sample_rng)
result = preprocessor(example, rng=sample_rng)

For stochastic datasets like coconut_pos_neg, this ensures identical category sampling for a given (seed, idx, epoch).

---

Optimizer & LR Schedule

Single param group (lm_head tied to embed_tokens):

• Optimizer: AdamW, betas=(0.9, 0.95), eps=1e-6, weight_decay=0.0
• Schedule: Linear warmup (0 → peak over warmup_steps) → Cosine decay (peak → alpha_f * peak)
• Gradient clipping: max_grad_norm=1.0 (configurable)

---

Checkpointing

Saves at every --save_interval steps (before eval, so eval bugs can't lose checkpoints):

checkpoints/step{N}/
├── model.pt          # raw_model.state_dict() (rank 0 only)
├── optimizer.pt      # optimizer.state_dict() (rank 0 only)
└── train/
    ├── rank0.pt      # training state + RNG (each rank saves its own)
    └── rank1.pt

Training state includes: global_step, total_tokens, global_train_examples_seen, world_size, args, config, and full RNG state (python, numpy, torch, cuda).

Auto-resume: Detects latest step{N}/ in --save_dir and resumes model, optimizer, RNG states, and data iterator position.

---

CLI Arguments

Arg	Default	Description
Model
--pretrained_path	/weka-mm/royg/ground_qwen3_5/Qwen3.5-0.8B	Pretrained model path
--freeze_vision	True	Freeze ViT + PatchMerger
Data
--mode	debug	Dataset mixture: debug, full, debug_single
--batch_size	1	Per-GPU batch size
--num_workers	4	DataLoader workers
--max_poly_points	64	Polygon subsampling (0=no limit)
--max_seq_len	8192	Truncate sequences exceeding this
Training
--max_steps	10000	Total optimizer steps
--lr	1e-5	Learning rate
--warmup_steps	200	Linear warmup steps
--alpha_f	0.1	Final LR ratio
--max_grad_norm	1.0	Gradient clipping (0=disabled)
--grad_accum_steps	1	Gradient accumulation
--min_pixels	3136	Image min pixels
--max_pixels	1003520	Image max pixels
--seed	42	Random seed
Checkpointing
--save_interval	1000	Save every N steps
--save_dir	/weka-mm/royg/ft_qwen3_5/checkpoints	Checkpoint directory
--resume_from	None	Specific checkpoint to resume from
--log_interval	10	Console log every N steps
Eval
--eval_interval	0	Eval every N steps (0=disabled)
--eval_max_examples	1024	Max samples per eval dataset
Wandb
--wandb	off	Enable wandb
--wandb_project	ft_qwen35	Wandb project
--wandb_entity	ai2-roygao-allen-institute-for-ai	Wandb entity
--wandb_name	auto	Run name
--wandb_group	None	Group runs
--wandb_log_interval	20	Wandb log frequency

---

Wandb Metrics

Category	Metrics
train/	loss
optim/	backbone_lr, backbone_grad_norm
throughput/	tokens_per_second, batches_per_second, total_tokens
batch/	seq_len, num_target_tokens_mean
system/	peak_gpu_memory_mb
eval/{ds}/	loss, mIoU, oIoU, P@{t} (RefCOCO); mSpanIoU, mMaskIoU, Text/Mask/Joint P/R/F1@{t} (baselines)

---

DDP Safety

1. Frozen ViT: requires_grad=False → excluded from DDP buckets → find_unused_parameters=False
2. All trainable params in graph: Cross-entropy uses lm_head (tied to embed_tokens) → all LLM layers participate
3. No data-dependent branching: Every batch computes loss identically
4. Gradient accumulation: model.no_sync() for non-final micro-steps

---

File Structure

ft_qwen3_5/
├── ft_qwen35/
│   ├── __init__.py
│   ├── config.py                # FTQwen35Config dataclass
│   ├── train.py                 # Training entry point (DDP, optimizer, loop, eval, ckpt)
│   ├── eval_metrics.py          # Mask IoU, span IoU, RefCOCO/caption baseline metrics
│   └── data/
│       ├── __init__.py
│       ├── dataset.py           # GroundingSourceDataset, DeterministicFTDataset, EvalFTDataset
│       ├── iterable_mixture.py  # IterableFTMixture (Molmo2-style non-packing mixture)
│       ├── preprocessor.py      # FTExamplePreprocessor (image → polygon text → chat tokens)
│       ├── collator.py          # FTCollator (pad batches)
│       └── formatter.py         # RLE→polygon, target text format, chat sequence construction
├── scripts/
│   ├── convert_dataset.py       # Convert one dataset → JSONL (offline, inspection utility)
│   ├── convert_all.py           # Batch convert all datasets (parallel)
│   ├── inspect_data.py          # Print stats from JSONL files
│   └── visualize_preprocessing.py  # Visualize polygons + char spans on images
└── data/                        # Output dir for JSONL files and visualizations

---

Quick Start

cd /weka-mm/royg/ft_qwen3_5

# Debug mode — single GPU, 3 steps
/opt/conda/bin/python -m ft_qwen35.train \
    --mode debug --batch_size 1 --max_steps 3 --log_interval 1

# Multi-GPU training (8 GPUs, full mode)
TRITON_CACHE_DIR=/tmp/.triton/cache \
/opt/conda/bin/torchrun --nproc-per-node 8 -m ft_qwen35.train \
    --mode full --batch_size 4 --grad_accum_steps 4 \
    --max_steps 10000 --eval_interval 500 --wandb

# 2-GPU smoke test with eval
TRITON_CACHE_DIR=/tmp/.triton/cache \
/opt/conda/bin/torchrun --nproc-per-node 2 -m ft_qwen35.train \
    --mode full --batch_size 1 --max_steps 3 \
    --eval_interval 1 --eval_max_examples 8 --log_interval 1

# Visualize preprocessing (polygon + char span check)
/opt/conda/bin/python scripts/visualize_preprocessing.py \
    --dataset coconut_pancap --split val --num_samples 5 \
    --max_poly_points 32 --output_dir data/vis

Note: Always set TRITON_CACHE_DIR=/tmp/.triton/cache for multi-GPU runs to avoid Triton cache race conditions on WekaFS.

---

Sample Validation & Retry

Every sample passes through _validate_sample():

1. pixel_values must not be None (image loaded)
2. input_ids must exist and be non-empty
3. labels must exist
4. At least one label is not -100 (model has something to learn)

DeterministicFTDataset retries up to 50 times with deterministic random replacement on validation failure.

_has_segments() pre-filters empty-segment samples at dataset construction time.