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/safetensorspair here, sofrom_pretrained("royguw/ft_qwen3_5")will not work. Use it to reproduce or inspect the baseline; train withft_qwen35/train.pyagainst the released data (royguw/grounding-dataandroyguw/mm-olmo-images).Contents:
ft_qwen35/(the package),scripts/,ft_eval_results*/(metrics JSON),ft_eval_cache*/(cached eval tensors,.pt), andwandb/(offline run logs). Paths inside the code and logs are absolute cluster paths and will need rewriting.
WANDB_API_KEYwas removed fromft_qwen35/train.pyfor 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:
- Split by
\n→ one string per instance - Split each line by
;→(polygon_part, spans_part) - Split
polygon_partby|→ list of segment strings - Split each segment by spaces, pair up →
[(x1,y1), (x2,y2), ...] - Split
spans_partby 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_catpos_neg(0.2): coconut_pos_neg, ade20k_pos_neg, lvis_pos_negpanoptic_caption(0.3): coconut_pancapreference(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
- Frozen ViT:
requires_grad=False→ excluded from DDP buckets →find_unused_parameters=False - All trainable params in graph: Cross-entropy uses lm_head (tied to embed_tokens) → all LLM layers participate
- No data-dependent branching: Every batch computes loss identically
- 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():
pixel_valuesmust not be None (image loaded)input_idsmust exist and be non-emptylabelsmust exist- 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.