OOM error while training Rasa/LaBSE

[aamra_work (aamra)]

I want to use the Language Model Featurizer specified in the docs while training.

My config looks like:

# The config recipe.
# https://rasa.com/docs/rasa/model-configuration/
recipe: default.v1

# Configuration for Rasa NLU.
# https://rasa.com/docs/rasa/nlu/components/
language: en

pipeline:
# # No configuration for the NLU pipeline was provided. The following default pipeline was used to train your model.
# # If you'd like to customize it, uncomment and adjust the pipeline.
# # See https://rasa.com/docs/rasa/tuning-your-model for more information.
   - name: WhitespaceTokenizer
#   - name: RegexFeaturizer
#   - name: LexicalSyntacticFeaturizer
   - name: LanguageModelFeaturizer
     # Name of the language model to use
     model_name: "bert"
     # Pre-Trained weights to be loaded
     model_weights: "rasa/LaBSE"
     cache_dir: null
   - name: CountVectorsFeaturizer
   - name: CountVectorsFeaturizer
     analyzer: char_wb
     min_ngram: 1
     max_ngram: 4
   - name: DIETClassifier
     epochs: 100
     constrain_similarities: true
     batch_size: 8
   - name: EntitySynonymMapper
   - name: ResponseSelector
     epochs: 100
     constrain_similarities: true
   - name: FallbackClassifier
     threshold: 0.3
     ambiguity_threshold: 0.1

Everything else is unchanged. But when I run rasa train, I get an Out of Memory error:

tensorflow.python.framework.errors_impl.ResourceExhaustedError: failed to allocate memory [Op:AddV2]

I am using GTX 1660 ti with 8GB Video memory. I did not change the default training data or stories.

I have not found anyone else with the same issues. My system requirements are:

rasa: 3.0.8
rasa-sdk: 3.0.5
OS: Ubuntu 18.04.6 LTS x86_64
Kernel: 5.4.0-113-generic
CUDA Version: 11.4
Driver Version: 470.57.02
-----------------------------------------------
TENSORFLOW:
tensorboard             2.8.0
tensorboard-data-server 0.6.1
tensorboard-plugin-wit  1.8.1
tensorflow              2.6.1
tensorflow-addons       0.14.0
tensorflow-estimator    2.6.0
tensorflow-hub          0.12.0
tensorflow-probability  0.13.0
tensorflow-text         2.6.0

[kunal-kumar-chaudhary (kunal kumar)]

try adjusting the batch_size, transformer_size, number_of_transformer_layers. make changes to these under the DIETClassifier section in the config.yml file. the problem here is “out of memory”, I resolved it using tweaking the above-mentioned parameters. The below link will help you with the tweak!

Components