# VILA训练脚本方案分析

## 1. ShareGPT4V vs VILA 格式差异

### ShareGPT4V格式（当前支持）
- **数据格式**: JSON文件 + 图像文件夹（分离存储）
- **数据集路径**: `DATASET_JSON` 指向JSON文件
- **图像路径**: `IMAGE_FOLDER` 指向图像文件夹
- **JSON结构**: 
  ```json
  [
    {
      "image": "relative/path/to/image.jpg",
      "conversations": [
        {"from": "human", "value": "..."},
        {"from": "gpt", "value": "..."}
      ]
    }
  ]
  ```
- **后端**: `CustomMultiModalDataset` (读取JSON + 图像文件夹)

### VILA WebDataset格式
- **数据格式**: WebDataset (.tar文件) + 元数据文件 (.wdsmeta.json)
- **元数据文件**: `/lustre/fsw/portfolios/nvr/users/ligengz/datasets/vila-webds-meta/*.wdsmeta.json`
- **元数据结构**:
  ```json
  {
    "name": "coyo-dev",
    "__kind__": "VILA-WebDataset",
    "wids_version": 1,
    "shardlist": [
      {
        "url": "/path/to/coyo25m-0-000000.tar",
        "nsamples": 10000,
        "filesize": 259962880
      }
    ]
  }
  ```
- **WebDataset特点**: 
  - 图像和文本数据打包在tar文件中
  - 每个tar文件包含多个样本
  - 需要WebDataset库来读取

## 2. 关键差异总结

| 项目 | ShareGPT4V | VILA WebDataset |
|------|------------|-----------------|
| **存储方式** | JSON + 图像文件夹（分离） | Tar文件（打包） |
| **数据加载** | 直接读取JSON | 需要WebDataset库 |
| **图像路径** | 相对路径，需要IMAGE_FOLDER | 打包在tar中 |
| **当前支持** | ✅ 已支持 | ❌ 不支持 |
| **后端类** | `CustomMultiModalDataset` | 需要新增或适配 |

## 3. 实现方案

### 方案A: 添加WebDataset后端（推荐，但需要代码修改）

**优点**:
- 保持WebDataset格式，无需转换
- 支持大规模数据集的高效加载
- 符合VILA原始数据格式

**缺点**:
- 需要修改代码，添加新的后端
- 需要安装WebDataset库
- 需要实现WebDataset到messages格式的转换

**实现步骤**:
1. 安装WebDataset: `pip install webdataset`
2. 在 `src/lmflow/datasets/dataset.py` 中添加 `backend == "webdataset"` 分支
3. 创建 `CustomWebDataset` 类，类似 `CustomMultiModalDataset`
4. 实现从WebDataset格式到messages格式的转换
5. 修改训练脚本，使用新的后端

**代码修改点**:
- `src/lmflow/datasets/dataset.py`: 添加webdataset后端
- `src/lmflow/datasets/multi_modal_dataset.py`: 添加 `CustomWebDataset` 类
- 训练脚本: 修改 `backend="webdataset"` 和数据集路径

### 方案B: 转换WebDataset为JSON格式（简单，但需要预处理）

**优点**:
- 无需修改代码
- 直接复用现有ShareGPT4V脚本
- 实现简单快速

**缺点**:
- 需要预处理步骤（转换WebDataset到JSON）
- 需要额外的存储空间（JSON + 图像）
- 可能丢失WebDataset的某些优势（如流式加载）

**实现步骤**:
1. 编写转换脚本，读取WebDataset并转换为JSON格式
2. 提取图像到文件夹
3. 生成JSON文件，格式与ShareGPT4V一致
4. 直接使用现有的ShareGPT4V训练脚本

**转换脚本示例**:
```python
# convert_vila_webds_to_json.py
import webdataset as wds
import json
from pathlib import Path

def convert_wdsmeta_to_json(wdsmeta_path, output_json, output_image_folder):
    # 读取元数据
    with open(wdsmeta_path) as f:
        meta = json.load(f)
    
    # 读取所有shard
    dataset = []
    for shard in meta['shardlist']:
        tar_path = shard['url']
        # 使用webdataset读取
        ds = wds.WebDataset(tar_path)
        for sample in ds:
            # 转换为JSON格式
            item = {
                "image": f"{Path(sample['__key__']).stem}.jpg",  # 保存图像
                "conversations": [...]  # 转换conversations格式
            }
            dataset.append(item)
    
    # 保存JSON
    with open(output_json, 'w') as f:
        json.dump(dataset, f)
```

### 方案C: 混合方案（推荐用于快速验证）

**步骤**:
1. **短期**: 使用方案B快速转换少量数据验证训练流程
2. **长期**: 实现方案A，支持原生WebDataset格式

## 4. 训练脚本差异对比

### ShareGPT4V脚本关键部分:
```bash
DATASET_JSON="/path/to/sharegpt4v.json"
IMAGE_FOLDER="/path/to/images"
# 使用 backend="custom_multi_modal"
```

### VILA脚本需要修改的部分:

**如果使用方案A（WebDataset后端）**:
```bash
DATASET_WDSMETA="/lustre/fsw/portfolios/nvr/users/ligengz/datasets/vila-webds-meta/xxx.wdsmeta.json"
# 使用 backend="webdataset"
# 不需要IMAGE_FOLDER（图像在tar中）
```

**如果使用方案B（转换后）**:
```bash
DATASET_JSON="/path/to/vila_converted.json"
IMAGE_FOLDER="/path/to/vila_images"
# 使用 backend="custom_multi_modal"（与ShareGPT4V相同）
```

## 5. 推荐方案

### 快速验证（方案B）
1. 编写转换脚本，转换少量数据（如1个shard）
2. 使用现有ShareGPT4V脚本训练
3. 验证训练流程是否正常

### 完整实现（方案A）
1. 添加WebDataset后端支持
2. 实现 `CustomWebDataset` 类
3. 创建VILA专用训练脚本

## 6. 具体实施建议

### 第一步: 快速验证（1-2小时）
1. ✅ 编写WebDataset到JSON的转换脚本
2. ✅ 转换1-2个shard测试
3. ✅ 使用ShareGPT4V脚本训练验证

### 第二步: 完整实现（1-2天）
1. 添加WebDataset后端
2. 实现数据加载类
3. 创建VILA训练脚本
4. 测试完整流程

## 7. 注意事项

1. **WebDataset格式细节**: 需要确认VILA WebDataset中每个样本的具体字段名（如`image`, `text`, `json`等）
2. **图像格式**: WebDataset中的图像可能是PIL Image对象或字节流
3. **Conversations格式**: 需要确认VILA WebDataset中conversations的存储格式
4. **性能**: WebDataset支持流式加载，可能比JSON+文件夹更高效
5. **路径问题**: WebDataset的tar文件路径可能需要调整（绝对路径 vs 相对路径）

## 8. 下一步行动

**建议先做**:
1. 查看一个实际的WebDataset tar文件内容，确认数据格式
2. 编写转换脚本（方案B）进行快速验证
3. 如果验证成功，再考虑实现方案A

**需要确认的信息**:
- VILA WebDataset中每个样本的字段结构
- 图像在tar中的存储方式
- Conversations的格式（是否已经是标准格式）
