code/docs/training.md

# M4SEG training

PRIMA trains on [M⁴Seg](https://huggingface.co/datasets/PLAN-Lab/M4Seg) using one dataset adapter. Follow the [data guide](datasets.md), then set the model assets, optional projector initialization weights, dataset root, and training annotation path in `configs/train.yaml`.

For a single GPU:

```bash
python train.py --config configs/train.yaml
```

For multiple GPUs:

```bash
deepspeed --num_gpus 4 train.py --config configs/train.yaml
```

To further fine-tune the complete [released model](https://huggingface.co/PLAN-Lab/PRIMA), run:

```bash
python train.py --config configs/train.yaml --merged_ckpt_path PLAN-Lab/PRIMA
```

This initializes from the merged release and attaches fresh language-model LoRA branches; separate base, EVA, Q-Former, and projector initialization paths are ignored. The dataset paths are still required. The default training configuration uses separate initialization assets to train from the pretrained backbone.

The merged checkpoint includes Q-Former and its tokenizer, so further fine-tuning with `merged_ckpt_path` needs no separate Q-Former checkpoint. Only training from the separate backbone assets requires the Q-Former initialization paths in `configs/train.yaml`.

Training updates LLM query/value LoRA, SQuARE attention and query tokens, the text-to-mask projection, SAM mask decoder, and optionally the visual-to-language projector. The language backbone, EVA/Q-Former backbone, and SAM image/prompt encoders remain frozen. A pretrained visual-to-language projector can be supplied as initialization; this repository does not train a separate pretraining stage.

Effective batch size is per-device batch size multiplied by GPU count and accumulation. The example gives an effective batch of 160 on four GPUs. AdamW uses linear learning-rate warmup followed by cosine decay. The example is a starting configuration; asset-specific recipes and expected results are documented separately.

Use `--max_steps` for a short smoke run. Scalar losses are printed at optimizer steps. Non-finite losses or gradients stop training.

Recoverable state and portable trainable-state exports are saved under the configured output directory. Resume with `--resume /path/to/recoverable_checkpoint` or `--auto_resume`. Keep batch size and accumulation unchanged when resuming. Distributed runs resume their DeepSpeed state; single-GPU runs restore optimizer, scheduler, progress, and random states. Portable exports omit frozen backbones and require the same base assets at inference.