Different batch sizes per dataset/loss (CachedMNRL + CoSENT) in multi-dataset training

[ddofer]

Hi,
I’m training a model using the v3 multi-dataset training API to combine a cached contrastive loss (CachedMultipleNegativesRankingLoss, or cached gist, etc) on one dataset with a regression loss (CoSENTLoss, or triplet loss) on a second dataset.

Because of the different memory requirements and nature of these losses, I want to use different batch sizes for each dataset. For example, 2K batch size with the cached mnrl, and 3264 with the triplet or cosent losses (i.e the max effective batch size I can fit in per device memory).

Here is a simplified version of my setup:

from sentence_transformers import SentenceTransformerTrainer, losses
from sentence_transformers.training_args import SentenceTransformerTrainingArguments, MultiDatasetBatchSamplers

train_dataset = {
    "contrastive_ds": dataset_a, # e.g. positive pairs
    "regression_ds": dataset_b,  # e.g. pairs with similarity scores
}

loss_model = {
    "contrastive_ds": losses.CachedMultipleNegativesRankingLoss(model, mini_batch_size=32),
    "regression_ds": losses.CoSENTLoss(model),
}

args = SentenceTransformerTrainingArguments(
    per_device_train_batch_size=64,  # <-- Applies globally?
    multi_dataset_batch_sampler=MultiDatasetBatchSamplers.ROUND_ROBIN,
    output_dir="./output"
)

trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    loss=loss_model,
)

My questions:

1. Dataset-Specific Batch Sizes: Is there a supported way to specify different per_device_train_batch_size values per dataset in SentenceTransformerTrainingArguments? (Multi GPU/DDP setup)
2. Interaction with CachedMNRL: Since CachedMNRL accepts a mini_batch_size parameter, is the recommended pattern to set a globally large per_device_train_batch_size? If I do this, won't CoSENTLoss attempt to process the massive global batch size all at once and cause an OOM error?
3. Best Practices: What is the recommended pattern for decoupling the batch sizes when mixing cached and non-cached losses via ROUND_ROBIN? (I keep getting nvcc crashes/issues, maybe due to batch sizes not being cleany divisible in the multigpu setting and different losses)

Thanks in advance for any guidance!

[ddofer]

Same issue applies with other losses (e.g. tsdae) that need a "small" batch size, vs the "big one" allowed by cached methods.
If not supported already, naively best solution would be allowing a way to set the "mini_batch_size " across multitasks/multidatasets etc? (I'm hoping it's already implementable though; the best practice in the documentation does mention cached mnrl and cosent!)

[tomaarsen]

Hello!

Apologies for the delay. I've been thinking about this issue for a few days, but I'm not sure what's best. To be explicit, here are my answers:

1. There are not. There's a few things you can define per-dataset, like prompts or learning_rate_mapping, batch size is not one of them. Currently, the Trainer combines all datasets into a torch ConcatDataset with specific dataset-aware "multi-dataset" batch samplers (RoundRobinBatchSampler or ProportionalBatchSampler). This ConcatDataset is then fed into a torch DataLoader, which also gets some global information including the batch size. I'm unsure if it's possible to fully defer the batch sizes to the batch samplers instead of the Data Collator. If yes, then it might be possible, but if no, then I don't think it'll be possible without moving away from the DataLoader class, which I'd really rather not do.
2. Yes, for CMNRL the general recommendation is to set a large per_device_train_batch_size and a small mini_batch_size. This indeed sadly conflicts with also having a non-GradCache loss like CoSENTLoss.
3. I'm not sure if there's a clean fix at the moment. The only thing I can think of right now is the creation of a CachedCoSENTLoss. Normally this doesn't make sense, as CoSENTLoss doesn't take in-batch negatives into account (I think), so there's normally no reason to use a huge batch size with CoSENTLoss.

It seems that perhaps the best solution is actually to fully extract the Cached GradCache functionality into a new "loss wrapper", with the goal that this wrapper can be applied on any loss akin to Loss Modifiers (docs) like MatryoshkaLoss. Then the mini_batch_size can be set on any loss, and you can set the batch size high regardless of what datasets/losses you're using.

But for now, I think my simplest recommendation is to try only training with the CMNRL loss/dataset, and/or converting the CoSENTLoss dataset into triplets so it can also be consumed by CMNRL.

• Tom Aarsen

[ddofer]

Dang, I was hoping I was missing a functionality. (I can't think of a reasonable way to make the cosent task binary, the continous biological signal is very hard to reasonably discretize; and big batch sizes are a killer).
That also prevent a big batch with triplet loss right? (It's not like I could just set up gradient accumulation "just" for the triplet loss part to avoid the OOM ?)
Thanks anyway! :(

[ddofer]

Porting the cached (e.g. just doing "gradient accumulation" does sound like a great overall solution; it'd be "good enough" , especially since all those "other" losses are relatively insensitive to larger batch sizes in terms of their behavior (unlike mnrl)).
I don't know the library well enough for it though :)