huggingface / huggingface/setfit

OOM on large datasets due to O(n²) pair generation in contrastive learning

Open
#628 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
2.8k
Forks
267
Avg merge
36m
Merged PRs (30d)
5

Description

### Problem

SetFit training goes OOM (Out of Memory) on main memory (not GPU VRAM) for large datasets. The root cause is that pair generation for contrastive learning materializes O(n²) data in memory before training begins.

### Environment
- SetFit version: 1.1.3
- Python version: 3.10
- OS: Linux

### Reproduction

Any dataset with a sufficiently large number of samples using contrastive loss (e.g., `CosineSimilarityLoss`) will likely OOM on machines with a given RAM. I tried a 40,000 sample dataset with num_iterations=1 on a 32 GB machine.

```python
from setfit import SetFitModel, Trainer, TrainingArguments
from datasets import load_dataset

# Load a large dataset
dataset = load_dataset("some_large_dataset") # e.g., 100k+ samples

model = SetFitModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
trainer = Trainer(
model=model,
train_dataset=dataset["train"],
args=TrainingArguments(num_iterations=20),
)
trainer.train() # OOM before training starts
```

### Root Cause Analysis

There are **3 layers** of memory explosion:

**Layer 1: `shuffle_combinations()` in `src/setfit/sampler.py`**

```python
def shuffle_combinations(iterable: Iterable, replacement: bool = True) -> Generator:
n = len(iterable)
k = 1 if not replacement else 0
idxs = np.stack(np.triu_indices(n, k), axis=-1) # O(n²) memory!
for i in np.random.RandomState(seed=42).permutation(len(idxs)): # Another O(n²) array!
_idx, idx = idxs[i, :]
yield iterable[_idx], iterable[idx]
```

Despite being typed as a `Generator`, it allocates ALL n*(n-1)/2 pair indices upfront before yielding anything.

**Layer 2: `ContrastiveDataset.generate_pairs()`**

```python
def generate_pairs(self) -> None:
for (_text, _label), (text, label) in shuffle_combinations(self.sentence_labels):
if is_positive:
self.pos_pairs.append(...) # Stores all pairs in lists
else:
self.neg_pairs.append(...)
```

**Layer 3: `trainer.py` line 618**

```python
dataset = Dataset.from_list(list(data_sampler)) # Materializes iterator + creates copy
```

### Suggested Solution

The solution involves replacing eager pair generation with streaming:

1. **`ContrastiveDataset`**: Generate pairs on the fly and track uniqueness using a set.

2. **Trainers**: Use `IterableDataset.from_generator()` instead of `Dataset.from_list(list(...))`

3. **Memory after fix**: O(n) for label groups + O(num_pairs_sampled) for uniqueness set

I have created a draft PR for this, would be happy to discuss and contribute it here. Here is the draft PR: https://github.com/huggingface/setfit/pull/627

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.