Continuous VLM Post-Training

This example builds a continuous post-training loop for a vision-language model. When labeled images are appended to a dataset, Mixtrain starts a LoRA fine-tuning run. The completed run registers a checkpoint, which starts a second routine that serves and evaluates the fine-tuned model against the base model. The example's architecture is useful for any training process where new data should produce a versioned model candidate and an evaluation.

The complete project, including training code and sample images, is in continuous-vlm-post-training. Let's break it down.

Setup

The setup script creates three datasets and an Eval in the Mixtrain workspace:

vlm-post-training-data contains labeled images, prompts, and if the row is for training or evaluation.
vlm-checkpoints acts as a registry for checkpoints produced by training runs.
vlm-eval-results stores answers from the base and fine-tuned models.
The Eval presents those answers beside the image, prompt, and expected answer.

Two routines are used to create the continuous loop:

TrainOnNewData routine starts the training workflow when new training rows are appended to the vlm-post-training-data dataset.
EvaluateNewCheckpoints routine starts model deployment and evaluation when a new checkpoint is registered in the vlm-checkpoints.

Both transitions are driven by dataset appends, routines declare on_added_rows() to recive new data. You do not need a separate scheduler or pipeline service to coordinate the loop, they are automatically triggered by the platform on the appropriate events.

Write the training workflow

A workflow is a MixFlow subclass with a run() method. Its arguments are the inputs users can set for each run, and _sandbox specifies the environment and compute for it:

from mixtrain import Dataset, MixFlow, Sandbox


class TrainVLM(MixFlow):
    _sandbox = Sandbox(
        image="pytorch/pytorch:2.6.0-cuda12.4-cudnn9-runtime",
        gpu="A10G",
        gpu_per_node=2,
    )

    def run(
        self,
        training_data: Dataset,
        base_model: str = "HuggingFaceTB/SmolVLM-256M-Instruct",
        num_epochs: int = 3,
        learning_rate: float = 2e-4,
        checkpoint_dataset: str = "vlm-checkpoints",
    ):
        ...

One thing to note here, training_data arugment is a Mixtrain dataset. You can pass a dataset directly rather than serializing a path or downloading it before the run. In this example, the routine passes a specific dataset version, so the training run remains reproducible even after more rows are appended.

Here we are also declaring Sandbox with a two(gpu_per_node=2) A10G node for the workflow to run on. The example's training code uses both that for multi-GPU training for distributed LoRA fine-tuning. See training.py for the PyTorch and TRL implementation code. You can also override the sandbox with a custom one, per run if needed. So running the same workflow on different GPUs doesn't require any changes to the workflow code.

Select the training split

The source dataset contains both training and held-out evaluation examples. The workflow selects training rows before preparing the data:

sample_count = prepare_splits(
    training_data.filter("split == 'train'"),
    dataset_path,
)

Dataset.filter() here allows you to pass python or sql expressions to filter the dataset. Note that chaining multiple transformations like this (select(), shuffle(), limit()) doesn't change the original dataset, unless explictly saved. This makes working with large dataset easy and efficient. See Working with datasets for the other dataset capabilities.

Register the trained checkpoint

After training, the workflow selects the epoch with the lowest validation loss and uploads that directory as a typed checkpoint:

checkpoint = Checkpoint.from_dir(
    str(best_checkpoint),
    f"/data/vlm-post-training-checkpoints/{experiment_id}",
    task="image-text-to-text",
    source_run=experiment_id,
)

Dataset("vlm-checkpoints").append(
    [{
        "experiment_id": experiment_id,
        "base_model": base_model,
        "checkpoint": checkpoint,
        "best_validation_loss": best_validation_loss,
        "dataset_version": training_data.version,
    }]
)

Checkpoint.from_dir() uploads the checkpoint files and records how the artifact should be served. Storing the resulting Checkpoint in a dataset gives each experiment a queryable record alongside its base model, validation metric, and source-data version.

Automatic training on new data

TrainOnNewData routine declares its on_added_rows() to recive a dataset of rows appended since its previous run. The routine uses it to query if new_rows contains any training data.

from mixtrain import Dataset, MixRoutine, Workflow, on_added_rows


class TrainOnNewData(MixRoutine):
    def run(
        self,
        new_rows=on_added_rows("vlm-post-training-data"),
        training_workflow: str = "train-vlm",
    ) -> Workflow | None:
        dataset_version = new_rows.version

        if new_rows.filter("split == 'train'").head(1).collect().num_rows == 0:
            return None

        run = Workflow(training_workflow).submit(
            inputs={
                "training_data": Dataset(
                    "vlm-post-training-data",
                    version=dataset_version,
                ),
            }
        )
        return Workflow(training_workflow, run_number=run["run_number"])

This example retrains on all labeled data available at that point and hence passes the entire dataset to the workflow. If your training strategy is incremental, you can pass new_rows instead. Returning the submitted Workflow associates the routine run with the training run, making the relationship visible in Mixtrain.

Evaluate each new checkpoint

The second EvaluateNewCheckpoints routine watches the checkpoint registry:

class EvaluateNewCheckpoints(MixRoutine):
    def run(
        self,
        new_checkpoints=on_added_rows("vlm-checkpoints"),
    ):
        for row in new_checkpoints.select(["experiment_id"]):
            evaluate_experiment(row["experiment_id"])

For each new experiment, it reads the base model and fine-tuned checkpoint, then creates serving models:

candidates = {
    "base": Checkpoint(base_model, task="image-text-to-text"),
    "finetuned": checkpoint["checkpoint"],
}

models = {
    candidate: Model.create(
        f"vlm-{experiment_id}-{candidate}",
        checkpoint=candidate_checkpoint,
        sandbox={"gpu": "A10G"},
    )
    for candidate, candidate_checkpoint in candidates.items()
}

Model.create(..., checkpoint=...) turns each checkpoint into a callable Mixtrain model. The same API accepts a Hugging Face model identifier for the baseline and the uploaded LoRA checkpoint for the candidate.

The routine loads held-out rows from the exact dataset version used for training and sends every input to both models:

eval_data = Dataset(
    "vlm-post-training-data",
    version=checkpoint["dataset_version"],
).filter("split == 'eval'")

batch_rows = Model.batch(
    [model.name for model in models.values()],
    inference_inputs(eval_data),
    input_columns=["prompt", "images", "max_tokens", "temperature"],
    error_col="model_errors",
)

Using the pinned version prevents later data changes from silently altering an experiment's evaluation set. Model.batch() adds one output column per model, making the results straightforward to append to vlm-eval-results and compare in the existing Eval.

Run the complete example

Start from the complete project on GitHub, which includes the setup scripts, training implementation, sample data, and images.

After cloning the repository, run setup_resources.py to create the source, checkpoint, and result datasets and the Eval. You can use the mixtrain cli commands in the readme to create the resources or you can import each hosted component directly:

Once the initial loop completes, append the included second batch to run it again. Each iteration preserves the source-data version, training checkpoint, validation metrics, model outputs, and evaluation results, so you can compare experiments without losing their provenance.

Datasets — version, query, and append multimodal data
Workflows — package code and run it on managed compute
Routines — respond to dataset and workflow events
Models — create and invoke serving models
Evaluations — inspect model outputs side by side

On this page