Multimodal input misalignment after deploying to AWS EC2 from on-prem

3. Check tokenization batch size

# logs from TorchServe
2026-06-16 12:04:32,117 INFO  [model_worker_0] batch size mismatch: text=32, image=31
Traceback (most recent call last):
  File ".../torchserve/model_handler.py", line 212, in preprocess
    inputs = self.tokenizer(batch["text"], padding=True, truncation=True)
ValueError: Expected input batch size 32 but got 31

The discrepancy aligns with the DataLoader worker seed issue.

4. Inspect environment variables

$ echo $PYTHONHASHSEED
# on‑prem: 0
# EC2: (empty)

Unset hash seed leads to nondeterministic dictionary ordering inside the tokenizer.

5. Capture a packet trace of the model server request

# Using tcpdump to ensure request payload is intact
sudo tcpdump -i eth0 -s 0 -w request.pcap port 8080

Analysis confirmed that the JSON payload contains both text and image fields; the problem is purely server‑side preprocessing.

Resolution

All identified gaps are addressed in the following changes.

1. Pin identical image libraries

Install the same libjpeg and Pillow versions used on‑prem.

# Before (Amazon Linux default)
sudo yum install -y libjpeg-turbo pillow

# After (match on‑prem libjpeg‑8)
sudo yum remove -y libjpeg-turbo
sudo yum install -y libjpeg libjpeg-devel
pip uninstall -y pillow
pip install pillow==8.2.0 --no-binary :all:

Rebuilding Pillow against libjpeg‑8 restores identical decoding behavior.

2. Enforce deterministic tokenization

# Add to startup script (e.g., /etc/profile.d/ml_env.sh)
export PYTHONHASHSEED=0
export TOKENIZER_SEED=42

In the model code, set the seed explicitly:

import os, random, torch, numpy as np
seed = int(os.getenv("TOKENIZER_SEED", "42"))
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)

3. Align DataLoader worker seeds

# Existing DataLoader (on‑prem)
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=0)

# Updated for EC2
def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    worker_init_fn=seed_worker,
)

This guarantees each forked worker starts with the same seed as the main process.

4. Centralize preprocessing configuration

Store all preprocessing constants in AWS Systems Manager Parameter Store and load them at runtime.

# Parameter Store entry: /ml/preprocess/image_size = 224
# Parameter Store entry: /ml/preprocess/text_max_len = 128

import boto3
ssm = boto3.client('ssm')
def get_param(name):
    return ssm.get_parameter(Name=name, WithDecryption=True)['Parameter']['Value']

IMAGE_SIZE = int(get_param('/ml/preprocess/image_size'))
MAX_TEXT_LEN = int(get_param('/ml/preprocess/text_max_len'))

Ensures identical pipelines across environments.

5. Re‑build the Docker image

Update the Dockerfile to include the pinned libraries and environment variables.

FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu20.04

# Install system deps
RUN apt-get update && apt-get install -y \\
    libjpeg-dev libpng-dev && \\
    rm -rf /var/lib/apt/lists/*

# Python deps
COPY requirements.txt .
RUN pip install -r requirements.txt \\
    && pip install pillow==8.2.0 torchvision==0.12.0

# Set deterministic env vars
ENV PYTHONHASHSEED=0
ENV TOKENIZER_SEED=42

Validation

After redeploying, run the following sanity checks:

Batch size consistency

# test script
from model import multimodal_predict
batch = {"text": ["sample"]*32, "image": ["img1.jpg"]*32}
outputs = multimodal_predict(batch)
print(len(outputs))  # should be 32

Log output shows no size mismatch errors.

Feature vector shape equality

features = multimodal_predict(batch, return_features=True)
print(features["image"].shape)   # torch.Size([32, 768])
print(features["text"].shape)    # torch.Size([32, 768])

Both modalities now share the same embedding dimension.

End‑to‑end inference test

curl -X POST http://ec2-instance:8080/predictions/model \\
  -H "Content-Type: application/json" \\
  -d '{"text":"A cat on a sofa","image":"s3://bucket/cat.jpg"}'

Response contains a coherent caption and confidence score, matching on‑prem baseline.

Best Practices and Prevention

• Always pin third‑party library versions (Pillow, torchvision, OpenCV) in requirements.txt and verify they are available for the target OS.
• Store preprocessing hyper‑parameters (image size, token max length, normalization stats) in a central configuration service such as AWS Systems Manager Parameter Store.
• Set PYTHONHASHSEED and any framework‑specific seeds at container start‑up to guarantee deterministic tokenization.
• When using DataLoader with num_workers > 0, provide a worker_init_fn that seeds each worker.
• Add health‑check endpoints that verify batch size parity and embedding dimensions before routing traffic.
• Include a CI step that runs a small multimodal sanity test on the target AMI (e.g., amazonlinux:2) to catch library drift early.

Related Topic Hub: Cloud Infrastructure Troubleshooting Hub

FAQ

1. Why does the model work locally but fail on EC2? The EC2 AMI ships with newer image libraries (libjpeg‑turbo 2.0, Pillow 9.x) that decode JPEGs slightly differently, producing tensors of a different dtype or shape. Pinning the same library versions resolves the discrepancy.
2. Do I need to reinstall the NVIDIA driver on the EC2 instance? The driver version must be compatible with the CUDA toolkit used by PyTorch. Follow the AWS EC2 GPU Instances guide to install the recommended driver; mismatched driver versions typically cause CUDA errors, not preprocessing misalignment.
3. How can I verify that tokenization is deterministic? Set PYTHONHASHSEED=0 and a fixed torch.manual_seed. Then run the tokenizer on a known sentence multiple times; the token IDs should be identical across runs and across environments.
4. Is the DataLoader seed the only source of batch size mismatch? Not always. Corrupted images that fail to decode can also shrink the image batch. Ensure the image loader has skip_invalid=True or pre‑filter the dataset.
5. Can I avoid rebuilding Pillow by using a different base image? Yes. Using an Ubuntu‑based AMI that matches the on‑prem OS can simplify library parity, but you must still pin the exact versions to guarantee reproducibility.