Implementation:OpenGVLab InternVL LLaVA Custom Training Script
| Knowledge Sources | |
|---|---|
| Domains | Model Training, Multimodal Learning, InternVL, Multi-Dataset Training |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
This module extends the standard LLaVA training script with custom data loading from cloud storage, weighted multi-dataset training, and InternVL-specific vision encoder support.
Description
The train_custom.py file is the InternVL-specific training entry point that builds upon the standard LLaVA training pipeline with several important extensions:
- TCSLoader: A cloud storage image loader that uses the petrel_client library to fetch images from Tencent Cloud Storage (TCS) backends, falling back to standard PIL loading when the client is unavailable
- WeightedConcatDataset: Combines multiple training datasets using configurable sampling weights, computing weights proportional to the square root of dataset sizes for balanced multi-dataset training via WeightedRandomSampler
- Meta-file-driven data configuration: The training data is specified via a JSON meta file (data_args.meta_path) where each dataset entry specifies its annotation file path, root image directory, and repeat_time for oversampling
- Extended LazySupervisedDataset: Supports both local file paths and TCS cloud paths for image loading, handles both .json and .jsonl annotation formats
The train() function mirrors the base train.py but adds TCS loader initialization, meta-file-based dataset construction, and hard-codes the v1 conversation template for preprocessing. The data pipeline supports InternViT-6B position embedding interpolation for resolution scaling (224 to 336/448).
Usage
Use this script when training LLaVA models with InternVL vision encoders, especially when working with multiple datasets that need weighted sampling, or when image data is stored on cloud storage backends.
Code Reference
Source Location
- Repository: OpenGVLab_InternVL
- File: internvl_chat_llava/llava/train/train_custom.py
- Lines: 1-1067
Signature
class TCSLoader(object):
def __init__(self, conf_path): ...
def __call__(self, fn) -> Image: ...
class WeightedConcatDataset(ConcatDataset):
def __init__(self, datasets, weights): ...
class LazySupervisedDataset(Dataset):
def __init__(self, meta, tokenizer, tcs_loader, data_args): ...
def __getitem__(self, i) -> Dict[str, torch.Tensor]: ...
def make_supervised_data_module(tokenizer, data_args, tcs_loader) -> Dict: ...
def train(attn_implementation=None): ...
Import
from llava.train.train_custom import train, TCSLoader, WeightedConcatDataset
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| model_name_or_path | str | Yes | HuggingFace model path or local checkpoint |
| meta_path | str | Yes | Path to JSON meta file describing datasets with annotation paths, roots, and repeat times |
| vision_tower | str | No | Vision encoder model name/path (e.g., InternViT-6B) |
| output_dir | str | Yes | Directory to save trained model |
| lora_enable | bool | No | Enable LoRA adapter training (default: False) |
| bits | int | No | Quantization bits: 4, 8, or 16 (default: 16) |
Outputs
| Name | Type | Description |
|---|---|---|
| Model checkpoint | Directory | Saved model weights, config, and tokenizer |
| LoRA weights | .bin files | LoRA adapter state dict and non-LoRA trainables (when lora_enable=True) |
| mm_projector.bin | .bin file | Multimodal projector weights (when tune_mm_mlp_adapter=True) |
Usage Examples
Basic Usage
# Training with multi-dataset meta file
# deepspeed llava/train/train_custom.py \
# --model_name_or_path path/to/llava-model \
# --vision_tower path/to/InternViT-6B-448px-V1-2 \
# --meta_path ./data/meta.json \
# --output_dir ./checkpoints/internvl-custom \
# --bf16 True --deepspeed ./scripts/zero3.json
# Example meta.json format:
# {
# "dataset1": {
# "root": "/path/to/images",
# "annotation": "/path/to/train.json",
# "repeat_time": 1
# },
# "dataset2": {
# "root": "s3://bucket/images",
# "annotation": "/path/to/train.jsonl",
# "repeat_time": 2
# }
# }