Problems in training multi-modal models with large batch sizes

I like to train a multimodal embedding model (images and text) using a cached loss. Therby, I found several issues that make it impossible to train with a large batch size.

Problems

1. The entire batch lives on the GPU for the whole optimizer step. The trainer moves the full batch to the GPU. Having thousands of images in a batch leads to OOMs. In some cases, it even led to main memory issues. For text inputs, it works since the text do not consume that much memory. However, for image inputs, it is a problem.

2. The whole batch is preprocessed upfront. The whole batch is loaded and preprocessed upfront. There is not way to do lazy loading approaches. I.e., load images only the images of the next n mini-batches. This upfront loading can cause main memory issues.

3. Fixed sample-count mini-batches don't bound memory. mini_batch_size counts samples, but memory scales with padded sequence length. I.e. the longest sample in a batch determines how many samples fit in a mini-batch. A dynamic token budget, that would be filled could solve this. E.g., for Qwen3-VL-Embedding-2B, the max sequence length is 260k, if one sample in a batch has this length all samples will be padded to this length. However, only the samples in the mini-batch need to be padded to the longest sample in the mini-batch.

While issue 3 is a comparable simple fix that can be easily integrated into the existing cached loss implementation. Fixing issue 1 and 2 is more difficult. It requires changes to the trainer, collator, and loss. And the user has to configure everything in a consistent way so that it works. Making it more difficult for the user to select the right combination. I.e. lazy loading would require a collator supporting this, a trainer avoiding that everything is directly loaded to the device and a loss that is handling lazy loading. Is there any chance to integrate something like this into the Sentence Transformer library or is it out-of-scope?