## Installation
```bash
# python 3.10.16
pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu118

pip install -r requirements.txt
```


## Warmup (optional)

<details>
<summary>Recommended training configurations</summary>


| Model Family           | Model Name                     | GPUs | Accelerate Config    |
|------------------------|--------------------------------|------|----------------------|
| **Qwen2-VL Models**    |                                |      |                      |
|                        | Qwen2-VL-2B                    | 2    | `deepspeed_zero2`    |
|                        | Qwen2-VL-2B-Instruct           | 2    | `deepspeed_zero2`    |
|                        | Qwen2-VL-7B                    | 4    | `deepspeed_zero2`    |
|                        | Qwen2-VL-7B-Instruct           | 4    | `deepspeed_zero2`    |
|                        | Qwen2-VL-72B                   | 24   | `deepspeed_zero3`    |
|                        | Qwen2-VL-72B-Instruct          | 24   | `deepspeed_zero3`    |
|                        |                                |      |                      |
| **Qwen2.5-VL Models**  |                                |      |                      |
|                        | Qwen2.5-VL-3B-Instruct         | 2    | `deepspeed_zero2`    |
|                        | Qwen2.5-VL-7B-Instruct         | 4    | `deepspeed_zero2`    |
|                        | Qwen2.5-VL-32B-Instruct        | 16   | `deepspeed_zero3`    |
|                        | Qwen2.5-VL-72B-Instruct        | 24   | `deepspeed_zero3`    |
|                        |                                |      |                      |
| **Gemma3 Models**      |                                |      |                      |
|                        | gemma-3-4b-pt                  | 4    | `deepspeed_zero2`    |
|                        | gemma-3-4b-it                  | 4    | `deepspeed_zero2`    |
|                        | gemma-3-12b-pt                 | 8    | `deepspeed_zero2`    |
|                        | gemma-3-12b-it                 | 8    | `deepspeed_zero2`    |
|                        | gemma-3-27b-pt                 | 16   | `deepspeed_zero3`    |
|                        | gemma-3-27b-it                 | 16   | `deepspeed_zero3`    |
|                        |                                |      |                      |
| **Llama-3.2-Vision**   |                                |      |                      |
|                        | Llama-3.2-11B-Vision           | 8    | `deepspeed_zero2`    |
|                        | Llama-3.2-11B-Vision-Instruct  | 8    | `deepspeed_zero2`    |
|                        | Llama-3.2-90B-Vision           | 32   | `deepspeed_zero3`    |
|                        | Llama-3.2-90B-Vision-Instruct  | 32   | `deepspeed_zero3`    |
| **Llava-1.5**          |                                |      |                      |
|                        | llava-1.5-7b-hf                | 8    | `deepspeed_zero2`    |
|                        | llava-1.5-13b-hf               | 8    | `deepspeed_zero3`    |
| **Llava-1.6**          |                                |      |                      |
|                        | llava-v1.6-vicuna-7b-hf        | 8    | `deepspeed_zero2`    |
|                        | llava-v1.6-vicuna-13b-hf       | 8    | `deepspeed_zero3`    |
|                        | llava-v1.6-34b-hf              | 24   | `deepspeed_zero3`    |
| **InternVL3**          |                                |      |                      |
|                        | InternVL3-1B                   | 2    | `deepspeed_zero2`    |
|                        | InternVL3-8B                   | 8    | `deepspeed_zero2`    |
|                        | InternVL3-14B                  | 8    | `deepspeed_zero3`    |
|                        | InternVL3-38B                  | NaN  | `deepspeed_zero3`    |
|                        | InternVL3-78B                  | NaN  | `deepspeed_zero3`    |

</details>

<details>
<summary>Examples of training on images</summary>

```bash
# train 7B model on images (default with Qwen2-VL-7B, deepspeed_zero2, one node, 8 gpus)
sh scripts/train.sh --data_config "animals/config_image"

# train 72B model on images
sh scripts/train.sh \
--nodes 3 --accelerate_config deepspeed_zero3 \
--model_name_or_path /mnt/lustrenew/mllm_safety-shared/models/huggingface/Qwen/Qwen2.5-VL-72B-Instruct \
--data_config "animals/config_image"

# train 32B model on images
sh scripts/train.sh \
--nodes 2 --accelerate_config deepspeed_zero3 \
--model_name_or_path /mnt/lustrenew/mllm_safety-shared/models/huggingface/Qwen/Qwen2.5-VL-32B-Instruct \
--data_config "animals/config_image"
```
</details>

<details>
<summary>Examples of training on patches</summary>

```bash
# split images into patches
PYTHONPATH=. srun -p p-cpu-new --quotatype=reserved --cpus-per-task=8 --time=30000 \
python src/tools/patches_split.py \
--src_images_dir "data/animals/files" \
--tgt_patches_dir "tmp/data/animals/files/4x4"
# train on patches (default with Qwen2-VL-7B, deepspeed_zero2, one node, 8 gpus)
sh scripts/train.sh \
--data_config "animals/config_patch" \
--data_overwrite_args "data.train[0].patches_dirs[0]=tmp/data/animals/files/4x4" \
--data_output_field "animals/config_patch/4x4" \
--epochs 5

# plotting
sh scripts/plot.sh models
```
</details>

<details>
<summary>Examples of single-node synchronous training (useful for debugging)</summary>

```bash
 sh scripts/tmp/dev.sh \
 --model_name_or_path "/mnt/lustrenew/mllm_safety-shared/models/huggingface/google/gemma-3-4b-pt" \
 --data_config "animals/config_image" --data_output_field "tmp" --epochs 20
```

</details>

## Main exps

<details>
<summary>Step-by-step walkthrough (optional)</summary>

Data preprocessing (this is model-agnostic)
```bash
dataset="animals"
filter_prompt="Can you tell what animal is shown partially in the image?"

# data spliting
for patch_ratio in 2 4 8; do
    PYTHONPATH=. srun -p p-cpu-new --quotatype=reserved --cpus-per-task=8 --time=30000 \
    python src/tools/patches_split.py \
    --src_images_dir "data/${dataset}/files" \
    --tgt_patches_dir "tmp/data/${dataset}/files/${patch_ratio}x${patch_ratio}" \
    --patch_ratio "${patch_ratio}" &
    sleep 1
done
wait

# data filtering and stitching
for patch_ratio in 4 8; do
    (
        PYTHONPATH=. srun -p mllm_safety --quotatype=reserved --gres=gpu:1 --cpus-per-task=8 --time=30000 \
        python src/tools/patches_filter.py \
        --src_patches_dir "tmp/data/${dataset}/files/${patch_ratio}x${patch_ratio}" \
        --tgt_patches_dir "tmp/data/${dataset}/files/${patch_ratio}x${patch_ratio}_unrecognizable" \
        --filter_class_name "Qwen_VL_Filter" --filter_kwargs "--prompt '${filter_prompt}'"

        PYTHONPATH=. srun -p p-cpu-new --quotatype=reserved --cpus-per-task=8 --time=30000 \
        python src/tools/patches_stitch.py \
        --src_patches_dir "tmp/data/${dataset}/files/${patch_ratio}x${patch_ratio}_unrecognizable" \
        --tgt_images_dir "tmp/data/${dataset}/files/${patch_ratio}x${patch_ratio}_unrecognizable_stitched" \
        --patch_ratio "${patch_ratio}"
    ) &
    sleep 1
done
wait
```

Training

```bash
dataset="animals"
model_name_or_path="/mnt/lustrenew/mllm_safety-shared/models/huggingface/Qwen/Qwen2-VL-7B"
quotatype="reserved"
nodes=1
gpu=4
accelerate_config="deepspeed_zero2"
num_seeds=3 # TODO: increase to 5 for final results

# train on images
sh scripts/train.sh \
--accelerate_config "${accelerate_config}" --nodes ${nodes} --gpu ${gpu} --quotatype "${quotatype}" \
--num_seeds ${num_seeds} \
--model_name_or_path "${model_name_or_path}" \
--data_config "${dataset}/config_image" \
--epochs 15

# train on unfiltered patches
for patch_ratio in 2 4 8; do
    sh scripts/train.sh \
    --accelerate_config "${accelerate_config}" --nodes ${nodes} --gpu ${gpu} --quotatype "${quotatype}" \
    --num_seeds ${num_seeds} \
    --model_name_or_path "${model_name_or_path}" \
    --data_config "${dataset}/config_patch" \
    --data_overwrite_args "data.train[0].patches_dirs[0]=tmp/data/${dataset}/files/${patch_ratio}x${patch_ratio}" \
    --data_output_field "${dataset}/config_patch/${patch_ratio}x${patch_ratio}" \
    --epochs 5
done

# train on filtered patches
for patch_ratio in 4 8; do
    sh scripts/train.sh \
    --accelerate_config "${accelerate_config}" --nodes ${nodes} --gpu ${gpu} --quotatype "${quotatype}" \
    --num_seeds ${num_seeds} \
    --model_name_or_path "${model_name_or_path}" \
    --data_config "${dataset}/config_patch" \
    --data_overwrite_args "data.train[0].patches_dirs[0]=tmp/data/${dataset}/files/${patch_ratio}x${patch_ratio}_unrecognizable" \
    --data_output_field "${dataset}/config_patch/others/${patch_ratio}x${patch_ratio}_unrecognizable" \
    --epochs 5
done

# train on stitched patches after filtering
for patch_ratio in 4 8; do
    sh scripts/train.sh \
    --accelerate_config "${accelerate_config}" --nodes ${nodes} --gpu ${gpu} --quotatype "${quotatype}" \
    --num_seeds ${num_seeds} \
    --model_name_or_path "${model_name_or_path}" \
    --data_config "${dataset}/config_image" \
    --data_overwrite_args "data.train[0].images_dirs[0]=tmp/data/${dataset}/files/${patch_ratio}x${patch_ratio}_unrecognizable_stitched" \
    --data_output_field "${dataset}/config_patch/others/${patch_ratio}x${patch_ratio}_unrecognizable_stitched" \
    --epochs 15
done
```

Plotting
```bash
sh scripts/plot.sh models
```
</details>

<!-- <details>
<summary>Automate all exps with scripts</summary> -->
### Automate all exps with scripts

Data preprocessing (this is model-agnostic)

```bash
sh scripts/exps/main/split.sh
sh scripts/exps/main/sfilter.sh # (optional) needed for ablation
```

Training 

```bash
# for full runs
sh scripts/exps/main/train.sh --user "zhouzhanhui"

# for full ablations
sh scripts/exps/main/train.sh --user "zhouzhanhui" --ablation

# for running on selected models and datasets
sh scripts/exps/main/train.sh \
--user "zhouzhanhui" \
--models Qwen2-VL-2B,Qwen2-VL-2B-Instruct,Qwen2-VL-7B,Qwen2-VL-7B-Instruct \
--datasets food,landmarks
```
<!-- </details> -->

## Moderation
```bash
# data preparation for the moderation tasks
PYTHONPATH=. srun -p p-cpu-new --quotatype=reserved --cpus-per-task=8 --time=30000 \
    python src/tools/patches_split.py \
    --src_images_dir "data/moderation/files" \
    --tgt_patches_dir "tmp/data/moderation/files/4x4" \
    --patch_ratio 4
PYTHONPATH=. srun -p mllm_safety --quotatype=reserved --gres=gpu:1 --cpus-per-task=8 --time=30000 \
    python src/tools/patches_filter.py \
    --src_patches_dir "tmp/data/moderation/files/4x4" \
    --tgt_patches_dir "tmp/data/moderation/files/4x4_safe" \
    --filter_class_name "Llama_Guard_3_11b_Vision_Filter"
PYTHONPATH=. srun -p p-cpu-new --quotatype=reserved --cpus-per-task=8 --time=30000 \
    python src/tools/patches_stitch.py \
    --src_patches_dir "tmp/data/moderation/files/4x4_safe" \
    --tgt_images_dir "tmp/data/moderation/files/4x4_safe_stitched"

# train
sh scripts/train.sh \
--data_config "moderation/config_patch" \
--data_overwrite_args "data.train[0].patches_dirs[0]=tmp/data/moderation/files/4x4_safe" \
--data_output_field "moderation/config_patch/4x4_safe" \
--epochs 5 \
--num_seeds 5 \
--gpu 4

# plotting
sh scripts/plot.sh models/moderation
```
