feat: Support for gradient_checkpointing for Sentence Transformers training

[sjrl]

Related Issues

• fixes N/A

Proposed Changes:

• Added support for gradient_checkpointing for Sentence Transformers training. This can greatly reduce the memory usage on a GPU allowing for much larger batch sizes at a smallish expense of training time. This is worth it for Multiple Negatives Ranking Loss (MNRL) training in Sentence Transformers because it has been shown that when using MNRL larger batch sizes (upwards of 128) can significantly boost retrieval metrics.
• More details on gradient checkpointing can be found here

How did you test it?

• Added two new integration tests to test SentenceTransformer retriever training

Notes for the reviewer

Checklist

• I have read the contributors guidelines and the code of conduct
• I have updated the related issue with new insights and changes
• I added tests that demonstrate the correct behavior of the change
• I've used one of the conventional commit types for my PR title: fix:, feat:, build:, chore:, ci:, docs:, style:, refactor:, perf:, test:.
• I documented my code
• I ran pre-commit hooks and fixed any issue

[coveralls]

Pull Request Test Coverage Report for Build 5256397223

• 0 of 0 changed or added relevant lines in 0 files are covered.
• 133 unchanged lines in 4 files lost coverage.
• Overall coverage increased (+0.07%) to 42.0%

---

Files with Coverage Reduction	New Missed Lines	%
nodes/prompt/invocation_layer/cohere.py	4	75.61%
nodes/prompt/invocation_layer/hugging_face.py	7	87.2%
nodes/retriever/dense.py	55	25.98%
nodes/retriever/_embedding_encoder.py	67	36.47%

Totals
Change from base Build 5253128824:	0.07%
Covered Lines:	9421
Relevant Lines:	22431

---

💛 - Coveralls

[ZanSara]

A bit puzzled by one bit of code specifically, let's talk about it before moving forward.

In addition, please next time open an issue before the PR, so we can have a discussion on the issue about the feature's design, and on the PR about the implementation details. I'm curious where the demand for this specific feature comes from 🙂

[ZanSara]

One more small thing to fix and this is ready to merge

[ZanSara]

I'm really sorry I didn't notice it earlier 🙈 but is this initialization loading the model? If yes, can we mock the model as well so that we don't actually download and load the weights from hf?

If we can do that we'll be able to add the @pytest.mark.unit marker to this test. Otherwise it won't run in CI.

[ZanSara]

As agreed with @sjrl we can close this for now and re-implement it in v2 later, if still relevant