Reducing GPU utilization while working with CLAREAugmenter

[dkoterwa]

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))

[tsalmonking]

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

tf.config.experimental.set_memory_growth(gpu, True)

[qiyanjun]

@tsalmonking thanks for looking into this, but I don't think the TensorFlow memory-growth theory holds up for the repro you posted: CLAREAugmenter only loads a PyTorch/HF distilroberta-base model (WordSwapMaskedLM / WordInsertionMaskedLM / WordMergeMaskedLM via transformers.AutoModelForCausalLM) — it never touches TensorFlow. The set_memory_growth call in attacker.py only matters for constraints like UniversalSentenceEncoder that load TF, and only in the --parallel multiprocessing 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: CLAREAugmenter builds two full copies of the masked-LM (shared_masked_lm is loaded once but referenced across a CompositeTransformation of 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 with pct_words_to_swap lower / on a truncated input) that would help narrow it down further. Leaving this open — this needs actual profiling, not a one-line fix.