Guide

Multimodal

Requires pytrio package version >= 0.2.7

PyTRIO supports image inputs for multimodal inference and training.

Input Processing

Both image and text inputs are wrapped in ModelInput in PyTRIO.

The main difference is that image inputs must be constructed with chunks, while text-only inputs can usually use the from_ints() method directly.

For multimodal inputs, a ModelInput consists of one or more chunks. Each chunk represents one part of the prompt, and the model reads them in the order in which they appear in the chunks list.

For example, the image-description task in this guide has three parts:

  1. Text tokens before the image, including the start markers for the user message and vision input.
  2. The image data.
  3. Text tokens after the image, including the end marker for the vision input, the user's question, and the start of the assistant response.

The corresponding structure is:

prompt = trio.ModelInput(
    chunks=[
        trio.types.EncodedTextChunk(tokens=...),
        trio.ImageChunk(data=Path(), format="png"),
        trio.types.EncodedTextChunk(tokens=...),
    ]
)

The chunk order is the exact order in which the model receives the content, so the special tokens before and after the image must not be rearranged.

Text Chunks (EncodedTextChunk)

EncodedTextChunk stores text tokens that have already been encoded by the tokenizer. It can also be understood as a TextChunk: convert the text into token IDs first, then pass them to the tokens field.

tokenizer = sampler.get_tokenizer()
tokens = tokenizer.encode("Describe this image.", add_special_tokens=False)

text_chunk = trio.types.EncodedTextChunk(tokens=tokens)

The ModelInput.from_ints(...) method commonly used for text-only inference essentially creates a ModelInput containing a single EncodedTextChunk. The following two forms represent the same text-only input:

input_ids = tokenizer.encode("Hello")

# Shorthand
prompt = trio.ModelInput.from_ints(input_ids)

# Explicit text chunk
prompt = trio.ModelInput(
    chunks=[trio.types.EncodedTextChunk(tokens=input_ids)]
)

When the prompt also contains an image, use the second form to arrange different chunk types explicitly.

Image Chunks (ImageChunk)

ImageChunk carries the binary contents of an image. It has two main fields:

  • data: the binary image data, which can be read with Path.read_bytes().
  • format: the image format. This guide uses png or jpeg.
from pathlib import Path

image_path = Path("example.png")
image_chunk = trio.ImageChunk(
    data=image_path.read_bytes(),
    format="png",
)

Pass the original image bytes directly. You do not need to convert the image into tokens or encode it as Base64. The format must match the actual file contents; both .jpg and .jpeg files use jpeg.

Multimodal Inference

To use the OpenAI Python SDK for image chat, see the OpenAI API image chat example.

First, create a sampling client with a vision-language model and get its tokenizer:

import pytrio as trio

sampler = trio.ServiceClient().create_sampling_client(
    base_model="Qwen/Qwen3.5-4B"
)
tokenizer = sampler.get_tokenizer()
encode = lambda text: tokenizer.encode(text, add_special_tokens=False)

Next, assemble the text and image into a ModelInput in the order required by the model:

prompt = trio.ModelInput(
    chunks=[
        trio.types.EncodedTextChunk(
            tokens=encode("<|im_start|>user\n<|vision_start|>")
        ),
        trio.ImageChunk(
            data=image_path.read_bytes(),
            format=image_format,
        ),
        trio.types.EncodedTextChunk(
            tokens=encode(
                "<|vision_end|>Describe this image.<|im_end|>\n"
                "<|im_start|>assistant\n<think>\n\n</think>\n\n"
            )
        ),
    ]
)

<|vision_start|> and <|vision_end|> mark the position of the image in the conversation. <|im_start|> and <|im_end|> mark message boundaries. These special tokens are part of the model's prompt format. Other models may use different templates, so follow the requirements of the model you use.

Finally, call sample just as you would for text inference, then call .result() to retrieve the remote inference result:

response = sampler.sample(
    prompt=prompt,
    num_samples=1,
    sampling_params=trio.SamplingParams(
        max_tokens=512,
        temperature=0.5,
        stop="<|im_end|>",
    ),
).result()

print(response.sequences[0].text)

Full Example

Save the following code as vlm_sample.py and pass it a local image:

python vlm_sample.py /path/to/image.png

The complete code is:

"""Usage: python vlm_sample.py /path/to/image.png"""

import sys
from pathlib import Path

import pytrio as trio


if len(sys.argv) != 2:
    raise SystemExit(f"Usage: python {Path(__file__).name} /path/to/image.png")

image_path = Path(sys.argv[1]).expanduser()
image_format = {".png": "png", ".jpg": "jpeg", ".jpeg": "jpeg"}.get(
    image_path.suffix.lower()
)
if image_format is None:
    raise SystemExit("Only PNG and JPEG images are supported")

sampler = trio.ServiceClient().create_sampling_client(
    base_model="Qwen/Qwen3.5-4B"
)
tokenizer = sampler.get_tokenizer()
encode = lambda text: tokenizer.encode(text, add_special_tokens=False)

prompt = trio.ModelInput(
    chunks=[
        trio.types.EncodedTextChunk(
            tokens=encode("<|im_start|>user\n<|vision_start|>")
        ),
        trio.ImageChunk(
            data=image_path.read_bytes(),
            format=image_format,
        ),
        trio.types.EncodedTextChunk(
            tokens=encode(
                "<|vision_end|>Describe this image.<|im_end|>\n"
                "<|im_start|>assistant\n<think>\n\n</think>\n\n"
            )
        ),
    ]
)

response = sampler.sample(
    prompt=prompt,
    num_samples=1,
    sampling_params=trio.SamplingParams(
        max_tokens=512,
        temperature=0.5,
        stop="<|im_end|>",
    ),
).result()

print(response.sequences[0].text)

Notes

  • base_model must support image input. Text-only models cannot process ImageChunk.
  • The image format must match the format argument of ImageChunk.
  • The special tokens surrounding the image depend on the model template. Update them when switching models.
  • EncodedTextChunk accepts token IDs, not unencoded strings.
  • Multimodal inference uses the same sampling parameters and response structure as text inference. The generated text is still available at response.sequences[0].text.

Multimodal Training

This section uses SFT as the main example for clarity.

Building a Datum

The only code-level difference between multimodal training and text-only training is how the Datum is constructed.

Both use the same cross_entropy loss function, and both use weights so that the model learns only from the assistant response. In text-only SFT, all tokens can be placed in one EncodedTextChunk, or from_ints() can be used directly. Multimodal SFT must preserve both text chunks and image chunks in the ModelInput.

A multimodal SFT Datum still consists of the following three parts:

  • model_input: the image, prompt text, and shifted completion.
  • target_tokens: the token that the model should predict at each input position.
  • weights: the loss weight for each position. Prompt and image positions use 0, while completion positions use 1.

Building an ImageChunk

For inference, passing the image data and format is enough. When constructing training data, you must also use expected_tokens to tell PyTRIO how many tokens the vision encoder will produce for the image. This allows PyTRIO to calculate the ModelInput length correctly and align target_tokens and weights with the input.

For Qwen3.5-4B, use the corresponding image processor to calculate the number of image patches:

import io

from PIL import Image
from transformers import AutoImageProcessor


processor_source = getattr(tokenizer, "name_or_path", "Qwen/Qwen3.5-4B")
image_processor = AutoImageProcessor.from_pretrained(
    processor_source,
    use_fast=False,
)


def encode_image(image: Image.Image, processor) -> trio.ImageChunk:
    image = image.convert("RGB")
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")

    patches = processor.get_number_of_image_patches(
        image.height,
        image.width,
        images_kwargs={},
    )
    expected_tokens = patches // processor.merge_size**2

    return trio.ImageChunk(
        data=buffer.getvalue(),
        format="png",
        expected_tokens=expected_tokens,
    )

expected_tokens depends on the vision preprocessing behavior of the model. When switching models, calculate it with the image processor that matches the remote model instead of hard-coding it.

Building Prompt Chunks

First use the tokenizer's chat template to generate the complete prompt required by the model, then replace the image placeholder with the actual ImageChunk.

For example, Qwen3.5 uses <|image_pad|> as its image placeholder. For a sample containing one image, split the generated prompt into the text before and after the image:

IMAGE_PAD = "<|image_pad|>"
PROMPT = (
    "Transcribe the mathematical formula in this image into LaTeX. "
    "Output only the LaTeX."
)

messages = [
    {
        "role": "user",
        "content": [
            {"type": "text", "text": PROMPT},
            {"type": "image", "image": "formula"},
        ],
    }
]
prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=False,
)

parts = prompt.split(IMAGE_PAD)
if len(parts) != 2:
    raise ValueError(f"Expected 1 image placeholder, got {len(parts) - 1}")

before_image, after_image = parts
prompt_chunks = [
    trio.types.EncodedTextChunk(
        tokens=tokenizer.encode(before_image, add_special_tokens=False)
    ),
    image_chunk,
    trio.types.EncodedTextChunk(
        tokens=tokenizer.encode(after_image, add_special_tokens=False)
    ),
]

The chat template has already inserted the special tokens required for the conversation and image. Set add_special_tokens=False when encoding the text again to avoid inserting duplicate special tokens.

Autoregressive Shift and Loss Mask

Suppose the answer tokens that the model should learn are stored in completion. Append only completion[:-1] to model_input. The last prompt position predicts the first answer token, and every subsequent answer token is predicted from the preceding answer token.

Use 0 as the placeholder in target_tokens for positions that do not participate in training, rather than the -100 value commonly used by HuggingFace. The weights array determines whether each position contributes to the loss.

The complete process_example function is shown below. The dataset's image field contains a PIL image, and its text field contains the target LaTeX:

import numpy as np


def process_example(example, tokenizer, processor) -> trio.Datum:
    image_chunk = encode_image(example["image"], processor)
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": PROMPT},
                {"type": "image", "image": "formula"},
            ],
        }
    ]
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )

    parts = prompt.split(IMAGE_PAD)
    if len(parts) != 2:
        raise ValueError(f"Expected 1 image placeholder, got {len(parts) - 1}")

    before_image, after_image = parts
    prompt_chunks = [
        trio.types.EncodedTextChunk(
            tokens=tokenizer.encode(before_image, add_special_tokens=False)
        ),
        image_chunk,
        trio.types.EncodedTextChunk(
            tokens=tokenizer.encode(after_image, add_special_tokens=False)
        ),
    ]

    prompt_length = len(trio.ModelInput(chunks=prompt_chunks))
    completion = tokenizer.encode(
        str(example["text"]).strip() + "<|im_end|>",
        add_special_tokens=False,
    )

    model_input = trio.ModelInput(
        chunks=[
            *prompt_chunks,
            trio.types.EncodedTextChunk(tokens=completion[:-1]),
        ]
    )
    target_tokens = np.zeros(len(model_input), dtype=np.int64)
    weights = np.zeros(len(model_input), dtype=np.float32)

    start = prompt_length - 1
    target_tokens[start : start + len(completion)] = completion
    weights[start : start + len(completion)] = 1.0

    return trio.Datum(
        model_input=model_input,
        loss_fn_inputs={
            "target_tokens": target_tokens,
            "weights": weights,
        },
    )

The completion includes <|im_end|> with a weight of 1, so the model also learns when to stop generating the answer.

After constructing the Datum objects, training is identical to text-only SFT:

processed_examples = [
    process_example(example, tokenizer, image_processor)
    for example in dataset
]

fwdbwd = training_client.forward_backward(
    processed_examples,
    loss_fn="cross_entropy",
)
optim = training_client.optim_step(
    trio.AdamParams(learning_rate=1e-4)
)

result = fwdbwd.result()
optim.result()

Full Example

The following example fine-tunes Qwen3.5-4B with multimodal LoRA SFT on the LaTeX_OCR dataset:

"""Run multimodal LoRA SFT on Qwen3.5-4B with LaTeX_OCR/small using PyTRIO.

Environment：
    pip install pytrio numpy datasets torch torchvision pillow

Run:
    python train.py --epochs 3 --batch-size 2
"""

from __future__ import annotations

import argparse
import io
from pathlib import Path

import numpy as np
import pytrio as trio
from datasets import load_dataset
from huggingface_hub import snapshot_download
from PIL import Image
from transformers import AutoImageProcessor


DATASET_ID = "linxy/LaTeX_OCR"
IMAGE_PAD = "<|image_pad|>"
PROMPT = (
    "Transcribe the mathematical formula in this image into LaTeX. "
    "Output only the LaTeX."
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="Qwen/Qwen3.5-4B")
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--learning-rate", type=float, default=1e-4)
    parser.add_argument("--rank", type=int, default=32)
    parser.add_argument("--max-samples", type=int, default=0, help="0 means all")
    parser.add_argument("--max-length", type=int, default=8192)
    parser.add_argument("--dataset-dir", default="data/LaTeX_OCR")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--checkpoint-name", default="latex-ocr-small-sft")
    return parser.parse_args()


def load_train_dataset(dataset_dir: str):
    """Download the small subset on first use, then read the local Parquet files."""
    local_dir = Path(dataset_dir).expanduser().resolve()
    files = sorted((local_dir / "small").glob("train-*.parquet"))
    if not files:
        snapshot_download(
            repo_id=DATASET_ID,
            repo_type="dataset",
            local_dir=local_dir,
            allow_patterns=["README.md", "small/*.parquet"],
        )
        files = sorted((local_dir / "small").glob("train-*.parquet"))
    if not files:
        raise FileNotFoundError(f"No train Parquet files found in {local_dir}")

    return load_dataset(
        "parquet",
        data_files={"train": [str(path) for path in files]},
        split="train",
    )


def encode_image(image: Image.Image, processor) -> trio.ImageChunk:
    """Convert a PIL image to a multimodal PyTRIO ImageChunk."""
    image = image.convert("RGB")
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")

    # TRIO needs to know how many tokens the vision encoder will produce.
    patches = processor.get_number_of_image_patches(
        image.height,
        image.width,
        images_kwargs={},
    )
    return trio.ImageChunk(
        data=buffer.getvalue(),
        format="png",
        expected_tokens=patches // processor.merge_size**2,
    )


def process_example(example, tokenizer, processor) -> trio.Datum:
    """Convert one image/text example into an answer-only SFT Datum."""
    image_chunk = encode_image(example["image"], processor)
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": PROMPT},
                {"type": "image", "image": "formula"},
            ],
        }
    ]
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )

    # apply_chat_template inserts an image placeholder; replace it with ImageChunk.
    parts = prompt.split(IMAGE_PAD)
    if len(parts) != 2:
        raise ValueError(f"Expected 1 image placeholder, got {len(parts) - 1}")
    before_image, after_image = parts
    prompt_chunks = [
        trio.types.EncodedTextChunk(
            tokens=tokenizer.encode(before_image, add_special_tokens=False)
        ),
        image_chunk,
        trio.types.EncodedTextChunk(
            tokens=tokenizer.encode(after_image, add_special_tokens=False)
        ),
    ]

    prompt_length = len(trio.ModelInput(chunks=prompt_chunks))
    completion = tokenizer.encode(
        str(example["text"]).strip() + "<|im_end|>",
        add_special_tokens=False,
    )

    # Autoregressive shift: predict the completion from the final prompt position.
    model_input = trio.ModelInput(
        chunks=[
            *prompt_chunks,
            trio.types.EncodedTextChunk(tokens=completion[:-1]),
        ]
    )
    target_tokens = np.zeros(len(model_input), dtype=np.int64)
    weights = np.zeros(len(model_input), dtype=np.float32)
    start = prompt_length - 1
    target_tokens[start : start + len(completion)] = completion
    weights[start : start + len(completion)] = 1.0

    return trio.Datum(
        model_input=model_input,
        loss_fn_inputs={"target_tokens": target_tokens, "weights": weights},
    )


def loss_per_token(result, batch: list[trio.Datum]) -> float:
    """Compute the mean NLL over supervised tokens from returned logprobs."""
    logprobs = np.concatenate(
        [output["logprobs"].tolist() for output in result.loss_fn_outputs]
    )
    weights = np.concatenate(
        [datum.loss_fn_inputs["weights"].tolist() for datum in batch]
    )
    return float(-np.dot(logprobs, weights) / weights.sum())


def main() -> None:
    args = parse_args()
    if args.batch_size < 1:
        raise ValueError("batch-size must be greater than 0")

    # 1. Download and load the small/train dataset locally.
    dataset = load_train_dataset(args.dataset_dir)
    if args.max_samples > 0:
        dataset = dataset.select(range(min(args.max_samples, len(dataset))))

    # 2. Connect to TRIO and create a LoRA training client.
    service_client = trio.ServiceClient()
    training_client = service_client.create_lora_training_client(
        base_model=args.model,
        rank=args.rank,
        seed=args.seed,
    )

    # 3. Get the tokenizer and image processor that match the remote model.
    tokenizer = training_client.get_tokenizer()
    processor_source = getattr(tokenizer, "name_or_path", args.model)
    image_processor = AutoImageProcessor.from_pretrained(
        processor_source,
        use_fast=False,
    )

    # 4. Convert the small dataset once to avoid repeated encoding each epoch.
    processed_examples = [
        process_example(example, tokenizer, image_processor) for example in dataset
    ]
    processed_examples = [
        datum
        for datum in processed_examples
        if len(datum.model_input) <= args.max_length
    ]
    if not processed_examples:
        raise RuntimeError("No training samples remain; check max-samples/max-length")

    # 5. Shuffle Datum objects each epoch and train in batch_size slices.
    step = 0
    for epoch in range(args.epochs):
        indices = np.random.default_rng(args.seed + epoch).permutation(
            len(processed_examples)
        )
        for start in range(0, len(indices), args.batch_size):
            batch_indices = indices[start : start + args.batch_size]
            batch = [processed_examples[index] for index in batch_indices]
            fwdbwd = training_client.forward_backward(batch, "cross_entropy")
            optim = training_client.optim_step(
                trio.AdamParams(learning_rate=args.learning_rate)
            )
            result = fwdbwd.result()
            optim.result()

            step += 1
            loss = loss_per_token(result, batch)
            print(
                f"epoch={epoch + 1} step={step} loss={loss:.4f}",
                flush=True,
            )

    # 6. Save weights that can be passed directly to SamplingClient.
    saved = training_client.save_weights_for_sampler(
        name=args.checkpoint_name
    ).result()
    print(f"saved_weights={saved.path}")


if __name__ == "__main__":
    main()

Inference

Previous Page

Save Weights

Next Page

Was this documentation helpful?

On this page

Input Processing
Text Chunks (EncodedTextChunk)
Image Chunks (ImageChunk)
Multimodal Inference
Full Example
Notes
Multimodal Training
Building a Datum
Building an ImageChunk
Building Prompt Chunks
Autoregressive Shift and Loss Mask
Full Example