cdpierse / cdpierse/transformers-interpret
[RobertaForSequenceClassification] RAM memory leaks during retrieving word attributions
- Dominant language
- Jupyter Notebook
- Stars
- 1.4k
- Forks
- 99
- PR merge metrics
- No merged PRs in 30d
Description
Hi folks,
I'm experiencing some memory leaks when using the Transformers Interpret (TI) library which triggers out of memory errors and kills the model service process.
## Setup
Here is how the TI lib is being using (roughly):
```python
class NLPModel:
"""
NLP Classification Model
"""
def __init__(self):
self.label_encoder: LabelEncoder = joblib.load(
...
)
self.tokenizer = AutoTokenizer.from_pretrained(
...
padding="max_length",
truncation=True,
)
self.model: "RobertaForSequenceClassification" = AutoModelForSequenceClassification.from_pretrained(
..., num_labels=self.num_classes
)
self.classifier = pipeline(
"text-classification", model=self.model, tokenizer=self.tokenizer, return_all_scores=True
)
self.explainer = SequenceClassificationExplainer(self.model, self.tokenizer) # the TI library
def get_word_attributions(self, agent_notes: str) -> List[WordAttribution]:
"""
Retrieves attributions for each word based on the model explainer
"""
with torch.no_grad():
raw_word_attributions: List[RawWordAttribution] = self.explainer(agent_notes)[1:-1]
# some post processing of the raw_word_attributions
# return processed word attributions
model = NLPModel()
def get_model() -> NLPModel:
return model
```
The code is running as a FastAPI view with one server worker:
```python
app = FastAPI()
router = APIRouter(prefix=URL_PREFIX)
# some views
@router.post("/predict/")
def get_predictions(payload: PredictionPayload, model: NLPModel = Depends(get_model)):
samples = payload.samples
predictions = model.predict(...) # regular forward pass on the roberta classification model
response: List[dict] = []
for sample, prediction in zip(samples, predictions):
word_attributions = model.get_word_attributions(sample.text)
response.append(
{
# .... some other information
"predictions": prediction,
"word_attributions": [attribution.dict() for attribution in word_attributions],
}
)
return JSONResponse(content=jsonable_encoder(response))
app.include_router(router)
```
```bash
uvicorn main:app --workers 1 --host 0.0.0.0 --port 8080 --log-level debug
```
The whole service is running on CPU/RAM, no GPU/CUDA is available.
## Problem
Now when I start to send sequential requests to the service, it allocates more and more memory after each of the request. Eventually this leads to OOM errors and the server gets killed by the system.
Memory allocation roughly looks like this (these statistics I have collected on my local docker environment where I have 6 GB RAM limit):

Using empirical experiments, I was able to define that problem lays in the following line:
```python
raw_word_attributions: List[RawWordAttribution] = self.explainer(agent_notes)[1:-1]
```
When I disabled the line, the service used not more than 500-700MB and the memory consumption almost stayed the same.
Now from what I understand the TI library calculates gradients in order to identify word attributions. I suspect that this is the reason of the issue. However, using zero_grad() on the whole model did not help me to clean up RAM. I have tried more tricks like forcing GC collection, removing the `explainer` and `model` instances but non of the things really helped.
Do you have any ideas how could I clean up RAM that the service uses after running the Transformers Interpret library?
Appreciate your help 🙏
Contributor guide
Assessment
This issue has not been assessed yet.