# VLM Training Guide

Guide for training Vision-Language Models with MDM (Masked Diffusion Model) and AR (Autoregressive) approaches.

## Training Modes

| Mode | Infrastructure | Framework | Use Case |
|------|---------------|-----------|----------|
| Single-Node | 1 node, 8 GPUs | DeepSpeed ZeRO-2/3 | Development, small datasets |
| Multi-Node | 2-16 nodes | FSDP + Slurm | Large-scale training |

## Directory Structure

```
training/
├── README.md                 # This file
├── single_node/              # Single-node training scripts
│   ├── train_j63_mdm.sh      # MDM training with J63 dataset
│   ├── train_j63_ar.sh       # AR training with J63 dataset
│   ├── train_sharegpt4v_vlm.sh
│   └── train_sharegpt4v_vlm_ar_baseline.sh
└── multinode/                # Multi-node training scripts
    ├── ar/                   # AR training configs
    │   ├── train_j63_ar_2node.sh
    │   └── train_j63_ar_8node.sh
    └── mdm/                  # MDM training configs
        ├── train_j63_mdm_2node.sh
        ├── train_j63_mdm_4node.sh
        ├── train_j63_mdm_8node.sh
        └── train_j63_mdm_16node.sh
```

---

## Quick Start

### Single-Node Training

**MDM Training (VILA J63 Dataset):**
```bash
bash chengyue_script/dvlm/training/single_node/train_j63_mdm.sh
```

**AR Baseline Training (ShareGPT4V):**
```bash
bash chengyue_script/dvlm/training/single_node/train_sharegpt4v_vlm_ar_baseline.sh
```

### Multi-Node Training

**MDM 8-Node Training:**
```bash
sbatch chengyue_script/dvlm/training/multinode/mdm/train_j63_mdm_8node.sh
```

**AR 8-Node Training:**
```bash
sbatch chengyue_script/dvlm/training/multinode/ar/train_j63_ar_8node.sh
```

---

## Configuration Parameters

### Model Parameters

| Parameter | Description | Default |
|-----------|-------------|---------|
| `MODEL_PATH` | Path to base model | `src/models/Qwen2.5-VL-3B-Instruct-mdm` |
| `TOKENIZER_NAME` | Tokenizer name/path | `Qwen/Qwen2.5-VL-3B-Instruct` |
| `TEMPLATE` | Conversation template | `qwen2_5_no_reasoning` |

### Training Parameters

| Parameter | Description | Default |
|-----------|-------------|---------|
| `EPOCHS` | Number of training epochs | 1 |
| `BATCH` | Per-device batch size | 1-2 |
| `GRAD_ACC` | Gradient accumulation steps | 4-8 |
| `LR` | Learning rate | 2e-5 |
| `SAVE_STEPS` | Save checkpoint every N steps | 500 |
| `SAVE_TOTAL_LIMIT` | Maximum checkpoints to keep | 20 |

### MDM-Specific Parameters

| Parameter | Description | Default |
|-----------|-------------|---------|
| `MDM` | Enable MDM training mode | 1 |
| `BD_SIZE` | Block size for diffusion | 8 |
| `LEARN_PADDING` | Learn on padding tokens | 1 |
| `CONPLE_MASK` | Use complementary masking | 1 |

### Image Processing Parameters

| Parameter | Description | Default |
|-----------|-------------|---------|
| `MIN_PIXELS` | Minimum image resolution | 200704 |
| `MAX_PIXELS` | Maximum image resolution | 200704 |

> **Note**: 200704 = 448×448, aligned with 14×14 patch size. Adjust based on memory constraints.

---

## Dataset Configuration

### ShareGPT4V Dataset

```bash
DATASET_JSON="/path/to/sharegpt4v_mix665k.json"
IMAGE_FOLDER="/path/to/ShareGPT4V/data"
export DATASET_BACKEND="custom_multi_modal"
```

### VILA J63 Mixture

```bash
DATASET_MIXTURE="j63_recipe"  # or j63_no_sft, j63_small, j63_tiny
export VILA_DATASETS="cs-oci-ord"
export DATASET_BACKEND="webdataset_multi_modal"
```

---

## Auto-Resume

Training scripts support automatic resume from the latest checkpoint:

```bash
# Auto-resume is enabled by default
AUTO_RESUME="${AUTO_RESUME:-1}"

# The script automatically finds the latest checkpoint
LATEST_CHECKPOINT=$(find "${OUTPUT_DIR}" -maxdepth 1 -type d -name "checkpoint-*" | sort -V | tail -n 1)
```

To force restart from scratch:
```bash
AUTO_RESUME=0 bash chengyue_script/dvlm/training/single_node/train_j63_mdm.sh
```

---

## Slurm Configuration (Multi-Node)

### Job Parameters

```bash
#SBATCH -A nvr_elm_llm                    # Account
#SBATCH -p batch_block1,batch_block3      # Partition
#SBATCH -t 04:00:00                       # Time limit
#SBATCH -N 8                              # Number of nodes
#SBATCH --gpus-per-node 8                 # GPUs per node
#SBATCH --array=1-5%1                     # Job array for requeue
#SBATCH --requeue                         # Enable requeue
#SBATCH --signal=B:USR1@120               # Signal handler for timeout
```

### Signal Handler for Graceful Shutdown

Scripts include a signal handler that catches Slurm timeout warnings (USR1) and saves checkpoints before requeue:

```bash
handle_timeout() {
    echo "Received USR1 signal (SLURM timeout warning)"
    echo "Gracefully stopping training for checkpoint save..."
    kill -TERM $TRAIN_PID 2>/dev/null || true
    wait $TRAIN_PID 2>/dev/null || true
    scontrol requeue $SLURM_JOB_ID
    exit 0
}
trap handle_timeout USR1
```

---

## Framework Selection

### Single-Node: DeepSpeed

```bash
deepspeed --master_port=23456 \
    examples/finetune_qwenvl.py \
    --deepspeed configs/ds_config_zero2_no_offload.json \
    ...
```

DeepSpeed configs:
- `ds_config_zero2_no_offload.json`: ZeRO-2 without CPU offload
- `ds_config_zero3.json`: ZeRO-3 with offload (for larger models)

### Multi-Node: FSDP

```bash
srun --nodes=${NNODES} --ntasks-per-node=1 \
    bash -c 'torchrun \
        --nnodes=${NNODES} \
        --nproc_per_node=${NGPUS_PER_NODE} \
        --node_rank=${SLURM_NODEID} \
        --master_addr=${MASTER_ADDR} \
        --master_port=${MASTER_PORT} \
        examples/finetune_qwenvl.py \
        --fsdp "full_shard" \
        ...'
```

> **Note**: Use `--fsdp "full_shard"` without `auto_wrap` to avoid NCCL hangs. The model's `_no_split_modules` attribute handles layer wrapping.

---

## Memory Optimization

### Activation Checkpointing

```bash
ACTIVATION_CHECKPOINTING=1
# Adds --gradient_checkpointing to training command
```

### CUDA Memory Configuration

```bash
export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True,max_split_size_mb:64,garbage_collection_threshold:0.9"
```

### Batch Size Tuning

| Model Size | GPUs | Suggested Batch | Grad Acc | Effective Batch |
|------------|------|-----------------|----------|-----------------|
| 3B | 8 | 1 | 8 | 64 |
| 3B | 64 | 1 | 4 | 256 |
| 7B | 8 | 1 | 16 | 128 |

---

## Monitoring

### GPU Monitoring

Scripts automatically log GPU usage:
```bash
nvidia-smi --query-gpu=timestamp,memory.used,memory.total,utilization.gpu,pstate \
           --format=csv,noheader,nounits -l 60 > ${log_dir}/gpu_monitor.log 2>&1 &
```

### WandB Integration

Training is logged to WandB by default:
```bash
wandb login
# Training will auto-log to your WandB project
```

To disable:
```bash
export WANDB_MODE=disabled
```

---

## Output Structure

```
output_models/
└── {RUN_ID}/
    ├── checkpoint-{step}/
    │   ├── config.json           # Model config
    │   ├── model.safetensors     # Model weights
    │   ├── optimizer.pt          # Optimizer state
    │   ├── scheduler.pt          # LR scheduler state
    │   ├── trainer_state.json    # Training state
    │   └── training_args.bin     # Training arguments
    └── ...
```

---

## Troubleshooting

### NCCL Timeout

If you see `_ALLGATHER_BASE` timeout errors:
1. Don't use `--fsdp "full_shard auto_wrap"`
2. Use `--fsdp "full_shard"` only
3. Increase `--ddp_timeout`

### Out of Memory

1. Reduce `BATCH` size
2. Increase `GRAD_ACC` to maintain effective batch size
3. Enable activation checkpointing: `ACTIVATION_CHECKPOINTING=1`
4. Reduce `MAX_PIXELS` for smaller images

### Dataset Not Found

For VILA datasets:
```bash
export VILA_DATASETS="cs-oci-ord"  # Correct dataset registry
```

For ShareGPT4V:
- Verify `DATASET_JSON` and `IMAGE_FOLDER` paths exist
- Check image paths in JSON are relative to `IMAGE_FOLDER`

