Repository: UKPLab/sentence-transformers

Issue:
Cannot restart training of model that necessitates `trust_remote_code=True` flag
## What is the issue

If your training is interrupted and you want to restart it with flag `resume_from_checkpoint`, the `Trainer` tries to reload the model via the `_load_from_checkpoint` method. However, this method does not take as input any keyword arguments. In the case where you are finetuning a model that necessitates the flag `trust_remote_code=True`, we are currently unable to reload the model to restart training.

https://github.com/UKPLab/sentence-transformers/blob/master/sentence_transformers/trainer.py#L734-L738



Discussion:

Patch:
diff --git a/sentence_transformers/SentenceTransformer.py b/sentence_transformers/SentenceTransformer.py
index 2bb55df..cb1a952 100644
--- a/sentence_transformers/SentenceTransformer.py
+++ b/sentence_transformers/SentenceTransformer.py
@@ -167,6 +167,7 @@ class SentenceTransformer(nn.Sequential, FitMixin):
         self.prompts = prompts or {}
         self.default_prompt_name = default_prompt_name
         self.similarity_fn_name = similarity_fn_name
+        self.trust_remote_code = trust_remote_code
         self.truncate_dim = truncate_dim
         self.model_card_data = model_card_data or SentenceTransformerModelCardData()
         self._model_card_vars = {}
diff --git a/sentence_transformers/losses/SoftmaxLoss.py b/sentence_transformers/losses/SoftmaxLoss.py
index 887356e..48a30c4 100644
--- a/sentence_transformers/losses/SoftmaxLoss.py
+++ b/sentence_transformers/losses/SoftmaxLoss.py
@@ -4,6 +4,8 @@ import logging
 from typing import Callable, Iterable
 
 import torch
+import transformers
+from packaging import version
 from torch import Tensor, nn
 
 from sentence_transformers.SentenceTransformer import SentenceTransformer
@@ -103,6 +105,13 @@ class SoftmaxLoss(nn.Module):
         )
         self.loss_fct = loss_fct
 
+        if version.parse(transformers.__version__) < version.parse("4.43.0"):
+            logger.warning(
+                "SoftmaxLoss requires transformers >= 4.43.0 to work correctly. "
+                "Otherwise, the classifier layer that maps embeddings to the labels cannot be updated. "
+                "Consider updating transformers with `pip install transformers>=4.43.0`."
+            )
+
     def forward(
         self, sentence_features: Iterable[dict[str, Tensor]], labels: Tensor
     ) -> Tensor | tuple[Tensor, Tensor]:
diff --git a/sentence_transformers/sampler.py b/sentence_transformers/sampler.py
index ad1a1d1..350f800 100644
--- a/sentence_transformers/sampler.py
+++ b/sentence_transformers/sampler.py
@@ -246,7 +246,10 @@ class ProportionalBatchSampler(SetEpochMixin, BatchSampler):
         batch_samplers = [iter(sampler) for sampler in self.batch_samplers]
         for dataset_idx in dataset_idx_sampler:
             sample_offset = sample_offsets[dataset_idx]
-            yield [idx + sample_offset for idx in next(batch_samplers[dataset_idx])]
+            try:
+                yield [idx + sample_offset for idx in next(batch_samplers[dataset_idx])]
+            except StopIteration:
+                continue
 
     def __len__(self) -> int:
         return sum([len(sampler) for sampler in self.batch_samplers])
diff --git a/sentence_transformers/trainer.py b/sentence_transformers/trainer.py
index 7677589..8cb3da3 100644
--- a/sentence_transformers/trainer.py
+++ b/sentence_transformers/trainer.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import logging
 import os
 import warnings
+from collections import OrderedDict
 from contextlib import nullcontext
 from typing import TYPE_CHECKING, Any, Callable
 
@@ -14,7 +15,6 @@ from transformers.data.data_collator import DataCollator
 from transformers.integrations import WandbCallback
 from transformers.trainer import TRAINING_ARGS_NAME
 from transformers.trainer_utils import EvalLoopOutput
-from transformers.training_args import ParallelMode
 
 from sentence_transformers.data_collator import SentenceTransformerDataCollator
 from sentence_transformers.evaluation import SentenceEvaluator, SequentialEvaluator
@@ -317,12 +317,13 @@ class SentenceTransformerTrainer(Trainer):
         if isinstance(loss_fn, dict) and dataset_name:
             loss_fn = loss_fn[dataset_name]
 
-        # Hackishly insert the distributed model into the loss function, if the loss stores the model
-        # Only called once per process
+        # Insert the wrapped (e.g. distributed or compiled) model into the loss function,
+        # if the loss stores the model. Only called once per process
         if (
-            self.args.parallel_mode != ParallelMode.NOT_PARALLEL
-            and hasattr(model, "module")
-            and hasattr(loss_fn, "model")
+            model == self.model_wrapped
+            and model != self.model  # Only if the model is wrapped
+            and hasattr(loss_fn, "model")  # Only if the loss stores the model
+            and loss_fn.model != model  # Only if the wrapped model is not already stored
         ):
             loss_fn = self.override_model_in_loss(loss_fn, model)
         loss = loss_fn(features, labels)
@@ -734,7 +735,7 @@ class SentenceTransformerTrainer(Trainer):
     def _load_from_checkpoint(self, checkpoint_path: str) -> None:
         from sentence_transformers import SentenceTransformer
 
-        loaded_model = SentenceTransformer(checkpoint_path)
+        loaded_model = SentenceTransformer(checkpoint_path, trust_remote_code=self.model.trust_remote_code)
         self.model.load_state_dict(loaded_model.state_dict())
 
     def create_model_card(
@@ -761,3 +762,41 @@ class SentenceTransformerTrainer(Trainer):
             self.model.model_card_data.add_tags(tags)
 
         self.model._create_model_card(self.args.output_dir, model_name=model_name)
+
+    def get_optimizer_cls_and_kwargs(
+        self, args: SentenceTransformerTrainingArguments, model: SentenceTransformer | None = None
+    ) -> tuple[Any, Any]:
+        """
+        We have to override the optimizer_grouped_parameters because the Trainer superclass bases it on the `model`
+        itself, but the SentenceTransformer losses can have weights that should be updated as well, e.g.
+        SoftmaxLoss (see #2872).
+
+        This method requires `transformers` >= 4.43.0.
+        """
+
+        if isinstance(self.loss, dict):
+            loss_model = nn.Sequential(OrderedDict(self.loss))
+        else:
+            loss_model = self.loss
+        optimizer_cls, optimizer_kwargs = super().get_optimizer_cls_and_kwargs(args, loss_model)
+
+        # If the kwargs were not overridden by the super() call, then we should override them here so that the potential
+        # weights in the loss(es) can also be updated.
+        if not {"params", "model", "optimizer_dict"} & set(optimizer_kwargs.keys()):
+            decay_parameters = self.get_decay_parameter_names(loss_model)
+            optimizer_kwargs["optimizer_dict"] = [
+                {
+                    "params": [
+                        p for n, p in loss_model.named_parameters() if (n in decay_parameters and p.requires_grad)
+                    ],
+                    "weight_decay": self.args.weight_decay,
+                },
+                {
+                    "params": [
+                        p for n, p in loss_model.named_parameters() if (n not in decay_parameters and p.requires_grad)
+                    ],
+                    "weight_decay": 0.0,
+                },
+            ]
+
+        return optimizer_cls, optimizer_kwargs


Test patch:
diff --git a/tests/samplers/test_no_duplicates_batch_sampler.py b/tests/samplers/test_no_duplicates_batch_sampler.py
index ad93ff0..4fb62d5 100644
--- a/tests/samplers/test_no_duplicates_batch_sampler.py
+++ b/tests/samplers/test_no_duplicates_batch_sampler.py
@@ -3,13 +3,15 @@ from __future__ import annotations
 import random
 
 import pytest
+import torch
 from datasets import Dataset
+from torch.utils.data import ConcatDataset
 
-from sentence_transformers.sampler import NoDuplicatesBatchSampler
+from sentence_transformers.sampler import NoDuplicatesBatchSampler, ProportionalBatchSampler
 
 
 @pytest.fixture
-def dummy_dataset():
+def dummy_dataset() -> Dataset:
     """
     Dummy dataset for testing purposes. The dataset looks as follows:
     {
@@ -25,7 +27,22 @@ def dummy_dataset():
     return Dataset.from_dict(data)
 
 
-def test_group_by_label_batch_sampler_label_a(dummy_dataset):
+@pytest.fixture
+def dummy_duplicates_dataset() -> Dataset:
+    """
+    Dummy dataset for testing purposes. The dataset looks as follows:
+    {
+        "anchor": ["anchor_1", "anchor_1", "anchor_1", ... "anchor_2", "anchor_2"],
+        "positive": ["positive_1", "positive_1", "positive_1", ... "positive_2", "positive_2"],
+    }
+    """
+    values = [{"anchor": "anchor_1", "positive": "positive_1"}] * 10 + [
+        {"anchor": "anchor_2", "positive": "positive_2"}
+    ] * 8
+    return Dataset.from_list(values)
+
+
+def test_group_by_label_batch_sampler_label_a(dummy_dataset: Dataset) -> None:
     batch_size = 10
 
     sampler = NoDuplicatesBatchSampler(
@@ -41,3 +58,36 @@ def test_group_by_label_batch_sampler_label_a(dummy_dataset):
     for batch in batches:
         batch_values = [dummy_dataset[i]["data"] for i in batch]
         assert len(batch_values) == len(set(batch_values)), f"Batch {batch} contains duplicate values: {batch_values}"
+
+
+@pytest.mark.parametrize("drop_last", [True, False])
+def test_proportional_no_duplicates(dummy_duplicates_dataset: Dataset, drop_last: bool) -> None:
+    batch_size = 2
+    sampler_1 = NoDuplicatesBatchSampler(
+        dataset=dummy_duplicates_dataset, batch_size=batch_size, drop_last=drop_last, valid_label_columns=["anchor"]
+    )
+    sampler_2 = NoDuplicatesBatchSampler(
+        dataset=dummy_duplicates_dataset, batch_size=batch_size, drop_last=drop_last, valid_label_columns=["positive"]
+    )
+
+    concat_dataset = ConcatDataset([dummy_duplicates_dataset, dummy_duplicates_dataset])
+
+    batch_sampler = ProportionalBatchSampler(
+        concat_dataset, [sampler_1, sampler_2], generator=torch.Generator(), seed=12
+    )
+    batches = list(iter(batch_sampler))
+
+    if drop_last:
+        # If we drop the last batch (i.e. incomplete batches), we should have 16 batches out of the 18 possible,
+        # because of the duplicates being skipped by the NoDuplicatesBatchSampler.
+        # Notably, we should not crash like reported in #2816.
+        assert len(batches) == 16
+        # All batches are the same size: 2
+        assert all(len(batch) == batch_size for batch in batches)
+        assert len(sum(batches, [])) == 32
+    else:
+        # If we don't drop incomplete batches, we should be able to do 18 batches, and get more data.
+        # Note: we don't get all data, because the NoDuplicatesBatchSampler will estimate the number of batches
+        # and it would require more (non-complete) batches to get all data.
+        assert len(batches) == 18
+        assert len(sum(batches, [])) == 34