Repository navigation
Reducing GPU utilization while working with CLAREAugmenter #812
Description
Hi,
I was trying to use CLAREAugmenter from your library and was surprised that it consumes more than 40GB of VRAM during the augmentation process.
I have some longer texts (~800 tokens) but did not expect that a method loading a simple distilroberta underneath will cause any memory issues.
Is there any way to reduce the memory consumption in this augmenter?
The code that I am using is really simple:
from textattack.augmentation.recipes import CLAREAugmenter
augmenter = CLAREAugmenter(pct_words_to_swap=0.1)
augmented_texts= []
for text in texts:
augmented_texts.append(augmenter.augment(text))
Activity
tsalmonking commentedon Dec 11, 2025
I found that what causes this is the fact that recipes like CLARE that use TensorFlow end up allocating the entire memory when not using parallelization, as memory growth is only set then.
Simply applying this before loading the augmentation should fix it:
TextAttack/textattack/attacker.py
Line 551 in 0d0929c
qiyanjun commentedon Aug 13, 2026
@tsalmonking thanks for looking into this, but I don't think the TensorFlow memory-growth theory holds up for the repro you posted:
CLAREAugmenteronly loads a PyTorch/HFdistilroberta-basemodel (WordSwapMaskedLM/WordInsertionMaskedLM/WordMergeMaskedLMviatransformers.AutoModelForCausalLM) — it never touches TensorFlow. Theset_memory_growthcall inattacker.pyonly matters for constraints likeUniversalSentenceEncoderthat load TF, and only in the--parallelmultiprocessing path, neither of which is in play in your snippet. Applying that fix wouldn't change anything for this repro.More likely culprits for the ~40GB:
CLAREAugmenterbuilds two full copies of the masked-LM (shared_masked_lmis loaded once but referenced across aCompositeTransformationof three sub-transformations, each generating candidate sets), and with ~800-token inputs the candidate-generation matrices scale with sequence length. If you can share GPU memory usage broken down by call (or run withpct_words_to_swaplower / on a truncated input) that would help narrow it down further. Leaving this open — this needs actual profiling, not a one-line fix.