Repository: oumi-ai/oumi
File: scripts/benchmarks/minimal_multimodal_training.py (Python)

Code before:
        return "test"
    elif dataset_name in (DatasetName.MERVE_VQAV2_SMALL,):
        return "validation"
    return "train"


def test_multimodal_trainer(
    model_name: ModelName = ModelName.BLIP2,
    dataset_name: DatasetName = DatasetName.COCO,
    batch_size: int = 2,
    max_steps: int = 20,
    optimizer: str = "sgd",
    logging_steps: int = 5,
    split: Optional[str] = None,
    test_inference: bool = False,
    test_save_state: bool = False,
    test_fsdp: bool = False,
):
    """Minimal multi-modal training loop."""
    if is_distributed():
        print("Initializing distributed process group")
        init_distributed()
    else:
        print("Not initializing distributed process group")

    if model_name in [ModelName.QWEN2_VL, ModelName.QWEN2_5_3B_VL] and batch_size != 1:
        print(
            f"Using batch size 1 for {model_name.value} (original: bs={batch_size}). "
            "The model only supports bs=1 because of variable-size image encodings."
        )
        batch_size = 1

    if not split:
        split = _get_default_dataset_split(dataset_name)

    #
    # Init model, processor, and dataset
    #
    model_params = ModelParams(
        model_name=model_name.value,
        torch_dtype_str="bfloat16",
        trust_remote_code=True,
        chat_template=_get_chat_template(model_name),
        freeze_layers=_get_freeze_layers(model_name),  # TODO: fix freeze + fsdp
    )
    if is_local_process_zero():
        print(f"ModelParams:\n{pformat(model_params)}")

    model = build_model(model_params)
    tokenizer: BaseTokenizer = build_tokenizer(model_params)
    processor: BaseProcessor = build_processor(

Review comment (suggestion):
nit: let's use tuple instead of array  ` ... in (ModelName.QWEN2_VL, ModelName.QWEN2_5_3B_VL)`
It's arguably cleaner because the list is supposed to be  immutable 

Code after:
        return "test"
    elif dataset_name in (DatasetName.MERVE_VQAV2_SMALL,):
        return "validation"
    return "train"


def test_multimodal_trainer(
    model_name: ModelName = ModelName.BLIP2,
    dataset_name: DatasetName = DatasetName.COCO,
    batch_size: int = 2,
    max_steps: int = 20,
    optimizer: str = "sgd",
    logging_steps: int = 5,
    split: Optional[str] = None,
    test_inference: bool = False,
    test_save_state: bool = False,
    test_fsdp: bool = False,
):
    """Minimal multi-modal training loop."""
    if is_distributed():
        print("Initializing distributed process group")
        init_distributed()
    else:
        print("Not initializing distributed process group")

    if model_name in (ModelName.QWEN2_VL, ModelName.QWEN2_5_3B_VL) and batch_size != 1:
        print(
            f"Using batch size 1 for {model_name.value} (original: bs={batch_size}). "
            "The model only supports bs=1 because of variable-size image encodings."
        )
        batch_size = 1

    if not split:
        split = _get_default_dataset_split(dataset_name)

    #
    # Init model, processor, and dataset
    #
    model_params = ModelParams(
        model_name=model_name.value,
        torch_dtype_str="bfloat16",
        trust_remote_code=True,
        chat_template=_get_chat_template(model_name),
        freeze_layers=_get_freeze_layers(model_name),  # TODO: fix freeze + fsdp
    )
    if is_local_process_zero():
        print(f"ModelParams:\n{pformat(model_params)}")

    model = build_model(model_params)
    tokenizer: BaseTokenizer = build_tokenizer(model_params)
    processor: BaseProcessor = build_processor(