# Diffusion Vision-Language Model (dVLM)

This directory contains a complete pipeline for training, evaluating, and running inference on Vision-Language Models using **Masked Diffusion Models (MDM)** and **Autoregressive (AR)** approaches based on Qwen2.5-VL.

## Architecture Overview

```
┌─────────────────────────────────────────────────────────────────────┐
│                           Data Pipeline                             │
├─────────────────────────────────────────────────────────────────────┤
│  ShareGPT4V / VILA J63  →  MultiModalDataset  →  DataCollatorQwenVL │
└───────────────────────────────────┬─────────────────────────────────┘
                                    ▼
┌─────────────────────────────────────────────────────────────────────┐
│                         Training Pipeline                            │
├─────────────────────────────────────────────────────────────────────┤
│  finetune_qwenvl.py  →  DeepSpeed (1-node) / FSDP+Slurm (N-nodes)  │
└───────────────────────────────────┬─────────────────────────────────┘
                                    ▼
┌─────────────────────────────────────────────────────────────────────┐
│                          Model Checkpoint                            │
├─────────────────────────────────────────────────────────────────────┤
│         Fast_dLLM_Qwen2_5_VLForConditionalGeneration                │
│                    ┌─────────┬─────────┐                            │
│                    │   AR    │   MDM   │                            │
│                    └─────────┴─────────┘                            │
└───────────────────────────────────┬─────────────────────────────────┘
                                    ▼
┌─────────────────────────────────────────────────────────────────────┐
│                    Evaluation / Inference                            │
├─────────────────────────────────────────────────────────────────────┤
│  eval_custom_model.py │ VLMEvalKit │ chatbot_mdm.py                 │
└─────────────────────────────────────────────────────────────────────┘
```

## Directory Structure

```
dvlm/
├── README.md                 # This file
├── training/                 # Training scripts
│   ├── README.md            # Training documentation
│   ├── single_node/         # Single-node training (DeepSpeed)
│   │   ├── train_j63_mdm.sh
│   │   ├── train_j63_ar.sh
│   │   └── train_sharegpt4v_vlm_ar_baseline.sh
│   └── multinode/           # Multi-node training (FSDP + Slurm)
│       ├── ar/              # AR training configs
│       └── mdm/             # MDM training configs
├── evaluation/              # Evaluation tools
│   ├── README.md            # Evaluation documentation
│   ├── eval_custom_model.py # Main evaluation script
│   ├── aggregate_results.py # Result aggregation
│   ├── monitor_checkpoints.py # Auto-eval monitor
│   ├── lmms_eval/           # lmms-eval wrapper
│   └── vlmeval/             # VLMEvalKit wrapper
├── inference/               # Inference and chatbots
│   ├── README.md            # Inference documentation
│   ├── chatbot.py           # AR chatbot
│   └── chatbot_mdm.py       # MDM chatbot
├── models/                  # Model utilities
│   ├── README.md            # Model documentation
│   └── generation_functions.py # MDM sampling
├── utils/                   # Utility scripts
└── demo_images/             # Demo images for testing
```

## Key Concepts

### Model Types

| Type | Description | Training | Generation |
|------|-------------|----------|------------|
| **AR** | Autoregressive | Standard causal LM loss | Token-by-token |
| **MDM** | Masked Diffusion | Block-wise denoising loss | Parallel within blocks |

### MDM Training

MDM (Masked Diffusion Model) uses block-wise denoising:

1. **Block Size (`bd_size`)**: Tokens are processed in blocks (default: 8)
2. **Masking**: Random tokens within each block are masked with `|<MASK>|`
3. **Denoising**: Model learns to predict masked tokens given context
4. **Generation**: Parallel sampling within blocks, sequential across blocks

Key training arguments:
- `--mdm 1`: Enable MDM training
- `--bd_size 8`: Block size for masking
- `--learn_padding 1`: Learn on padding tokens
- `--conplemenrary_mask 1`: Use complementary masking strategy

### Data Flow

```
Input: {"messages": [{"role": "user", "content": [{"type": "image"}, {"type": "text"}]}, ...]}
  │
  ▼
DataCollatorForQwenVL
  ├── Apply chat template
  ├── Process vision info (resize, normalize)
  ├── Tokenize with Qwen2.5-VL processor
  ├── Build labels (mask non-assistant tokens)
  └── Pad to bd_size multiple (if MDM)
  │
  ▼
Output: {input_ids, attention_mask, labels, pixel_values, image_grid_thw}
```

## Quick Start

### 0. 环境准备

请先按照根目录 `README.md` 的 **Setup** 完成 `dvlm-shiyi` 环境安装。
VLMEvalKit 的可选依赖也统一在根目录 `README.md` 的 Optional 部分说明。

### 1. Training

**Single-Node MDM:**
```bash
bash chengyue_script/dvlm/training/single_node/train_j63_mdm.sh
```

**Multi-Node MDM (8 nodes):**
```bash
sbatch chengyue_script/dvlm/training/multinode/mdm/train_j63_mdm_8node.sh
```

### 2. Evaluation

**Custom Checkpoint:**
```bash
python chengyue_script/dvlm/evaluation/eval_custom_model.py \
  --model_path output_models/j63_mdm/checkpoint-2000 \
  --tasks gqa,pope,textvqa,mmmu \
  --num_gpus 8
```

### 3. Inference

**MDM Chatbot:**
```bash
python chengyue_script/dvlm/inference/chatbot_mdm.py \
  --model-name output_models/j63_mdm/checkpoint-2000 \
  --demo
```

## Supported Benchmarks

| Benchmark | Type | Metric | Tool |
|-----------|------|--------|------|
| GQA | VQA | Exact match | eval_custom_model.py |
| POPE | Y/N | F1 score | eval_custom_model.py |
| TextVQA | OCR VQA | VQA accuracy | eval_custom_model.py |
| MMMU | Multiple choice | Accuracy | eval_custom_model.py |
| MMBench | Multiple choice | Accuracy | eval_custom_model.py |
| MathVista | Math reasoning | Accuracy | VLMEvalKit (GPT judge) |

## Model Checkpoints

Checkpoints are saved to `output_models/` with the following structure:
```
output_models/
└── {run_id}/
    ├── checkpoint-{step}/
    │   ├── config.json
    │   ├── model.safetensors
    │   ├── trainer_state.json
    │   └── eval_vlmeval/    # Evaluation results (if run)
    └── checkpoint-{step}/
```

## Environment Variables

| Variable | Description | Default |
|----------|-------------|---------|
| `MIN_PIXELS` | Minimum image resolution | 200704 |
| `MAX_PIXELS` | Maximum image resolution | 200704 |
| `VILA_DATASETS` | VILA dataset registry | cs-oci-ord |
| `DATASET_BACKEND` | Dataset backend type | custom_multi_modal |

## Related Files

| File | Location | Purpose |
|------|----------|---------|
| `finetune_qwenvl.py` | `examples/` | Main training script |
| `multi_modal_dataset.py` | `src/lmflow/datasets/` | Dataset and collator |
| `modeling_dfm.py` | `chengyue_script/` | DFM model implementation |
| `Fast_dLLM_*` | `src/models/Qwen2_5VL_MDM/` | Custom model architecture |

