integrate mock vision backbone into model

[epwalsh]

Adds an abstraction for vision backbones, along with a mock dataset to test it. You can run a test training job on 2 GPUs with the mock dataset like so:

torchrun --nproc-per-node=2 scripts/train.py configs/mm-tiny.yaml

[epwalsh]

This took some thinking, but you can validate it with this little example:

import torch

B = 2
S = 8
D = 16
P = 3  # num patches (max across instances)
x = torch.zeros(B, S, D)

img_emb = torch.rand(B, P, D)

# use -1 for padding
image_offsets = torch.tensor([[1, 5, 6], [3, -1, -1]])
assert image_offsets.shape == (B, P)

image_offsets_mask = image_offsets > 0
batch_idx = torch.arange(0, B).repeat_interleave(image_offsets_mask.sum(dim=-1))
x.index_put_((batch_idx, image_offsets[image_offsets_mask]), img_emb[image_offsets_mask])

[notion-workspace]

Integrate mock vision backbone into OLMo

[notion-workspace]

Figure out how to put the image data into the sequences